Add sphere and cylinder internal wrapping for spatial tendons to MJX.

PiperOrigin-RevId: 720507688
Change-Id: I49fda2839e68623ff5e1f3c46bdc3d77acfc36ef

doc/changelog.rst

@@ -13,6 +13,10 @@ General
   attaching can be restored by setting the deep copy flag to 1.
 - Added :ref:`potential<sensor-e_potential>` and :ref:`kinetic<sensor-e_kinetic>` energy sensors.
 
+MJX
+^^^
+- Added support for spatial tendons with internal sphere and cylinder wrapping.
+
 Version 3.2.7 (Jan 14, 2025)
 ----------------------------
 

doc/mjx.rst

@@ -220,7 +220,7 @@ The following features are **fully supported** in MJX:
    * - :ref:`Actuator Bias <mjtBias>`
      - ``NONE``, ``AFFINE``, ``MUSCLE``
    * - :ref:`Tendon Wrapping <mjtWrap>`
-     - ``JOINT``, ``SITE``, ``PULLEY``
+     - ``JOINT``, ``SITE``, ``PULLEY``, ``SPHERE``, ``CYLINDER``
    * - :ref:`Geom <mjtGeom>`
      - ``PLANE``, ``HFIELD``, ``SPHERE``, ``CAPSULE``, ``BOX``, ``MESH`` are fully implemented. ``ELLIPSOID`` and
        ``CYLINDER`` are implemented but only collide with other primitives, note that ``BOX`` is implemented as a mesh.
@@ -239,7 +239,7 @@ The following features are **fully supported** in MJX:
    * - Fluid Model
      - :ref:`flInertia`
    * - :ref:`Tendons <tendon>`
-     - :ref:`Fixed <tendon-fixed>`
+     - :ref:`Fixed <tendon-fixed>`, :ref:`Spatial <tendon-spatial>`
    * - :ref:`Sensors <mjtSensor>`
      - ``MAGNETOMETER``, ``CAMPROJECTION``, ``RANGEFINDER``, ``JOINTPOS``, ``TENDONPOS``, ``ACTUATORPOS``, ``BALLQUAT``,
        ``FRAMEPOS``, ``FRAMEXAXIS``, ``FRAMEYAXIS``, ``FRAMEZAXIS``, ``FRAMEQUAT``, ``SUBTREECOM``, ``CLOCK``,
@@ -265,12 +265,8 @@ The following features are **in development** and coming soon:
      - ``IMPLICIT``
    * - Dynamics
      - :ref:`Inverse <mj_inverse>`
-   * - :ref:`Tendon Wrapping <mjtWrap>`
-     - ``SPHERE``, ``CYLINDER`` (external wrapping is supported)
    * - Fluid Model
      - :ref:`flEllipsoid`
-   * - :ref:`Tendons <tendon>`
-     - :ref:`Spatial <tendon-spatial>`
    * - :ref:`Sensors <mjtSensor>`
      - All except ``PLUGIN``, ``USER``
    * - Lights

mjx/mujoco/mjx/_src/io.py

