Remove dense actuator_moment for _put_data_c.

PiperOrigin-RevId: 779234984
Change-Id: I19ac2849aa2b00b1688781491c9230558c24c7a2

doc/changelog.rst

@@ -25,6 +25,7 @@ Bug fixes
 ^^^^^^^^^
 - Inverse dynamics were not being computed correctly when :ref:`tendon armature<tendon-spatial-armature>` was present,
   now fixed.
+- Fix bug in ``mjx.put_data`` where ``actuator_moment`` was not being copied correctly for the C implementation.
 
 Documentation
 ^^^^^^^^^^^^^

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.
+  # TODO(stunya): remove reliance on JAX _put_contact.
-  # 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.
   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(
+          if d_i.impl == types.Impl.JAX:
-              actuator_moment,
+            mujoco.mju_dense2sparse(
-              d_i._impl.actuator_moment,
+                actuator_moment,
-              moment_rownnz,
+                d_i._impl.actuator_moment,
-              moment_rowadr,
+                moment_rownnz,
-              moment_colind,
+                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

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))