From 4e5005ef9ebea1461a3cabc3f0c2b00d19d65970 Mon Sep 17 00:00:00 2001 From: Sheng Zhong Date: Wed, 21 Aug 2024 17:21:13 -0400 Subject: [PATCH] Add IK animation and visualize IK for parallel goals and configs --- README.md | 2 + tests/test_inverse_kinematics.py | 98 +++++++++++++++++++------------- 2 files changed, 62 insertions(+), 38 deletions(-) diff --git a/README.md b/README.md index eefb0d6..7a668cb 100644 --- a/README.md +++ b/README.md @@ -174,6 +174,8 @@ Compared to other IK libraries, these are the typical advantages over them: - batched in both goal specification and retries from different starting configurations - goal orientation in addition to goal position +![IK](https://i.imgur.com/QgaUME9.gif) + See `tests/test_inverse_kinematics.py` for usage, but generally what you need is below: ```python full_urdf = os.path.join(search_path, urdf) diff --git a/tests/test_inverse_kinematics.py b/tests/test_inverse_kinematics.py index 7fd6547..d522d4e 100644 --- a/tests/test_inverse_kinematics.py +++ b/tests/test_inverse_kinematics.py @@ -1,6 +1,7 @@ import os from timeit import default_timer as timer +import numpy as np import torch import pytorch_kinematics as pk @@ -55,7 +56,8 @@ def test_jacobian_follower(): goal_pos = goal[..., :3, 3] goal_rot = pk.matrix_to_euler_angles(goal[..., :3, :3], "XYZ") - ik = pk.PseudoInverseIK(chain, max_iterations=30, num_retries=10, + num_retries = 10 + ik = pk.PseudoInverseIK(chain, max_iterations=30, num_retries=num_retries, joint_limits=lim.T, early_stopping_any_converged=True, early_stopping_no_improvement="all", @@ -85,53 +87,73 @@ def test_jacobian_follower(): p.setAdditionalSearchPath(search_path) yaw = 90 - pitch = -55 - dist = 1. - target = goal_pos[0].cpu().numpy() + pitch = -65 + # dist = 1. + dist = 2.4 + target = np.array([2., 1.5, 0]) p.resetDebugVisualizerCamera(dist, yaw, pitch, target) - p.loadURDF("plane.urdf", [0, 0, 0], useFixedBase=True) + plane_id = p.loadURDF("plane.urdf", [0, 0, 0], useFixedBase=True) + p.changeVisualShape(plane_id, -1, rgbaColor=[0.3, 0.3, 0.3, 1]) + + # make 1 per retry with positional offsets + robots = [] + num_robots = 16 + # 4x4 grid position offset + offset = 1.0 m = rob_tf.get_matrix() pos = m[0, :3, 3] rot = m[0, :3, :3] quat = pk.matrix_to_quaternion(rot) pos = pos.cpu().numpy() rot = pk.wxyz_to_xyzw(quat).cpu().numpy() - armId = p.loadURDF(urdf, basePosition=pos, baseOrientation=rot, useFixedBase=True) - _make_robot_translucent(armId, alpha=0.6) - # p.resetBasePositionAndOrientation(armId, [0, 0, 0], [0, 0, 0, 1]) - # draw goal - # place a translucent sphere at the goal + for i in range(num_robots): + this_offset = np.array([i % 4 * offset, i // 4 * offset, 0]) + armId = p.loadURDF(urdf, basePosition=pos + this_offset, baseOrientation=rot, useFixedBase=True) + # _make_robot_translucent(armId, alpha=0.6) + robots.append({"id": armId, "offset": this_offset, "pos": pos}) + show_max_num_retries_per_goal = 10 - for goal_num in range(M): - # draw cone to indicate pose instead of sphere - visId = p.createVisualShape(p.GEOM_MESH, fileName="meshes/cone.obj", meshScale=1.0, - rgbaColor=[0., 1., 0., 0.5]) - # visId = p.createVisualShape(p.GEOM_SPHERE, radius=0.05, rgbaColor=[0., 1., 0., 0.5]) - r = goal_rot[goal_num] - xyzw = pk.wxyz_to_xyzw(pk.matrix_to_quaternion(pk.euler_angles_to_matrix(r, "XYZ"))) - goalId = p.createMultiBody(baseMass=0, baseVisualShapeIndex=visId, - basePosition=goal_pos[goal_num].cpu().numpy(), - baseOrientation=xyzw.cpu().numpy()) - - solutions = sol.solutions[goal_num] - # sort based on if they converged - converged = sol.converged[goal_num] - idx = torch.argsort(converged.to(int), descending=True) - solutions = solutions[idx] - - # print how many retries converged for this one - print("Goal %d converged %d / %d" % (goal_num, converged.sum(), converged.numel())) - - for i, q in enumerate(solutions): - if i > show_max_num_retries_per_goal: - break - for dof in range(q.shape[0]): - p.resetJointState(armId, dof, q[dof]) - input("Press enter to continue") - - p.removeBody(goalId) + + goals = [] + # draw cone to indicate pose instead of sphere + visId = p.createVisualShape(p.GEOM_MESH, fileName="meshes/cone.obj", meshScale=1.0, + rgbaColor=[0., 1., 0., 0.5]) + for i in range(num_robots): + goals.append(p.createMultiBody(baseMass=0, baseVisualShapeIndex=visId)) + + try: + import window_recorder + with window_recorder.WindowRecorder(save_dir="."): + # batch over goals with num_robots + for j in range(0, M, num_robots): + this_selection = slice(j, j + num_robots) + r = goal_rot[this_selection] + xyzw = pk.wxyz_to_xyzw(pk.matrix_to_quaternion(pk.euler_angles_to_matrix(r, "XYZ"))) + + solutions = sol.solutions[this_selection, :, :] + converged = sol.converged[this_selection, :] + + # print how many retries converged for this one + print("Goal %d to %d converged %d / %d" % (j, j + num_robots, converged.sum(), converged.numel())) + + # outer loop over retries, inner loop over goals (for each robot shown in parallel) + for ii in range(num_retries): + if ii > show_max_num_retries_per_goal: + break + for jj in range(num_robots): + p.resetBasePositionAndOrientation(goals[jj], + goal_pos[j + jj].cpu().numpy() + robots[jj]["offset"], + xyzw[jj].cpu().numpy()) + armId = robots[jj]["id"] + q = solutions[jj, ii, :] + for dof in range(q.shape[0]): + p.resetJointState(armId, dof, q[dof]) + + input("Press enter to continue") + except ImportError: + print("pip install window_recorder") while True: p.stepSimulation()