diff --git a/doc/changelog.rst b/doc/changelog.rst index 55cf4d2d..14608cee 100644 --- a/doc/changelog.rst +++ b/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) ----------------------------------- diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index 682cc88e..c6bbdbc1 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/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 - dh = d.replace(qM=d.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)) + if support.is_sparse(m): + 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) diff --git a/mjx/mujoco/mjx/_src/forward_test.py b/mjx/mujoco/mjx/_src/forward_test.py index cc767ae8..0a777797 100644 --- a/mjx/mujoco/mjx/_src/forward_test.py +++ b/mjx/mujoco/mjx/_src/forward_test.py @@ -27,14 +27,14 @@ import numpy as np _TOLERANCE = 1e-5 -def _assert_eq(a, b, name): - tol = _TOLERANCE * 10 # avoid test noise +def _assert_eq(a, b, name, tol=_TOLERANCE): + 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): - _assert_eq(getattr(a, attr), getattr(b, attr), attr) +def _assert_attr_eq(a, b, attr, tol=_TOLERANCE): + _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)(mx, mjx.put_data(m, d)) + dx = jax.jit(mjx.step)(mjx.put_model(m), 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)(mx, mjx.put_data(m, d)) + dx = jax.jit(mjx.rungekutta4)(mjx.put_model(m), 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)(mx, mjx.put_data(m, d)) + dx = jax.jit(mjx.euler)(mjx.put_model(m), mjx.put_data(m, d)) np.testing.assert_allclose(dx.qvel, 1 + m.opt.timestep) diff --git a/mjx/mujoco/mjx/_src/solver.py b/mjx/mujoco/mjx/_src/solver.py index cd0cc63c..dae5a82b 100644 --- a/mjx/mujoco/mjx/_src/solver.py +++ b/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) - h = d.qM + support.make_m(m, d.efc_J.T * d.efc_D * active, d.efc_J.T) - dh = d.replace(qM=h) - dh = smooth.factor_m(m, dh) - mgrad = smooth.solve_m(m, dh, grad) + active = (ctx.Jaref < 0).at[: ne + nf].set(True) + h = (d.efc_J.T * d.efc_D * active) @ d.efc_J + h = support.full_m(m, d) + h + h_ = jax.scipy.linalg.cho_factor(h) + 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) diff --git a/mjx/mujoco/mjx/_src/solver_test.py b/mjx/mujoco/mjx/_src/solver_test.py index 692f6619..f4fdf732 100644 --- a/mjx/mujoco/mjx/_src/solver_test.py +++ b/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): - _assert_eq(getattr(a, attr), getattr(b, attr), attr) +def _assert_attr_eq(a, b, attr, tol=_TOLERANCE): + _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_forward(m, d) - mx = mjx.put_model(m) - dx = jax.jit(mjx.solve)(mx, mjx.put_data(m, d)) + mujoco.mj_step(m, d, 20) # significant constraint forces at 20 steps + + # 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') 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 - 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)) + mujoco.mj_step(m, d, 20) # significant constraint forces at 20 steps + + # CG does not converge as quickly as Newton but is cheaper to calculate + m.opt.solver = mujoco.mjtSolver.mjSOL_CG + 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): - """Test solver works with dense mass matrices.""" + def test_sparse(self): + """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_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)) + mujoco.mj_step(m, d, 20) # significant constraint forces at 20 steps + + # 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') nnz = dx.efc_J.any(axis=1) diff --git a/mjx/mujoco/mjx/test_data/constraints.xml b/mjx/mujoco/mjx/test_data/constraints.xml index 52eb95d8..c7a9cd04 100644 --- a/mjx/mujoco/mjx/test_data/constraints.xml +++ b/mjx/mujoco/mjx/test_data/constraints.xml @@ -5,7 +5,7 @@ * solref, solimp --> -