Switch MJX DataJAX from legacy qM format to CSR M format

Switch MJX DataJAX to use the CSR-format `M` as its primary sparse inertia representation, matching MuJoCo C `mjData.M`. This allows both `mjData.qM` and `mjModel.mapM2M` to be deleted in the future.

PiperOrigin-RevId: 932433465
Change-Id: I194ffbfdba0cdd9c718277ee0d5de7ac20f2b06f

mjx/mujoco/mjx/_src/forward.py

@@ -352,10 +352,11 @@ def euler(m: Model, d: Data) -> Data:
   qacc = d.qacc
   if not m.opt.disableflags & DisableBit.EULERDAMP:
     if support.is_sparse(m):
-      qM = d._impl.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)
+      diag_adr = m.M_rowadr + m.M_rownnz - 1
+      M = d._impl.M.at[diag_adr].add(m.opt.timestep * m.dof_damping)
     else:
-      qM = d._impl.qM + jp.diag(m.opt.timestep * m.dof_damping)
+      M = d._impl.M + jp.diag(m.opt.timestep * m.dof_damping)
-    dh = d.tree_replace({'_impl.qM': qM})
+    dh = d.tree_replace({'_impl.M': M})
     dh = smooth.factor_m(m, dh)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
     qacc = smooth.solve_m(m, dh, qfrc)
@@ -418,7 +419,7 @@ def implicit(m: Model, d: Data) -> Data:
   qacc = d.qacc
   if qderiv is not None:
     # TODO(robotics-simulation): use smooth.factor_m / solve_m here:
-    qm = support.full_m(m, d) if support.is_sparse(m) else d._impl.qM
+    qm = support.full_m(m, d) if support.is_sparse(m) else d._impl.M
     qm -= m.opt.timestep * qderiv
     qh, _ = jax.scipy.linalg.cho_factor(qm)
     qfrc = d.qfrc_smooth + d.qfrc_constraint

mjx/mujoco/mjx/_src/io.py

