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

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

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

PiperOrigin-RevId: 600696483
Change-Id: If69bb9a2e21ba8dad6ca23f093ce7b7ceae644ff
This commit is contained in:
Erik Frey
2024-01-23 00:31:01 -08:00
committed by Copybara-Service
parent aceb52bd09
commit 0a7be1732c
14 changed files with 324 additions and 143 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/changelog.rst
+10 -8
View File
@@ -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
code will require modification.
10. Improved the implmentation of the :ref:`rollout<PySample>` module. Note the changes below are breaking, dependent
code will require modification.
- 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.
- Outputs are no longer squeezed and always have dim=3.
- 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.
- 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
:github:issue:`1270`.
10. Fixed a bug in the :ref:`muscle model<CMuscle>` that led to non-zero values outside the lower
11. Fixed a bug that prevented the use of pins with plugins if flexes are not in the worldbody. Fixes
:github:issue:`1270`.
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
+8
View File
@@ -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
+3 -1
View File
@@ -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
+2
View File
@@ -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
+3 -3
View File
@@ -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 = smooth.factor_m(m, d, mh)
dh = d.replace(qM=d.qM.at[m.dof_Madr].add(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)
return _advance(m, d, d.act_dot, qacc)
mjx/mujoco/mjx/_src/io.py
+40 -4
View File
@@ -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,
qLD=zero_nm,
qLDiagInv=zero_nv,
qLDiagSqrtInv=zero_nv,
qM=zero_nm if support.is_sparse(m) else zero_nv_nv,
qLD=zero_nm if support.is_sparse(m) else zero_nv_nv,
qLDiagInv=zero_nv if support.is_sparse(m) else zero_0,
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 -9
View File
@@ -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.qLD.shape, (nm,))
self.assertEqual(d.qLDiagInv.shape, (nv,))
self.assertEqual(d.qLDiagSqrtInv.shape, (nv,))
self.assertEqual(d.qM.shape, (nv, nv))
self.assertEqual(d.qLD.shape, (nv, nv))
self.assertEqual(d.qLDiagInv.shape, (0,))
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)
np.testing.assert_allclose(dx_from_sparse.efc_J, dx.efc_J, atol=1e-8)
dx_sparse = mjx.put_data(m, d)
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
+17 -75
View File
@@ -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]))]
return jp.concatenate([inert, off * mass, jp.expand_dims(mass, 0)])
inert = inert[([0, 1, 2, 0, 0, 1], [0, 1, 2, 1, 2, 2])]
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)
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)
qm = support.make_m(m, crb_cdof, d.cdof, m.dof_armature)
d = d.replace(qM=qm)
return d
def factor_m(
m: Model,
d: Data,
qM: jax.Array, # pylint:disable=invalid-name
) -> Data:
"""Gets sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd."""
def factor_m(m: Model, d: Data) -> Data:
"""Gets factorizaton of inertia-like matrix M, assumed spd."""
if not support.is_sparse(m):
qh, _ = jax.scipy.linalg.cho_factor(d.qM)
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
+3 -18
View File
@@ -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, dx, dx.qM)
dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
_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
+7 -12
View File
@@ -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 = smooth.mul_m(m, d, d.qacc) if M is None else M @ d.qacc
ma = support.mul_m(m, d, 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 = smooth.dense_m(m, d) + h
h_ = jax.scipy.linalg.cho_factor(h)
mgrad = jax.scipy.linalg.cho_solve(h_, grad)
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)
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
+34 -6
View File
@@ -39,37 +39,65 @@ def _assert_attr_eq(a, b, attr):
class SolverTest(absltest.TestCase):
def test_solver(self):
"""Test solver."""
def test_newton(self):
"""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
m.opt.solver = mujoco.mjtSolver.mjSOL_NEWTON
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
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
+108 -1
View File
@@ -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
+40
View File
@@ -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
+24 -6
View File
@@ -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,)
qLD: L'*D*L factorization of M (sparse) (nM,)
qLDiagInv: 1/diag(D) (nv,)
qLDiagSqrtInv: 1/sqrt(diag(D)) (nv,)
qM: total inertia if sparse: (nM,)
if dense: (nv, nv)
qLD: L'*D*L (or Cholesky) factorization of M. if sparse: (nM,)
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