Inverse kinematics

.gitignore

@@ -1,4 +1,5 @@
 .idea
+*.mp4
 *.so
 *.pkl
 *.egg-info

pyproject.toml

@@ -46,10 +46,16 @@ dependencies = [
     'numpy',
     'pyyaml',
     'torch',
+    'matplotlib',
+    'pytorch_seed',
 ]
 
 [project.optional-dependencies]
-test = ["pytest"]
+test = [
+    "pytest",
+    "pybullet",
+    "jax",
+]
 
 [project.urls]
 "Homepage" = "https://github.com/UM-ARM-Lab/pytorch_kinematics"

src/pytorch_kinematics/__init__.py

@@ -3,3 +3,4 @@
 from pytorch_kinematics.mjcf import *
 from pytorch_kinematics.transforms import *
 from pytorch_kinematics.chain import *
+from pytorch_kinematics.ik import *

src/pytorch_kinematics/chain.py

@@ -458,10 +458,10 @@ def _generate_serial_chain_recurse(root_frame, end_frame_name):
                     return [child] + frames
         return None
 
-    def jacobian(self, th, locations=None):
+    def jacobian(self, th, locations=None, **kwargs):
         if locations is not None:
             locations = tf.Transform3d(pos=locations)
-        return jacobian.calc_jacobian(self, th, tool=locations)
+        return jacobian.calc_jacobian(self, th, tool=locations, **kwargs)
 
     def forward_kinematics(self, th, end_only: bool = True):
         """ Like the base class, except `th` only needs to contain the joints in the SerialChain, not all joints. """

src/pytorch_kinematics/jacobian.py

@@ -3,7 +3,7 @@
 from pytorch_kinematics import transforms
 
 
