From 8084b06a35a20a1bf58b3ad830273e60d9fe78fa Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Tue, 17 Jun 2025 12:50:56 -0700
Subject: [PATCH] Add tendon armature with sparse qM to MJX. Fixes #2681.

PiperOrigin-RevId: 772583693
Change-Id: I827d61aacada33f53500bf54613ccb8145bc854b
---
 mjx/mujoco/mjx/_src/smooth.py      | 46 +++++++++++++++---------------
 mjx/mujoco/mjx/_src/smooth_test.py | 14 +++++----
 2 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py
index bef176d6..169e5a25 100644
--- a/mjx/mujoco/mjx/_src/smooth.py
+++ b/mjx/mujoco/mjx/_src/smooth.py
@@ -1194,19 +1194,23 @@ def tendon_armature(m: Model, d: Data) -> Data:
   if not m.ntendon:
     return d
 
-  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.'
-    )
+  # TODO(taylorhowell): if sparse, compute sparse JTAJ
+  JTAJ = d._impl.ten_J.T @ jax.vmap(jp.multiply)(
+      d._impl.ten_J, m.tendon_armature
+  )
+
+  if support.is_sparse(m):
+    ij = []
+    for i in range(m.nv):
+      j = i
+      while j > -1:
+        ij.append((i, j))
+        j = m.dof_parentid[j]
+
+    i, j = (jp.array(x) for x in zip(*ij))
+    JTAJ = JTAJ[(i, j)]
+
+  return d.tree_replace({'_impl.qM': d._impl.qM + JTAJ})
 
 
 def tendon_dot(m: Model, d: Data) -> jax.Array:
@@ -1329,13 +1333,9 @@ def tendon_bias(m: Model, d: Data) -> Data:
   # 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.')
+  return d.tree_replace({
+      'qfrc_bias': (
+          d.qfrc_bias
+          + jp.sum(jax.vmap(jp.multiply)(d._impl.ten_J, coef), axis=0)
+      )
+  })
diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py
index 6c038522..506d8109 100644
--- a/mjx/mujoco/mjx/_src/smooth_test.py
+++ b/mjx/mujoco/mjx/_src/smooth_test.py
@@ -22,6 +22,7 @@ import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
 from mujoco.mjx._src.types import ConeType  # pylint: disable=g-importing-member
+from mujoco.mjx._src.types import JacobianType  # pylint: disable=g-importing-member
 import numpy as np
 
 # tolerance for difference between MuJoCo and MJX smooth calculations - mostly
@@ -317,7 +318,8 @@ 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):
+  @parameterized.parameters(JacobianType.DENSE, JacobianType.SPARSE)
+  def test_tendon_armature(self, jacobian):
     """Tests MJX tendon armature matches MuJoCo."""
     m = mujoco.MjModel.from_xml_string("""
         <mujoco>
@@ -345,14 +347,11 @@ class TendonTest(parameterized.TestCase):
           </keyframe>
         </mujoco>
         """)
-
+    m.opt.jacobian = jacobian
     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)
 
@@ -363,6 +362,11 @@ class TendonTest(parameterized.TestCase):
     dx = mjx.crb(mx, dx)
     dx = mjx.tendon_armature(mx, dx)
 
+    if jacobian == JacobianType.DENSE:
+      qM = np.zeros((m.nv, m.nv))  # pylint: disable=invalid-name
+      mujoco.mj_fullM(m, qM, d.qM)
+    else:
+      qM = d.qM  # pylint: disable=invalid-name
     _assert_eq(dx._impl.qM, qM, 'qM')
 
     dx = mjx.rne(mx, dx)