Add tendon armature to MJX.

PiperOrigin-RevId: 769294284
Change-Id: Idfbd4f355eb26fba5035973ab989a9d3af05c85f

doc/changelog.rst

@@ -36,6 +36,10 @@ Python bindings
 
 - Added examples of procedural terrain generation to the Model Editing tutorial: |mjspec_colab|
 
+MJX
+^^^
+- Added tendon armature.
+
 Version 3.3.2 (April 28, 2025)
 ------------------------------
 

mjx/mujoco/mjx/__init__.py

@@ -48,6 +48,8 @@ from mujoco.mjx._src.smooth import rne
 from mujoco.mjx._src.smooth import rne_postconstraint
 from mujoco.mjx._src.smooth import subtree_vel
 from mujoco.mjx._src.smooth import tendon
+from mujoco.mjx._src.smooth import tendon_armature
+from mujoco.mjx._src.smooth import tendon_bias
 from mujoco.mjx._src.smooth import transmission
 from mujoco.mjx._src.solver import solve
 from mujoco.mjx._src.support import apply_ft

mjx/mujoco/mjx/_src/forward.py

@@ -73,6 +73,7 @@ def fwd_position(m: Model, d: Data) -> Data:
   d = smooth.camlight(m, d)
   d = smooth.tendon(m, d)
   d = smooth.crb(m, d)
+  d = smooth.tendon_armature(m, d)
   d = smooth.factor_m(m, d)
   d = collision_driver.collision(m, d)
   d = constraint.make_constraint(m, d)
@@ -93,6 +94,7 @@ def fwd_velocity(m: Model, d: Data) -> Data:
   d = smooth.com_vel(m, d)
   d = passive.passive(m, d)
   d = smooth.rne(m, d)
+  d = smooth.tendon_bias(m, d)
   return d
 
 

mjx/mujoco/mjx/_src/smooth.py

@@ -1184,3 +1184,155 @@ def transmission(m: Model, d: Data) -> Data:
       {'_impl.actuator_length': length, '_impl.actuator_moment': moment}
   )
   return d