@@ -669,8 +669,7 @@ def _make_data_jax(
       'wrap_xpos': (m.nwrap, 6, float_),
       'actuator_moment': (m.nu, m.nv, float_),
       'crb': (m.nbody, 10, float_),
-      'qM': (m.nM, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
+      'M': (m.nC, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
-      'M': (m.nC, float_),
       'qLD': (m.nC, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
       'qLDiagInv': (m.nv, float_) if support.is_sparse(m) else (0, float_),
       'ten_velocity': (m.ntendon, float_),
@@ -1045,11 +1044,11 @@ def _put_data_jax(
 
     impl_fields[fname] = value
 
-  # convert qM and qLD if jacobian is dense
+  # convert M and qLD if jacobian is dense
   if not support.is_sparse(m):
-    impl_fields['qM'] = np.zeros((m.nv, m.nv))
+    impl_fields['M'] = np.zeros((m.nv, m.nv))
     mujoco.mju_sym2dense(
-        impl_fields['qM'],
+        impl_fields['M'],
         d.M,
         m.M_rownnz,
         m.M_rowadr,
@@ -1057,9 +1056,9 @@ def _put_data_jax(
     )
     # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing
     try:
-      impl_fields['qLD'], _ = scipy.linalg.cho_factor(impl_fields['qM'])
+      impl_fields['qLD'], _ = scipy.linalg.cho_factor(impl_fields['M'])
     except scipy.linalg.LinAlgError:
-      # this happens when qM is empty or unstable simulation
+      # this happens when M is empty or unstable simulation
       impl_fields['qLD'] = np.zeros((m.nv, m.nv))
     impl_fields['qLDiagInv'] = np.zeros(0)
 
@@ -1333,15 +1332,6 @@ def _get_data_into(
   d = jax.device_get(d)
   batch_size = d.qpos.shape[0] if batched else 1
 
-  dof_i, dof_j = [], []
-  if d.impl == types.Impl.JAX:
-    for i in range(m.nv):
-      j = i
-      while j > -1:
-        dof_i.append(i)
-        dof_j.append(j)
-        j = m.dof_parentid[j]
-
   for i in range(batch_size):
     d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
     result_i = result[i] if batched else result
@@ -1446,8 +1436,14 @@ def _get_data_into(
       elif field.name.startswith('efc_'):
         value = value[efc_active]
       if d.impl == types.Impl.JAX:
-        if field.name == 'qM' and not support.is_sparse(m):
+        if field.name == 'M' and not support.is_sparse(m):
-          value = value[dof_i, dof_j]
+          M_csr = np.zeros(m.nC)
+          for i in range(m.nv):
+            adr = m.M_rowadr[i]
+            for k in range(m.M_rownnz[i]):
+              col = m.M_colind[adr + k]
+              M_csr[adr + k] = value[i, col]
+          value = M_csr
         elif field.name == 'qLD':
           value = np.zeros(m.nC)
         elif field.name == 'qLDiagInv' and not support.is_sparse(m):
@@ -1464,9 +1460,9 @@ def _get_data_into(
       else:
         setattr(result_i, field.name, value)
 
-    # TODO(taylorhowell): remove mapping once qM is deprecated
+    if hasattr(result_i, 'qM'):
-    # map inertia (sparse) to reduced inertia (compressed sparse) representation
+      result_i.qM.fill(0.0)
-    result_i.M[:] = result_i.qM[m.mapM2M]
+      result_i.qM[m.mapM2M] = result_i.M
 
     # recalculate qLD and qLDiagInv as MJX and MuJoCo have different
     # representations of the Cholesky decomposition.

mjx/mujoco/mjx/_src/io_test.py

@@ -476,20 +476,18 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d._impl.efc_force.shape, (nefc,))
 
     if impl == 'jax':
-      self.assertEqual(d._impl.qM.shape, (nv, nv))
+      self.assertEqual(d._impl.M.shape, (nv, nv))
       self.assertEqual(d._impl.qLD.shape, (nv, nv))
       self.assertEqual(d._impl.qLDiagInv.shape, (0,))
 
-
     # test sparse
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mjx.make_data(m, impl=impl)
-    self.assertEqual(d._impl.qM.shape, (nm,))
+    self.assertEqual(d._impl.M.shape, (m.nC,))
-    self.assertEqual(d._impl.qLD.shape, (nm,))
+    self.assertEqual(d._impl.qLD.shape, (m.nC,))
     self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
 
 
-
   @mock.patch.dict(os.environ, {'MJX_GPU_DEFAULT_WARP': 'true'})
   def test_make_data_warp(self):
     if not mjxw.WARP_INSTALLED:
@@ -605,7 +603,7 @@ class DataIOTest(parameterized.TestCase):
     )
 
     # check sparse mass matrices are correct
-    np.testing.assert_allclose(dx_sparse._impl.qM, d.qM, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse._impl.M, d.M, atol=1e-8)
     np.testing.assert_allclose(dx_sparse._impl.qLD, d.qLD, atol=1e-8)
     np.testing.assert_allclose(
         dx_sparse._impl.qLDiagInv, d.qLDiagInv, atol=1e-8
@@ -619,7 +617,7 @@ class DataIOTest(parameterized.TestCase):
     if impl == 'jax':
       qm = np.zeros((m.nv, m.nv))
       mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
-      np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8)
+      np.testing.assert_allclose(dx_from_dense._impl.M, qm, atol=1e-8)
 
 
   def test_put_data_warp_ndim(self):
@@ -865,8 +863,8 @@ class DataIOTest(parameterized.TestCase):
       mjx.make_data(m)
 
   @parameterized.parameters(JacobianType.DENSE, JacobianType.SPARSE)
-  def test_qm_mapm2m(self, jacobian):
+  def test_m_mapm2m(self, jacobian):
-    """Test that qM is mapped to M."""
+    """Test that M matches MuJoCo."""
     m = test_util.load_test_file('humanoid/humanoid.xml')
     m.opt.jacobian = jacobian
     d = mujoco.MjData(m)

mjx/mujoco/mjx/_src/smooth.py

@@ -307,7 +307,7 @@ def crb(m: Model, d: Data) -> Data:
   crb_dof = jp.take(crb_body, jp.array(m.dof_bodyid), axis=0)
   crb_cdof = jax.vmap(math.inert_mul)(crb_dof, d.cdof)
   qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
-  d = d.tree_replace({'_impl.qM': qm})
+  d = d.tree_replace({'_impl.M': qm})
   return d
 
 
@@ -317,56 +317,55 @@ def factor_m(m: Model, d: Data) -> Data:
     raise ValueError('factor_m requires JAX backend implementation.')
 
   if not support.is_sparse(m):
-    qh, _ = jax.scipy.linalg.cho_factor(d._impl.qM)
+    qh, _ = jax.scipy.linalg.cho_factor(d._impl.M)
     d = d.tree_replace({'_impl.qLD': qh})
     return d
 
-  # build up indices for where we will do backwards updates over qLD
   depth = []
   for i in range(m.nv):
     depth.append(depth[m.dof_parentid[i]] + 1 if m.dof_parentid[i] != -1 else 0)
   updates = {}
-  madr_ds = []
+  diag_ds = []
   for i in range(m.nv):
-    madr_d = madr_ij = m.dof_Madr[i]
+    diag_i = m.M_rowadr[i] + m.M_rownnz[i] - 1
-    j = i
+    start_i = m.M_rowadr[i]
-    while True:
+    for k in range(m.M_rownnz[i]):
-      madr_ds.append(madr_d)
+      diag_ds.append(diag_i)
-      madr_ij, j = madr_ij + 1, m.dof_parentid[j]
+    for k in range(m.M_rownnz[i] - 1):
-      if j == -1:
+      adr = start_i + k
-        break
+      j = m.M_colind[adr]
-      out_beg, out_end = tuple(m.dof_Madr[j : j + 2])
+      out_beg = m.M_rowadr[j]
+      width = m.M_rownnz[j]
+      out_end = out_beg + width
       updates.setdefault(depth[j], []).append(
-          (out_beg, out_end, madr_d, madr_ij)
+          (out_beg, out_end, diag_i, adr, start_i)
       )
 
-  qld = d._impl.qM
+  qld = d._impl.M
 
   for _, updates in sorted(updates.items(), reverse=True):
-    # combine the updates into one update batch (per depth level)
     rows = []
     madr_ijs = []
     pivots = []
     out = []
 
-    for b, e, madr_d, madr_ij in updates:
+    for b, e, piv_i, madr_ij, start_i in updates:
       width = e - b
-      rows.append(np.arange(madr_ij, madr_ij + width))
+      rows.append(np.arange(start_i, start_i + width))
       madr_ijs.append(np.full((width,), madr_ij))
-      pivots.append(np.full((width,), madr_d))
+      pivots.append(np.full((width,), piv_i))
       out.append(np.arange(b, e))
-    rows = np.concatenate(rows)
+    if out:
-    madr_ijs = np.concatenate(madr_ijs)
+      rows = np.concatenate(rows)
-    pivots = np.concatenate(pivots)
+      madr_ijs = np.concatenate(madr_ijs)
-    out = np.concatenate(out)
+      pivots = np.concatenate(pivots)
+      out = np.concatenate(out)
 
-    # apply the update batch
+      qld = qld.at[out].add(-(qld[madr_ijs] / qld[pivots]) * qld[rows])
-    qld = qld.at[out].add(-(qld[madr_ijs] / qld[pivots]) * qld[rows])
-    # TODO(erikfrey): determine if this minimum value guarding is necessary:
-    # qld = qld.at[dof_madr].set(jp.maximum(qld[dof_madr], _MJ_MINVAL))
 
-  qld_diag = qld[m.dof_Madr]
+  diag_adr = m.M_rowadr + m.M_rownnz - 1
-  qld = (qld / qld[jp.array(madr_ds)]).at[m.dof_Madr].set(qld_diag)
+  qld_diag = qld[diag_adr]
+  qld = (qld / qld[np.array(diag_ds)]).at[diag_adr].set(qld_diag)
 
   d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag})
   return d
@@ -386,13 +385,12 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
 
   updates_i, updates_j = {}, {}
   for i in range(m.nv):
-    madr_ij, j = m.dof_Madr[i], i
+    start_i = m.M_rowadr[i]
-    while True:
+    for k in range(m.M_rownnz[i] - 1):
-      madr_ij, j = madr_ij + 1, m.dof_parentid[j]
+      adr = start_i + k
-      if j == -1:
+      j = m.M_colind[adr]
-        break
+      updates_i.setdefault(depth[i], []).append((i, adr, j))
-      updates_i.setdefault(depth[i], []).append((i, madr_ij, j))
+      updates_j.setdefault(depth[j], []).append((j, adr, i))
-      updates_j.setdefault(depth[j], []).append((j, madr_ij, i))
 
   # x <- inv(L') * x
   for _, vals in sorted(updates_j.items(), reverse=True):
@@ -1317,7 +1315,7 @@ def transmission(m: Model, d: Data) -> Data:
 
 
 def tendon_armature(m: Model, d: Data) -> Data:
-  """Add tendon armature to qM."""
+  """Add tendon armature to M."""
   if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
     raise ValueError('tendon_armature requires JAX backend implementation.')
 
@@ -1330,17 +1328,11 @@ def tendon_armature(m: Model, d: Data) -> Data:
   )
 
   if support.is_sparse(m):
-    ij = []
+    i = np.repeat(np.arange(m.nv), m.M_rownnz)
-    for i in range(m.nv):
+    j = m.M_colind
-      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})
+  return d.tree_replace({'_impl.M': d._impl.M + JTAJ})
 
 
 def tendon_dot(m: Model, d: Data) -> jax.Array:

mjx/mujoco/mjx/_src/smooth_test.py

@@ -89,13 +89,10 @@ class SmoothTest(absltest.TestCase):
     # crb
     dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d))
     _assert_attr_eq(d, dx._impl, 'crb')
