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

Switch MJX DataJAX from legacy qM format to CSR M format

Browse Source 
Switch MJX DataJAX to use the CSR-format `M` as its primary sparse inertia representation, matching MuJoCo C `mjData.M`. This allows both `mjData.qM` and `mjModel.mapM2M` to be deleted in the future.

PiperOrigin-RevId: 932433465
Change-Id: I194ffbfdba0cdd9c718277ee0d5de7ac20f2b06f
This commit is contained in:
Yuval Tassa
2026-06-15 06:52:12 -07:00
committed by Copybara-Service
parent 3317921999
commit 60b5cac066
7 changed files with 96 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/forward.py
+5 -4
View File
@@ -352,10 +352,11 @@ def euler(m: Model, d: Data) -> Data:
qacc = d.qacc
if not m.opt.disableflags & DisableBit.EULERDAMP:
if support.is_sparse(m):
qM = d._impl.qM.at[m.dof_Madr].add(m.opt.timestep * m.dof_damping)
diag_adr = m.M_rowadr + m.M_rownnz - 1
M = d._impl.M.at[diag_adr].add(m.opt.timestep * m.dof_damping)
else:
qM = d._impl.qM + jp.diag(m.opt.timestep * m.dof_damping)
dh = d.tree_replace({'_impl.qM': qM})
M = d._impl.M + jp.diag(m.opt.timestep * m.dof_damping)
dh = d.tree_replace({'_impl.M': M})
dh = smooth.factor_m(m, dh)
qfrc = d.qfrc_smooth + d.qfrc_constraint
qacc = smooth.solve_m(m, dh, qfrc)
@@ -418,7 +419,7 @@ def implicit(m: Model, d: Data) -> Data:
qacc = d.qacc
if qderiv is not None:
# TODO(robotics-simulation): use smooth.factor_m / solve_m here:
qm = support.full_m(m, d) if support.is_sparse(m) else d._impl.qM
qm = support.full_m(m, d) if support.is_sparse(m) else d._impl.M
qm -= m.opt.timestep * qderiv
qh, _ = jax.scipy.linalg.cho_factor(qm)
qfrc = d.qfrc_smooth + d.qfrc_constraint
mjx/mujoco/mjx/_src/io.py
+17 -21
View File
@@ -669,8 +669,7 @@ def _make_data_jax(
'wrap_xpos': (m.nwrap, 6, float_),
'actuator_moment': (m.nu, m.nv, float_),
'crb': (m.nbody, 10, float_),
'qM': (m.nM, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
'M': (m.nC, float_),
'M': (m.nC, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
'qLD': (m.nC, float_) if support.is_sparse(m) else (m.nv, m.nv, float_),
'qLDiagInv': (m.nv, float_) if support.is_sparse(m) else (0, float_),
'ten_velocity': (m.ntendon, float_),
@@ -1045,11 +1044,11 @@ def _put_data_jax(
impl_fields[fname] = value
# convert qM and qLD if jacobian is dense
# convert M and qLD if jacobian is dense
if not support.is_sparse(m):
impl_fields['qM'] = np.zeros((m.nv, m.nv))
impl_fields['M'] = np.zeros((m.nv, m.nv))
mujoco.mju_sym2dense(
impl_fields['qM'],
impl_fields['M'],
d.M,
m.M_rownnz,
m.M_rowadr,
@@ -1057,9 +1056,9 @@ def _put_data_jax(
)
# TODO(erikfrey): derive L*L' from L'*D*L instead of recomputing
try:
impl_fields['qLD'], _ = scipy.linalg.cho_factor(impl_fields['qM'])
impl_fields['qLD'], _ = scipy.linalg.cho_factor(impl_fields['M'])
except scipy.linalg.LinAlgError:
# this happens when qM is empty or unstable simulation
# this happens when M is empty or unstable simulation
impl_fields['qLD'] = np.zeros((m.nv, m.nv))
impl_fields['qLDiagInv'] = np.zeros(0)
@@ -1333,15 +1332,6 @@ def _get_data_into(
d = jax.device_get(d)
batch_size = d.qpos.shape[0] if batched else 1
dof_i, dof_j = [], []
if d.impl == types.Impl.JAX:
for i in range(m.nv):
j = i
while j > -1:
dof_i.append(i)
dof_j.append(j)
j = m.dof_parentid[j]
for i in range(batch_size):
d_i = jax.tree_util.tree_map(lambda x, i=i: x[i], d) if batched else d
result_i = result[i] if batched else result
@@ -1446,8 +1436,14 @@ def _get_data_into(
elif field.name.startswith('efc_'):
value = value[efc_active]
if d.impl == types.Impl.JAX:
if field.name == 'qM' and not support.is_sparse(m):
value = value[dof_i, dof_j]
if field.name == 'M' and not support.is_sparse(m):
M_csr = np.zeros(m.nC)
for i in range(m.nv):
adr = m.M_rowadr[i]
for k in range(m.M_rownnz[i]):
col = m.M_colind[adr + k]
M_csr[adr + k] = value[i, col]
value = M_csr
elif field.name == 'qLD':
value = np.zeros(m.nC)
elif field.name == 'qLDiagInv' and not support.is_sparse(m):
@@ -1464,9 +1460,9 @@ def _get_data_into(
else:
setattr(result_i, field.name, value)
# TODO(taylorhowell): remove mapping once qM is deprecated
# map inertia (sparse) to reduced inertia (compressed sparse) representation
result_i.M[:] = result_i.qM[m.mapM2M]
if hasattr(result_i, 'qM'):
result_i.qM.fill(0.0)
result_i.qM[m.mapM2M] = result_i.M
# recalculate qLD and qLDiagInv as MJX and MuJoCo have different
# representations of the Cholesky decomposition.
mjx/mujoco/mjx/_src/io_test.py
+7 -9
View File
@@ -476,20 +476,18 @@ class DataIOTest(parameterized.TestCase):
self.assertEqual(d._impl.efc_force.shape, (nefc,))
if impl == 'jax':
self.assertEqual(d._impl.qM.shape, (nv, nv))
self.assertEqual(d._impl.M.shape, (nv, nv))
self.assertEqual(d._impl.qLD.shape, (nv, nv))
self.assertEqual(d._impl.qLDiagInv.shape, (0,))
# test sparse
m.opt.jacobian = mujoco.mjtJacobian.mjJAC_SPARSE
d = mjx.make_data(m, impl=impl)
self.assertEqual(d._impl.qM.shape, (nm,))
self.assertEqual(d._impl.qLD.shape, (nm,))
self.assertEqual(d._impl.M.shape, (m.nC,))
self.assertEqual(d._impl.qLD.shape, (m.nC,))
self.assertEqual(d._impl.qLDiagInv.shape, (nv,))
@mock.patch.dict(os.environ, {'MJX_GPU_DEFAULT_WARP': 'true'})
def test_make_data_warp(self):
if not mjxw.WARP_INSTALLED:
@@ -605,7 +603,7 @@ class DataIOTest(parameterized.TestCase):
)
# check sparse mass matrices are correct
np.testing.assert_allclose(dx_sparse._impl.qM, d.qM, atol=1e-8)
np.testing.assert_allclose(dx_sparse._impl.M, d.M, atol=1e-8)
np.testing.assert_allclose(dx_sparse._impl.qLD, d.qLD, atol=1e-8)
np.testing.assert_allclose(
dx_sparse._impl.qLDiagInv, d.qLDiagInv, atol=1e-8
@@ -619,7 +617,7 @@ class DataIOTest(parameterized.TestCase):
if impl == 'jax':
qm = np.zeros((m.nv, m.nv))
mujoco.mju_sym2dense(qm, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
np.testing.assert_allclose(dx_from_dense._impl.qM, qm, atol=1e-8)
np.testing.assert_allclose(dx_from_dense._impl.M, qm, atol=1e-8)
def test_put_data_warp_ndim(self):
@@ -865,8 +863,8 @@ class DataIOTest(parameterized.TestCase):
mjx.make_data(m)
@parameterized.parameters(JacobianType.DENSE, JacobianType.SPARSE)
def test_qm_mapm2m(self, jacobian):
"""Test that qM is mapped to M."""
def test_m_mapm2m(self, jacobian):
"""Test that M matches MuJoCo."""
m = test_util.load_test_file('humanoid/humanoid.xml')
m.opt.jacobian = jacobian
d = mujoco.MjData(m)
mjx/mujoco/mjx/_src/smooth.py
+37 -45
View File
@@ -307,7 +307,7 @@ def crb(m: Model, d: Data) -> Data:
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)
d = d.tree_replace({'_impl.qM': qm})
d = d.tree_replace({'_impl.M': qm})
return d
@@ -317,56 +317,55 @@ def factor_m(m: Model, d: Data) -> Data:
raise ValueError('factor_m requires JAX backend implementation.')
if not support.is_sparse(m):
qh, _ = jax.scipy.linalg.cho_factor(d._impl.qM)
qh, _ = jax.scipy.linalg.cho_factor(d._impl.M)
d = d.tree_replace({'_impl.qLD': qh})
return d
# build up indices for where we will do backwards updates over qLD
depth = []
for i in range(m.nv):
depth.append(depth[m.dof_parentid[i]] + 1 if m.dof_parentid[i] != -1 else 0)
updates = {}
madr_ds = []
diag_ds = []
for i in range(m.nv):
madr_d = madr_ij = m.dof_Madr[i]
j = i
while True:
madr_ds.append(madr_d)
madr_ij, j = madr_ij + 1, m.dof_parentid[j]
if j == -1:
break
out_beg, out_end = tuple(m.dof_Madr[j : j + 2])
diag_i = m.M_rowadr[i] + m.M_rownnz[i] - 1
start_i = m.M_rowadr[i]
for k in range(m.M_rownnz[i]):
diag_ds.append(diag_i)
for k in range(m.M_rownnz[i] - 1):
adr = start_i + k
j = m.M_colind[adr]
out_beg = m.M_rowadr[j]
width = m.M_rownnz[j]
out_end = out_beg + width
updates.setdefault(depth[j], []).append(
(out_beg, out_end, madr_d, madr_ij)
(out_beg, out_end, diag_i, adr, start_i)
)
qld = d._impl.qM
qld = d._impl.M
for _, updates in sorted(updates.items(), reverse=True):
# combine the updates into one update batch (per depth level)
rows = []
madr_ijs = []
pivots = []
out = []
for b, e, madr_d, madr_ij in updates:
for b, e, piv_i, madr_ij, start_i in updates:
width = e - b
rows.append(np.arange(madr_ij, madr_ij + width))
rows.append(np.arange(start_i, start_i + width))
madr_ijs.append(np.full((width,), madr_ij))
pivots.append(np.full((width,), madr_d))
pivots.append(np.full((width,), piv_i))
out.append(np.arange(b, e))
rows = np.concatenate(rows)
madr_ijs = np.concatenate(madr_ijs)
pivots = np.concatenate(pivots)
out = np.concatenate(out)
if out:
rows = np.concatenate(rows)
madr_ijs = np.concatenate(madr_ijs)
pivots = np.concatenate(pivots)
out = np.concatenate(out)
# apply the update batch
qld = qld.at[out].add(-(qld[madr_ijs] / qld[pivots]) * qld[rows])
# TODO(erikfrey): determine if this minimum value guarding is necessary:
# qld = qld.at[dof_madr].set(jp.maximum(qld[dof_madr], _MJ_MINVAL))
qld = qld.at[out].add(-(qld[madr_ijs] / qld[pivots]) * qld[rows])
qld_diag = qld[m.dof_Madr]
qld = (qld / qld[jp.array(madr_ds)]).at[m.dof_Madr].set(qld_diag)
diag_adr = m.M_rowadr + m.M_rownnz - 1
qld_diag = qld[diag_adr]
qld = (qld / qld[np.array(diag_ds)]).at[diag_adr].set(qld_diag)
d = d.tree_replace({'_impl.qLD': qld, '_impl.qLDiagInv': 1 / qld_diag})
return d
@@ -386,13 +385,12 @@ def solve_m(m: Model, d: Data, x: jax.Array) -> jax.Array:
updates_i, updates_j = {}, {}
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
updates_i.setdefault(depth[i], []).append((i, madr_ij, j))
updates_j.setdefault(depth[j], []).append((j, madr_ij, i))
start_i = m.M_rowadr[i]
for k in range(m.M_rownnz[i] - 1):
adr = start_i + k
j = m.M_colind[adr]
updates_i.setdefault(depth[i], []).append((i, adr, j))
updates_j.setdefault(depth[j], []).append((j, adr, i))
# x <- inv(L') * x
for _, vals in sorted(updates_j.items(), reverse=True):
@@ -1317,7 +1315,7 @@ def transmission(m: Model, d: Data) -> Data:
def tendon_armature(m: Model, d: Data) -> Data:
"""Add tendon armature to qM."""
"""Add tendon armature to M."""
if not isinstance(m._impl, ModelJAX) or not isinstance(d._impl, DataJAX):
raise ValueError('tendon_armature requires JAX backend implementation.')
@@ -1330,17 +1328,11 @@ def tendon_armature(m: Model, d: Data) -> Data:
)
if support.is_sparse(m):
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))
i = np.repeat(np.arange(m.nv), m.M_rownnz)
j = m.M_colind
JTAJ = JTAJ[(i, j)]
return d.tree_replace({'_impl.qM': d._impl.qM + JTAJ})
return d.tree_replace({'_impl.M': d._impl.M + JTAJ})
def tendon_dot(m: Model, d: Data) -> jax.Array:
mjx/mujoco/mjx/_src/smooth_test.py
+7 -10
View File
@@ -89,13 +89,10 @@ class SmoothTest(absltest.TestCase):
# crb
dx = jax.jit(mjx.crb)(mx, mjx.put_data(m, d))
_assert_attr_eq(d, dx._impl, 'crb')
_assert_attr_eq(d, dx._impl, 'qM')
_assert_attr_eq(d, dx._impl, 'M')
# factor_m
dx = jax.jit(mjx.factor_m)(mx, mjx.put_data(m, d))
qLDLegacy = np.zeros(mx.nM) # pylint:disable=invalid-name
for i in range(m.nC):
qLDLegacy[m.mapM2M[i]] = d.qLD[i]
_assert_eq(qLDLegacy, dx._impl.qLD, 'qLD')
_assert_attr_eq(d, dx._impl, 'qLD')
_assert_attr_eq(d, dx._impl, 'qLDiagInv')
# com_vel
dx = jax.jit(mjx.com_vel)(mx, mjx.put_data(m, d))
@@ -430,18 +427,18 @@ class TendonTest(parameterized.TestCase):
dx = mjx.put_data(m, d)
dx = dx.tree_replace(
{'_impl.qM': jp.zeros((m.nv, m.nv)), 'qfrc_bias': jp.zeros(m.nv)}
{'_impl.M': jp.zeros((m.nv, m.nv)), 'qfrc_bias': jp.zeros(m.nv)}
)
dx = mjx.crb(mx, dx)
dx = mjx.tendon_armature(mx, dx)
if jacobian == JacobianType.DENSE:
qM = np.zeros((m.nv, m.nv)) # pylint: disable=invalid-name
mujoco.mju_sym2dense(qM, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
M = np.zeros((m.nv, m.nv)) # pylint: disable=invalid-name
mujoco.mju_sym2dense(M, d.M, m.M_rownnz, m.M_rowadr, m.M_colind)
else:
qM = d.qM # pylint: disable=invalid-name
_assert_eq(dx._impl.qM, qM, 'qM')
M = d.M # pylint: disable=invalid-name
_assert_eq(dx._impl.M, M, 'M')
dx = mjx.rne(mx, dx)
dx = mjx.tendon_bias(mx, dx)
mjx/mujoco/mjx/_src/support.py
+20 -33
View File
@@ -60,14 +60,8 @@ def make_m(
) -> 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))
i = np.repeat(np.arange(m.nv), m.M_rownnz)
j = m.M_colind
if not is_sparse(m):
qm = a @ b.T
@@ -82,29 +76,23 @@ def make_m(
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)
diag_adr = m.M_rowadr + m.M_rownnz - 1
qm = qm.at[diag_adr].add(d)
return qm
def full_m(m: Model, d: Data) -> jax.Array:
"""Reconstitute dense mass matrix from qM."""
"""Reconstitute dense mass matrix from M."""
if not is_sparse(m):
return d._impl.qM # pytype: disable=attribute-error
return d._impl.M # pytype: disable=attribute-error
ij = []
for i in range(m.nv):
j = i
while j > -1:
ij.append((i, j))
j = m.dof_parentid[j]
i = np.repeat(np.arange(m.nv), m.M_rownnz)
j = m.M_colind
i, j = (jp.array(x) for x in zip(*ij))
mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d._impl.qM) # pytype: disable=attribute-error
mat = jp.zeros((m.nv, m.nv)).at[(i, j)].set(d._impl.M) # pytype: disable=attribute-error
# also set upper triangular
mat = mat + jp.tril(mat, -1).T
@@ -116,24 +104,23 @@ def mul_m(m: Model, d: Data, vec: jax.Array) -> jax.Array:
"""Multiply vector by inertia matrix."""
if not is_sparse(m):
return d._impl.qM @ vec # pytype: disable=attribute-error
return d._impl.M @ vec # pytype: disable=attribute-error
diag_mul = d._impl.qM[jp.array(m.dof_Madr)] * vec # pytype: disable=attribute-error
diag_adr = m.M_rowadr + m.M_rownnz - 1
diag_mul = d._impl.M[diag_adr] * vec # pytype: disable=attribute-error
is_, js, madr_ijs = [], [], []
for i in range(m.nv):
madr_ij, j = m.dof_Madr[i], i
adr = m.M_rowadr[i]
for k in range(m.M_rownnz[i] - 1):
is_.append(i)
js.append(m.M_colind[adr + k])
madr_ijs.append(adr + k)
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 = (np.array(x, dtype=np.int32) for x in (is_, js, madr_ijs))
i, j, madr_ij = (jp.array(x, dtype=jp.int32) for x in (is_, js, madr_ijs))
out = diag_mul.at[i].add(d._impl.qM[madr_ij] * vec[j]) # pytype: disable=attribute-error
out = out.at[j].add(d._impl.qM[madr_ij] * vec[i]) # pytype: disable=attribute-error
out = diag_mul.at[i].add(d._impl.M[madr_ij] * vec[j]) # pytype: disable=attribute-error
out = out.at[j].add(d._impl.M[madr_ij] * vec[i]) # pytype: disable=attribute-error
return out
mjx/mujoco/mjx/_src/types.py
+3 -1
View File
@@ -696,6 +696,9 @@ class Model(PyTreeNode):
dof_treeid: np.ndarray
dof_Madr: np.ndarray # pylint:disable=invalid-name
dof_simplenum: np.ndarray
M_rowadr: np.ndarray # pylint:disable=invalid-name
M_rownnz: np.ndarray # pylint:disable=invalid-name
M_colind: np.ndarray # pylint:disable=invalid-name
dof_solref: jax.Array
dof_solimp: jax.Array
dof_frictionloss: jax.Array
@@ -996,7 +999,6 @@ class DataJAX(PyTreeNode):
wrap_xpos: jax.Array
actuator_moment: jax.Array
crb: jax.Array
qM: jax.Array # pylint:disable=invalid-name
M: jax.Array # pylint:disable=invalid-name
qLD: jax.Array # pylint:disable=invalid-name
qLDiagInv: jax.Array # pylint:disable=invalid-name