From df6fc041d4a6778936ce722ef16f2b801fc47021 Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Sat, 1 Mar 2025 22:52:02 -0800 Subject: [PATCH 1/3] Add ability to restrict IK to certain dofs. --- CHANGELOG.md | 1 + mink/__init__.py | 2 ++ mink/configuration.py | 29 ++++++++++++++++++++++++++--- mink/exceptions.py | 12 ++++++++++++ tests/test_configuration.py | 28 ++++++++++++++++++++++++++++ tests/test_posture_task.py | 2 +- 6 files changed, 70 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 26a5d58a..9b4647a7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ All notable changes to this project will be documented in this file. - Added a new humanoid example using the Apollo robot. - Added `utils.get_subtree_body_ids` to get all bodies belonging to the subtree starting at a given body. - Example usage of this function can be found in the [ALOHA example script](examples/arm_aloha.py) where it is used to selectively apply gravity compensation torques to the left and right arm bodies. +- Added `joint_names` parameter to `Configuration` to allow performing IK on a subset of joints. ### Changed diff --git a/mink/__init__.py b/mink/__init__.py index ff1dae56..83e35672 100644 --- a/mink/__init__.py +++ b/mink/__init__.py @@ -15,6 +15,7 @@ MinkError, NotWithinConfigurationLimits, UnsupportedFrame, + InvalidJointName, ) from .lie import SE3, SO3, MatrixLieGroup from .limits import ( @@ -70,6 +71,7 @@ "NotWithinConfigurationLimits", "TargetNotSet", "InvalidMocapBody", + "InvalidJointName", "SUPPORTED_FRAMES", "FRAME_TO_ENUM", "FRAME_TO_JAC_FUNC", diff --git a/mink/configuration.py b/mink/configuration.py index b82e3bfc..7b944423 100644 --- a/mink/configuration.py +++ b/mink/configuration.py @@ -8,7 +8,7 @@ """ import logging -from typing import Optional +from typing import Optional, Sequence import mujoco import numpy as np @@ -38,6 +38,7 @@ def __init__( self, model: mujoco.MjModel, q: Optional[np.ndarray] = None, + joint_names: Optional[Sequence[str]] = None, ): """Constructor. @@ -45,11 +46,27 @@ def __init__( model: Mujoco model. q: Configuration to initialize from. If None, the configuration is initialized to the default configuration `qpos0`. + joint_names: List of joints to be controlled. If None (default), all + joints are used """ self.model = model self.data = mujoco.MjData(model) self.update(q=q) + if joint_names is None: + dof_ids = np.arange(model.nv) + else: + dof_ids = [] + for joint_name in joint_names: + joint_id = mujoco.mj_name2id( + model, mujoco.mjtObj.mjOBJ_JOINT, joint_name + ) + if joint_id == -1: + raise exceptions.InvalidJointName(joint_name, model) + dof_ids.append(model.jnt_dofadr[joint_id]) + dof_ids = np.array(dof_ids) + self.dof_ids = dof_ids + def update(self, q: Optional[np.ndarray] = None) -> None: """Run forward kinematics. @@ -85,9 +102,11 @@ def check_limits(self, tol: float = 1e-6, safety_break: bool = True) -> None: """ for jnt in range(self.model.njnt): jnt_type = self.model.jnt_type[jnt] + jnt_id = self.model.jnt_dofadr[jnt] if ( jnt_type == mujoco.mjtJoint.mjJNT_FREE or not self.model.jnt_limited[jnt] + or jnt_id not in self.dof_ids ): continue padr = self.model.jnt_qposadr[jnt] @@ -222,7 +241,9 @@ def integrate(self, velocity: np.ndarray, dt: float) -> np.ndarray: The new configuration after integration. """ q = self.data.qpos.copy() - mujoco.mj_integratePos(self.model, q, velocity, dt) + qvel = np.zeros((self.nv,)) + qvel[self.dof_ids] = velocity[self.dof_ids] + mujoco.mj_integratePos(self.model, q, qvel, dt) return q def integrate_inplace(self, velocity: np.ndarray, dt: float) -> None: @@ -232,7 +253,9 @@ def integrate_inplace(self, velocity: np.ndarray, dt: float) -> None: velocity: The velocity in tangent space. dt: Integration duration in [s]. """ - mujoco.mj_integratePos(self.model, self.data.qpos, velocity, dt) + qvel = np.zeros((self.nv,)) + qvel[self.dof_ids] = velocity[self.dof_ids] + mujoco.mj_integratePos(self.model, self.data.qpos, qvel, dt) self.update() # Aliases. diff --git a/mink/exceptions.py b/mink/exceptions.py index f9aaf0a2..5ad0e43e 100644 --- a/mink/exceptions.py +++ b/mink/exceptions.py @@ -96,3 +96,15 @@ def __init__( f"{lower} <= {value} <= {upper}" ) super().__init__(message) + + +class InvalidJointName(MinkError): + """Exception raised when a joint name is not found in the robot model.""" + + def __init__(self, joint_name: str, model: mujoco.MjModel): + joint_names = [model.joint(i).name for i in range(model.njnt)] + message = ( + f"Joint '{joint_name}' does not exist in the model. Available " + f"joint names: {joint_names}" + ) + super().__init__(message) diff --git a/tests/test_configuration.py b/tests/test_configuration.py index 415f5a9f..35ff0f93 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -117,6 +117,34 @@ def test_check_limits_freejoint(self): configuration.update(q=q) configuration.check_limits(safety_break=True) # Should not raise. + def test_raises_error_if_joint_name_is_invalid(self): + joint_names = ["invalid_joint_name"] + with self.assertRaises(mink.InvalidJointName): + mink.Configuration(self.model, joint_names=joint_names) + + def test_dof_subset(self): + joint_names = ["elbow_joint", "shoulder_pan_joint"] + configuration = mink.Configuration(self.model, joint_names=joint_names) + + # Check that the DOF IDs are set correctly. + expected_dof_ids = np.array( + [self.model.joint(name).dofadr[0] for name in joint_names] + ) + np.testing.assert_equal(configuration.dof_ids, expected_dof_ids) + + # Integrate a velocity and check that only the specified joints move. + configuration.update(self.q_ref) + qvel = np.random.randn(self.model.nv) * 0.1 + dt = 1e-1 + configuration.integrate_inplace(qvel, dt) + qpos = configuration.q.copy() + + expected_qpos = qpos.copy() + expected_qpos[expected_dof_ids] = ( + self.q_ref[expected_dof_ids] + qvel[expected_dof_ids] * dt + ) + np.testing.assert_almost_equal(qpos, expected_qpos) + if __name__ == "__main__": absltest.main() diff --git a/tests/test_posture_task.py b/tests/test_posture_task.py index 59b2fd5e..7b8b41c1 100644 --- a/tests/test_posture_task.py +++ b/tests/test_posture_task.py @@ -18,7 +18,7 @@ class TestPostureTask(absltest.TestCase): @classmethod def setUpClass(cls): - cls.model = load_robot_description("stretch_mj_description") + cls.model = load_robot_description("talos_mj_description") def setUp(self): self.configuration = Configuration(self.model) From f3cba8f2cd80e53a196be06ef643fb2c39a3de4a Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Sun, 2 Mar 2025 10:47:17 -0800 Subject: [PATCH 2/3] Fix dof IDs. --- mink/configuration.py | 16 ++++------------ mink/limits/velocity_limit.py | 6 +++--- mink/utils.py | 28 ++++++++++++++++++++++++++-- tests/test_utils.py | 29 ++++++++++++++++++++++++++++- 4 files changed, 61 insertions(+), 18 deletions(-) diff --git a/mink/configuration.py b/mink/configuration.py index 7b944423..72917e2f 100644 --- a/mink/configuration.py +++ b/mink/configuration.py @@ -16,6 +16,7 @@ from . import constants as consts from . import exceptions from .lie import SE3, SO3 +from .utils import get_dof_ids class Configuration: @@ -47,25 +48,16 @@ def __init__( q: Configuration to initialize from. If None, the configuration is initialized to the default configuration `qpos0`. joint_names: List of joints to be controlled. If None (default), all - joints are used + joints are used. """ self.model = model self.data = mujoco.MjData(model) self.update(q=q) if joint_names is None: - dof_ids = np.arange(model.nv) + self.dof_ids = np.arange(model.nv) else: - dof_ids = [] - for joint_name in joint_names: - joint_id = mujoco.mj_name2id( - model, mujoco.mjtObj.mjOBJ_JOINT, joint_name - ) - if joint_id == -1: - raise exceptions.InvalidJointName(joint_name, model) - dof_ids.append(model.jnt_dofadr[joint_id]) - dof_ids = np.array(dof_ids) - self.dof_ids = dof_ids + self.dof_ids = np.asarray(get_dof_ids(model, joint_names)) def update(self, q: Optional[np.ndarray] = None) -> None: """Run forward kinematics. diff --git a/mink/limits/velocity_limit.py b/mink/limits/velocity_limit.py index 3848b986..ea011bd1 100644 --- a/mink/limits/velocity_limit.py +++ b/mink/limits/velocity_limit.py @@ -1,6 +1,6 @@ """Joint velocity limit.""" -from typing import Mapping +from typing import Mapping, Optional import mujoco import numpy as np @@ -28,7 +28,7 @@ class VelocityLimit(Limit): indices: np.ndarray limit: np.ndarray - projection_matrix: np.ndarray + projection_matrix: Optional[np.ndarray] def __init__( self, @@ -58,7 +58,7 @@ def __init__( f"Got: {max_vel.shape}" ) index_list.extend(range(vadr, vadr + vdim)) - limit_list.extend(max_vel.tolist()) + limit_list.extend(float(val) for val in max_vel) self.indices = np.array(index_list) self.indices.setflags(write=False) diff --git a/mink/utils.py b/mink/utils.py index fc9a3bc8..d0592eed 100644 --- a/mink/utils.py +++ b/mink/utils.py @@ -1,10 +1,10 @@ -from typing import Optional +from typing import Optional, Sequence import mujoco import numpy as np from . import constants as consts -from .exceptions import InvalidKeyframe, InvalidMocapBody +from .exceptions import InvalidKeyframe, InvalidMocapBody, InvalidJointName def move_mocap_to_frame( @@ -56,6 +56,30 @@ def get_freejoint_dims(model: mujoco.MjModel) -> tuple[list[int], list[int]]: return q_ids, v_ids +def get_dof_ids(model: mujoco.MjModel, joint_names: Sequence[str]) -> list[int]: + """Get DOF indices for a list of joint names. + + Args: + model: Mujoco model. + joint_names: List of joint names. + + Returns: + List of DOF indices. + + Raises: + InvalidJointName: If any joint name is invalid. + """ + dof_ids: list[int] = [] + for joint_name in joint_names: + jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, joint_name) + if jid == -1: + raise InvalidJointName(joint_name, model) + vadr = model.jnt_dofadr[jid] + vdim = consts.dof_width(model.jnt_type[jid]) + dof_ids.extend(range(vadr, vadr + vdim)) + return dof_ids + + def custom_configuration_vector( model: mujoco.MjModel, key_name: Optional[str] = None, diff --git a/tests/test_utils.py b/tests/test_utils.py index 440adfd4..d27e90dc 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -6,7 +6,7 @@ from robot_descriptions.loaders.mujoco import load_robot_description from mink import utils -from mink.exceptions import InvalidKeyframe, InvalidMocapBody +from mink.exceptions import InvalidKeyframe, InvalidMocapBody, InvalidJointName class TestUtils(absltest.TestCase): @@ -108,6 +108,33 @@ def test_get_freejoint_dims(self): np.asarray(list(range(0, 6))), ) + def test_get_dof_ids_throws_error_if_joint_name_invalid(self): + with self.assertRaises(InvalidJointName): + utils.get_dof_ids(self.model, ["invalid_joint_name"]) + + def test_get_dof_ids(self): + xml_str = """ + + + + + + + + + + + + + + + + + """ + model = mujoco.MjModel.from_xml_string(xml_str) + dof_ids = utils.get_dof_ids(model, ["ball", "hinge2"]) + np.testing.assert_allclose(dof_ids, [0, 1, 2, 4]) + def test_get_subtree_geom_ids(self): xml_str = """ From 4df2e055d2ebfb091b019fec467ca52de48f5e1c Mon Sep 17 00:00:00 2001 From: Kevin Zakka Date: Sun, 2 Mar 2025 10:52:44 -0800 Subject: [PATCH 3/3] Combine tests. --- tests/test_utils.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index d27e90dc..c82d9e54 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -108,10 +108,6 @@ def test_get_freejoint_dims(self): np.asarray(list(range(0, 6))), ) - def test_get_dof_ids_throws_error_if_joint_name_invalid(self): - with self.assertRaises(InvalidJointName): - utils.get_dof_ids(self.model, ["invalid_joint_name"]) - def test_get_dof_ids(self): xml_str = """ @@ -123,8 +119,8 @@ def test_get_dof_ids(self): - - + + @@ -134,6 +130,8 @@ def test_get_dof_ids(self): model = mujoco.MjModel.from_xml_string(xml_str) dof_ids = utils.get_dof_ids(model, ["ball", "hinge2"]) np.testing.assert_allclose(dof_ids, [0, 1, 2, 4]) + with self.assertRaises(InvalidJointName): + utils.get_dof_ids(model, ["invalid_joint_name"]) def test_get_subtree_geom_ids(self): xml_str = """