-def calc_jacobian(serial_chain, th, tool=None):
+def calc_jacobian(serial_chain, th, tool=None, ret_eef_pose=False):
     """
     Return robot Jacobian J in base frame (N,6,DOF) where dot{x} = J dot{q}
     The first 3 rows relate the translational velocities and the
@@ -57,4 +57,6 @@ def calc_jacobian(serial_chain, th, tool=None):
     j_tr[:, :3, :3] = rotation
     j_tr[:, 3:, 3:] = rotation
     j_w = j_tr @ j_eef
+    if ret_eef_pose:
+        return j_w, pose
     return j_w

tests/meshes/cone.mtl

@@ -0,0 +1,2 @@
+# Blender 3.6.4 MTL File: 'ycb.blend'
+# www.blender.org

tests/meshes/cone.obj

@@ -0,0 +1,169 @@
+# Blender 3.6.4
+# www.blender.org
+mtllib cone.mtl
+o Cone
+v 0.097681 0.000000 -0.097681
+v 0.095804 -0.019057 -0.097681
+v 0.090246 -0.037381 -0.097681
+v 0.081219 -0.054269 -0.097681
+v 0.069071 -0.069071 -0.097681
+v 0.054269 -0.081219 -0.097681
+v 0.037381 -0.090246 -0.097681
+v 0.019057 -0.095804 -0.097681
+v 0.000000 -0.097681 -0.097681
+v -0.019057 -0.095804 -0.097681
+v -0.037381 -0.090246 -0.097681
+v -0.054269 -0.081219 -0.097681
+v -0.069071 -0.069071 -0.097681
+v -0.081219 -0.054269 -0.097681
+v -0.090246 -0.037381 -0.097681
+v -0.095804 -0.019057 -0.097681
+v -0.097681 0.000000 -0.097681
+v -0.095804 0.019057 -0.097681
+v -0.090246 0.037381 -0.097681
+v -0.081219 0.054269 -0.097681
+v -0.069071 0.069071 -0.097681
+v -0.054269 0.081219 -0.097681
+v -0.037381 0.090246 -0.097681
+v -0.019057 0.095804 -0.097681
+v 0.000000 0.097681 -0.097681
+v 0.019057 0.095804 -0.097681
+v 0.037381 0.090246 -0.097681
+v 0.054269 0.081219 -0.097681
+v 0.069071 0.069071 -0.097681
+v 0.081219 0.054269 -0.097681
+v 0.090246 0.037381 -0.097681
+v 0.095804 0.019057 -0.097681
+v 0.000000 0.000000 0.097681
+vn 0.8910 -0.0878 0.4455
+vn 0.8567 -0.2599 0.4455
+vn 0.7896 -0.4220 0.4455
+vn 0.6921 -0.5680 0.4455
+vn 0.5680 -0.6921 0.4455
+vn 0.4220 -0.7896 0.4455
+vn 0.2599 -0.8567 0.4455
+vn 0.0878 -0.8910 0.4455
+vn -0.0878 -0.8910 0.4455
+vn -0.2599 -0.8567 0.4455
+vn -0.4220 -0.7896 0.4455
+vn -0.5680 -0.6921 0.4455
+vn -0.6921 -0.5680 0.4455
+vn -0.7896 -0.4220 0.4455
+vn -0.8567 -0.2599 0.4455
+vn -0.8910 -0.0878 0.4455
+vn -0.8910 0.0878 0.4455
+vn -0.8567 0.2599 0.4455
+vn -0.7896 0.4220 0.4455
+vn -0.6921 0.5680 0.4455
+vn -0.5680 0.6921 0.4455
+vn -0.4220 0.7896 0.4455
+vn -0.2599 0.8567 0.4455
+vn -0.0878 0.8910 0.4455
+vn 0.0878 0.8910 0.4455
+vn 0.2599 0.8567 0.4455
+vn 0.4220 0.7896 0.4455
+vn 0.5680 0.6921 0.4455
+vn 0.6921 0.5680 0.4455
+vn 0.7896 0.4220 0.4455
+vn -0.0000 -0.0000 -1.0000
+vn 0.8567 0.2599 0.4455
+vn 0.8910 0.0878 0.4455
+vt 0.250000 0.490000
+vt 0.250000 0.250000
+vt 0.296822 0.485388
+vt 0.341844 0.471731
+vt 0.383337 0.449553
+vt 0.419706 0.419706
+vt 0.449553 0.383337
+vt 0.471731 0.341844
+vt 0.485388 0.296822
+vt 0.490000 0.250000
+vt 0.485388 0.203178
+vt 0.471731 0.158156
+vt 0.449553 0.116663
+vt 0.419706 0.080294
+vt 0.383337 0.050447
+vt 0.341844 0.028269
+vt 0.296822 0.014612
+vt 0.250000 0.010000
+vt 0.203178 0.014612
+vt 0.158156 0.028269
+vt 0.116663 0.050447
+vt 0.080294 0.080294
+vt 0.050447 0.116663
+vt 0.028269 0.158156
+vt 0.014612 0.203178
+vt 0.010000 0.250000
+vt 0.014612 0.296822
+vt 0.028269 0.341844
+vt 0.050447 0.383337
+vt 0.080294 0.419706
+vt 0.116663 0.449553
+vt 0.158156 0.471731
+vt 0.750000 0.490000
+vt 0.796822 0.485388
+vt 0.841844 0.471731
+vt 0.883337 0.449553
+vt 0.919706 0.419706
+vt 0.949553 0.383337
+vt 0.971731 0.341844
+vt 0.985388 0.296822
+vt 0.990000 0.250000
+vt 0.985388 0.203178
+vt 0.971731 0.158156
+vt 0.949553 0.116663
+vt 0.919706 0.080294
+vt 0.883337 0.050447
+vt 0.841844 0.028269
+vt 0.796822 0.014612
+vt 0.750000 0.010000
+vt 0.703178 0.014612
+vt 0.658156 0.028269
+vt 0.616663 0.050447
+vt 0.580294 0.080294
+vt 0.550447 0.116663
+vt 0.528269 0.158156
+vt 0.514612 0.203178
+vt 0.510000 0.250000
+vt 0.514612 0.296822
+vt 0.528269 0.341844
+vt 0.550447 0.383337
+vt 0.580294 0.419706
+vt 0.616663 0.449553
+vt 0.658156 0.471731
+vt 0.703178 0.485388
+vt 0.203178 0.485388
+s 0
+f 1/1/1 33/2/1 2/3/1
+f 2/3/2 33/2/2 3/4/2
+f 3/4/3 33/2/3 4/5/3
+f 4/5/4 33/2/4 5/6/4
+f 5/6/5 33/2/5 6/7/5
+f 6/7/6 33/2/6 7/8/6
+f 7/8/7 33/2/7 8/9/7
+f 8/9/8 33/2/8 9/10/8
+f 9/10/9 33/2/9 10/11/9
+f 10/11/10 33/2/10 11/12/10
+f 11/12/11 33/2/11 12/13/11
+f 12/13/12 33/2/12 13/14/12
+f 13/14/13 33/2/13 14/15/13
+f 14/15/14 33/2/14 15/16/14
+f 15/16/15 33/2/15 16/17/15
+f 16/17/16 33/2/16 17/18/16
+f 17/18/17 33/2/17 18/19/17
+f 18/19/18 33/2/18 19/20/18
+f 19/20/19 33/2/19 20/21/19
+f 20/21/20 33/2/20 21/22/20
+f 21/22/21 33/2/21 22/23/21
+f 22/23/22 33/2/22 23/24/22
+f 23/24/23 33/2/23 24/25/23
+f 24/25/24 33/2/24 25/26/24
+f 25/26/25 33/2/25 26/27/25
+f 26/27/26 33/2/26 27/28/26
+f 27/28/27 33/2/27 28/29/27
+f 28/29/28 33/2/28 29/30/28
+f 29/30/29 33/2/29 30/31/29
+f 30/31/30 33/2/30 31/32/30
+f 1/33/31 2/34/31 3/35/31 4/36/31 5/37/31 6/38/31 7/39/31 8/40/31 9/41/31 10/42/31 11/43/31 12/44/31 13/45/31 14/46/31 15/47/31 16/48/31 17/49/31 18/50/31 19/51/31 20/52/31 21/53/31 22/54/31 23/55/31 24/56/31 25/57/31 26/58/31 27/59/31 28/60/31 29/61/31 30/62/31 31/63/31 32/64/31
+f 31/32/32 33/2/32 32/65/32
+f 32/65/33 33/2/33 1/1/33

tests/test_inverse_kinematics.py

@@ -0,0 +1,136 @@
+import os
+from timeit import default_timer as timer
+
+import torch
+
+import pytorch_kinematics as pk
+import pytorch_seed
+
+import pybullet as p
+import pybullet_data
+
+visualize = True
+
+
+def _make_robot_translucent(robot_id, alpha=0.4):
+    def make_transparent(link):
+        link_id = link[1]
+        rgba = list(link[7])
+        rgba[3] = alpha
+        p.changeVisualShape(robot_id, link_id, rgbaColor=rgba)
+
+    visual_data = p.getVisualShapeData(robot_id)
+    for link in visual_data:
+        make_transparent(link)
+
+
+def test_jacobian_follower():
+    pytorch_seed.seed(2)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    # device = "cpu"
+    urdf = "kuka_iiwa/model.urdf"
+    search_path = pybullet_data.getDataPath()
+    full_urdf = os.path.join(search_path, urdf)
+    chain = pk.build_serial_chain_from_urdf(open(full_urdf).read(), "lbr_iiwa_link_7")
+    chain = chain.to(device=device)
+
+    # robot frame
+    pos = torch.tensor([0.0, 0.0, 0.0], device=device)
+    rot = torch.tensor([0.0, 0.0, 0.0], device=device)
+    rob_tf = pk.Transform3d(pos=pos, rot=rot, device=device)
+
+    # world frame goal
+    M = 1000
+    # generate random goal joint angles (so these are all achievable)
+    # use the joint limits to generate random joint angles
+    lim = torch.tensor(chain.get_joint_limits(), device=device)
+    goal_q = torch.rand(M, 7, device=device) * (lim[1] - lim[0]) + lim[0]
+
+    # get ee pose (in robot frame)
+    goal_in_rob_frame_tf = chain.forward_kinematics(goal_q)
+
+    # transform to world frame for visualization
+    goal_tf = rob_tf.compose(goal_in_rob_frame_tf)
+    goal = goal_tf.get_matrix()
+    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,
+                            joint_limits=lim.T,
+                            early_stopping_any_converged=True,
+                            early_stopping_no_improvement="all",
+                            # line_search=pk.BacktrackingLineSearch(max_lr=0.2),
+                            debug=False,
+                            lr=0.2)
+
+    # do IK
+    timer_start = timer()
+    sol = ik.solve(goal_in_rob_frame_tf)
+    timer_end = timer()
+    print("IK took %f seconds" % (timer_end - timer_start))
+    print("IK converged number: %d / %d" % (sol.converged.sum(), sol.converged.numel()))
+    print("IK took %d iterations" % sol.iterations)
+    print("IK solved %d / %d goals" % (sol.converged_any.sum(), M))
+
+    # visualize everything
+    if visualize:
+        p.connect(p.GUI)
+        p.setRealTimeSimulation(False)
+        p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
+        p.setAdditionalSearchPath(search_path)
+
+        yaw = 90
+        pitch = -55
+        dist = 1.
+        target = goal_pos[0].cpu().numpy()
+        p.resetDebugVisualizerCamera(dist, yaw, pitch, target)
+
+        p.loadURDF("plane.urdf", [0, 0, 0], useFixedBase=True)
+        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
+        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)
+
+        while True:
+            p.stepSimulation()
+
+
+if __name__ == "__main__":
+    test_jacobian_follower()