From 38bbd041bcdf787d490cf335a7f13bf5e830911e Mon Sep 17 00:00:00 2001 From: vittorione94 Date: Thu, 18 May 2023 15:56:54 +0300 Subject: [PATCH 1/2] update inverse kinematics to support multiple sites --- dm_control/utils/CMU-CMU-02-02_04_poses.xml | 5 + dm_control/utils/common.xml | 31 +++ dm_control/utils/humanoid.xml | 231 ++++++++++++++++++++ dm_control/utils/inverse_kinematics.py | 92 ++++++-- dm_control/utils/inverse_kinematics_test.py | 61 +++++- dm_control/utils/task.xml | 65 ++++++ dm_control/utils/visualize.py | 65 ++++++ 7 files changed, 527 insertions(+), 23 deletions(-) create mode 100644 dm_control/utils/CMU-CMU-02-02_04_poses.xml create mode 100644 dm_control/utils/common.xml create mode 100644 dm_control/utils/humanoid.xml create mode 100644 dm_control/utils/task.xml create mode 100644 dm_control/utils/visualize.py diff --git a/dm_control/utils/CMU-CMU-02-02_04_poses.xml b/dm_control/utils/CMU-CMU-02-02_04_poses.xml new file mode 100644 index 00000000..bd906724 --- /dev/null +++ b/dm_control/utils/CMU-CMU-02-02_04_poses.xml @@ -0,0 +1,5 @@ + + + + + diff --git a/dm_control/utils/common.xml b/dm_control/utils/common.xml new file mode 100644 index 00000000..8194b72e --- /dev/null +++ b/dm_control/utils/common.xml @@ -0,0 +1,31 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/dm_control/utils/humanoid.xml b/dm_control/utils/humanoid.xml new file mode 100644 index 00000000..ecb11b5c --- /dev/null +++ b/dm_control/utils/humanoid.xml @@ -0,0 +1,231 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/dm_control/utils/inverse_kinematics.py b/dm_control/utils/inverse_kinematics.py index 623d2ea9..fa0048f1 100644 --- a/dm_control/utils/inverse_kinematics.py +++ b/dm_control/utils/inverse_kinematics.py @@ -35,7 +35,7 @@ def qpos_from_site_pose(physics, - site_name, + sites_names, target_pos=None, target_quat=None, joint_names=None, @@ -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): """Find joint positions that satisfy a target site position and/or rotation. Args: physics: A `mujoco.Physics` instance. - site_name: A string specifying the name of the target site. + sites_names: A list of strings specifying the names of the target sites. target_pos: A (3,) numpy array specifying the desired Cartesian position of the site, or None if the position should be unconstrained (default). One or both of `target_pos` or `target_quat` must be specified. @@ -79,7 +80,8 @@ def qpos_from_site_pose(physics, max_steps: (optional) The maximum number of iterations to perform. inplace: (optional) If True, `physics.data` will be modified in place. Default value is False, i.e. a copy of `physics.data` will be made. - + null_space_method: (optional) If True uses the null space method to find the + update norm for the joint angles, otherwise uses the damped least squares. Returns: An `IKResult` namedtuple with the following fields: qpos: An (nq,) numpy array of joint positions. @@ -102,8 +104,12 @@ def qpos_from_site_pose(physics, jac_pos, jac_rot = jac[:3], jac[3:] err_pos, err_rot = err[:3], err[3:] else: - jac = np.empty((3, physics.model.nv), dtype=dtype) - err = np.empty(3, dtype=dtype) + if len(sites_names) > 1: + jac = np.empty((len(sites_names), 3, physics.model.nv), dtype=dtype) + err = np.empty((len(sites_names), 3), dtype=dtype) + else: + jac = np.empty((3, physics.model.nv), dtype=dtype) + err = np.empty(3, dtype=dtype) if target_pos is not None: jac_pos, jac_rot = jac, None err_pos, err_rot = err, None @@ -127,12 +133,13 @@ def qpos_from_site_pose(physics, mjlib.mj_fwdPosition(physics.model.ptr, physics.data.ptr) # Convert site name to index. - site_id = physics.model.name2id(site_name, 'site') + site_ids = [physics.model.name2id(site_name, 'site') + for site_name in sites_names] # These are views onto the underlying MuJoCo buffers. mj_fwdPosition will # update them in place, so we can avoid indexing overhead in the main loop. - site_xpos = physics.named.data.site_xpos[site_name] - site_xmat = physics.named.data.site_xmat[site_name] + site_xpos = np.squeeze(physics.named.data.site_xpos[sites_names]) + site_xmat = np.squeeze(physics.named.data.site_xmat[sites_names]) # This is an index into the rows of `update` and the columns of `jac` # that selects DOFs associated with joints that we are allowed to manipulate. @@ -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 + if err_norm < tol: logging.debug('Converged after %i steps: err_norm=%3g', steps, err_norm) success = True break else: # TODO(b/112141670): Generalize this to other entities besides sites. - mjlib.mj_jacSite( - physics.model.ptr, physics.data.ptr, jac_pos, jac_rot, site_id) - jac_joints = jac[:, dof_indices] - - # 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)): + site_id = site_ids[idx] + mjlib.mj_jacSite( + physics.model.ptr, physics.data.ptr, jac_pos[idx], jac_rot, site_id) + else: + mjlib.mj_jacSite( + physics.model.ptr, physics.data.ptr, jac_pos, jac_rot, site_ids[0]) + + if null_space_method: + jac_joints = jac[:, dof_indices] + + # TODO(b/112141592): This does not take joint limits into consideration. + reg_strength = ( regularization_strength if err_norm > regularization_threshold else 0.0) - update_joints = nullspace_method( + update_joints = nullspace_method( jac_joints, err, regularization_strength=reg_strength) - update_norm = np.linalg.norm(update_joints) + update_norm = np.linalg.norm(update_joints) + else: + update_joints = np.empty((len(site_ids), physics.model.nv)) + for idx in range(len(site_ids)): + update_joints[idx] = damped_least_squares(jac[idx], err[idx], + regularization_strength=regularization_strength) + + update_joints = np.mean(update_joints, axis=0) + update_norm = np.linalg.norm(update_joints) # Check whether we are still making enough progress, and halt if not. progress_criterion = err_norm / update_norm @@ -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: + joint_names = physics.named.data.qpos.axes.row.names + limited_joints = joint_names[1:] # ignore root joint + lower, upper = physics.named.model.jnt_range[limited_joints].T + physics.named.data.qpos[limited_joints] = np.clip(physics.named.data.qpos[limited_joints], + lower, upper) + # 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]) + site_xmat = np.squeeze(physics.named.data.site_xmat[sites_names]) logging.debug('Step %2i: err_norm=%-10.3g update_norm=%-10.3g', steps, err_norm, update_norm) @@ -258,3 +292,27 @@ def nullspace_method(jac_joints, delta, regularization_strength=0.0): return np.linalg.solve(hess_approx, joint_delta) else: return np.linalg.lstsq(hess_approx, joint_delta, rcond=-1)[0] + +def damped_least_squares(jac_joints, delta, regularization_strength): + """Calculates the joint velocities to achieve a specified end effector delta. + + Args: + jac_joints: The Jacobian of the end effector with respect to the joints. A + numpy array of shape `(ndelta, nv)`, where `ndelta` is the size of `delta` + and `nv` is the number of degrees of freedom. + delta: The desired end-effector delta. A numpy array of shape `(3,)` or + `(6,)` containing either position deltas, rotation deltas, or both. + regularization_strength: (optional) Coefficient of the quadratic penalty + on joint movements. Default is zero, i.e. no regularization. + + Returns: + An `(nv,)` numpy array of joint velocities. + + Reference: + Buss, S. R. S. (2004). Introduction to inverse kinematics with jacobian + transpose, pseudoinverse and damped least squares methods. + https://www.math.ucsd.edu/~sbuss/ResearchWeb/ikmethods/iksurvey.pdf + """ + 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 diff --git a/dm_control/utils/inverse_kinematics_test.py b/dm_control/utils/inverse_kinematics_test.py index 4c8898c8..fac25801 100644 --- a/dm_control/utils/inverse_kinematics_test.py +++ b/dm_control/utils/inverse_kinematics_test.py @@ -24,6 +24,7 @@ from dm_control.mujoco.wrapper import mjbindings from dm_control.utils import inverse_kinematics as ik import numpy as np +import os mjlib = mjbindings.mjlib @@ -80,6 +81,54 @@ def __call__(self, physics): class InverseKinematicsTest(parameterized.TestCase): + def testQposFromMultipleSitesPose(self): + dir_path = os.path.dirname(os.path.realpath(__file__)) + model_dir = os.path.join(dir_path, "./task.xml") + physics = mujoco.Physics.from_xml_path(model_dir) + + target_pos = physics.model.key_mpos[0] + target_pos = target_pos.reshape((-1, 3)) + target_quat = None + + _SITES_NAMES = [] + + body_names = [ + "pelvis", "head", "ltoe", "rtoe", "lheel", "rheel", + "lknee", "rknee", "lhand", "rhand", "lelbow", "relbow", + "lshoulder", "rshoulder", "lhip", "rhip", + ] + + for name in body_names: + _SITES_NAMES.append("tracking[" + name + "]") + + _MAX_STEPS = 5000 + result = ik.qpos_from_site_pose( + physics=physics, + sites_names=_SITES_NAMES, + target_pos=target_pos, + target_quat=target_quat, + joint_names=None, + tol=1e-14, + regularization_threshold=0.5, + regularization_strength=1e-2, + max_update_norm=2.0, + progress_thresh=5000.0, + max_steps=_MAX_STEPS, + inplace=False, + null_space_method=False + ) + + self.assertLessEqual(result.steps, _MAX_STEPS) + physics.data.qpos[:] = result.qpos + + save_path = os.path.join(dir_path, "./result_qpos") + np.save(save_path, result.qpos) + mjlib.mj_fwdPosition(physics.model.ptr, physics.data.ptr) + + pos = physics.named.data.site_xpos[_SITES_NAMES] + err_norm = np.linalg.norm(target_pos - pos) + self.assertLessEqual(err_norm, 0.11) + @parameterized.parameters(itertools.product(_TARGETS, _INPLACE)) def testQposFromSitePose(self, target, inplace): physics = mujoco.Physics.from_xml_string(_ARM_XML) @@ -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], target_pos=target_pos, target_quat=target_quat, joint_names=_JOINTS, @@ -133,7 +182,7 @@ def testNamedJointsWithMultipleDOFs(self): target_pos = (0.05, 0.05, 0) result = ik.qpos_from_site_pose( physics=physics, - site_name=site_name, + sites_names=[site_name], target_pos=target_pos, joint_names=joint_names, tol=_TOL, @@ -150,7 +199,7 @@ def testNamedJointsWithMultipleDOFs(self): physics.reset() result = ik.qpos_from_site_pose( physics=physics, - site_name=site_name, + sites_names=[site_name], target_pos=target_pos, joint_names=joint_names[:1], tol=_TOL, @@ -170,7 +219,7 @@ def testAllowedJointNameTypes(self, joint_names): target_pos = (0.05, 0.05, 0) ik.qpos_from_site_pose( physics=physics, - site_name=site_name, + sites_names=[site_name], target_pos=target_pos, joint_names=joint_names, tol=_TOL, @@ -192,7 +241,7 @@ def testDisallowedJointNameTypes(self, joint_names): with self.assertRaisesWithLiteralMatch(ValueError, expected_message): ik.qpos_from_site_pose( physics=physics, - site_name=site_name, + sites_names=[site_name], target_pos=target_pos, joint_names=joint_names, tol=_TOL, @@ -206,7 +255,7 @@ def testNoTargetPosOrQuat(self): ValueError, ik._REQUIRE_TARGET_POS_OR_QUAT): ik.qpos_from_site_pose( physics=physics, - site_name=site_name, + sites_names=[site_name], tol=_TOL, max_steps=_MAX_STEPS, inplace=True) diff --git a/dm_control/utils/task.xml b/dm_control/utils/task.xml new file mode 100644 index 00000000..5c92bce4 --- /dev/null +++ b/dm_control/utils/task.xml @@ -0,0 +1,65 @@ + + + + + + + + + + + diff --git a/dm_control/utils/visualize.py b/dm_control/utils/visualize.py new file mode 100644 index 00000000..6fb0a4cb --- /dev/null +++ b/dm_control/utils/visualize.py @@ -0,0 +1,65 @@ +from dm_control import mujoco +from dm_control import suite +from dm_control import viewer +from dm_control.rl import control +from dm_control.suite import common +import numpy as np +import os +import math +from scipy.signal import savgol_filter +import re + +def get_model_and_assets(model_name): + with open(model_name, mode='rb') as f: + return f.read(), common.ASSETS + + +class Task(suite.base.Task): + def __init__(self, mocap_qpos, random=None): + super().__init__(random=random) + self.mocap_qpos = mocap_qpos + self.init = False + + def initialize_episode(self, physics): + super().initialize_episode(physics) + physics.forward() + physics.data.qpos[:] = self.mocap_qpos[0] + physics.forward() + + def get_observation(self, physics): + return {} + + def get_reward(self, physics): + return 0 + + def after_step(self, physics): + target_pos = physics.model.key_mpos[0] + target_pos = target_pos.reshape((-1, 3)) + physics.data.mocap_pos[:] = target_pos + physics.data.qpos[:] = self.mocap_qpos[0] + + physics.data.qacc[:] = 0 + physics.data.qvel[:] = 0 + + + def before_step(self, action, physics): + physics.forward() + + +if __name__ == '__main__': + dir_path = os.path.dirname(os.path.realpath(__file__)) + model_dir = os.path.join(dir_path, "./task.xml") + physics = mujoco.Physics.from_xml_path(model_dir) + + qpos_dir = os.path.join(dir_path, "./result_qpos.npy") + mocap_qpos = np.load(qpos_dir) + + 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, + # 0.00806733, -0.0921235, -0.00874875, -0.0488637, -0.0894064, + # 0.0219203, -0.432015, -0.374747, -0.0965783, 0.0391758, -0.0760225, + # 0.132042, -0.256406, -0.350065, -0.494566, 0.714612, -0.744779, + # -0.922223, 0.818614, -0.698744, -1.00263]])) + + viewer.launch(env) From a14bacee7985ecc0726017a289850feff80eb163 Mon Sep 17 00:00:00 2001 From: vittorione94 Date: Fri, 19 May 2023 19:00:20 +0300 Subject: [PATCH 2/2] created directory for test assets --- dm_control/utils/inverse_kinematics_test.py | 4 ++-- .../utils/{ => testing/assets}/CMU-CMU-02-02_04_poses.xml | 0 dm_control/utils/{ => testing/assets}/common.xml | 0 dm_control/utils/{ => testing/assets}/humanoid.xml | 0 dm_control/utils/{ => testing/assets}/task.xml | 0 dm_control/utils/visualize.py | 4 ++-- 6 files changed, 4 insertions(+), 4 deletions(-) rename dm_control/utils/{ => testing/assets}/CMU-CMU-02-02_04_poses.xml (100%) rename dm_control/utils/{ => testing/assets}/common.xml (100%) rename dm_control/utils/{ => testing/assets}/humanoid.xml (100%) rename dm_control/utils/{ => testing/assets}/task.xml (100%) diff --git a/dm_control/utils/inverse_kinematics_test.py b/dm_control/utils/inverse_kinematics_test.py index fac25801..6947f6c1 100644 --- a/dm_control/utils/inverse_kinematics_test.py +++ b/dm_control/utils/inverse_kinematics_test.py @@ -83,7 +83,7 @@ class InverseKinematicsTest(parameterized.TestCase): def testQposFromMultipleSitesPose(self): dir_path = os.path.dirname(os.path.realpath(__file__)) - model_dir = os.path.join(dir_path, "./task.xml") + model_dir = os.path.join(dir_path, "./testing/assets/task.xml") physics = mujoco.Physics.from_xml_path(model_dir) target_pos = physics.model.key_mpos[0] @@ -121,7 +121,7 @@ def testQposFromMultipleSitesPose(self): self.assertLessEqual(result.steps, _MAX_STEPS) physics.data.qpos[:] = result.qpos - save_path = os.path.join(dir_path, "./result_qpos") + save_path = os.path.join(dir_path, "./testing/assets/result_qpos") np.save(save_path, result.qpos) mjlib.mj_fwdPosition(physics.model.ptr, physics.data.ptr) diff --git a/dm_control/utils/CMU-CMU-02-02_04_poses.xml b/dm_control/utils/testing/assets/CMU-CMU-02-02_04_poses.xml similarity index 100% rename from dm_control/utils/CMU-CMU-02-02_04_poses.xml rename to dm_control/utils/testing/assets/CMU-CMU-02-02_04_poses.xml diff --git a/dm_control/utils/common.xml b/dm_control/utils/testing/assets/common.xml similarity index 100% rename from dm_control/utils/common.xml rename to dm_control/utils/testing/assets/common.xml diff --git a/dm_control/utils/humanoid.xml b/dm_control/utils/testing/assets/humanoid.xml similarity index 100% rename from dm_control/utils/humanoid.xml rename to dm_control/utils/testing/assets/humanoid.xml diff --git a/dm_control/utils/task.xml b/dm_control/utils/testing/assets/task.xml similarity index 100% rename from dm_control/utils/task.xml rename to dm_control/utils/testing/assets/task.xml diff --git a/dm_control/utils/visualize.py b/dm_control/utils/visualize.py index 6fb0a4cb..117a3565 100644 --- a/dm_control/utils/visualize.py +++ b/dm_control/utils/visualize.py @@ -48,10 +48,10 @@ def before_step(self, action, physics): if __name__ == '__main__': dir_path = os.path.dirname(os.path.realpath(__file__)) - model_dir = os.path.join(dir_path, "./task.xml") + model_dir = os.path.join(dir_path, "./testing/assets/task.xml") physics = mujoco.Physics.from_xml_path(model_dir) - qpos_dir = os.path.join(dir_path, "./result_qpos.npy") + qpos_dir = os.path.join(dir_path, "./testing/assets/result_qpos.npy") mocap_qpos = np.load(qpos_dir) env = control.Environment(physics, Task(mocap_qpos=[mocap_qpos]))