Add spatial tendons with site wrapping to MJX.

PiperOrigin-RevId: 673783473
Change-Id: I09b13787e3e0914c8de21d81db0347c4638b2f1d

doc/changelog.rst

@@ -19,6 +19,7 @@ General
      for any object type.
    - The experimental use of 2D/3D elasticity plugins with :ref:`composite<body-composite>` has been removed. Users
      should instead use :ref:`flexcomp<body-flexcomp>`, which provides the correct collision behavior.
+
 - Added the :ref:`nativeccd<option-flag-nativeccd>` flag. When this flag is enabled, general convex collision
   detection is handled natively, as opposed to using `libccd <https://github.com/danfis/libccd>`__. This feature is in
   early stages of testing.
@@ -63,6 +64,7 @@ MJX
 - Fixed a bug where ``qLDiagInv`` had the wrong size for sparse mass matrices.
 - Added support for joint and tendon :ref:`frictionloss <coFriction>`.
 - Added support for :ref:`connect<equality-connect>` equality constraints using two sites.
+- Added support for :ref:`spatial tendons <tendon-spatial>` with site wrapping.
 
 Bug fixes
 ^^^^^^^^^
@@ -73,7 +75,7 @@ Bug fixes
 - Fixed a bug in tendon wrapping around spheres. Before this fix, tendons that wrapped around spheres with an
   externally-placed :ref:`sidesite<spatial-geom-sidesite>` could jump inside the sphere instead of wrapping around it.
 - Fixed a bug that caused :at:`meshdir` and :at:`texturedir` to be overwritten during model
