Adds support for explicit dense/sparse mass matrices to MJX.

This increases performance, particularly for the Newton solver on TPU.

PiperOrigin-RevId: 600696483
Change-Id: If69bb9a2e21ba8dad6ca23f093ce7b7ceae644ff

doc/changelog.rst

@@ -26,21 +26,23 @@ MJX
 5. Added :at:`site` transmission.
 6. Updated MJX colab tutorial with more stable quadruped environment.
 7. Added ``mjx.ray`` which mirrors :ref:`mj_ray` for planes, spheres, capsules, boxes, and meshes.
+8. Added ``mjx.is_sparse`` which mirrors :ref:`mj_isSparse` and ``mjx.full_m`` which mirrors :ref:`mj_fullM`.
+9. Added support for specifying sparse or dense mass matrices via :ref:`option-jacobian`.
 
 Python bindings
 ^^^^^^^^^^^^^^^
-8. Improved the implmentation of the :ref:`rollout<PySample>` module. Note the changes below are breaking, dependent
+10. Improved the implmentation of the :ref:`rollout<PySample>` module. Note the changes below are breaking, dependent
-   code will require modification.
+    code will require modification.
 
-   - Uses :ref:`mjSTATE_FULLPHYSICS<geFullPhysics>` as state spec, enabling divergence detection by inspecting time.
+    - Uses :ref:`mjSTATE_FULLPHYSICS<geFullPhysics>` as state spec, enabling divergence detection by inspecting time.
-   - Allows user-defined control spec for any combination of :ref:`user input<geInput>` fields as controls.
+    - Allows user-defined control spec for any combination of :ref:`user input<geInput>` fields as controls.
-   - Outputs are no longer squeezed and always have dim=3.
+    - Outputs are no longer squeezed and always have dim=3.
 
 Bug fixes
 ^^^^^^^^^
-9. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
+11. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
-   :github:issue:`1270`.
+    :github:issue:`1270`.
-10. Fixed a bug in the :ref:`muscle model<CMuscle>` that led to non-zero values outside the lower
+12. Fixed a bug in the :ref:`muscle model<CMuscle>` that led to non-zero values outside the lower
     bound of the length range. Fixes :github:issue:`1342`.
 
 

doc/mjx.rst

@@ -349,3 +349,11 @@ For MJX to perform well, some configuration parameters should be adjusted from t
 
 :ref:`option-flag` element
   Disabling ``eulerdamp`` can help performance and is often not needed for stability.
+
+:ref:`option-jacobian` element
+  Explicitly setting "dense" or "sparse" may speed up simulation depending on your device. Modern TPUs have specialized
+  hardware for rapidly operating over sparse matrices, whereas GPUs tend to be faster with dense matrices as long as
+  they fit onto the device. As such, the behavior in MJX for the default "auto" setting is sparse if ``nv`` is 60 or
+  greater, or if MJX detects a TPU as the default backend, otherwise "dense". For TPU, using "sparse" with the
+  Newton solver can speed up simulation by 2x to 3x. For GPU, choosing "dense" may impart a more modest speedup of 10%
+  to 20%, as long as the dense matrices can fit on the device.

mjx/mujoco/mjx/__init__.py

@@ -39,8 +39,10 @@ from mujoco.mjx._src.smooth import com_vel
 from mujoco.mjx._src.smooth import crb
 from mujoco.mjx._src.smooth import factor_m
 from mujoco.mjx._src.smooth import kinematics
-from mujoco.mjx._src.smooth import mul_m
 from mujoco.mjx._src.smooth import rne
 from mujoco.mjx._src.smooth import transmission
 from mujoco.mjx._src.solver import solve
+from mujoco.mjx._src.support import is_sparse
+from mujoco.mjx._src.support import full_m
+from mujoco.mjx._src.support import mul_m
 from mujoco.mjx._src.types import *

mjx/mujoco/mjx/_src/device_test.py

@@ -75,6 +75,7 @@ class DeviceTest(parameterized.TestCase):
   def testdevice_get(self, fname):
     """Test getting MjData from a device."""
     m = test_util.load_test_file(fname)
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE  # force sparse for testing
     mx = device.device_put(m)
     dx = mjx.make_data(mx)
     d = mujoco.MjData(m)
