diff --git a/doc/changelog.rst b/doc/changelog.rst index 401cf60e..7e2001fb 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -25,6 +25,7 @@ Bug fixes ^^^^^^^^^ - Inverse dynamics were not being computed correctly when :ref:`tendon armature` was present, now fixed. +- Fix bug in ``mjx.put_data`` where ``actuator_moment`` was not being copied correctly for the C implementation. Documentation ^^^^^^^^^^^^^ diff --git a/mjx/mujoco/mjx/_src/io.py b/mjx/mujoco/mjx/_src/io.py index ae747e85..e980afe5 100644 --- a/mjx/mujoco/mjx/_src/io.py +++ b/mjx/mujoco/mjx/_src/io.py @@ -600,7 +600,7 @@ def _make_data_c( 'moment_rownnz': (m.nu, np.int32), 'moment_rowadr': (m.nu, np.int32), 'moment_colind': (m.nJmom, np.int32), - 'actuator_moment': (m.nu, m.nv, float_), + 'actuator_moment': (m.nJmom, float_), 'bvh_aabb_dyn': (nbvhdynamic, 6, float_), 'bvh_active': (nbvh, np.uint8), 'flexedge_velocity': (nflexedge, float_), @@ -940,22 +940,10 @@ def _put_data_c( # TODO(stunya): support islanding via C impl. impl_fields['solver_niter'] = impl_fields['solver_niter'][0] - # TODO(btaba): remove dense actuator moment. - # convert sparse representation of actuator_moment to dense matrix - moment = np.zeros((m.nu, m.nv)) - mujoco.mju_sparse2dense( - moment, - d.actuator_moment, - d.moment_rownnz, - d.moment_rowadr, - d.moment_colind, - ) - impl_fields['actuator_moment'] = moment - - # TODO(btaba): remove reliance on JAX _put_contact. + # TODO(stunya): remove reliance on JAX _put_contact. contact, contact_map = _put_contact(d.contact, dim, efc_address) - # TODO(btaba): remove reliance on dense efc_J. + # TODO(stunya): remove reliance on dense efc_J. if mujoco.mj_isSparse(m): efc_j = np.zeros((d.efc_J_rownnz.shape[0], m.nv)) mujoco.mju_sparse2dense( @@ -1130,13 +1118,16 @@ def _get_data_into( moment_colind = np.zeros(m.nJmom, dtype=np.int32) actuator_moment = np.zeros(m.nJmom) if m.nu: - mujoco.mju_dense2sparse( - actuator_moment, - d_i._impl.actuator_moment, - moment_rownnz, - moment_rowadr, - moment_colind, - ) + if d_i.impl == types.Impl.JAX: + mujoco.mju_dense2sparse( + actuator_moment, + d_i._impl.actuator_moment, + moment_rownnz, + moment_rowadr, + moment_colind, + ) + else: + actuator_moment = d_i._impl.actuator_moment result_i.moment_rownnz[:] = moment_rownnz result_i.moment_rowadr[:] = moment_rowadr result_i.moment_colind[:] = moment_colind diff --git a/mjx/mujoco/mjx/_src/io_test.py b/mjx/mujoco/mjx/_src/io_test.py index f8d56d09..706d7e5e 100644 --- a/mjx/mujoco/mjx/_src/io_test.py +++ b/mjx/mujoco/mjx/_src/io_test.py @@ -289,7 +289,10 @@ class DataIOTest(parameterized.TestCase): self.assertEqual(d._impl.cinert.shape, (nbody, 10)) self.assertEqual(d._impl.crb.shape, (nbody, 10)) self.assertEqual(d._impl.actuator_length.shape, (1,)) - self.assertEqual(d._impl.actuator_moment.shape, (1, nv)) + if impl == 'jax': + self.assertEqual(d._impl.actuator_moment.shape, (1, nv)) + elif impl == 'c': + self.assertEqual(d._impl.actuator_moment.shape, (m.nJmom,)) self.assertEqual(d._impl.contact.dist.shape, (ncon,)) self.assertEqual(d._impl.contact.pos.shape, (ncon, 3)) self.assertEqual(d._impl.contact.frame.shape, (ncon, 3, 3))