More unit tests.

examples/main.py

@@ -7,7 +7,7 @@
 
 import mink
 from mink.lie import SE3, SO3
-from mink.utils import set_mocap_pose_from_site
+from mink.utils import custom_configuration_vector, set_mocap_pose_from_site
 
 _HERE = Path(__file__).parent
 _XML = _HERE / "universal_robots_ur5e" / "scene.xml"
@@ -45,16 +45,16 @@
         # ),
     ]
 
-    # max_velocities = {
-    #     "shoulder_pan": np.pi,
-    #     "shoulder_lift": np.pi,
-    #     "elbow": np.pi,
-    #     "wrist_1": np.pi,
-    #     "wrist_2": np.pi,
-    #     "wrist_3": np.pi,
-    # }
-    # velocity_limit = mink.VelocityLimit(model, max_velocities)
-    # limits.append(velocity_limit)
+    max_velocities = {
+        "shoulder_pan": np.pi,
+        "shoulder_lift": np.pi,
+        "elbow": np.pi,
+        "wrist_1": np.pi,
+        "wrist_2": np.pi,
+        "wrist_3": np.pi,
+    }
+    velocity_limit = mink.VelocityLimit(model, max_velocities)
+    limits.append(velocity_limit)
 
     ## =================== ##
 
@@ -71,7 +71,14 @@
     ) as viewer:
         mujoco.mjv_defaultFreeCamera(model, viewer.cam)
 
-        mujoco.mj_resetDataKeyframe(model, data, model.key("home").id)
+        # mujoco.mj_resetDataKeyframe(model, data, model.key("home").id)
+        q_ref = custom_configuration_vector(
+            model,
+            "home",
+            shoulder_pan=(0.0, 0.1),
+        )
+        data.qpos = q_ref
+
         configuration.update(data.qpos)
         mujoco.mj_forward(model, data)
 

mink/configuration.py

@@ -6,7 +6,6 @@
 from . import exceptions
 from .lie import SE3, SO3
 
