Add ability to restrict IK to certain dofs.

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
 

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",

mink/configuration.py

@@ -8,14 +8,15 @@
 """
 
 import logging
-from typing import Optional
+from typing import Optional, Sequence
 
 import mujoco
 import numpy as np
 
 from . import constants as consts
 from . import exceptions
 from .lie import SE3, SO3
+from .utils import get_dof_ids
 
 
 class Configuration:
@@ -38,18 +39,26 @@ def __init__(
         self,
         model: mujoco.MjModel,
         q: Optional[np.ndarray] = None,
+        joint_names: Optional[Sequence[str]] = None,
     ):
         """Constructor.
 
         Args:
             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:
+            self.dof_ids = np.arange(model.nv)
+        else:
+            self.dof_ids = np.asarray(get_dof_ids(model, joint_names))
+
     def update(self, q: Optional[np.ndarray] = None) -> None:
         """Run forward kinematics.
 
@@ -85,9 +94,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 +233,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 +245,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.

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)

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)

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,

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()

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)

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,31 @@ def test_get_freejoint_dims(self):
             np.asarray(list(range(0, 6))),
         )
 
+    def test_get_dof_ids(self):
+        xml_str = """
+        <mujoco>
+          <worldbody>
+            <body>
+              <joint type="ball" name="ball"/>
+              <geom type="sphere" size=".1" mass=".1"/>
+              <body>
+                <joint type="hinge" name="hinge1" range="0 1.57"/>
+                <geom type="sphere" size=".1" mass=".1"/>
+                <body>
+                  <joint type="hinge" name="hinge2" range="0 1.57"/>
+                  <geom type="sphere" size=".1" mass=".1"/>
+                </body>
+              </body>
+            </body>
+          </worldbody>
+        </mujoco>
+        """
+        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 = """
         <mujoco>