+
+
+def tendon_armature(m: Model, d: Data) -> Data:
+  """Add tendon armature to qM."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('tendon_armature requires JAX backend implementation.')
+
+  if not support.is_sparse(m):
+    return d.tree_replace({
+        '_impl.qM': (
+            d._impl.qM
+            + d._impl.ten_J.T
+            @ jax.vmap(jp.multiply)(d._impl.ten_J, m.tendon_armature)
+        )
+    })
+  else:
+    # TODO(taylorhowell): implement tendon armature with sparse qM
+    raise NotImplementedError(
+        'Tendon armature with sparse qM is not implemented.'
+    )
+
+
+def tendon_dot(m: Model, d: Data) -> jax.Array:
+  """Compute time derivative of dense tendon Jacobian for one tendon."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('tendon_dot requires JAX backend implementation.')
+
+  ten_Jdot = jp.zeros((m.ntendon, m.nv))  # pylint: disable=invalid-name
+
+  if not m.ntendon:
+    return ten_Jdot
+
+  # process pulleys
+  (wrap_id_pulley,) = np.nonzero(m.wrap_type == WrapType.PULLEY)
+
+  divisor = np.ones(m.nwrap)
+  for adr, num in zip(m.tendon_adr, m.tendon_num):
+    for id_pulley in wrap_id_pulley:
+      if adr <= id_pulley < adr + num:
+        divisor[id_pulley : adr + num] = np.maximum(
+            mujoco.mjMINVAL, m.wrap_prm[id_pulley]
+        )
+
+  # process spatial tendon sites
+  (wrap_id_site,) = np.nonzero(m.wrap_type == WrapType.SITE)
+
+  # 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)
+  wrap_objid_site0 = m.wrap_objid[wrap_id_site_pair]
+  wrap_objid_site1 = m.wrap_objid[wrap_id_site_pair + 1]
+  site_bodyid0 = m.site_bodyid[wrap_objid_site0]
+  site_bodyid1 = m.site_bodyid[wrap_objid_site1]
+  site_xpos0 = d.site_xpos[wrap_objid_site0]
+  site_xpos1 = d.site_xpos[wrap_objid_site1]
+  subtree_com0 = d.subtree_com[m.body_rootid[site_bodyid0]]
+  subtree_com1 = d.subtree_com[m.body_rootid[site_bodyid1]]
+  site_vel0 = jax.vmap(lambda a, b: a[3:] - jp.cross(b, a[:3]))(
+      d.cvel[site_bodyid0], site_xpos0 - subtree_com0
+  )
+  site_vel1 = jax.vmap(lambda a, b: a[3:] - jp.cross(b, a[:3]))(
+      d.cvel[site_bodyid1], site_xpos1 - subtree_com1
+  )
+
+  @jax.vmap
+  def _momentdot(wpnt0, wpnt1, wvel0, wvel1, body0, body1):
+    # dpnt = 3D position difference, normalize
+    dpnt = wpnt1 - wpnt0
+    norm = math.norm(dpnt)
+    dpnt = jp.where(
+        norm < mujoco.mjMINVAL, jp.array([1.0, 0.0, 0.0]), dpnt / norm
+    )
+
+    # dvel = d / dt(dpnt)
+    dvel = wvel1 - wvel0
+    dot = jp.dot(dpnt, dvel)
+    dvel += dpnt * -dot
+    dvel = jp.where(norm > mujoco.mjMINVAL, dvel / norm, 0.0)
+
+    # get endpoint JacobianDots, subtract
+    jacp1, _ = support.jac_dot(m, d, wpnt0, body0)
+    jacp2, _ = support.jac_dot(m, d, wpnt1, body1)
+    jacdif = jacp2 - jacp1
+
+    # chain rule, first term: Jdot += d / dt(jac2 - jac1) * dpnt
+    tmp0 = jacdif @ dpnt
+
+    # get endpoint Jacobians, subtract
+    jacp1, _ = support.jac(m, d, wpnt0, body0)
+    jacp2, _ = support.jac(m, d, wpnt1, body1)
+    jacdif = jacp2 - jacp1
+
+    # chain rule, second term: Jdot += (jac2 - jac1) * d/dt (dpnt)
+    tmp1 = jacdif @ dvel
+
+    return jp.where(body0 != body1, tmp0 + tmp1, jp.zeros(m.nv))
+
+  momentdots = _momentdot(
+      site_xpos0,
+      site_xpos1,
+      site_vel0,
+      site_vel1,
+      site_bodyid0,
+      site_bodyid1,
+  )
+
+  if wrap_id_site_pair.size:
+    divisor_site_pair = divisor[wrap_id_site_pair]
+    momentdots /= divisor_site_pair[:, None]
+
+  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_has_site = tendon_nsite > 0
+  (tendon_id_site,) = np.nonzero(tendon_has_site)
+  tendon_nsite = tendon_nsite[tendon_has_site]
+  tendon_with_site = tendon_nsite.size
+  ten_site_id = np.repeat(np.arange(tendon_with_site), tendon_nsite)
+
+  momentdot = jax.ops.segment_sum(momentdots, ten_site_id, tendon_with_site)
+  ten_Jdot = ten_Jdot.at[tendon_id_site].set(momentdot)  # pylint: disable=invalid-name
+
+  # TODO(taylorhowell): time derivatives for geoms
+
+  return ten_Jdot
+
+
+def tendon_bias(m: Model, d: Data) -> Data:
+  """Add bias force due to tendon armature."""
+  if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
+    raise ValueError('tendon_bias requires JAX backend implementation.')
+
+  if not m.ntendon:
+    return d
+
+  # get dense d/dt(tendon Jacobian) for each tendon
+  ten_Jdot = tendon_dot(m, d)  # pylint: disable=invalid-name
+
+  # add bias term: qfrc += ten_J * armature * ten_Jdot @ qvel
+  coef = m.tendon_armature * jp.dot(ten_Jdot, d.qvel)
+
+  if not support.is_sparse(m):
+    return d.tree_replace({
+        'qfrc_bias': (
+            d.qfrc_bias
+            + jp.sum(jax.vmap(jp.multiply)(d._impl.ten_J, coef), axis=0)
+        )
+    })
+  else:
+    # TODO(taylorhowell): implement tendon bias with sparse qM
+    raise NotImplementedError('Tendon bias with sparse qM is not implemented.')

mjx/mujoco/mjx/_src/smooth_test.py

@@ -17,6 +17,7 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
+from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
@@ -316,6 +317,58 @@ class TendonTest(parameterized.TestCase):
     _assert_eq(d.wrap_obj, dx._impl.wrap_obj, 'wrap_obj')
     _assert_eq(d.wrap_xpos, dx._impl.wrap_xpos, 'wrap_xpos')
 