@@ -85,6 +86,7 @@ class DeviceTest(parameterized.TestCase):
   def testdevice_get_batched(self, fname):
     """Test getting MjData from a device."""
     m = test_util.load_test_file(fname)
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE  # force sparse for testing
     mx = device.device_put(m)
     batch_size = 32
 

mjx/mujoco/mjx/_src/forward.py

@@ -66,7 +66,7 @@ def fwd_position(m: Model, d: Data) -> Data:
   d = smooth.kinematics(m, d)
   d = smooth.com_pos(m, d)
   d = smooth.crb(m, d)
-  d = smooth.factor_m(m, d, d.qM)
+  d = smooth.factor_m(m, d)
   d = collision_driver.collision(m, d)
   d = constraint.make_constraint(m, d)
   d = smooth.transmission(m, d)
@@ -288,8 +288,8 @@ def euler(m: Model, d: Data) -> Data:
   qacc = d.qacc
   if not m.opt.disableflags & DisableBit.EULERDAMP:
     # TODO(robotics-simulation): can this be done with a smaller perf hit
-    mh = d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)
+    dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping))
-    dh = smooth.factor_m(m, d, mh)
+    dh = smooth.factor_m(m, dh)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
     qacc = smooth.solve_m(m, dh, qfrc)
   return _advance(m, d, d.act_dot, qacc)

mjx/mujoco/mjx/_src/io.py

@@ -23,8 +23,10 @@ import mujoco
 from mujoco.mjx._src import collision_driver
 from mujoco.mjx._src import constraint
 from mujoco.mjx._src import mesh
+from mujoco.mjx._src import support
 from mujoco.mjx._src import types
 import numpy as np
+import scipy
 
 
 def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
@@ -35,6 +37,9 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
   if o.cone not in set(types.ConeType):
     raise NotImplementedError(f'{mujoco.mjtCone(o.cone)}')
 
+  if o.jacobian not in set(types.JacobianType):
+    raise NotImplementedError(f'{mujoco.mjtJacobian(o.jacobian)}')
+
   if o.solver not in set(types.SolverType):
     raise NotImplementedError(f'{mujoco.mjtSolver(o.solver)}')
 
@@ -49,6 +54,7 @@ def _put_option(o: mujoco.MjOption, device=None) -> types.Option:
   }
   static_fields['integrator'] = types.IntegratorType(o.integrator)
   static_fields['cone'] = types.ConeType(o.cone)
+  static_fields['jacobian'] = types.JacobianType(o.jacobian)
   static_fields['solver'] = types.SolverType(o.solver)
   static_fields['disableflags'] = types.DisableBit(o.disableflags)
 
@@ -137,8 +143,10 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
   ne, nf, nl, nc = constraint.count_constraints(m)
   nefc = ne + nf + nl + nc
 
+  zero_0 = jp.zeros(0, dtype=jp.float32)
   zero_nv = jp.zeros(m.nv, dtype=jp.float32)
   zero_nv_6 = jp.zeros((m.nv, 6), dtype=jp.float32)
+  zero_nv_nv = jp.zeros((m.nv, m.nv), dtype=jp.float32)
   zero_nbody_3 = jp.zeros((m.nbody, 3), dtype=jp.float32)
   zero_nbody_6 = jp.zeros((m.nbody, 6), dtype=jp.float32)
   zero_nbody_10 = jp.zeros((m.nbody, 10), dtype=jp.float32)
@@ -180,10 +188,9 @@ def make_data(m: Union[types.Model, mujoco.MjModel]) -> types.Data:
       actuator_length=zero_nu,
       actuator_moment=jp.zeros((m.nu, m.nv), dtype=jp.float32),
       crb=zero_nbody_10,
-      qM=zero_nm,
+      qM=zero_nm if support.is_sparse(m) else zero_nv_nv,
-      qLD=zero_nm,
+      qLD=zero_nm if support.is_sparse(m) else zero_nv_nv,
-      qLDiagInv=zero_nv,
+      qLDiagInv=zero_nv if support.is_sparse(m) else zero_0,
-      qLDiagSqrtInv=zero_nv,
       contact=types.Contact.zero(ncon),
       efc_J=jp.zeros((nefc, m.nv), dtype=jp.float32),
       efc_frictionloss=zero_nefc,
@@ -237,6 +244,14 @@ def get_data(
       mujoco.mjtConstraint.mjCNSTR_CONTACT_PYRAMIDAL,
   ]).repeat([ne, nf, nl, nc])
 
