Update (#32)

* Docstring fixes. * Tune aloha example.
kevinzakka · Sep 29, 2024 · 4a6803a · 4a6803a
1 parent 0978316
commit 4a6803a
Show file tree

Hide file tree

Showing 10 changed files with 91 additions and 47 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file.
 
 ## Unreleased
 
+### Changed
+
+- Improved ALOHA example script.
+
 ## [0.0.5] - 2024-09-27
 
 ### Changed

diff --git a/examples/aloha/aloha.xml b/examples/aloha/aloha.xml
@@ -137,6 +137,7 @@
                   <inertial pos="0.0395662 -2.56311e-07 0.00400649" quat="0.62033 0.619916 -0.339682 0.339869"
                     mass="0.251652" diaginertia="0.000689546 0.000650316 0.000468142"/>
                   <joint name="left/wrist_rotate" class="wrist_rotate"/>
+                  <!-- <geom name="left_ee_sphere" type="sphere" size=".07" pos=".1 0 0" class="collision" rgba=".5 .5 .5 .3"/> -->
                   <site name="left/gripper" pos="0.13 0 -.003" group="5"/>
                   <body name="left/gripper_base" euler="0 1.57 -1.57" pos="0.035 0 0">
                     <inertial pos="0.000182154 -0.0341589 -0.0106026" quat="0.435286 0.557074 -0.551539 0.442718"
@@ -148,7 +149,8 @@
                     <geom class="visual" pos="0 -0.03525 -0.0227" quat="0 -1 0 -1" type="mesh" mesh="vx300s_7_gripper_wrist_mount"/>
                     <!-- <geom class="collision" pos="0 -0.03525 -0.0227" quat="0 -1 0 -1" type="capsule" mesh="vx300s_7_gripper_wrist_mount"/> -->
                     <geom class="visual" pos="0 -0.0824748 -0.0095955" quat="0 0 -0.21644 -0.976296" type="mesh" mesh="d405_solid"/>
-                    <geom class="collision" pos="0 -0.0824748 -0.0095955" quat="0 0 -0.21644 -0.976296" type="box" mesh="d405_solid"/>
+                    <!-- <geom class="collision" pos="0 -0.0824748 -0.0095955" quat="0 0 -0.21644 -0.976296" type="box" mesh="d405_solid"/> -->
+                    <geom name="left_camera" class="collision" pos="0 -0.08 -0.02" quat="0 0 -0.21644 -0.976296" type="sphere" size=".035"/>
                     <camera name="wrist_cam_left" pos="0 -0.0824748 -0.0095955" mode="fixed" euler="2.70525955359 0 0"
                             focal="1.93e-3 1.93e-3" resolution="1280 720" sensorsize="3896e-6 2140e-6"/>
                    <body name="left/left_finger_link" pos="0.0191 -0.0141637 0.0211727" quat="1 -1 -1 1">
@@ -158,7 +160,7 @@
                       <geom pos="0.0141637 0.0211727 0.06" class="visual" quat="1 1 1 -1" type="mesh"
                         mesh="vx300s_8_custom_finger_left"/>
                       <!-- <geom pos="0.0141637 0.0211727 0.06" class="collision" quat="1 1 1 -1" type="mesh" mesh="vx300s_8_custom_finger_left"/> -->
-                      <geom class="collision" size="0.0137306 0.0460639" pos="0.0141512 -0.0347039 0.0119043" quat="0.544446 0.455511 0.456118 -0.536698"/>
+                      <geom name="left_left_finger_coll" class="collision" size="0.0137306 0.0460639" pos="0.0141512 -0.0347039 0.0119043" quat="0.544446 0.455511 0.456118 -0.536698"/>
                       <geom name="left/left_g0" pos="0.013  -0.0892 0.0268" class="sphere_collision"/>
                       <geom name="left/left_g1" pos="0.0222 -0.0892 0.0268" class="sphere_collision"/>
                       <geom name="left/left_g2" pos="0.0182 -0.0845 0.0266" class="sphere_collision"/>
@@ -171,7 +173,7 @@
                       <geom pos="0.0141637 -0.0211727 0.0597067" class="visual" quat="1 -1 -1 -1" type="mesh"
                         mesh="vx300s_8_custom_finger_right"/>
                       <!-- <geom pos="0.0141637 -0.0211727 0.0597067" class="collision" quat="1 -1 -1 -1" type="mesh" mesh="vx300s_8_custom_finger_right"/> -->
-                      <geom class="collision" size="0.013729 0.0460644" pos="0.0141715 0.0347097 0.011759" quat="0.538959 -0.458738 -0.452903 -0.542183"/>
+                      <geom name="left_right_finger_coll" class="collision" size="0.013729 0.0460644" pos="0.0141715 0.0347097 0.011759" quat="0.538959 -0.458738 -0.452903 -0.542183"/>
                       <geom name="left/right_g0" pos="0.013  0.0892 0.0268" class="sphere_collision"/>
                       <geom name="left/right_g1" pos="0.0222 0.0892 0.0268" class="sphere_collision"/>
                       <geom name="left/right_g2" pos="0.0182 0.0845 0.0266" class="sphere_collision"/>
@@ -225,6 +227,7 @@
                   <inertial pos="0.0395662 -2.56311e-07 0.00400649" quat="0.62033 0.619916 -0.339682 0.339869"
                     mass="0.251652" diaginertia="0.000689546 0.000650316 0.000468142"/>
                   <joint name="right/wrist_rotate" class="wrist_rotate"/>
+                  <!-- <geom name="right_ee_sphere" type="sphere" size=".07" pos=".1 0 0" class="collision" rgba=".5 .5 .5 .3"/> -->
                   <site name="right/gripper" pos="0.13 0 -.003" group="5"/>
                   <body name="right/gripper_base" euler="0 1.57 -1.57" pos="0.035 0 0">
                     <inertial pos="0.000182154 -0.0341589 -0.0106026" quat="0.435286 0.557074 -0.551539 0.442718"
@@ -236,7 +239,8 @@
                     <geom class="visual" pos="0 -0.03525 -0.0227" quat="0 -0.707107 0 -0.707107" type="mesh" mesh="vx300s_7_gripper_wrist_mount"/>
                     <!-- <geom class="collision" pos="0 -0.03525 -0.0227" quat="0 -0.707107 0 -0.707107" type="capsule" mesh="vx300s_7_gripper_wrist_mount"/> -->
                     <geom class="visual" pos="0 -0.0824748 -0.0095955" quat="0 0 -0.21644 -0.976296" type="mesh" mesh="d405_solid"/>
-                    <geom class="collision" pos="0 -0.0824748 -0.0095955" quat="0 0 -0.21644 -0.976296" type="box" mesh="d405_solid"/>
+                    <!-- <geom class="collision" pos="0 -0.0824748 -0.0095955" quat="0 0 -0.21644 -0.976296" type="box" mesh="d405_solid"/> -->
+                    <geom name="right_camera" class="collision" pos="0 -0.08 -0.02" quat="0 0 -0.21644 -0.976296" type="sphere" size=".035"/>
                     <camera name="wrist_cam_right" pos="0 -0.0824748 -0.0095955" mode="fixed" euler="2.70525955359 0 0"
                             focal="1.93e-3 1.93e-3" resolution="1280 720" sensorsize="3896e-6 2140e-6"/>
                     <body name="right/left_finger_link" pos="0.0191 -0.0141637 0.0211727" quat="1 -1 -1 1">
@@ -246,7 +250,7 @@
                       <geom pos="0.0141637 0.0211727 0.06" class="visual" quat="1 1 1 -1" type="mesh"
                         mesh="vx300s_8_custom_finger_left"/>
                       <!-- <geom pos="0.0141637 0.0211727 0.06" class="collision" quat="1 1 1 -1" type="mesh" mesh="vx300s_8_custom_finger_left"/> -->
-                      <geom class="collision" size="0.0137306 0.0460639" pos="0.0141512 -0.0347039 0.0119043" quat="0.544446 0.455511 0.456118 -0.536698"/>
+                      <geom name="right_left_finger_coll" class="collision" size="0.0137306 0.0460639" pos="0.0141512 -0.0347039 0.0119043" quat="0.544446 0.455511 0.456118 -0.536698"/>
                       <geom name="right/left_g0" pos="0.013  -0.0892 0.0268" class="sphere_collision"/>
                       <geom name="right/left_g1" pos="0.0222 -0.0892 0.0268" class="sphere_collision"/>
                       <geom name="right/left_g2" pos="0.0182 -0.0845 0.0266" class="sphere_collision"/>
@@ -259,7 +263,7 @@
                       <geom pos="0.0141637 -0.0211727 0.0597067" class="visual" quat="1 -1 -1 -1" type="mesh"
                         mesh="vx300s_8_custom_finger_right"/>
                       <!-- <geom pos="0.0141637 -0.0211727 0.0597067" class="collision" quat="1 -1 -1 -1" type="mesh" mesh="vx300s_8_custom_finger_right"/> -->
-                      <geom class="collision" size="0.013729 0.0460644" pos="0.0141715 0.0347097 0.011759" quat="0.538959 -0.458738 -0.452903 -0.542183"/>
+                      <geom name="right_right_finger_coll" class="collision" size="0.013729 0.0460644" pos="0.0141715 0.0347097 0.011759" quat="0.538959 -0.458738 -0.452903 -0.542183"/>
                       <geom name="right/right_g0" pos="0.013  0.0892 0.0268" class="sphere_collision"/>
                       <geom name="right/right_g1" pos="0.0222 0.0892 0.0268" class="sphere_collision"/>
                       <geom name="right/right_g2" pos="0.0182 0.0845 0.0266" class="sphere_collision"/>
@@ -287,4 +291,15 @@
 
   <include file="joint_position_actuators.xml"/>
   <include file="keyframe_ctrl.xml"/>
+
+  <!-- <sensor>
+    <fromto geom1="left_left_finger_coll" geom2="right_left_finger_coll" cutoff="1"/>
+    <fromto geom1="left_left_finger_coll" geom2="right_right_finger_coll" cutoff="1"/>
+    <fromto geom1="left_right_finger_coll" geom2="right_left_finger_coll" cutoff="1"/>
+    <fromto geom1="left_right_finger_coll" geom2="right_right_finger_coll" cutoff="1"/>
+    <fromto geom1="left_left_finger_coll" geom2="table" cutoff="1"/>
+    <fromto geom1="left_right_finger_coll" geom2="table" cutoff="1"/>
+    <fromto geom1="right_left_finger_coll" geom2="table" cutoff="1"/>
+    <fromto geom1="right_right_finger_coll" geom2="table" cutoff="1"/>
+  </sensor> -->
 </mujoco>
diff --git a/examples/aloha/keyframe_ctrl.xml b/examples/aloha/keyframe_ctrl.xml
@@ -1,10 +1,10 @@
 <mujoco>
   <keyframe>
     <key name="neutral_pose" qpos="
-    0 -0.96 1.16 0 -0.3 0 0.0084 0.0084
-    0 -0.96 1.16 0 -0.3 0 0.0084 0.0084"
+    0 -0.96 1.16 0 -0.3 0 0.037 0.037
+    0 -0.96 1.16 0 -0.3 0 0.037 0.037"
       ctrl="
-    0 -0.96 1.16 0 -0.3 0 0.0084
-    0 -0.96 1.16 0 -0.3 0 0.0084"/>
+    0 -0.96 1.16 0 -0.3 0 0.037
+    0 -0.96 1.16 0 -0.3 0 0.037"/>
   </keyframe>
 </mujoco>
diff --git a/examples/aloha/scene.xml b/examples/aloha/scene.xml
@@ -126,7 +126,7 @@
       <geom type="box" size=".05 .05 .05" contype="0" conaffinity="0" rgba=".6 .3 .3 .2"/>
     </body>
     <body name="right/target" pos="0.5 0 .5" quat="0 1 0 0" mocap="true">
-      <geom type="box" size=".05 .05 .05" contype="0" conaffinity="0" rgba=".6 .3 .3 .2"/>
+      <geom type="box" size=".05 .05 .05" contype="0" conaffinity="0" rgba=".3 .3 .6 .2"/>
     </body>
   </worldbody>
 </mujoco>
diff --git a/examples/arm_aloha.py b/examples/arm_aloha.py
@@ -57,6 +57,7 @@
             orientation_cost=1.0,
             lm_damping=1.0,
         ),