-    _assert_attr_eq(d, dx._impl, 'qM')
+    _assert_attr_eq(d, dx._impl, 'M')
     # factor_m
     dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
-    qLDLegacy = np.zeros(mx.nM)  # pylint:disable=invalid-name
+    _assert_attr_eq(d, dx._impl, 'qLD')
-    for i in range(m.nC):
-      qLDLegacy[m.mapM2M[i]] = d.qLD[i]
-    _assert_eq(qLDLegacy, dx._impl.qLD, 'qLD')
     _assert_attr_eq(d, dx._impl, 'qLDiagInv')
     # com_vel
     dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
@@ -430,18 +427,18 @@ class TendonTest(parameterized.TestCase):
     dx = mjx.put_data(m, d)
 
     dx = dx.tree_replace(
-        {'_impl.qM': jp.zeros((m.nv, m.nv)), 'qfrc_bias': jp.zeros(m.nv)}
+        {'_impl.M': jp.zeros((m.nv, m.nv)), 'qfrc_bias': jp.zeros(m.nv)}
     )
 
     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
+      M = np.zeros((m.nv, m.nv))  # pylint: disable=invalid-name
-      mujoco.mju_sym2dense(qM, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
+      mujoco.mju_sym2dense(M, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
     else:
-      qM = d.qM  # pylint: disable=invalid-name
+      M = d.M  # pylint: disable=invalid-name
-    _assert_eq(dx._impl.qM, qM, 'qM')
+    _assert_eq(dx._impl.M, M, 'M')
 
     dx = mjx.rne(mx, dx)
     dx = mjx.tendon_bias(mx, dx)

mjx/mujoco/mjx/_src/support.py

@@ -60,14 +60,8 @@ def make_m(
 ) -> jax.Array:
   """Computes M = a @ b.T + diag(d)."""
 
-  ij = []
+  i = np.repeat(np.arange(m.nv), m.M_rownnz)
-  for i in range(m.nv):
+  j = m.M_colind
-    j = i
-    while j > -1:
-      ij.append((i, j))
-      j = m.dof_parentid[j]
-
-  i, j = (jp.array(x) for x in zip(*ij))
 
   if not is_sparse(m):
     qm = a @ b.T
@@ -82,29 +76,23 @@ def make_m(
   b_j = jp.take(b, j, axis=0)
   qm = jax.vmap(jp.dot)(a_i, b_j)
 
-  # add diagonal
   if d is not None:
-    qm = qm.at[m.dof_Madr].add(d)
+    diag_adr = m.M_rowadr + m.M_rownnz - 1
+    qm = qm.at[diag_adr].add(d)
 
   return qm
 
 
 def full_m(m: Model, d: Data) -> jax.Array:
-  """Reconstitute dense mass matrix from qM."""
+  """Reconstitute dense mass matrix from M."""
 
   if not is_sparse(m):
-    return d._impl.qM  # pytype: disable=attribute-error
+    return d._impl.M  # pytype: disable=attribute-error
 
-  ij = []
+  i = np.repeat(np.arange(m.nv), m.M_rownnz)
-  for i in range(m.nv):
+  j = m.M_colind
-    j = i
-    while j > -1:
-      ij.append((i, j))
-      j = m.dof_parentid[j]
 
-  i, j = (jp.array(x) for x in zip(*ij))
+  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d._impl.M)  # pytype: disable=attribute-error
-
-  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d._impl.qM)  # pytype: disable=attribute-error
 
   # also set upper triangular
   mat = mat + jp.tril(mat, -1).T
@@ -116,24 +104,23 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
   """Multiply vector by inertia matrix."""
 
   if not is_sparse(m):
-    return d._impl.qM @ vec  # pytype: disable=attribute-error
+    return d._impl.M @ vec  # pytype: disable=attribute-error
 
-  diag_mul = d._impl.qM[jp.array(m.dof_Madr)] * vec  # pytype: disable=attribute-error
+  diag_adr = m.M_rowadr + m.M_rownnz - 1
+  diag_mul = d._impl.M[diag_adr] * vec  # pytype: disable=attribute-error
 
   is_, js, madr_ijs = [], [], []
   for i in range(m.nv):
-    madr_ij, j = m.dof_Madr[i], i
+    adr = m.M_rowadr[i]
+    for k in range(m.M_rownnz[i] - 1):
+      is_.append(i)
+      js.append(m.M_colind[adr + k])
+      madr_ijs.append(adr + k)
 
-    while True:
+  i, j, madr_ij = (np.array(x, dtype=np.int32) for x in (is_, js, madr_ijs))
-      madr_ij, j = madr_ij + 1, m.dof_parentid[j]
-      if j == -1:
-        break
-      is_, js, madr_ijs = is_ + [i], js + [j], madr_ijs + [madr_ij]
 
-  i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs))
+  out = diag_mul.at[i].add(d._impl.M[madr_ij] * vec[j])  # pytype: disable=attribute-error
-
+  out = out.at[j].add(d._impl.M[madr_ij] * vec[i])  # pytype: disable=attribute-error
-  out = diag_mul.at[i].add(d._impl.qM[madr_ij] * vec[j])  # pytype: disable=attribute-error
-  out = out.at[j].add(d._impl.qM[madr_ij] * vec[i])  # pytype: disable=attribute-error
 
   return out
 

mjx/mujoco/mjx/_src/types.py

@@ -696,6 +696,9 @@ class Model(PyTreeNode):
   dof_treeid: np.ndarray
   dof_Madr: np.ndarray  # pylint:disable=invalid-name
   dof_simplenum: np.ndarray
+  M_rowadr: np.ndarray  # pylint:disable=invalid-name
+  M_rownnz: np.ndarray  # pylint:disable=invalid-name
+  M_colind: np.ndarray  # pylint:disable=invalid-name
   dof_solref: jax.Array
   dof_solimp: jax.Array
   dof_frictionloss: jax.Array
@@ -996,7 +999,6 @@ class DataJAX(PyTreeNode):
   wrap_xpos: jax.Array
   actuator_moment: jax.Array
   crb: jax.Array
-  qM: jax.Array  # pylint:disable=invalid-name
   M: jax.Array  # pylint:disable=invalid-name
   qLD: jax.Array  # pylint:disable=invalid-name
   qLDiagInv: jax.Array  # pylint:disable=invalid-name