+  dof_i, dof_j = [], []
+  for i in range(m.nv):
+    j = i
+    while j > -1:
+      dof_i.append(i)
+      dof_j.append(j)
+      j = m.dof_parentid[j]
+
   ds = []
   for i in range(batch_size):
     dx_i = jax.tree_map(lambda x, i=i: x[i], dx) if batched else d
@@ -267,6 +282,15 @@ def get_data(
       if field.name == 'efc_J':
         value = value[efc_active].reshape(-1)
 
+      if field.name == 'qM' and not support.is_sparse(m):
+        value = value[dof_i, dof_j]
+
+      if field.name == 'qLD' and not support.is_sparse(m):
+        value = value[dof_i, dof_j]
+
+      if field.name == 'qLDiagInv' and not support.is_sparse(m):
+        value = np.ones(m.nv)
+
       if value.shape:
         getattr(d_i, field.name)[:] = value
       else:
@@ -346,6 +370,18 @@ def put_data(m: mujoco.MjModel, d: mujoco.MjData, device=None) -> types.Data:
       value[value_beg:value_beg+size] = fields[fname][d_beg:d_beg+size]
     fields[fname] = value
 
+  # convert qM and qLD if jacobian is dense
+  if not support.is_sparse(m):
+    fields['qM'] = np.zeros((m.nv, m.nv))
+    mujoco.mj_fullM(m, fields['qM'], d.qM)
+    # TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing
+    try:
+      fields['qLD'], _ = scipy.linalg.cho_factor(fields['qM'])
+    except scipy.linalg.LinAlgError:
+      # this happens when qM is empty or unstable simulation
+      fields['qLD'] = np.zeros((m.nv, m.nv))
+    fields['qLDiagInv'] = np.zeros(0)
+
   fields = jax.device_put(fields, device=device)
   fields['contact'] = _put_contact(d.contact, ncon, device=device)
 

mjx/mujoco/mjx/_src/io_test.py

@@ -25,7 +25,7 @@ import numpy as np
 
 _MULTIPLE_CONVEX_OBJECTS = """
   <mujoco>
-    <option timestep="0.001"/>
+    <option timestep="0.001" jacobian="dense"/>
     <default>
       <geom solref=".006 1"/>
     </default>
@@ -224,7 +224,6 @@ class DataIOTest(parameterized.TestCase):
     ncon = 46
     nv = 19
     nefc = 185
-    nm = 64
 
     self.assertEqual(d.qpos.shape, (nq,))
     self.assertEqual(d.qvel.shape, (nv,))
@@ -251,10 +250,9 @@ class DataIOTest(parameterized.TestCase):
     self.assertEqual(d.crb.shape, (nbody, 10))
     self.assertEqual(d.actuator_length.shape, (1,))
     self.assertEqual(d.actuator_moment.shape, (1, nv))
-    self.assertEqual(d.qM.shape, (nm,))
+    self.assertEqual(d.qM.shape, (nv, nv))
-    self.assertEqual(d.qLD.shape, (nm,))
+    self.assertEqual(d.qLD.shape, (nv, nv))
-    self.assertEqual(d.qLDiagInv.shape, (nv,))
+    self.assertEqual(d.qLDiagInv.shape, (0,))
-    self.assertEqual(d.qLDiagSqrtInv.shape, (nv,))
     self.assertEqual(d.contact.dist.shape, (ncon,))
     self.assertEqual(d.contact.pos.shape, (ncon, 3))
     self.assertEqual(d.contact.frame.shape, (ncon, 3, 3))
@@ -291,7 +289,11 @@ class DataIOTest(parameterized.TestCase):
     np.testing.assert_allclose(dx.xpos, d.xpos)
     np.testing.assert_allclose(dx.cvel, d.cvel)
     np.testing.assert_allclose(dx.cdof_dot, d.cdof_dot)
-    np.testing.assert_allclose(dx.qM, d.qM)
+
+    # check that qM is transformed properly
+    qm = np.zeros((m.nv, m.nv), dtype=np.float64)
+    mujoco.mj_fullM(m, qm, d.qM)
+    np.testing.assert_allclose(qm, mjx.full_m(mjx.put_model(m), dx))
 
     # 4 contacts, 2 for each capsule against the plane
     self.assertEqual(dx.contact.dist.shape, (4,))
@@ -335,8 +337,22 @@ class DataIOTest(parameterized.TestCase):
     m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 2)
-    dx_from_sparse = mjx.put_data(m, d)
+    dx_sparse = mjx.put_data(m, d)
-    np.testing.assert_allclose(dx_from_sparse.efc_J, dx.efc_J, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse.efc_J, dx.efc_J, atol=1e-8)
+
+    # check sparse mass matrices are correct
+    np.testing.assert_allclose(dx_sparse.qM, d.qM, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse.qLD, d.qLD, atol=1e-8)
+    np.testing.assert_allclose(dx_sparse.qLDiagInv, d.qLDiagInv, atol=1e-8)
+
+    # check dense mass matrices are correct
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 2)
+    dx_from_dense = mjx.put_data(m, d)
+    qm = np.zeros((m.nv, m.nv))
+    mujoco.mj_fullM(m, qm, d.qM)
+    np.testing.assert_allclose(dx_from_dense.qM, qm, atol=1e-8)
 
   def test_get_data(self):
     """Test that get_data makes correct MjData."""

mjx/mujoco/mjx/_src/smooth.py

@@ -136,7 +136,7 @@ def com_pos(m: Model, d: Data) -> Data:
   pos, mass = scan.body_tree(
       m, subtree_sum, 'bb', 'bb', d.xipos, m.body_mass, reverse=True
   )
-  cond = jp.tile(mass < jp.array(mujoco.mjMINVAL), (3, 1)).T
+  cond = jp.tile(mass < mujoco.mjMINVAL, (3, 1)).T
   subtree_com = jp.where(cond, d.xipos, jax.vmap(jp.divide)(pos, mass))
   d = d.replace(subtree_com=subtree_com)
 
@@ -144,12 +144,12 @@ def com_pos(m: Model, d: Data) -> Data:
   @jax.vmap
   def inert_com(inert, ximat, off, mass):
     h = jp.cross(off, -jp.eye(3))
-    inert = ximat @ jp.diag(inert) @ ximat.T + h @ h.T * mass
+    inert = (ximat * inert) @ ximat.T + h @ h.T * mass
     # cinert is triu(inert), mass * off, mass
-    inert = inert[(jp.array([0, 1, 2, 0, 0, 1]), jp.array([0, 1, 2, 1, 2, 2]))]
+    inert = inert[([0, 1, 2, 0, 0, 1], [0, 1, 2, 1, 2, 2])]
-    return jp.concatenate([inert, off * mass, jp.expand_dims(mass, 0)])
+    return jp.concatenate([inert, off * mass, mass[None]])
 
-  root_com = subtree_com[jp.array(m.body_rootid)]
+  root_com = subtree_com[m.body_rootid]
   offset = d.xipos - root_com
   cinert = inert_com(m.body_inertia, d.ximat, offset, m.body_mass)
   d = d.replace(cinert=cinert)
@@ -208,36 +208,21 @@ def crb(m: Model, d: Data) -> Data:
   crb_body = crb_body.at[0].set(0.0)
   d = d.replace(crb=crb_body)
 
-  # TODO(erikfrey): do centralized take fn?
   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)
-  dof_i, dof_j, diag = [], [], []
-  for i in range(m.nv):
-    diag.append(len(dof_i))
-    j = i
-    while j > -1:
-      dof_i, dof_j = dof_i + [i], dof_j + [j]
-      j = m.dof_parentid[j]
-
-  crb_codf_i = jp.take(crb_cdof, jp.array(dof_i), axis=0)
-  cdof_j = jp.take(d.cdof, jp.array(dof_j), axis=0)
-  qm = jax.vmap(jp.dot)(crb_codf_i, cdof_j)
-
-  # add armature to diagonal
-  qm = qm.at[jp.array(diag)].add(m.dof_armature)
-
   d = d.replace(qM=qm)
 
   return d
 
 
-def factor_m(
+def factor_m(m: Model, d: Data) -> Data:
-    m: Model,
+  """Gets factorizaton of inertia-like matrix M, assumed spd."""
-    d: Data,
+
-    qM: jax.Array,  # pylint:disable=invalid-name
+  if not support.is_sparse(m):
-) -> Data:
+    qh, _ = jax.scipy.linalg.cho_factor(d.qM)
-  """Gets sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd."""
+    d = d.replace(qLD=qh)
+    return d
 
   # build up indices for where we will do backwards updates over qLD
   # TODO(erikfrey): do fewer updates by combining non-overlapping ranges
@@ -255,7 +240,7 @@ def factor_m(
       madr_j_range = tuple(m.dof_Madr[j : j + 2])
       updates.setdefault(madr_j_range, []).append((madr_d, madr_ij))
 
-  qld = qM
+  qld = d.qM
 
   for (out_beg, out_end), vals in sorted(updates.items(), reverse=True):
     madr_d, madr_ij = jp.array(vals).T
@@ -281,6 +266,9 @@ def factor_m(
 def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
   """Computes sparse backsubstitution:  x = inv(L'*D*L)*y ."""
 
+  if not support.is_sparse(m):
+    return jax.scipy.linalg.cho_solve((d.qLD, False), x)
+
   updates_i, updates_j = {}, {}
   for i in range(m.nv):
     madr_ij, j = m.dof_Madr[i], i
@@ -307,52 +295,6 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
   return x
 
 
-def dense_m(m: Model, d: Data) -> jax.Array:
-  """Reconstitute dense mass matrix from qM."""
-
-  is_, js, madr_ijs = [], [], []
-  for i in range(m.nv):
-    madr_ij, j = m.dof_Madr[i], i
-
-    while True:
-      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))
-
-  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM[madr_ij])
-
-  # diagonal, upper triangular, lower triangular
-  mat = jp.diag(d.qM[jp.array(m.dof_Madr)]) + mat + mat.T
-
-  return mat
-
-
-def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
-  """Multiply vector by inertia matrix."""
-
-  diag_mul = d.qM[jp.array(m.dof_Madr)] * vec
-
-  is_, js, madr_ijs = [], [], []
-  for i in range(m.nv):
-    madr_ij, j = m.dof_Madr[i], i
-
-    while True:
-      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.qM[madr_ij] * vec[j])
-  out = out.at[j].add(d.qM[madr_ij] * vec[i])
-
-  return out
-
-
 def com_vel(m: Model, d: Data) -> Data:
   """Computes cvel, cdof_dot."""
 

