From 5ef465fd6b300b29c9b998fa5358f6fe6b8f7ec1 Mon Sep 17 00:00:00 2001 From: Johnson Wang Date: Wed, 1 Jan 2025 11:46:26 -0500 Subject: [PATCH] add custom ik chain --- examples/rigid/ik_custom_chain.py | 89 +++++++++++++++++++ .../entities/rigid_entity/rigid_entity.py | 66 +++++++++++--- 2 files changed, 144 insertions(+), 11 deletions(-) create mode 100644 examples/rigid/ik_custom_chain.py diff --git a/examples/rigid/ik_custom_chain.py b/examples/rigid/ik_custom_chain.py new file mode 100644 index 00000000..67d18e3d --- /dev/null +++ b/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() diff --git a/genesis/engine/entities/rigid_entity/rigid_entity.py b/genesis/engine/entities/rigid_entity/rigid_entity.py index 955af57c..f4f5dfa4 100644 --- a/genesis/engine/entities/rigid_entity/rigid_entity.py +++ b/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 = (