Merge pull request #430 from zswang666/feature/user_defined_ik_chain

[feature] support custom ik chain

Author: YilingQiao

examples/rigid/ik_custom_chain.py

@@ -0,0 +1,89 @@
+import argparse
+
+import numpy as np
+
+import genesis as gs
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-v", "--vis", action="store_true", default=False)
+    args = parser.parse_args()
+
+    ########################## init ##########################
+    gs.init(seed=0, precision="32", logging_level="debug")
+
+    ########################## create a scene ##########################
+    scene = gs.Scene(
+        viewer_options=gs.options.ViewerOptions(
+            camera_pos=(2.5, 0.0, 1.5),
+            camera_lookat=(0.0, 0.0, 0.5),
+            camera_fov=40,
+        ),
+        show_viewer=args.vis,
+        rigid_options=gs.options.RigidOptions(
+            gravity=(0, 0, 0),
+            enable_collision=False,
+            enable_joint_limit=False,
+        ),
+    )
+
+    target_1 = scene.add_entity(
+        gs.morphs.Mesh(
+            file="meshes/axis.obj",
+            scale=0.05,
+        ),
+        surface=gs.surfaces.Default(color=(1, 0.5, 0.5, 1)),
+    )
+
+    ########################## entities ##########################
+    robot = scene.add_entity(
+        morph=gs.morphs.URDF(
+            scale=1.0,
+            file="urdf/shadow_hand/shadow_hand.urdf",
+        ),
+        surface=gs.surfaces.Reflective(color=(0.4, 0.4, 0.4)),
+    )
+
+    ########################## build ##########################
+    scene.build()
+    scene.reset()
+
+    target_quat = np.array([1, 0, 0, 0])
+    index_finger_distal = robot.get_link("index_finger_distal")
+
+    dofs_idx_local = []
+    for v in robot.joints:
+        if v.name in [
+            "wrist_joint",
+            "index_finger_joint1",
+            "index_finger_join2",
+            "index_finger_joint3",
+        ]:
+            dof_idx_local_v = v.dof_idx_local
+            if isinstance(dof_idx_local_v, list):
+                dofs_idx_local.extend(dof_idx_local_v)
+            else:
+                assert isinstance(dof_idx_local_v, int)
+                dofs_idx_local.append(dof_idx_local_v)
+
+    center = np.array([0.033, -0.01, 0.42])
+    r1 = 0.05
+
+    for i in range(2000):
+        index_finger_pos = center + np.array([0, np.cos(i / 90 * np.pi) - 1.0, np.sin(i / 90 * np.pi) - 1.0]) * r1
+
+        target_1.set_qpos(np.concatenate([index_finger_pos, target_quat]))
+
+        qpos = robot.inverse_kinematics_multilink(
+            links=[index_finger_distal],  # IK targets
+            poss=[index_finger_pos],
+            dofs_idx_local=dofs_idx_local,  # IK wrt these dofs
+        )
+
+        robot.set_qpos(qpos)
+        scene.step()
+
+
+if __name__ == "__main__":
+    main()

genesis/engine/entities/rigid_entity/rigid_entity.py

@@ -738,6 +738,7 @@ def inverse_kinematics(
         pos_mask=[True, True, True],
         rot_mask=[True, True, True],
         max_step_size=0.5,
+        dofs_idx_local=None,
         return_error=False,
     ):
         """
@@ -771,6 +772,8 @@ def inverse_kinematics(
             Mask for rotation axis alignment. Defaults to [True, True, True]. E.g.: If you only want the link's Z-axis to be aligned with the Z-axis in the given quat, you can set it to [False, False, True].
         max_step_size : float, optional
             Maximum step size in q space for each IK solver step. Defaults to 0.5.
+        dofs_idx_local : None | array_like, optional
+            The indices of the dofs to set. If None, all dofs will be set. Note that here this uses the local `q_idx`, not the scene-level one. Defaults to None. This is used to specify which dofs the IK is applied to.
         return_error : bool, optional
             Whether to return the final errorqpos. Defaults to False.
 
@@ -803,6 +806,7 @@ def inverse_kinematics(
             pos_mask=pos_mask,
             rot_mask=rot_mask,
             max_step_size=max_step_size,
+            dofs_idx_local=dofs_idx_local,
             return_error=return_error,
         )
 
@@ -830,6 +834,7 @@ def inverse_kinematics_multilink(
         pos_mask=[True, True, True],
         rot_mask=[True, True, True],
         max_step_size=0.5,
+        dofs_idx_local=None,
         return_error=False,
     ):
         """