@@ -123,29 +123,6 @@ def put_model(
       if t == mujoco.mjtGeom.mjGEOM_MESH:
         mesh_geomid.add(g)
 
-  # check for spatial tendon internal geom wrapping
-  if m.ntendon:
-    # find sphere or cylinder geoms (if any exist)
-    (wrap_id_geom,) = np.nonzero(
-        (m.wrap_type == mujoco.mjtWrap.mjWRAP_SPHERE)
-        | (m.wrap_type == mujoco.mjtWrap.mjWRAP_CYLINDER)
-    )
-    wrap_objid_geom = m.wrap_objid[wrap_id_geom]
-    geom_pos = m.geom_pos[wrap_objid_geom]
-    geom_size = m.geom_size[wrap_objid_geom, 0]
-
-    # find sidesites (if any exist)
-    side_id = np.round(m.wrap_prm[wrap_id_geom]).astype(int)
-    side = m.site_pos[side_id]
-
-    # check for sidesite inside geom
-    if np.any(
-        (np.linalg.norm(side - geom_pos, axis=1) < geom_size) & (side_id >= 0)
-    ):
-      raise NotImplementedError(
-          'Internal wrapping with sphere and cylinder geoms is not'
-          ' implemented for spatial tendons.'
-      )
 
   # check for unsupported sensor and equality constraint combinations
   sensor_rne_postconstraint = (
@@ -199,6 +176,30 @@ def put_model(
   fields['opt'] = _make_option(m.opt, _full_compat=_full_compat)
   fields['stat'] = _make_statistic(m.stat)
 
+  # spatial tendon wrap inside
+  fields['wrap_inside_maxiter'] = 5
+  fields['wrap_inside_tolerance'] = 1.0e-4
+  fields['wrap_inside_z_init'] = 1.0 - 1.0e-5
+  fields['is_wrap_inside'] = np.zeros(0, dtype=bool)
+  if m.nsite:
+    # find sphere or cylinder geoms (if any exist)
+    (wrap_id_geom,) = np.nonzero(
+        (m.wrap_type == mujoco.mjtWrap.mjWRAP_SPHERE)
+        | (m.wrap_type == mujoco.mjtWrap.mjWRAP_CYLINDER)
+    )
+    wrap_objid_geom = m.wrap_objid[wrap_id_geom]
+    geom_pos = m.geom_pos[wrap_objid_geom]
+    geom_size = m.geom_size[wrap_objid_geom, 0]
+
+    # find sidesites (if any exist)
+    side_id = np.round(m.wrap_prm[wrap_id_geom]).astype(int)
+    side = m.site_pos[side_id]
+
+    # wrap inside flag
+    fields['is_wrap_inside'] = np.array(
+        (np.linalg.norm(side - geom_pos, axis=1) < geom_size) & (side_id >= 0)
+    )
+
   # Pre-compile meshes for MJX collisions.
   fields['mesh_convex'] = [None] * m.nmesh
   if not _full_compat:

mjx/mujoco/mjx/_src/io_test.py

@@ -20,6 +20,7 @@
 from jax import numpy as jp
 import mujoco
 from mujoco import mjx
+from mujoco.mjx._src import test_util
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import ConeType
 # pylint: enable=g-importing-member
@@ -171,33 +172,6 @@ def test_pgs_not_implemented(self):
           )
       )
 
-  def test_spatial_tendon_not_implemented(self):
-    with self.assertRaises(NotImplementedError):
-      mjx.put_model(mujoco.MjModel.from_xml_string("""
-        <mujoco>
-          <worldbody>
-            <body name="arm">
-              <joint name="arm" axis="0 1 0"/>
-              <geom name="shoulder" type="sphere" size=".05"/>
-              <site name="arm" pos="-.1 0 .05"/>
-              <site name="sidesite" pos="0 0 0"/>
-            </body>
-            <body name="slider" pos=".05 0 -.2">
-              <joint name="slider" type="slide" damping="1"/>
-              <geom name="slider" type="box" size=".01 .01 .01"/>
-              <site name="slider" pos="0 0 .01"/>
-            </body>
-          </worldbody>
-
-          <tendon>
-            <spatial name="rope" range="0 .35">
-              <site site="slider"/>
-              <geom geom="shoulder" sidesite="sidesite"/>
-              <site site="arm"/>
-            </spatial>
-          </tendon>
-        </mujoco>"""))
-
   def test_margin_gap_mesh_not_implemented(self):
     with self.assertRaises(NotImplementedError):
       mjx.put_model(mujoco.MjModel.from_xml_string("""
@@ -225,6 +199,20 @@ def test_implicitfast_fluid_not_implemented(self):
           <worldbody/>
         </mujoco>"""))
 
+  def test_wrap_inside(self):
+    m = test_util.load_test_file('tendon/wrap_sidesite.xml')
+    mx0 = mjx.put_model(m)
+    np.testing.assert_equal(
+        mx0.is_wrap_inside,
+        np.array([1, 0, 1, 0, 1, 1, 0]),
+    )
+    m.site_pos[2] = m.site_pos[1]
+    mx1 = mjx.put_model(m)
+    np.testing.assert_equal(
+        mx1.is_wrap_inside,
+        np.array([0, 0, 1, 0, 1, 0, 0]),
+    )
+
 
 class DataIOTest(parameterized.TestCase):
   """IO tests for mjx.Data."""

mjx/mujoco/mjx/_src/smooth.py

@@ -812,6 +812,7 @@ def _length_moment(pnt0, pnt1, body0, body1):
   geom_xmat = d.geom_xmat[wrap_objid_geom]
   geom_size = m.geom_size[wrap_objid_geom, 0]
   geom_type = m.wrap_type[wrap_id_geom]
+  is_sphere = geom_type == WrapType.SPHERE
 
   # get body ids for site-geom-site instances
   body_id_site0 = m.site_bodyid[wrap_objid_site0]
@@ -823,17 +824,52 @@ def _length_moment(pnt0, pnt1, body0, body1):
   side = d.site_xpos[side_id]
   has_sidesite = np.expand_dims(np.array(side_id >= 0), -1)
 
+  # wrap inside
+  # TODO(taylorhowell): check that is_wrap_inside is consistent with
+  # site and geom relative positions
+  (wrap_inside_id,) = np.nonzero(m.is_wrap_inside)
+  (wrap_outside_id,) = np.nonzero(~m.is_wrap_inside)
+
   # compute geom wrap length and connect points (if wrap occurs)
-  lengths_geomgeom, geom_pnt0, geom_pnt1 = jax.vmap(support.wrap)(
-      site_pnt0,
-      site_pnt1,
-      geom_xpos,
-      geom_xmat,
-      geom_size,
-      side,
-      has_sidesite,
-      geom_type == WrapType.SPHERE,
+  v_wrap = jax.vmap(
+      support.wrap, in_axes=(0, 0, 0, 0, 0, 0, 0, 0, None, None, None, None)
+  )
+  lengths_inside, pnt0_inside, pnt1_inside = v_wrap(
+      site_pnt0[wrap_inside_id],
+      site_pnt1[wrap_inside_id],
+      geom_xpos[wrap_inside_id],
+      geom_xmat[wrap_inside_id],
+      geom_size[wrap_inside_id],
+      side[wrap_inside_id],
+      has_sidesite[wrap_inside_id],
+      is_sphere[wrap_inside_id],
+      True,
+      m.wrap_inside_maxiter,
+      m.wrap_inside_tolerance,
+      m.wrap_inside_z_init,
+  )
+
+  lengths_outside, pnt0_outside, pnt1_outside = v_wrap(
+      site_pnt0[wrap_outside_id],
+      site_pnt1[wrap_outside_id],
+      geom_xpos[wrap_outside_id],
+      geom_xmat[wrap_outside_id],
+      geom_size[wrap_outside_id],
+      side[wrap_outside_id],
+      has_sidesite[wrap_outside_id],
+      is_sphere[wrap_outside_id],
+      False,
+      m.wrap_inside_maxiter,
+      m.wrap_inside_tolerance,
+      m.wrap_inside_z_init,
   )
+
+  wrap_id = np.argsort(np.concatenate([wrap_inside_id, wrap_outside_id]))
+  vstack_ = lambda x, y: jp.vstack([x, y])[wrap_id]
+  lengths_geomgeom = vstack_(lengths_inside, lengths_outside)
+  geom_pnt0 = vstack_(pnt0_inside, pnt0_outside)
+  geom_pnt1 = vstack_(pnt1_inside, pnt1_outside)
+
   lengths_geomgeom = lengths_geomgeom.reshape(-1)
 
   # identify geoms where wrap does not occur