Fixes to the MJX Euler integrator and Newton solver, and test improvements to catch future bugs like this.

PiperOrigin-RevId: 606389094
Change-Id: I50d0143f338fe29a823c930dba82995ce8a64ce2

doc/changelog.rst

@@ -16,6 +16,8 @@ MJX
    - Avoid calling ``mjx.ncon`` in ``mjx.get_data_into`` when ``nc`` can be derived from ``mjx.Data``.
 2. Fixed a bug in ``mjx-viewer`` that prevented it from running.  Updated ``mjx-viewer`` to use newer
    ``mjx.get_data_into`` function call.
+3. Fixed a bug in ``mjx.euler`` that applied incorrect damping when using dense mass matrices.
+4. Fixed a bug in ``mjx.solve`` that was causing slow convergence when using ``mjSOL_NEWTON`` in :ref:`mjtSolver`.
 
 Version 3.1.2 (February 05, 2024)
 -----------------------------------

mjx/mujoco/mjx/_src/forward.py

@@ -287,8 +287,10 @@ def euler(m: Model, d: Data) -> Data:
   # integrate damping implicitly
   qacc = d.qacc
   if not m.opt.disableflags & DisableBit.EULERDAMP:
-    # TODO(robotics-simulation): can this be done with a smaller perf hit
+    if support.is_sparse(m):
-    dh = d.replace(qM=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))
+    else:
+      dh = d.replace(qM=d.qM + jp.diag(m.opt.timestep * m.dof_damping))
     dh = smooth.factor_m(m, dh)
     qfrc = d.qfrc_smooth + d.qfrc_constraint
     qacc = smooth.solve_m(m, dh, qfrc)

mjx/mujoco/mjx/_src/forward_test.py

@@ -27,14 +27,14 @@ import numpy as np
 _TOLERANCE = 1e-5
 
 
-def _assert_eq(a, b, name):
+def _assert_eq(a, b, name, tol=_TOLERANCE):
-  tol = _TOLERANCE * 10  # avoid test noise
+  tol = tol * 10  # avoid test noise
   err_msg = f'mismatch: {name}'
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
 
-def _assert_attr_eq(a, b, attr):
+def _assert_attr_eq(a, b, attr, tol=_TOLERANCE):
-  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr, tol=tol)
 
 
 class ForwardTest(absltest.TestCase):
@@ -46,7 +46,7 @@ class ForwardTest(absltest.TestCase):
     d.ctrl = np.array([-18, 0.59, 0.47])
     d.xfrc_applied[0, 2] = 0.1  # torque
     d.xfrc_applied[1, 4] = 0.3  # linear force
-    mujoco.mj_step(m, d, 100)  # get some dynamics going
+    mujoco.mj_step(m, d, 20)  # get some dynamics going
     mujoco.mj_forward(m, d)
 
     mx = mjx.put_model(m)
@@ -75,14 +75,13 @@ class ForwardTest(absltest.TestCase):
     d.ctrl = np.array([-18, 0.59, 0.47])
     d.xfrc_applied[0, 2] = 0.1  # torque
     d.xfrc_applied[1, 4] = 0.3  # linear force
-    mujoco.mj_step(m, d, 100)  # get some dynamics going
+    mujoco.mj_step(m, d, 20)  # get some dynamics going
 
-    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.step)(mjx.put_model(m), mjx.put_data(m, d))
-    dx = jax.jit(mjx.step)(mx, mjx.put_data(m, d))
     mujoco.mj_step(m, d)
     _assert_attr_eq(d, dx, 'act')
     _assert_attr_eq(d, dx, 'time')
-    _assert_attr_eq(d, dx, 'qvel')
+    _assert_attr_eq(d, dx, 'qvel', tol=5e-4)
     _assert_attr_eq(d, dx, 'qpos')
 
   def test_rk4(self):
@@ -111,8 +110,7 @@ class ForwardTest(absltest.TestCase):
     mujoco.mj_step(m, d, 10)  # let dynamics get state significantly non-zero
     mujoco.mj_forward(m, d)
 
