-
Notifications
You must be signed in to change notification settings - Fork 668
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
update inverse kinematics to support multiple sites #399
base: main
Are you sure you want to change the base?
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Do we need the additional index here? I think np.squeeze
is a view.
@@ -207,8 +231,18 @@ def qpos_from_site_pose(physics, | |||
# Update `physics.qpos`, taking quaternions into account. | |||
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Please remove this comment block.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you for your contribution. I left a couple of comments.
@@ -35,7 +35,7 @@ | |||
|
|||
|
|||
def qpos_from_site_pose(physics, | |||
site_name, | |||
sites_names, |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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
.
@@ -170,24 +177,41 @@ def qpos_from_site_pose(physics, | |||
mjlib.mju_quat2Vel(err_rot, err_rot_quat, 1) | |||
err_norm += np.linalg.norm(err_rot) * rot_weight | |||
|
|||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Remove
""" | ||
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) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
indent 2 more spaces. Same for other line continuations below.
@@ -46,12 +46,13 @@ def qpos_from_site_pose(physics, | |||
max_update_norm=2.0, | |||
progress_thresh=20.0, | |||
max_steps=100, | |||
inplace=False): | |||
inplace=False, | |||
null_space_method=True): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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 |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
max_steps
@@ -90,7 +139,7 @@ def testQposFromSitePose(self, target, inplace): | |||
while True: | |||
result = ik.qpos_from_site_pose( | |||
physics=physics2, | |||
site_name=_SITE_NAME, | |||
sites_names=[_SITE_NAME], |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.