mjx/mujoco/mjx/_src/smooth_test.py

@@ -16,7 +16,6 @@
 
 from absl.testing import absltest
 import jax
-from jax import numpy as jp
 import mujoco
 from mujoco import mjx
 from mujoco.mjx._src import test_util
@@ -49,6 +48,8 @@ class SmoothTest(absltest.TestCase):
     """Tests MJX smooth functions match MuJoCo smooth functions."""
 
     m = test_util.load_test_file('pendula.xml')
+    # # force MJX sparse for testing:
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mujoco.MjData(m)
     # give the system a little kick to ensure we have non-identity rotations
     d.qvel = np.random.random(m.nv)
@@ -79,8 +80,7 @@ class SmoothTest(absltest.TestCase):
     _assert_attr_eq(d, dx, 'crb')
     _assert_attr_eq(d, dx, 'qM')
     # factor_m
-    dx = mjx.put_data(m, d)
+    dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
-    dx = jax.jit(mjx.factor_m)(mx, dx, dx.qM)
     _assert_attr_eq(d, dx, 'qLD')
     _assert_attr_eq(d, dx, 'qLDiagInv')
     # com_vel
@@ -95,21 +95,6 @@ class SmoothTest(absltest.TestCase):
     _assert_attr_eq(d, dx, 'actuator_length')
     _assert_attr_eq(d, dx, 'actuator_moment')
 