-_SUPPORTED_JOINT_TYPES = {mujoco.mjtJoint.mjJNT_HINGE, mujoco.mjtJoint.mjJNT_SLIDE}
 _SUPPORTED_OBJ_TYPES = {"body", "geom", "site"}
 
 _TYPE_TO_ENUM = {
@@ -34,7 +33,7 @@
 
 class Configuration:
     """A struct that provides convenient access to kinematic quantities such as frame
-    transforms and frame jacobians.
+    transforms and frame jacobians. Frames can be defined at bodies, geoms or sites.
 
     The `update` function ensures the proper forward kinematics functions have been
     called, namely:
@@ -100,7 +99,7 @@ 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]
             if (
-                jnt_type not in _SUPPORTED_JOINT_TYPES
+                jnt_type == mujoco.mjtJoint.mjJNT_FREE
                 or not self.model.jnt_limited[jnt]
             ):
                 continue

mink/limits/configuration_limit.py

@@ -48,7 +48,7 @@ def __init__(
             qpos_dim = qpos_width(jnt_type)
             jnt_range = model.jnt_range[jnt]
             padr = model.jnt_qposadr[jnt]
-            if jnt_type == mujoco.mjtJoint.mjJNT_FREE:
+            if jnt_type == mujoco.mjtJoint.mjJNT_FREE or not model.jnt_limited[jnt]:
                 continue
             lower[padr : padr + qpos_dim] = jnt_range[0] + min_distance_from_limits
             upper[padr : padr + qpos_dim] = jnt_range[1] - min_distance_from_limits
@@ -64,7 +64,6 @@ def __init__(
         self.upper = upper
         self.model = model
         self.gain = gain
-        self.G = np.vstack([self.projection_matrix, -self.projection_matrix])
 
     def compute_qp_inequalities(
         self,
@@ -95,5 +94,6 @@ def compute_qp_inequalities(
 
         p_min = self.gain * delta_q_min[self.indices]
         p_max = self.gain * delta_q_max[self.indices]
+        self.G = np.vstack([self.projection_matrix, -self.projection_matrix])
         h = np.hstack([p_max, -p_min])
         return Constraint(G=self.G, h=h)

mink/utils.py

@@ -1,8 +1,37 @@
 import mujoco
 import numpy as np
 
+from .exceptions import InvalidKeyframe
 from .lie import SE3, SO3
 
+_TYPE_TO_ENUM = {
+    "body": mujoco.mjtObj.mjOBJ_BODY,
+    "geom": mujoco.mjtObj.mjOBJ_GEOM,
+    "site": mujoco.mjtObj.mjOBJ_SITE,
+}
+
+_TYPE_TO_XMAT_ATTRIBUTE = {
+    "body": "xmat",
+    "geom": "geom_xmat",
+    "site": "site_xmat",
+}
+
+_TYPE_TO_POS_ATTRIBUTE = {
+    "body": "xpos",
+    "geom": "geom_xpos",
+    "site": "site_xpos",
+}
+
+
+def dof_width(joint_type: int) -> int:
+    """Get the dimensionality of the joint in qvel."""
+    return {0: 6, 1: 3, 2: 1, 3: 1}[joint_type]
+
+
+def qpos_width(joint_type: int) -> int:
+    """Get the dimensionality of the joint in qpos."""
+    return {0: 7, 1: 4, 2: 1, 3: 1}[joint_type]
+
 
 def set_mocap_pose_from_site(
     model: mujoco.MjModel,
@@ -40,6 +69,22 @@ def set_mocap_pose_from_body(
     mujoco.mju_mat2Quat(data.mocap_quat[mocap_id], data.xmat[body_id])
 
 
+def set_mocap_pose_from_obj(
+    model: mujoco.MjModel,
+    data: mujoco.MjData,
+    mocap_name: str,
+    obj_name: str,
+    obj_type: str,
+):
+    mocap_id = model.body(mocap_name).mocapid[0]
+    obj_id = mujoco.mj_name2id(model, _TYPE_TO_ENUM[obj_name], obj_name)
+    data.mocap_pos[mocap_id] = getattr(data, _TYPE_TO_POS_ATTRIBUTE[obj_type])[obj_id]
+    mujoco.mju_mat2Quat(
+        data.mocap_quat[mocap_id],
+        getattr(data, _TYPE_TO_XMAT_ATTRIBUTE[obj_type])[obj_id],
+    )
+
+
 def pose_from_mocap(
     model: mujoco.MjModel,
     data: mujoco.MjData,
@@ -67,14 +112,41 @@ def get_freejoint_dims(model: mujoco.MjModel) -> tuple[list[int], list[int]]:
 
 def custom_configuration_vector(
     model: mujoco.MjModel,
+    key_name: str | None = None,
     **kwargs,
 ) -> np.ndarray:
-    raise NotImplementedError
-
-
-def dof_width(joint_type: int) -> int:
-    return {0: 6, 1: 3, 2: 1, 3: 1}[joint_type]
-
-
-def qpos_width(joint_type: int) -> int:
-    return {0: 7, 1: 4, 2: 1, 3: 1}[joint_type]
+    """Generate a configuration vector where named joints have specific values.
+
+    Args:
+        model: An MjModel instance.
+        key_name: Optional keyframe name to initialize the configuration vector from.
+            Otherwise, the default pose qpos0 is used.
+        kwargs: Custom values for joint coordinates.
+
+    Returns:
+        Configuration vector where named joints have the values specified in
+            keyword arguments, and other joints have their neutral value or value
+            defined in the keyframe if provided.
+    """
+    data = mujoco.MjData(model)
+    if key_name is not None:
+        key_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_KEY, key_name)
+        if key_id == -1:
+            raise InvalidKeyframe(key_name, model)
+        mujoco.mj_resetDataKeyframe(model, data, key_id)
+    else:
+        mujoco.mj_resetData(model, data)
+    q = data.qpos.copy()
+    for name, value in kwargs.items():
+        jid = model.joint(name).id
+        jnt_type = model.jnt_type[jid]
+        jnt_dim = qpos_width(jnt_type)
+        qid = model.jnt_dofadr[jid]
+        value = np.atleast_1d(value)
+        if value.shape != (jnt_dim,):
+            raise ValueError(
+                f"Joint {name} should have a qpos value of {jnt_dim,} but "
+                f"got {value.shape}"
+            )
+        q[qid : qid + jnt_dim] = value
+    return q

tests/test_configuration.py

@@ -1,3 +1,5 @@
+"""Tests for configuration.py."""
+
 import numpy as np
 from absl.testing import absltest
 from robot_descriptions.loaders.mujoco import load_robot_description
@@ -15,6 +17,11 @@ def setUpClass(cls):
     def setUp(self):
         self.q_ref = self.model.key("home").qpos
 
+    def test_nq_nv(self):
+        configuration = mink.Configuration(self.model)
+        self.assertEqual(configuration.nq, self.model.nq)
+        self.assertEqual(configuration.nv, self.model.nv)
+
     def test_initialize_from_keyframe(self):
         """Test that keyframe initialization correctly updates the configuration."""
         configuration = mink.Configuration(self.model)

tests/test_configuration_limit.py

@@ -1,9 +1,12 @@
 """Tests for configuration_limit.py."""
 
+import mujoco
+import numpy as np
 from absl.testing import absltest
 from robot_descriptions.loaders.mujoco import load_robot_description
 
 from mink import Configuration
+from mink.limits import ConfigurationLimit, VelocityLimit
 
 
 class TestConfigurationLimit(absltest.TestCase):
@@ -16,6 +19,97 @@ def setUpClass(cls):
     def setUp(self):
         self.configuration = Configuration(self.model)
         self.configuration.update_from_keyframe("home")
+        self.velocities = {
+            "shoulder_pan_joint": np.pi,
+            "shoulder_lift_joint": np.pi,
+            "elbow_joint": np.pi,
+            "wrist_1_joint": np.pi,
+            "wrist_2_joint": np.pi,
+            "wrist_3_joint": np.pi,
+        }
+        self.vel_limit = VelocityLimit(self.model, self.velocities)
+
+    def test_dimensions(self):
+        limit = ConfigurationLimit(self.model)
+        nv = self.configuration.nv
+        self.assertEqual(limit.projection_matrix.shape, (nv, nv))
+        self.assertEqual(len(limit.indices), nv)
+
+    def test_model_with_no_limit(self):
+        empty_model = mujoco.MjModel.from_xml_string("<mujoco></mujoco>")
+        empty_bounded = ConfigurationLimit(empty_model)
+        self.assertEqual(len(empty_bounded.indices), 0)
+        self.assertIsNone(empty_bounded.projection_matrix)
+
+    def test_model_with_subset_of_velocities_limited(self):
+        xml_str = """
+        <mujoco>
+          <worldbody>
+            <body>
+              <joint type="hinge" name="hinge_unlimited"/>
+              <geom type="sphere" size=".1" mass=".1"/>
+              <body>
+                <joint type="hinge" name="hinge_limited" limited="true" range="0 1.57"/>
+                <geom type="sphere" size=".1" mass=".1"/>
+              </body>
+            </body>
+          </worldbody>
+        </mujoco>
+        """
+        model = mujoco.MjModel.from_xml_string(xml_str)
+        limit = ConfigurationLimit(model)
+        nb = 1
+        nv = model.nv
+        self.assertEqual(limit.projection_matrix.shape, (nb, nv))
+        self.assertEqual(len(limit.indices), nb)
+
+    def test_freejoint_ignored(self):
+        xml_str = """
+        <mujoco>
+          <worldbody>
+            <body>
+              <joint type="free" name="floating"/>
+              <geom type="sphere" size=".1" mass=".1"/>
+              <body>
+                <joint type="hinge" name="hinge" range="0 1.57" limited="true"/>
+                <geom type="sphere" size=".1" mass=".1"/>
+              </body>
+            </body>
+          </worldbody>
+        </mujoco>
+        """
+        model = mujoco.MjModel.from_xml_string(xml_str)
+        limit = ConfigurationLimit(model)
+        nb = 1
+        nv = model.nv
+        self.assertEqual(limit.projection_matrix.shape, (nb, nv))
+        self.assertEqual(len(limit.indices), nb)
+
+    def test_far_from_limit(self, tol=1e-10):
+        """Limit has no effect when the configuration is far away."""
+        dt = 1e-3  # [s]
+        limit = ConfigurationLimit(self.model)
+        G, h = limit.compute_qp_inequalities(self.configuration, dt=dt)
+        # When we are far away from configuration limits, the velocity limit is
+        # simply the configuration-agnostic one from the robot.
+        self.assertLess(np.max(+G @ self.vel_limit.limit * dt - h), -tol)
+        self.assertLess(np.max(-G @ self.vel_limit.limit * dt - h), -tol)
+
+    def test_configuration_limit_repulsion(self, tol=1e-10):
+        """Velocities are scaled down when close to a configuration limit."""
+        dt = 1e-3  # [s]
+        slack_vel = 5e-4  # [rad] / [s]
+        limit = ConfigurationLimit(self.model, gain=0.5)
+        # Override configuration limits to `q +/- slack_vel * dt`.
+        limit.lower = self.configuration.integrate(
+            -slack_vel * np.ones((self.configuration.nv,)), dt
+        )
+        limit.upper = self.configuration.integrate(
+            +slack_vel * np.ones((self.configuration.nv,)), dt
+        )
+        _, h = limit.compute_qp_inequalities(self.configuration, dt)
+        self.assertLess(np.max(h), slack_vel * dt + tol)
+        self.assertGreater(np.min(h), -slack_vel * dt - tol)
 
 
 if __name__ == "__main__":

tests/test_jacobians.py

@@ -1,3 +1,5 @@
+"""Test task jacobian matrices against finite differences."""
+
 import mujoco
 import numpy as np
 from absl.testing import absltest
@@ -74,7 +76,7 @@ def J(q) -> np.ndarray:
 
     def test_frame_task(self):
         frame_task = mink.FrameTask(
-            frame_name="attachment_site",
+            frame_name="left_foot",
             frame_type="site",
             position_cost=1.0,
             orientation_cost=1.0,