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

Use sparse (uncompressed) actuator_moment in mj_transmission.

Browse Source 
PiperOrigin-RevId: 692179704
Change-Id: Ic30ac5a98dc13de2028e378df65dc88ba3912bf5
This commit is contained in:
Taylor Howell
2024-11-01 08:05:14 -07:00
committed by Copybara-Service
parent 1d58576d28
commit a51f346059
13 changed files with 230 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mjx/mujoco/mjx/_src/io.py
+41 -6
View File
@@ -274,6 +274,9 @@ def make_data(
'wrap_obj': (m.nwrap, 2, jp.int32),
'wrap_xpos': (m.nwrap, 6, float),
'actuator_length': (m.nu, float),
'moment_rownnz': (m.nu, jp.int32),
'moment_rowadr': (m.nu, jp.int32),
'moment_colind': (m.nu, m.nv, jp.int32),
'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),
@@ -427,6 +430,25 @@ def get_data_into(
result_i.contact.efc_address[:] = efc_map[result_i.contact.efc_address]
continue
# MuJoCo actuator_moment is sparse, MJX uses a dense representation.
if field.name == 'actuator_moment' and m.nu:
moment_rownnz = np.zeros(m.nu, dtype=int)
moment_rowadr = np.zeros(m.nu, dtype=int)
moment_colind = np.zeros(m.nu * m.nv, dtype=int)
actuator_moment = np.zeros(m.nu * m.nv)
mujoco.mju_dense2sparse(
actuator_moment,
d.actuator_moment,
moment_rownnz,
moment_rowadr,
moment_colind,
)
result_i.moment_rownnz[:] = moment_rownnz
result_i.moment_rowadr[:] = moment_rowadr
result_i.moment_colind[:] = moment_colind.reshape((m.nu, m.nv))
result_i.actuator_moment[:] = actuator_moment.reshape((m.nu, m.nv))
continue
value = getattr(d_i, field.name)
if field.name in ('nefc', 'ncon'):
@@ -532,6 +554,17 @@ def put_data(
# MJX does not support islanding, so only transfer the first solver_niter
fields['solver_niter'] = fields['solver_niter'][0]
# convert sparse representation of actuator_moment to dense matrix
moment = np.zeros((m.nu, m.nv))
mujoco.mju_sparse2dense(
moment,
d.actuator_moment.reshape(-1),
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind.reshape(-1),
)
fields['actuator_moment'] = moment
contact, contact_map = _make_contact(d.contact, dim, efc_address)
# pad efc fields: MuJoCo efc arrays are sparse for inactive constraints.
@@ -539,12 +572,14 @@ def put_data(
# neither: it contains zeros for inactive constraints, and efc_J is always
# (nefc, nv). this may change in the future.
if mujoco.mj_isSparse(m):
nr = d.efc_J_rownnz.shape[0]
efc_j = np.zeros((nr, m.nv))
for i in range(nr):
rowadr = d.efc_J_rowadr[i]
for j in range(d.efc_J_rownnz[i]):
efc_j[i, d.efc_J_colind[rowadr + j]] = fields['efc_J'][rowadr + j]
efc_j = np.zeros((d.efc_J_rownnz.shape[0], m.nv))
mujoco.mju_sparse2dense(
efc_j,
fields['efc_J'],
d.efc_J_rownnz,
d.efc_J_rowadr,
d.efc_J_colind,
)
fields['efc_J'] = efc_j
else:
fields['efc_J'] = fields['efc_J'].reshape((-1 if m.nv else 0, m.nv))
mjx/mujoco/mjx/_src/smooth_test.py
+22 -2
View File
@@ -117,7 +117,17 @@ class SmoothTest(absltest.TestCase):
# transmission
dx = jax.jit(mjx.transmission)(mx, dx)
_assert_attr_eq(d, dx, 'actuator_length')
_assert_attr_eq(d, dx, 'actuator_moment')
# convert sparse actuator_moment to dense representation
moment = np.zeros((m.nu, m.nv))
mujoco.mju_sparse2dense(
moment,
d.actuator_moment.reshape(-1),
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind.reshape(-1),
)
_assert_eq(moment, dx.actuator_moment, 'actuator_moment')
def test_disable_gravity(self):
m = mujoco.MjModel.from_xml_string("""
@@ -178,7 +188,17 @@ class SmoothTest(absltest.TestCase):
mujoco.mj_transmission(m, d)
dx = jax.jit(mjx.transmission)(mx, dx)
_assert_attr_eq(d, dx, 'actuator_length')
_assert_attr_eq(d, dx, 'actuator_moment')
# convert sparse actuator_moment to dense representation
moment = np.zeros((m.nu, m.nv))
mujoco.mju_sparse2dense(
moment,
d.actuator_moment.reshape(-1),
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind.reshape(-1),
)
_assert_eq(moment, dx.actuator_moment, 'actuator_moment')
def test_subtree_vel(self):
"""Tests MJX subtree_vel function matches MuJoCo mj_subtreeVel."""
mjx/mujoco/mjx/_src/types.py
+6
View File
@@ -1228,6 +1228,9 @@ class Data(PyTreeNode):
wrap_obj: geom id; -1: site; -2: pulley (nwrap*2,)
wrap_xpos: Cartesian 3D points in all path (nwrap*2, 3)
actuator_length: actuator lengths (nu,)
moment_rownnz: number of non-zeros in actuator_moment row (nu,)
moment_rowadr: row start address in colind array (nu,)
moment_colind: column indices in sparse Jacobian (nu, nv)
actuator_moment: actuator moments (nu, nv)
crb: com-based composite inertia and mass (nbody, 10)
qM: total inertia if sparse: (nM,)
@@ -1350,6 +1353,9 @@ class Data(PyTreeNode):
wrap_obj: jax.Array
wrap_xpos: jax.Array
actuator_length: jax.Array
moment_rownnz: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
moment_rowadr: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
moment_colind: jax.Array = _restricted_to('mujoco') # pylint:disable=invalid-name
actuator_moment: jax.Array
crb: jax.Array
qM: jax.Array # pylint:disable=invalid-name
mjx/mujoco/mjx/integration_test/smooth_test.py
+21 -4
View File
@@ -68,10 +68,27 @@ class TransmissionIntegrationTest(parameterized.TestCase):
mujoco.mj_transmission(m, d)
dx = transmission_jit_fn(mx, dx)
for field in ['actuator_length', 'actuator_moment']:
_assert_attr_eq(
d, dx, field, seed, f'transmission{seed}', atol=1e-4
)
_assert_attr_eq(
d, dx, 'actuator_length', seed, f'transmission{seed}', atol=1e-4
)
# convert sparse actuator_moment to dense representation
moment = np.zeros((m.nu, m.nv))
mujoco.mju_sparse2dense(
moment,
d.actuator_moment.reshape(-1),
d.moment_rownnz,
d.moment_rowadr,
d.moment_colind.reshape(-1),
)
_assert_eq(
moment,
dx.actuator_moment,
'actuator_moment',
seed,
f'transmission{seed}',
atol=1e-4,
)
if __name__ == '__main__':