+  def test_tendon_armature(self):
+    """Tests MJX tendon armature matches MuJoCo."""
+    m = mujoco.MjModel.from_xml_string("""
+        <mujoco>
+          <worldbody>
+            <site name="site0" pos="1 0 1"/>
+            <body>
+              <joint type="slide" axis="0 0 1"/>
+              <joint type="hinge" axis="0 1 0"/>
+              <geom type="box" size="0.1 0.1 0.1" mass="1" pos="1 0 0"/>
+              <site name="site1"/>
+            </body>
+          </worldbody>
+          <tendon>
+            <spatial armature="123">
+              <site site="site0"/>
+              <site site="site1"/>
+            </spatial>
+            <spatial armature="456">
+              <site site="site0"/>
+              <site site="site1"/>
+            </spatial>
+          </tendon>
+          <keyframe>
+            <key qpos="1.2345 1.2345" qvel="1.2345 1.2345"/>
+          </keyframe>
+        </mujoco>
+        """)
+
+    d = mujoco.MjData(m)
+    mujoco.mj_resetDataKeyframe(m, d, 0)
+    mujoco.mj_forward(m, d)
+
+    qM = np.zeros((m.nv, m.nv))  # pylint: disable=invalid-name
+    mujoco.mj_fullM(m, qM, d.qM)
+
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+
+    dx = dx.tree_replace(
+        {'_impl.qM': jp.zeros((m.nv, m.nv)), 'qfrc_bias': jp.zeros(m.nv)}
+    )
+
+    dx = mjx.crb(mx, dx)
+    dx = mjx.tendon_armature(mx, dx)
+
+    _assert_eq(dx._impl.qM, qM, 'qM')
+
+    dx = mjx.rne(mx, dx)
+    dx = mjx.tendon_bias(mx, dx)
+    _assert_eq(dx.qfrc_bias, d.qfrc_bias, 'qfrc_bias')
+
 
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/support.py

@@ -142,6 +142,7 @@ def jac(
     m: Model, d: Data, point: jax.Array, body_id: jax.Array
 ) -> Tuple[jax.Array, jax.Array]:
   """Compute pair of (NV, 3) Jacobians of global point attached to body."""
+  # TODO(taylorhowell): statically construct mask
   fn = lambda carry, b: b if carry is None else b + carry
   mask = (jp.arange(m.nbody) == body_id) * 1
   mask = scan.body_tree(m, fn, 'b', 'b', mask, reverse=True)
@@ -155,6 +156,42 @@ def jac(
   return jacp, jacr
 
 
+def jac_dot(
+    m: Model, d: Data, point: jax.Array, body_id: jax.Array
+) -> Tuple[jax.Array, jax.Array]:
+  """Compute pair of (NV, 3) Jacobian time derivatives of global point attached to body."""
+  # TODO(taylorhowell): statically construct mask
+  fn = lambda carry, b: b if carry is None else b + carry
+  mask = (jp.arange(m.nbody) == body_id) * 1
+  mask = scan.body_tree(m, fn, 'b', 'b', mask, reverse=True)
+  mask = mask[jp.array(m.dof_bodyid)] > 0
+
+  offset = point - d.subtree_com[jp.array(m.body_rootid)[body_id]]
+  pvel_lin = d.cvel[body_id][3:] - jp.cross(offset, d.cvel[body_id][:3])
+
+  cdof = d._impl.cdof
+  cdof_dot = d._impl.cdof_dot
+
+  # check for quaternion
+  jnt_type = m.jnt_type[m.dof_jntid]
+  dof_adr = m.jnt_dofadr[m.dof_jntid]
+  is_quat = (jnt_type == JointType.BALL) | (
+      jnt_type == JointType.FREE & (np.arange(m.nv) >= dof_adr + 3)
+  )
+
+  # compute cdof_dot for quaternion (use current body cvel)
+  cdof_dot_quat = jax.vmap(math.motion_cross)(d.cvel[m.dof_bodyid], cdof)
+  cdof_dot = jp.where(is_quat[:, None], cdof_dot_quat, cdof_dot)
+
+  jacp = jax.vmap(
+      lambda a, b: a[3:] + jp.cross(a[:3], offset) + jp.cross(b[:3], pvel_lin)
+  )(cdof_dot, cdof)
+  jacp = jax.vmap(jp.multiply)(jacp, mask)
+  jacr = jax.vmap(jp.multiply)(cdof_dot[:, :3], mask)  # pytype: disable=attribute-error
+
+  return jacp, jacr
+
+
 def apply_ft(
     m: Model,
     d: Data,