-    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.rungekutta4)(mjx.put_model(m), mjx.put_data(m, d))
-    dx = jax.jit(mjx.rungekutta4)(mx, mjx.put_data(m, d))
     mujoco.mj_RungeKutta(m, d, 4)
 
     _assert_attr_eq(d, dx, 'qvel')
@@ -120,6 +118,30 @@ class ForwardTest(absltest.TestCase):
     _assert_attr_eq(d, dx, 'act')
     _assert_attr_eq(d, dx, 'time')
 
+  def test_eulerdamp(self):
+    m = test_util.load_test_file('pendula.xml')
+    self.assertTrue((m.dof_damping > 0).any())
+
+    d = mujoco.MjData(m)
+    d.qvel[:] = 1.0
+    d.qacc[:] = 1.0
+    mujoco.mj_forward(m, d)
+    dx = jax.jit(mjx.euler)(mjx.put_model(m), mjx.put_data(m, d))
+    mujoco.mj_Euler(m, d)
+
+    _assert_attr_eq(d, dx, 'qpos')
+
+    # also test sparse
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
+    d = mujoco.MjData(m)
+    d.qvel[:] = 1.0
+    d.qacc[:] = 1.0
+    mujoco.mj_forward(m, d)
+    dx = jax.jit(mjx.euler)(mjx.put_model(m), mjx.put_data(m, d))
+    mujoco.mj_Euler(m, d)
+
+    _assert_attr_eq(d, dx, 'qpos')
+
   def test_disable_eulerdamp(self):
     m = test_util.load_test_file('pendula.xml')
     self.assertTrue((m.dof_damping > 0).any())
@@ -128,8 +150,7 @@ class ForwardTest(absltest.TestCase):
     d = mujoco.MjData(m)
     d.qvel[:] = 1.0
     d.qacc[:] = 1.0
-    mx = mjx.put_model(m)
+    dx = jax.jit(mjx.euler)(mjx.put_model(m), mjx.put_data(m, d))
-    dx = jax.jit(mjx.euler)(mx, mjx.put_data(m, d))
 
     np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep)
 

mjx/mujoco/mjx/_src/solver.py

@@ -218,13 +218,13 @@ def _update_gradient(m: Model, d: Data, ctx: _Context) -> _Context:
     mgrad = smooth.solve_m(m, d, grad)
   elif m.opt.solver == SolverType.NEWTON:
     ne, nf, *_ = constraint.count_constraints(m)
-    active = (ctx.Jaref < 0).at[:ne + nf].set(True)
+    active = (ctx.Jaref < 0).at[: ne + nf].set(True)
-    h = d.qM + support.make_m(m, d.efc_J.T * d.efc_D * active, d.efc_J.T)
+    h = (d.efc_J.T * d.efc_D * active) @ d.efc_J
-    dh = d.replace(qM=h)
+    h = support.full_m(m, d) + h
-    dh = smooth.factor_m(m, dh)
+    h_ = jax.scipy.linalg.cho_factor(h)
-    mgrad = smooth.solve_m(m, dh, grad)
+    mgrad = jax.scipy.linalg.cho_solve(h_, grad)
   else:
-    raise NotImplementedError(f"unsupported solver type: {m.opt.solver}")
+    raise NotImplementedError(f'unsupported solver type: {m.opt.solver}')
 
   ctx = ctx.replace(grad=grad, Mgrad=mgrad)
 

mjx/mujoco/mjx/_src/solver_test.py

@@ -33,8 +33,8 @@ def _assert_eq(a, b, name, tol=_TOLERANCE):
   np.testing.assert_allclose(a, b, err_msg=err_msg, atol=tol, rtol=tol)
 
 
-def _assert_attr_eq(a, b, attr):
+def _assert_attr_eq(a, b, attr, tol=_TOLERANCE):
-  _assert_eq(getattr(a, attr), getattr(b, attr), attr)
+  _assert_eq(getattr(a, attr), getattr(b, attr), attr, tol=tol)
 
 
 class SolverTest(absltest.TestCase):
