diff --git a/mjx/mujoco/mjx/_src/forward.py b/mjx/mujoco/mjx/_src/forward.py index 8e0de506..d2b12891 100644 --- a/mjx/mujoco/mjx/_src/forward.py +++ b/mjx/mujoco/mjx/_src/forward.py @@ -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 diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index 03c7af65..b896ccc7 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -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. diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index 1dec6057..ff014402 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -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) diff --git a/mjx/mujoco/mjx/_src/smooth.py b/mjx/mujoco/mjx/_src/smooth.py index 55437fbf..1e8b4df6 100644 --- a/mjx/mujoco/mjx/_src/smooth.py +++ b/mjx/mujoco/mjx/_src/smooth.py @@ -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: diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py index b421b68f..b6814745 100644 --- a/mjx/mujoco/mjx/_src/smooth_test.py +++ b/mjx/mujoco/mjx/_src/smooth_test.py @@ -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) diff --git a/mjx/mujoco/mjx/_src/support.py b/mjx/mujoco/mjx/_src/support.py index ee2e0e40..914a0ba9 100644 --- a/mjx/mujoco/mjx/_src/support.py +++ b/mjx/mujoco/mjx/_src/support.py @@ -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 diff --git a/mjx/mujoco/mjx/_src/types.py b/mjx/mujoco/mjx/_src/types.py index 0d9a29ed..73cf2573 100644 --- a/mjx/mujoco/mjx/_src/types.py +++ b/mjx/mujoco/mjx/_src/types.py @@ -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