@@ -863,6 +868,8 @@ def inverse_kinematics_multilink(
             Mask for rotation axis alignment. Defaults to [True, True, True]. E.g.: If you only want the link's Z-axis to be aligned with the Z-axis in the given quat, you can set it to [False, False, True].
         max_step_size : float, optional
             Maximum step size in q space for each IK solver step. Defaults to 0.5.
+        dofs_idx_local : None | array_like, optional
+            The indices of the dofs to set. If None, all dofs will be set. Note that here this uses the local `q_idx`, not the scene-level one. Defaults to None. This is used to specify which dofs the IK is applied to.
         return_error : bool, optional
             Whether to return the final errorqpos. Defaults to False.
 
@@ -949,11 +956,33 @@ def inverse_kinematics_multilink(
             [self._solver._process_dim(torch.as_tensor(quat, dtype=gs.tc_float, device=gs.device)) for quat in quats]
         )
 
+        dofs_idx = self._get_dofs_idx_local(dofs_idx_local)
+        n_dofs = len(dofs_idx)
+        if n_dofs == 0:
+            gs.raise_exception("Target dofs not provided.")
+        links_idx_by_dofs = []
+        for v in self.links:
+            links_idx_by_dof_at_v = v.joint.dof_idx_local
+            if links_idx_by_dof_at_v is None:
+                link_relevant = False
+            elif isinstance(links_idx_by_dof_at_v, list):
+                link_relevant = any([vv in dofs_idx for vv in links_idx_by_dof_at_v])
+            else:
+                link_relevant = links_idx_by_dof_at_v in dofs_idx
+            if link_relevant:
+                links_idx_by_dofs.append(v.idx_local)  # converted to global later
+        links_idx_by_dofs = self._get_ls_idx(links_idx_by_dofs)
+        n_links_by_dofs = len(links_idx_by_dofs)
+
         self._kernel_inverse_kinematics(
             links_idx,
             poss,
             quats,
             n_links,
+            dofs_idx,
+            n_dofs,
+            links_idx_by_dofs,
+            n_links_by_dofs,
             custom_init_qpos,
             init_qpos,
             max_samples,
@@ -990,6 +1019,10 @@ def _kernel_inverse_kinematics(
         poss: ti.types.ndarray(),
         quats: ti.types.ndarray(),
         n_links: ti.i32,
+        dofs_idx: ti.types.ndarray(),
+        n_dofs: ti.i32,
+        links_idx_by_dofs: ti.types.ndarray(),
+        n_links_by_dofs: ti.i32,
         custom_init_qpos: ti.i32,
         init_qpos: ti.types.ndarray(),
         max_samples: ti.i32,
@@ -1063,30 +1096,40 @@ def _kernel_inverse_kinematics(
                     for i_ee in range(n_links):
                         # update jacobian for ee link
                         i_l_ee = links_idx[i_ee]
-                        self._func_get_jacobian(i_l_ee, i_b, pos_mask, rot_mask)
+                        self._func_get_jacobian(
+                            i_l_ee, i_b, pos_mask, rot_mask
+                        )  # NOTE: we still compute jacobian for all dofs as we haven't found a clean way to implement this
 
-                        # copy to multi-link jacobian
-                        for i_error, i_dof in ti.ndrange(6, self.n_dofs):
+                        # copy to multi-link jacobian (only for the effective n_dofs instead of self.n_dofs)
+                        for i_error, i_dof in ti.ndrange(6, n_dofs):
                             i_row = i_ee * 6 + i_error
-                            self._IK_jacobian[i_row, i_dof, i_b] = self._jacobian[i_error, i_dof, i_b]
+                            i_dof_ = dofs_idx[i_dof]
+                            self._IK_jacobian[i_row, i_dof, i_b] = self._jacobian[i_error, i_dof_, i_b]
 
-                    # compute dq = jac.T @ inverse(jac @ jac.T + diag) @ error
-                    lu.mat_transpose(self._IK_jacobian, self._IK_jacobian_T, n_error_dims, self.n_dofs, i_b)
+                    # compute dq = jac.T @ inverse(jac @ jac.T + diag) @ error (only for the effective n_dofs instead of self.n_dofs)
+                    lu.mat_transpose(self._IK_jacobian, self._IK_jacobian_T, n_error_dims, n_dofs, i_b)
                     lu.mat_mul(
                         self._IK_jacobian,
                         self._IK_jacobian_T,
                         self._IK_mat,
                         n_error_dims,
-                        self.n_dofs,
+                        n_dofs,
                         n_error_dims,
                         i_b,
                     )
                     lu.mat_add_eye(self._IK_mat, damping**2, n_error_dims, i_b)
                     lu.mat_inverse(self._IK_mat, self._IK_L, self._IK_U, self._IK_y, self._IK_inv, n_error_dims, i_b)
                     lu.mat_mul_vec(self._IK_inv, self._IK_err_pose, self._IK_vec, n_error_dims, n_error_dims, i_b)
-                    lu.mat_mul_vec(
-                        self._IK_jacobian_T, self._IK_vec, self._IK_delta_qpos, self.n_dofs, n_error_dims, i_b
-                    )
+
+                    for i in range(self.n_dofs):  # IK_delta_qpos = IK_jacobian_T @ IK_vec
+                        self._IK_delta_qpos[i, i_b] = 0
+                    for i in range(n_dofs):
+                        for j in range(n_error_dims):
+                            i_ = dofs_idx[
+                                i
+                            ]  # NOTE: IK_delta_qpos uses the original indexing instead of the effective n_dofs
+                            self._IK_delta_qpos[i_, i_b] += self._IK_jacobian_T[i, j, i_b] * self._IK_vec[j, i_b]
+
                     for i in range(self.n_dofs):
                         self._IK_delta_qpos[i, i_b] = ti.math.clamp(
                             self._IK_delta_qpos[i, i_b], -max_step_size, max_step_size
@@ -1159,7 +1202,8 @@ def _kernel_inverse_kinematics(
 
                     # Resample init q
                     if respect_joint_limit and i_sample < max_samples - 1:
-                        for i_l in range(self.link_start, self.link_end):
+                        for _i_l in range(n_links_by_dofs):
+                            i_l = links_idx_by_dofs[_i_l]
                             I_l = [i_l, i_b] if ti.static(self.solver._options.batch_links_info) else i_l
                             l_info = self._solver.links_info[I_l]
                             I_dof_start = (