-  :ref:`attachment<meAttachment>.
+  :ref:`attachment<meAttachment>`.
 
 Python bindings
 ^^^^^^^^^^^^^^^

doc/mjx.rst

@@ -195,6 +195,8 @@ The following features are **fully supported** in MJX:
      - ``FIXED``, ``AFFINE``
    * - :ref:`Actuator Bias <mjtBias>`
      - ``NONE``, ``AFFINE``
+   * - :ref:`Tendon Wrapping <mjtWrap>`
+     - ``JOINT``, ``SITE``
    * - :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.
@@ -244,7 +246,7 @@ The following features are **in development** and coming soon:
    * - :ref:`Actuator Bias <mjtBias>`
      - ``MUSCLE``
    * - :ref:`Tendon Wrapping <mjtWrap>`
-     - ``NONE``, ``JOINT``, ``PULLEY``, ``SITE``, ``SPHERE``, ``CYLINDER``
+     - ``PULLEY``, ``SPHERE``, ``CYLINDER``
    * - Fluid Model
      - :ref:`flEllipsoid`
    * - :ref:`Tendons <tendon>`

mjx/mujoco/mjx/_src/io_test.py

@@ -168,6 +168,7 @@ class ModelIOTest(parameterized.TestCase):
           <tendon>
             <spatial name="rope" range="0 .35">
               <site site="slider"/>
+              <geom geom="shoulder"/>
               <site site="arm"/>
             </spatial>
           </tendon>

mjx/mujoco/mjx/_src/smooth.py

@@ -27,6 +27,7 @@ from mujoco.mjx._src.types import DisableBit
 from mujoco.mjx._src.types import JointType
 from mujoco.mjx._src.types import Model
 from mujoco.mjx._src.types import TrnType
+from mujoco.mjx._src.types import WrapType
 # pylint: enable=g-importing-member
 import numpy as np
 
@@ -686,13 +687,106 @@ def tendon(m: Model, d: Data) -> Data:
   if not m.ntendon:
     return d
 
-  ten_id = np.repeat(np.arange(m.ntendon), m.tendon_num)
+  # process joint tendons
-  length = m.wrap_prm * d.qpos[m.jnt_qposadr[m.wrap_objid]]
+  (wrap_id_jnt,) = np.nonzero(m.wrap_type == WrapType.JOINT)
-  ten_length = jax.ops.segment_sum(length, ten_id, m.ntendon)
+  (tendon_id_jnt,) = np.nonzero(np.isin(m.tendon_adr, wrap_id_jnt))
-  ten_j = jp.zeros((m.ntendon, m.nv))
-  ten_j = ten_j.at[ten_id, m.jnt_dofadr[m.wrap_objid]].set(m.wrap_prm)
 
-  return d.replace(ten_length=ten_length, ten_J=ten_j)
+  ntendon_jnt = tendon_id_jnt.size
+  wrap_objid_jnt = m.wrap_objid[wrap_id_jnt]
+  tendon_num_jnt = m.tendon_num[tendon_id_jnt]
+
+  moment_jnt = m.wrap_prm[wrap_id_jnt]
+  length_jnt = jax.ops.segment_sum(
+      moment_jnt * d.qpos[m.jnt_qposadr[wrap_objid_jnt]],
+      np.repeat(np.arange(ntendon_jnt), tendon_num_jnt),
+      ntendon_jnt,
+  )
+
+  adr_moment_jnt = np.repeat(tendon_id_jnt, tendon_num_jnt)
+  dofadr_moment_jnt = m.jnt_dofadr[wrap_objid_jnt]
+
+  # process spatial tendon sites
+  (wrap_id_site,) = np.nonzero(m.wrap_type == WrapType.SITE)
+  nwrap_site = wrap_id_site.size
+
+  # find consecutive sites, skipping tendon transitions
+  (pair_id,) = np.nonzero(np.diff(wrap_id_site) == 1)
+  wrap_id_site_pair = np.setdiff1d(wrap_id_site[pair_id], m.tendon_adr[1:] - 1)
+  (tendon_id_site,) = np.nonzero(np.isin(m.tendon_adr, wrap_id_site_pair))
+
+  id0 = m.wrap_objid[wrap_id_site_pair]
+  id1 = m.wrap_objid[wrap_id_site_pair + 1]
+
+  @jax.vmap
+  def _length_moment(pnt0, pnt1, body0, body1):
+    dif = pnt1 - pnt0
+    length = jp.linalg.norm(dif)
+    vec = jp.where(
+        length < mujoco.mjMINVAL, jp.array([1.0, 0.0, 0.0]), dif / length
+    )
+
+    jacp1, _ = support.jac(m, d, pnt0, body0)
+    jacp2, _ = support.jac(m, d, pnt1, body1)
+    jacdif = jacp2 - jacp1
+    moment = jp.where(body0 != body1, jacdif @ vec, jp.zeros(m.nv))
+
+    return length, moment
+
+  lengths_site, moments_site = _length_moment(
+      d.site_xpos[id0], d.site_xpos[id1], m.site_bodyid[id0], m.site_bodyid[id1]
+  )
+
+  tendon_nsite = np.array([
+      sum((wrap_id_site_pair >= adr) & (wrap_id_site_pair < adr + num))
+      for adr, num in zip(m.tendon_adr, m.tendon_num)
+  ])
+  tendon_nsite = tendon_nsite[tendon_nsite > 0]
+  tendon_wrapnum_site = tendon_nsite + 1
+  tendon_with_site = sum([s > 0 for s in tendon_nsite])
+  ten_site_id = np.repeat(np.arange(tendon_with_site), tendon_nsite)
+
+  length_site = jax.ops.segment_sum(lengths_site, ten_site_id, tendon_with_site)
+  moment_site = jax.ops.segment_sum(moments_site, ten_site_id, tendon_with_site)
+
+  # assemble length and moment
+  ten_length = (
+      jp.zeros_like(d.ten_length)
+      .at[np.concatenate([tendon_id_jnt, tendon_id_site])]
+      .set(jp.concatenate([length_jnt, length_site]))
+  )
+  ten_moment = (
+      jp.zeros_like(d.ten_J)
+      .at[adr_moment_jnt, dofadr_moment_jnt]
+      .set(moment_jnt)
+  )
+  ten_moment = ten_moment.at[tendon_id_site].set(moment_site)
+
+  # wrap
+  wrap_xpos = jp.concatenate([
+      d.site_xpos[m.wrap_objid[wrap_id_site]],
+      jp.zeros((2 * m.nwrap - nwrap_site, 3)),
+  ]).reshape((m.nwrap, 6))
+
+  ten_wrapnum = np.zeros(m.ntendon)
+  ten_wrapnum[tendon_id_site] = tendon_wrapnum_site
+
+  ten_wrapadr = [0]
+  for wn in ten_wrapnum[:-1]:
+    ten_wrapadr.append(ten_wrapadr[-1] + wn)
+  ten_wrapadr = np.array(ten_wrapadr).astype(int)
+
+  wrap_obj = np.zeros(m.nwrap * 2, dtype=int)
+  wrap_obj[:nwrap_site] = -1
+  wrap_obj = wrap_obj.reshape((-1, 2))
+
+  return d.replace(
+      ten_length=ten_length,
+      ten_J=ten_moment,
+      ten_wrapadr=jp.array(ten_wrapadr),
+      ten_wrapnum=jp.array(ten_wrapnum),
+      wrap_xpos=wrap_xpos,
+      wrap_obj=jp.array(wrap_obj),
+  )
 
 
 def _site_dof_mask(m: Model) -> np.ndarray:

mjx/mujoco/mjx/_src/smooth_test.py

@@ -15,6 +15,7 @@
 """Tests for smooth dynamics functions."""
 
 from absl.testing import absltest
+from absl.testing import parameterized
 import jax
 import mujoco
 from mujoco import mjx
@@ -230,5 +231,36 @@ class SmoothTest(absltest.TestCase):
     _assert_eq(d.cfrc_int, dx.cfrc_int, 'cfrc_int')
 
 
+class TendonTest(parameterized.TestCase):
+
+  @parameterized.parameters(
+      'tendon/fixed.xml',
+      'tendon/site.xml',
+      'tendon/fixed_site.xml',
+      'tendon/no_tendon.xml',
+  )
+  def test_tendon(self, filename):
+    """Tests MJX tendon function matches MuJoCo mj_tendon."""
+    m = test_util.load_test_file(filename)
+    d = mujoco.MjData(m)
+    # give the system a little kick to ensure we have non-identity rotations
+    d.qvel = np.random.random(m.nv)
+    mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    mujoco.mj_forward(m, d)
+    dx = jax.jit(mjx.forward)(mx, dx)
+
+    _assert_eq(d.ten_length, dx.ten_length, 'ten_length')
+    _assert_eq(d.ten_J, dx.ten_J, 'ten_J')
+    _assert_eq(d.ten_wrapnum, dx.ten_wrapnum, 'ten_wrapnum')
+    _assert_eq(d.ten_wrapadr, dx.ten_wrapadr, 'ten_wrapadr')
+    if d.wrap_obj.shape == dx.wrap_obj.shape:
+      _assert_eq(d.wrap_obj, dx.wrap_obj, 'wrap_obj')
+    if d.wrap_xpos.shape == dx.wrap_xpos.shape:
+      _assert_eq(d.wrap_xpos, dx.wrap_xpos, 'wrap_xpos')
+
+
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/types.py

@@ -197,9 +197,11 @@ class WrapType(enum.IntEnum):
 
   Members:
     JOINT: constant moment arm
+    SITE: pass through site
   """
   JOINT = mujoco.mjtWrap.mjWRAP_JOINT
-  # unsupported: NONE, PULLEY, SITE, SPHERE, CYLINDER
+  SITE = mujoco.mjtWrap.mjWRAP_SITE
+  # unsupported: NONE, PULLEY, SPHERE, CYLINDER
 
 
 class TrnType(enum.IntEnum):

mjx/mujoco/mjx/test_data/tendon/fixed.xml

@@ -0,0 +1,23 @@
+<mujoco model="fixed">
+  <worldbody>
+    <body>
+      <joint name="joint0" type="hinge" axis="0 1 0"/>
+      <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+      <body pos="0.5 0 0">
+        <joint name="joint1" type="hinge" axis="0 1 0"/>
+        <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+        <body pos="0.5 0 0">
+          <joint name="joint2" type="hinge" axis="0 1 0"/>
+          <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <tendon>
+    <fixed>
+      <joint joint="joint0" coef=".1"/>
+      <joint joint="joint1" coef=".2"/>
+      <joint joint="joint2" coef=".3"/>
+    </fixed>
+  </tendon>
+</mujoco>

mjx/mujoco/mjx/test_data/tendon/fixed_site.xml

@@ -0,0 +1,39 @@
+<mujoco model="fixed_site">
+  <worldbody>
+    <body>
+      <joint name="joint0" type="hinge" axis="0 1 0"/>
+      <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+      <site name="site0" pos="0.25 0 0.1" size="0.025"/>
+      <body pos="0.5 0 0">
+        <joint name="joint1" type="hinge" axis="0 1 0"/>
+        <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+        <site name="site1" pos="0.25 0 0.1" size="0.025"/>
+        <body pos="0.5 0 0">
+          <joint name="joint2" type="hinge" axis="0 1 0"/>
+          <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+          <site name="site2" pos="0.25 0 0.1" size="0.025"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <tendon>
+    <spatial width="0.0125">
+      <site site="site0"/>
+      <site site="site1"/>
+      <site site="site2"/>
+    </spatial>
+    <fixed>
+      <joint joint="joint0" coef=".1"/>
+      <joint joint="joint1" coef=".2"/>
+      <joint joint="joint2" coef=".3"/>
+    </fixed>
+    <fixed>
+      <joint joint="joint0" coef=".1"/>
+      <joint joint="joint2" coef=".3"/>
+    </fixed>
+    <spatial width="0.0125">
+      <site site="site1"/>
+      <site site="site2"/>
+    </spatial>
+  </tendon>
+</mujoco>

mjx/mujoco/mjx/test_data/tendon/no_tendon.xml

@@ -0,0 +1,16 @@
+<mujoco model="no_tendon">
+  <worldbody>
+    <body>
+      <joint type="hinge" axis="0 1 0"/>
+      <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+      <body pos="0.5 0 0">
+        <joint type="hinge" axis="0 1 0"/>
+        <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+        <body pos="0.5 0 0">
+          <joint type="hinge" axis="0 1 0"/>
+          <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+</mujoco>

mjx/mujoco/mjx/test_data/tendon/site.xml

@@ -0,0 +1,26 @@
+<mujoco model="site">
+  <worldbody>
+    <body>
+      <joint type="hinge" axis="0 1 0"/>
+      <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+      <site name="site0" pos="0.25 0 0.1" size="0.025"/>
+      <body pos="0.5 0 0">
+        <joint type="hinge" axis="0 1 0"/>
+        <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+        <site name="site1" pos="0.25 0 0.1" size="0.025"/>
+        <body pos="0.5 0 0">
+          <joint type="hinge" axis="0 1 0"/>
+          <geom type="capsule" size="0.05 0.5" fromto="0 0 0 0.5 0 0"/>
+          <site name="site2" pos="0.25 0 0.1" size="0.025"/>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <tendon>
+    <spatial width="0.0125">
+      <site site="site0"/>
+      <site site="site1"/>
+      <site site="site2"/>
+    </spatial>
+  </tendon>
+</mujoco>