chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
PiperOrigin-RevId: 606389094
Change-Id: I50d0143f338fe29a823c930dba82995ce8a64ce2
This commit is contained in:
Erik Frey
2024-02-12 15:45:25 -08:00
committed by Copybara-Service
parent 805d862d5f
commit 47bb4a828e
6 changed files with 89 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+2
View File
@@ -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
+4 -2
View File
@@ -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)
mjx/mujoco/mjx/_src/forward_test.py
+34 -13
View File
@@ -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)
mjx/mujoco/mjx/_src/solver.py
+6 -6
View File
@@ -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)
mjx/mujoco/mjx/_src/solver_test.py
+42 -24
View File
@@ -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)
mjx/mujoco/mjx/test_data/constraints.xml
+1 -1
View File
@@ -5,7 +5,7 @@
* solref, solimp
-->
<mujoco>
<option timestep="0.02"/>
<option timestep="0.015"/>
<default>
<default class="box">