@@ -42,13 +42,20 @@ class SolverTest(absltest.TestCase):
   def test_newton(self):
     """Test newton solver."""
     m = test_util.load_test_file('constraints.xml')
+    # it's critical that mgrad is optimally calculated, so lower iterations
+    # to be sure that MJX is converging as quickly as MuJoCo
+    m.opt.iterations = 1
     d = mujoco.MjData(m)
-    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    mujoco.mj_step(m, d, 20)  # significant constraint forces at 20 steps
-    mujoco.mj_forward(m, d)
+
-    mx = mjx.put_model(m)
+    # mj_forward overwrites qacc_warmstart, so let's restore it to what it was
-    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+    # at the beginning of the step so that MJX does not have a trivial solution
+    warmstart = d.qacc_warmstart.copy()
+    mujoco.mj_forward(m, d)
+    d.qacc_warmstart = warmstart
+
+    dx = jax.jit(mjx.solve)(mjx.put_model(m), 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)
@@ -58,23 +65,30 @@ class SolverTest(absltest.TestCase):
     """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
+    mujoco.mj_step(m, d, 20)  # significant constraint forces at 20 steps
-    m.opt.solver = mujoco.mjtSolver.mjSOL_CG
+
-    mujoco.mj_forward(m, d)
+    # CG does not converge as quickly as Newton but is cheaper to calculate
-    mx = mjx.put_model(m)
+    m.opt.solver = mujoco.mjtSolver.mjSOL_CG
-    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+    m.opt.iterations = 8
+
+    # mj_forward overwrites qacc_warmstart, so let's restore it to what it was
+    # at the beginning of the step so that MJX does not have a trivial solution
+    warmstart = d.qacc_warmstart.copy()
+    mujoco.mj_forward(m, d)
+    d.qacc_warmstart = warmstart
+
+    dx = jax.jit(mjx.solve)(mjx.put_model(m), 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')
+    _assert_attr_eq(d, dx, 'qfrc_constraint', tol=8e-4)
     nnz = dx.efc_J.any(axis=1)
-    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force')
+    _assert_eq(d.efc_force, dx.efc_force[nnz], 'efc_force', tol=5e-4)
 
   def test_no_warmstart(self):
     """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
+    mujoco.mj_step(m, d, 20)  # significant constraint forces at 20 steps
     m.opt.disableflags |= mujoco.mjtDisableBit.mjDSBL_WARMSTART
     mujoco.mj_forward(m, d)
     mx = mjx.put_model(m)
@@ -83,17 +97,21 @@ class SolverTest(absltest.TestCase):
     # 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):
+  def test_sparse(self):
-    """Test solver works with dense mass matrices."""
+    """Test solver works with sparse mass matrices."""
     m = test_util.load_test_file('constraints.xml')
+    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
     d = mujoco.MjData(m)
-    mujoco.mj_step(m, d, 100)  # at 100 steps mix of active/inactive constraints
+    mujoco.mj_step(m, d, 20)  # significant constraint forces at 20 steps
-    mujoco.mj_forward(m, d)
+
-    m.opt.jacobian = mujoco.mjtJacobian.mjJAC_DENSE
+    # mj_forward overwrites qacc_warmstart, so let's restore it to what it was
-    mx = mjx.put_model(m)
+    # at the beginning of the step so that MJX does not have a trivial solution
-    dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d))
+    warmstart = d.qacc_warmstart.copy()
+    mujoco.mj_forward(m, d)
+    d.qacc_warmstart = warmstart
+
+    dx = jax.jit(mjx.solve)(mjx.put_model(m), 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)

mjx/mujoco/mjx/test_data/constraints.xml

@@ -5,7 +5,7 @@
 * solref, solimp
 -->
 <mujoco>
-  <option timestep="0.02"/>
+  <option timestep="0.015"/>
 
   <default>
     <default class="box">