+        posture_task := mink.PostureTask(model, cost=1e-4),
     ]
 
     # Enable collision avoidance between the following geoms:
@@ -65,10 +66,12 @@
     # geoms starting at subtree "right wrist" - geoms starting at subtree "left wrist".
     l_wrist_geoms = mink.get_subtree_geom_ids(model, model.body("left/wrist_link").id)
     r_wrist_geoms = mink.get_subtree_geom_ids(model, model.body("right/wrist_link").id)
+    l_geoms = mink.get_subtree_geom_ids(model, model.body("left/upper_arm_link").id)
+    r_geoms = mink.get_subtree_geom_ids(model, model.body("right/upper_arm_link").id)
     frame_geoms = mink.get_body_geom_ids(model, model.body("metal_frame").id)
     collision_pairs = [
         (l_wrist_geoms, r_wrist_geoms),
-        (l_wrist_geoms + r_wrist_geoms, frame_geoms + ["table"]),
+        (l_geoms + r_geoms, frame_geoms + ["table"]),
     ]
     collision_avoidance_limit = mink.CollisionAvoidanceLimit(
         model=model,
@@ -86,9 +89,9 @@
     l_mid = model.body("left/target").mocapid[0]
     r_mid = model.body("right/target").mocapid[0]
     solver = "quadprog"
-    pos_threshold = 1e-4
-    ori_threshold = 1e-4
-    max_iters = 20
+    pos_threshold = 1e-2
+    ori_threshold = 1e-2
+    max_iters = 2
 
     with mujoco.viewer.launch_passive(
         model=model, data=data, show_left_ui=False, show_right_ui=False
@@ -99,6 +102,7 @@
         mujoco.mj_resetDataKeyframe(model, data, model.key("neutral_pose").id)
         configuration.update(data.qpos)
         mujoco.mj_forward(model, data)
+        posture_task.set_target_from_configuration(configuration)
 
         # Initialize mocap targets at the end-effector site.
         mink.move_mocap_to_frame(model, data, "left/target", "left/gripper", "site")
@@ -118,7 +122,7 @@
                     rate.dt,
                     solver,
                     limits=limits,
-                    damping=1e-3,
+                    damping=1e-5,
                 )
                 configuration.integrate_inplace(vel, rate.dt)
 

diff --git a/mink/configuration.py b/mink/configuration.py
@@ -43,8 +43,8 @@ def __init__(
 
         Args:
             model: Mujoco model.
-            q: Configuration to initialize from. If None, the configuration
-                is initialized to the default configuration `qpos0`.
+            q: Configuration to initialize from. If None, the configuration is
+                initialized to the default configuration `qpos0`.
         """
         self.model = model
         self.data = mujoco.MjData(model)

diff --git a/mink/limits/collision_avoidance_limit.py b/mink/limits/collision_avoidance_limit.py
@@ -19,6 +19,23 @@
 
 @dataclass(frozen=True)
 class Contact:
+    """Struct to store contact information between two geoms.
+
+    Attributes:
+        dist: Smallest signed distance between geom1 and geom2. If no collision of
+            distance smaller than distmax is found, this value is equal to distmax [1].
+        fromto: Segment connecting the closest points on geom1 and geom2. The first
+            three elements are the coordinates of the closest point on geom1, and the
+            last three elements are the coordinates of the closest point on geom2.
+        geom1: ID of geom1.
+        geom2: ID of geom2.
+        distmax: Maximum distance between geom1 and geom2.
+
+    References:
+        [1] MuJoCo API documentation. `mj_geomDistance` function.
+            https://mujoco.readthedocs.io/en/latest/APIreference/APIfunctions.html
+    """
+
     dist: float
     fromto: np.ndarray
     geom1: int
@@ -27,12 +44,16 @@ class Contact:
 
     @property
     def normal(self) -> np.ndarray:
+        """Contact normal pointing from geom1 to geom2."""
         normal = self.fromto[3:] - self.fromto[:3]
-        return normal / (np.linalg.norm(normal) + 1e-9)
+        mujoco.mju_normalize3(normal)
+        return normal
 
     @property
     def inactive(self) -> bool:
-        return self.dist == self.distmax and not self.fromto.any()
+        """Returns True if no distance smaller than distmax is detected between geom1
+        and geom2."""
+        return self.dist == self.distmax
 
 
 def compute_contact_normal_jacobian(

diff --git a/mink/limits/limit.py b/mink/limits/limit.py
@@ -15,9 +15,7 @@ class Constraint(NamedTuple):
     """
 
     G: Optional[np.ndarray] = None
-    """Shape (nv, nv)."""
     h: Optional[np.ndarray] = None
-    """Shape (nv,)."""
 
     @property
     def inactive(self) -> bool:

diff --git a/mink/limits/velocity_limit.py b/mink/limits/velocity_limit.py
@@ -18,11 +18,12 @@ class VelocityLimit(Limit):
     Floating base joints are ignored.
 
     Attributes:
-        indices: Tangent indices corresponding to velocity-limited joints.
+        indices: Tangent indices corresponding to velocity-limited joints. Shape (nb,).
         limit: Maximum allowed velocity magnitude for velocity-limited joints, in
-            [m]/[s] for slide joints and [rad]/[s] for hinge joints.
+            [m]/[s] for slide joints and [rad]/[s] for hinge joints. Shape (nb,).
         projection_matrix: Projection from tangent space to subspace with
-            velocity-limited joints.
+            velocity-limited joints. Shape (nb, nv) where nb is the dimension of the
+            velocity-limited subspace and nv is the dimension of the tangent space.
     """
 
     indices: np.ndarray
@@ -46,26 +47,26 @@ def __init__(
         for joint_name, max_vel in velocities.items():
             jid = model.joint(joint_name).id
             jnt_type = model.jnt_type[jid]
-            jnt_dim = dof_width(jnt_type)
-            jnt_id = model.jnt_dofadr[jid]
             if jnt_type == mujoco.mjtJoint.mjJNT_FREE:
                 raise LimitDefinitionError(f"Free joint {joint_name} is not supported")
+            vadr = model.jnt_dofadr[jid]
+            vdim = dof_width(jnt_type)
             max_vel = np.atleast_1d(max_vel)
-            if max_vel.shape != (jnt_dim,):
+            if max_vel.shape != (vdim,):
                 raise LimitDefinitionError(
-                    f"Joint {joint_name} must have a limit of shape ({jnt_dim},). "
+                    f"Joint {joint_name} must have a limit of shape ({vdim},). "
                     f"Got: {max_vel.shape}"
                 )
-            index_list.extend(range(jnt_id, jnt_id + jnt_dim))
+            index_list.extend(range(vadr, vadr + vdim))
             limit_list.extend(max_vel.tolist())
 
         self.indices = np.array(index_list)
         self.indices.setflags(write=False)
         self.limit = np.array(limit_list)
         self.limit.setflags(write=False)
 
-        dim = len(self.indices)
-        self.projection_matrix = np.eye(model.nv)[self.indices] if dim > 0 else None
+        nb = len(self.indices)
+        self.projection_matrix = np.eye(model.nv)[self.indices] if nb > 0 else None
 
     def compute_qp_inequalities(
         self, configuration: Configuration, dt: float
@@ -89,7 +90,8 @@ def compute_qp_inequalities(
 
         Returns:
             Pair :math:`(G, h)` representing the inequality constraint as
-            :math:`G \Delta q \leq h`, or ``None`` if there is no limit.
+            :math:`G \Delta q \leq h`, or ``None`` if there is no limit. G has
+            shape (2nb, nv) and h has shape (2nb,).
         """
         del configuration  # Unused.
         if self.projection_matrix is None:

diff --git a/tests/test_velocity_limit.py b/tests/test_velocity_limit.py
@@ -7,7 +7,6 @@
 
 from mink import Configuration
 from mink.limits import LimitDefinitionError, VelocityLimit
-from mink.utils import get_freejoint_dims
 
 
 class TestVelocityLimit(absltest.TestCase):
@@ -22,22 +21,23 @@ def setUp(self):
         self.configuration.update_from_keyframe("stand")
         # NOTE(kevin): These velocities are arbitrary and do not match real hardware.
         self.velocities = {
-            self.model.joint(i).name: 3.14 for i in range(1, self.model.njnt)
+            self.model.joint(i).name: np.pi for i in range(1, self.model.njnt)
         }
 
     def test_dimensions(self):
         limit = VelocityLimit(self.model, self.velocities)
         nv = self.configuration.nv
-        nb = nv - len(get_freejoint_dims(self.model)[1])
+        nb = nv - 6
         self.assertEqual(len(limit.indices), nb)
         self.assertEqual(limit.projection_matrix.shape, (nb, nv))
 
     def test_indices(self):
         limit = VelocityLimit(self.model, self.velocities)
-        expected = np.arange(6, self.model.nv)  # Freejoint (0-5) is not limited.
+        expected = np.arange(6, self.model.nv)
         self.assertTrue(np.allclose(limit.indices, expected))
 
     def test_model_with_no_limit(self):
+        """Tests that a velocity limit applied to a model with no limits is a no-op."""
         empty_model = mujoco.MjModel.from_xml_string("<mujoco></mujoco>")
         empty_bounded = VelocityLimit(empty_model)
         self.assertEqual(len(empty_bounded.indices), 0)
@@ -47,25 +47,23 @@ def test_model_with_no_limit(self):
         self.assertIsNone(h)
 
     def test_model_with_subset_of_velocities_limited(self):
-        partial_velocities = {}
+        limit_subset = {}
         for i, (key, value) in enumerate(self.velocities.items()):
             if i > 2:
                 break
-            partial_velocities[key] = value
-        limit = VelocityLimit(self.model, partial_velocities)
+            limit_subset[key] = value
+        limit = VelocityLimit(self.model, limit_subset)
         nb = 3
         nv = self.model.nv
         self.assertEqual(limit.projection_matrix.shape, (nb, nv))
         self.assertEqual(len(limit.indices), nb)
-        expected_limit = np.asarray(
-            [
-                3.14,
-            ]
-            * nb
-        )
-        np.testing.assert_allclose(limit.limit, expected_limit)
+        np.testing.assert_allclose(limit.limit, np.asarray([np.pi] * nb))
+        G, h = limit.compute_qp_inequalities(self.configuration, 1e-3)
+        self.assertEqual(G.shape, (2 * nb, nv))
+        self.assertEqual(h.shape, (2 * nb,))
 
     def test_model_with_ball_joint(self):
+        """Test that ball joints are handled correctly."""
         xml_str = """
         <mujoco>
           <worldbody>
@@ -91,6 +89,7 @@ def test_model_with_ball_joint(self):
         self.assertEqual(limit.projection_matrix.shape, (nb, model.nv))
 
     def test_ball_joint_invalid_limit_shape(self):
+        """Ball joints should have limits of shape (3,)."""
         xml_str = """
         <mujoco>
           <worldbody>
@@ -115,6 +114,7 @@ def test_ball_joint_invalid_limit_shape(self):
         self.assertEqual(str(cm.exception), expected_error_message)
 
     def test_that_freejoint_raises_error(self):
+        """Trying to apply a velocity limit to a freejoint should raise an error."""
         xml_str = """
         <mujoco>
           <worldbody>