-  def test_mul_m(self):
-    m = test_util.load_test_file('pendula.xml')
-    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
-    mujoco.mj_forward(m, d)
-    mx = mjx.put_model(m)
-    dx = mjx.put_data(m, d)
-    vec = np.random.random(m.nv)
-    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
-    mj_vec = np.zeros(m.nv)
-    mujoco.mj_mulM(m, d, mj_vec, vec)
-    _assert_eq(mj_vec, mjx_vec, 'mul_m')
-
   def test_disable_gravity(self):
     m = mujoco.MjModel.from_xml_string("""
         <mujoco>

mjx/mujoco/mjx/_src/solver.py

@@ -14,14 +14,13 @@
 # ==============================================================================
 """Constraint solvers."""
 
-from typing import Optional
-
 import jax
 from jax import numpy as jp
 import mujoco
 from mujoco.mjx._src import constraint
 from mujoco.mjx._src import math
 from mujoco.mjx._src import smooth
+from mujoco.mjx._src import support
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.dataclasses import PyTreeNode
 from mujoco.mjx._src.types import Data
@@ -39,7 +38,6 @@ class _Context(PyTreeNode):
     qfrc_constraint: constraint force (from Data)     (nv,)
     Jaref: Jac*qacc - aref                            (nefc,)
     efc_force: constraint force in constraint space   (nefc,)
-    M: dense mass matrix, populated for nv < 100      (nv, nv)
     Ma: M*qacc                                        (nv,)
     grad: gradient of master cost                     (nv,)
     Mgrad: M / grad                                   (nv,)
@@ -54,7 +52,6 @@ class _Context(PyTreeNode):
   qfrc_constraint: jax.Array
   Jaref: jax.Array  # pylint: disable=invalid-name
   efc_force: jax.Array
-  M: Optional[jax.Array]
   Ma: jax.Array  # pylint: disable=invalid-name
   grad: jax.Array
   Mgrad: jax.Array  # pylint: disable=invalid-name
@@ -68,15 +65,13 @@ class _Context(PyTreeNode):
   def create(cls, m: Model, d: Data, grad: bool = True) -> '_Context':
     jaref = d.efc_J @ d.qacc - d.efc_aref
     # TODO(robotics-team): determine nv at which sparse mul is faster
-    M = smooth.dense_m(m, d) if m.nv < 100 else None  # pylint: disable=invalid-name
+    ma = support.mul_m(m, d, d.qacc)
-    ma = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc
     nv_0 = jp.zeros(m.nv)
     ctx = _Context(
         qacc=d.qacc,
         qfrc_constraint=d.qfrc_constraint,
         Jaref=jaref,
         efc_force=d.efc_force,
-        M=M,
         Ma=ma,
         grad=nv_0,
         Mgrad=nv_0,
@@ -224,10 +219,10 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
   elif m.opt.solver == SolverType.NEWTON:
     ne, nf, *_ = constraint.count_constraints(m)
     active = (ctx.Jaref < 0).at[:ne + nf].set(True)
-    h = (d.efc_J.T * d.efc_D * active) @ d.efc_J
+    h = d.qM + support.make_m(m, d.efc_J.T * d.efc_D * active, d.efc_J.T)
-    h = smooth.dense_m(m, d) + h
+    dh = d.replace(qM=h)
-    h_ = jax.scipy.linalg.cho_factor(h)
+    dh = smooth.factor_m(m, dh)
-    mgrad = jax.scipy.linalg.cho_solve(h_, grad)
+    mgrad = smooth.solve_m(m, dh, grad)
   else:
     raise NotImplementedError(f"unsupported solver type: {m.opt.solver}")
 
@@ -255,7 +250,7 @@ def _linesearch(m: Model, d: Data, ctx: _Context) -> _Context:
   gtol = m.opt.tolerance * m.opt.ls_tolerance * smag
 
   # compute Mv, Jv
-  mv = smooth.mul_m(m, d, ctx.search) if ctx.M is None else ctx.M @ ctx.search
+  mv = support.mul_m(m, d, ctx.search)
   jv = d.efc_J @ ctx.search
 
   # prepare quadratics

mjx/mujoco/mjx/_src/solver_test.py

@@ -39,37 +39,65 @@ def _assert_attr_eq(a, b, attr):
 
 class SolverTest(absltest.TestCase):
 
-  def test_solver(self):
+  def test_newton(self):
-    """Test solver."""
+    """Test newton solver."""
     m = test_util.load_test_file('constraints.xml')
     d = mujoco.MjData(m)
     mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
     mujoco.mj_forward(m, d)
     mx = mjx.put_model(m)
-
     dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+
     _assert_attr_eq(d, dx, 'qacc_warmstart')
     _assert_attr_eq(d, dx, 'qacc')
     _assert_attr_eq(d, dx, 'qfrc_constraint')
     nnz = dx.efc_J.any(axis=1)
     _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
 
-    # also test normal CG
+  def test_cg(self):
+    """Test CG solver."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
     m.opt.solver = mujoco.mjtSolver.mjSOL_CG
     mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
     dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+
     _assert_attr_eq(d, dx, 'qacc_warmstart')
     _assert_attr_eq(d, dx, 'qacc')
     _assert_attr_eq(d, dx, 'qfrc_constraint')
+    nnz = dx.efc_J.any(axis=1)
     _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
 
-    # without warmstart, the solution is not as close
+  def test_no_warmstart(self):
-    m.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON
+    """Test no warmstart."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
     m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART
     mujoco.mj_forward(m, d)
     mx = mjx.put_model(m)
     dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+    nnz = dx.efc_J.any(axis=1)
+    # without warmstart, the solution is not as close
     _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=2e-2)
 
+  def test_dense(self):
+    """Test solver works with dense mass matrices."""
+    m = test_util.load_test_file('constraints.xml')
+    d = mujoco.MjData(m)
+    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    mujoco.mj_forward(m, d)
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+
+    _assert_attr_eq(d, dx, 'qacc_warmstart')
+    _assert_attr_eq(d, dx, 'qacc')
+    _assert_attr_eq(d, dx, 'qfrc_constraint')
+    nnz = dx.efc_J.any(axis=1)
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+
 if __name__ == '__main__':
   absltest.main()

mjx/mujoco/mjx/_src/support.py

@@ -14,17 +14,124 @@
 # ==============================================================================
 """Engine support functions."""
 
-from typing import Tuple
+from typing import Optional, Tuple, Union
 
 import jax
 from jax import numpy as jp
+import mujoco
 from mujoco.mjx._src import scan
 # pylint: disable=g-importing-member
 from mujoco.mjx._src.types import Data
+from mujoco.mjx._src.types import JacobianType
 from mujoco.mjx._src.types import Model
 # pylint: enable=g-importing-member
 
 
+def is_sparse(m: Union[mujoco.MjModel, Model]) -> bool:
+  """Return True if this model should create sparse mass matrices.
+
+  Args:
+    m: a MuJoCo or MJX model
+
+  Returns:
+    True if provided model should create sparse mass matrices
+
+  Modern TPUs have specialized hardware for rapidly operating over sparse
+  matrices, whereas GPUs tend to be faster with dense matrices as long as they
+  fit onto the device.  As such, the default behavior in MJX (via
+  ``JacobianType.AUTO``) is sparse if ``nv`` is >= 60 or MJX detects a TPU as
+  the default backend, otherwise dense.
+  """
+  # AUTO is a rough heuristic - you may see better performance for your workload
+  # and compute by explicitly setting jacobian to dense or sparse
+  if m.opt.jacobian == JacobianType.AUTO:
+    return m.nv >= 60 or jax.default_backend() == 'tpu'
+  return m.opt.jacobian == JacobianType.SPARSE
+
+
+def make_m(
+    m: Model, a: jax.Array, b: jax.Array, d: Optional[jax.Array] = None
+) -> jax.Array:
+  """Computes M = a @ b.T + diag(d)."""
+
+  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))
+
+  if not is_sparse(m):
+    qm = a @ b.T
+    if d is not None:
+      qm += jp.diag(d)
+    mask = jp.zeros((m.nv, m.nv), dtype=bool).at[(i, j)].set(True)
+    qm = qm * mask
+    qm = qm + jp.tril(qm, -1).T
+    return qm
+
+  a_i = jp.take(a, i, axis=0)
+  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)
+
+  return qm
+
+
+def full_m(m: Model, d: Data) -> jax.Array:
+  """Reconstitute dense mass matrix from qM."""
+
+  if not is_sparse(m):
+    return d.qM
+
+  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))
+
+  mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d.qM)
+
+  # also set upper triangular
+  mat = mat + jp.tril(mat, -1).T
+
+  return mat
+
+
+def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
+  """Multiply vector by inertia matrix."""
+
+  if not is_sparse(m):
+    return d.qM @ vec
+
+  diag_mul = d.qM[jp.array(m.dof_Madr)] * vec
+
+  is_, js, madr_ijs = [], [], []
+  for i in range(m.nv):
+    madr_ij, j = m.dof_Madr[i], i
+
+    while True:
+      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.qM[madr_ij] * vec[j])
+  out = out.at[j].add(d.qM[madr_ij] * vec[i])
+
+  return out
+
+
 def jac(
     m: Model, d: Data, point: jax.Array, body_id: jax.Array
 ) -> Tuple[jax.Array, jax.Array]:

mjx/mujoco/mjx/_src/support_test.py

@@ -27,6 +27,46 @@ import numpy as np
 
 class SupportTest(parameterized.TestCase):
 
+  def test_mul_m(self):
+    m = test_util.load_test_file('pendula.xml')
+    # first test sparse
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    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
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    vec = np.random.random(m.nv)
+    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
+    mj_vec = np.zeros(m.nv)
+    mujoco.mj_mulM(m, d, mj_vec, vec)
+    np.testing.assert_allclose(mjx_vec, mj_vec, atol=5e-5, rtol=5e-5)
+
+    # also check dense
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    mujoco.mj_forward(m, d)
+    mx = mjx.put_model(m)
+    dx = mjx.put_data(m, d)
+    mjx_vec = jax.jit(mjx.mul_m)(mx, dx, jp.array(vec))
+    np.testing.assert_allclose(mjx_vec, mj_vec, atol=5e-5, rtol=5e-5)
+
+  def test_full_m(self):
+    m = test_util.load_test_file('pendula.xml')
+    # for the model to be sparse to exercise MJX full_M
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    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)
+    mjx_full_m = jax.jit(support.full_m)(mx, dx)
+    mj_full_m = np.zeros((m.nv, m.nv), dtype=np.float64)
+    mujoco.mj_fullM(m, mj_full_m, d.qM)
+    np.testing.assert_allclose(mjx_full_m, mj_full_m, atol=5e-5, rtol=5e-5)
+
   @parameterized.parameters('constraints.xml', 'pendula.xml')
   def test_jac(self, fname):
     np.random.seed(0)

mjx/mujoco/mjx/_src/types.py

@@ -123,6 +123,19 @@ class ConeType(enum.IntEnum):
   # unsupported: ELLIPTIC
 
 
+class JacobianType(enum.IntEnum):
+  """Type of constraint Jacobian.
+
+  Attributes:
+    DENSE: dense
+    SPARSE: sparse
+    AUTO: sparse if nv>60 and device is TPU, dense otherwise
+  """
+  DENSE = mujoco.mjtJacobian.mjJAC_DENSE
+  SPARSE = mujoco.mjtJacobian.mjJAC_SPARSE
+  AUTO = mujoco.mjtJacobian.mjJAC_AUTO
+
+
 class SolverType(enum.IntEnum):
   """Constraint solver algorithm.
 
@@ -215,6 +228,10 @@ class Option(PyTreeNode):
       nonzero. Not used by mj
     integrator:       integration mode
     cone:             type of friction cone
+    jacobian:         matrix layout for mass matrices (dense or sparse)
+                      (note that this is different from MuJoCo, where jacobian
+                      specifies whether efc_J and its accompanying matrices
+                      are dense or sparse.
     solver:           solver algorithm
     iterations:       number of main solver iterations
     ls_iterations:    maximum number of CG/Newton linesearch iterations
@@ -232,7 +249,7 @@ class Option(PyTreeNode):
   # unsupported: magnetic, o_margin, o_solref, o_solimp
   integrator: IntegratorType
   cone: ConeType
-  # unsupported: jacobian
+  jacobian: JacobianType
   solver: SolverType
   iterations: int
   ls_iterations: int
@@ -600,10 +617,12 @@ class Data(PyTreeNode):
     actuator_length: actuator lengths                             (nu,)
     actuator_moment: actuator moments                             (nu, nv)
     crb: com-based composite inertia and mass                     (nbody, 10)
-    qM: total inertia (sparse)                                    (nM,)
+    qM: total inertia                                  if sparse: (nM,)
-    qLD: L'*D*L factorization of M (sparse)                       (nM,)
+                                                       if dense:  (nv, nv)
-    qLDiagInv: 1/diag(D)                                          (nv,)
+    qLD: L'*D*L (or Cholesky) factorization of M.      if sparse: (nM,)
-    qLDiagSqrtInv: 1/sqrt(diag(D))                                (nv,)
+                                                       if dense:  (nv, nv)
+    qLDiagInv: 1/diag(D)                               if sparse: (nv,)
+                                                       if dense:  (0,)
     contact: list of all detected contacts                        (ncon,)
     efc_J: constraint Jacobian                                    (nefc, nv)
     efc_frictionloss: frictionloss (friction)                     (nefc,)
@@ -660,7 +679,6 @@ class Data(PyTreeNode):
   qM: jax.Array  # pylint:disable=invalid-name
   qLD: jax.Array  # pylint:disable=invalid-name
   qLDiagInv: jax.Array  # pylint:disable=invalid-name
-  qLDiagSqrtInv: jax.Array  # pylint:disable=invalid-name
   contact: Contact
   efc_J: jax.Array  # pylint:disable=invalid-name
   efc_frictionloss: jax.Array