update inverse kinematics to support multiple sites#399
update inverse kinematics to support multiple sites#399vittorione94 wants to merge 2 commits intogoogle-deepmind:mainfrom
Conversation
|
Could you please put test assets in their own directory, similar to https://github.com/deepmind/dm_control/tree/main/dm_control/mujoco/testing/assets ? |
|
Unlike the original single-site version, looks like the multiple site generalization doesn't support simultaneous position and orientation fitting, does it? |
|
@vaxenburg you're right the multiple site generalization doesn't support simultaneous position and orientation fitting. |
|
I just tested this and can confirm that this works as expected at least with the default parameters taken from |
sbohez
left a comment
There was a problem hiding this comment.
Few small comments but looks good otherwise!
| # TODO(b/112141592): This does not take joint limits into consideration. | ||
| reg_strength = ( | ||
| if len(site_ids) > 1: | ||
| for idx in range(len(site_ids)): |
There was a problem hiding this comment.
Switch to
for idx, site_id in enumerate(site_ids):
...
| # clip joint angles to their respective limits | ||
| if len(sites_names) > 1: | ||
| joint_names = physics.named.data.qpos.axes.row.names | ||
| limited_joints = joint_names[1:] # ignore root joint |
There was a problem hiding this comment.
You probably want to explicitly check which joints are actually limited using jnt_limited.
|
|
||
| # Compute the new Cartesian position of the site. | ||
| mjlib.mj_fwdPosition(physics.model.ptr, physics.data.ptr) | ||
| site_xpos = np.squeeze(physics.named.data.site_xpos[sites_names]) |
There was a problem hiding this comment.
Do we need the additional index here? I think np.squeeze is a view.
| mjlib.mj_integratePos(physics.model.ptr, physics.data.qpos, update_nv, 1) | ||
|
|
||
| # clip joint angles to their respective limits | ||
| if len(sites_names) > 1: |
There was a problem hiding this comment.
Why only when there are multiple site names?
|
|
||
| env = control.Environment(physics, Task(mocap_qpos=[mocap_qpos])) | ||
|
|
||
| # env = control.Environment(physics, Task(mocap_qpos=[[-0.04934, -0.00198, 1.25512, 0.99691, 0.0161, -0.04859, -0.05959, |
There was a problem hiding this comment.
Please remove this comment block.
nimrod-gileadi
left a comment
There was a problem hiding this comment.
Thank you for your contribution. I left a couple of comments.
|
|
||
| def qpos_from_site_pose(physics, | ||
| site_name, | ||
| sites_names, |
There was a problem hiding this comment.
This is a widely used function, that we don't want to break API compatibility on.
Can you instead add a new function, qpos_from_site_poses which takes a parameter (called site_names rather than sites_names, BTW)?
The existing qpos_from_site_pose function could call down to the new qpos_from_site_poses.
| mjlib.mju_quat2Vel(err_rot, err_rot_quat, 1) | ||
| err_norm += np.linalg.norm(err_rot) * rot_weight | ||
|
|
||
|
|
| """ | ||
| JJ_t = jac_joints.dot(jac_joints.T) | ||
| JJ_t += np.eye(JJ_t.shape[0]) * regularization_strength | ||
| return jac_joints.T.dot(np.linalg.inv(JJ_t)).dot(delta) No newline at end of file |
There was a problem hiding this comment.
Add newline at the end of the file.
| for idx in range(len(site_ids)): | ||
| site_id = site_ids[idx] | ||
| mjlib.mj_jacSite( | ||
| physics.model.ptr, physics.data.ptr, jac_pos[idx], jac_rot, site_id) |
There was a problem hiding this comment.
indent 2 more spaces. Same for other line continuations below.
| max_steps=100, | ||
| inplace=False): | ||
| inplace=False, | ||
| null_space_method=True): |
There was a problem hiding this comment.
What do you think about not using a boolean variable, and instead having an enum for regularization_method that can be NULLSPACE_METHOD or DAMPED_LEAST_SQUARES? That would allow us to add more methosd in the future without breaking the API.
BTW: I think it's confusing that there's a function called nullspace_method and a variable called null_space_method. If we stick with the boolean, name the variable use_nullspace_method.
Same below.
| target_pos = target_pos.reshape((-1, 3)) | ||
| target_quat = None | ||
|
|
||
| _SITES_NAMES = [] |
There was a problem hiding this comment.
This naming is reserved for module-wide constants. Name it site_names.
| "lshoulder", "rshoulder", "lhip", "rhip", | ||
| ] | ||
|
|
||
| for name in body_names: |
There was a problem hiding this comment.
Maybe just use a list comprehension, e.g.:
site_names = [f"tracking[{name}]" for name in body_names]
| for name in body_names: | ||
| _SITES_NAMES.append("tracking[" + name + "]") | ||
|
|
||
| _MAX_STEPS = 5000 |
| result = ik.qpos_from_site_pose( | ||
| physics=physics2, | ||
| site_name=_SITE_NAME, | ||
| sites_names=[_SITE_NAME], |
There was a problem hiding this comment.
The tests should prove that the existing functions continue to work.
Maybe add a parameterized annotation that will run both qpos_from_site_pose and qpos_from_site_poses on the same input.
I have updated the inverse kinematics script to support multiple sites. It still supports the previous script and all the tests run correctly. There is also a new visualise script that just visualise the output position in the humanoid.