Add nJmom, the number of non-zeros in sparse actuator_moment matrix, to mjModel.

PiperOrigin-RevId: 698343986
Change-Id: I3d4a4dea1ec095d4e7b231c2fffc0e6241cee39f

mjx/mujoco/mjx/_src/io.py

@@ -276,7 +276,7 @@ def make_data(
         '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),
+        'moment_colind': (m.nJmom, 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),
@@ -431,22 +431,23 @@ def get_data_into(
         continue
 
       # MuJoCo actuator_moment is sparse, MJX uses a dense representation.
-      if field.name == 'actuator_moment' and m.nu:
+      if field.name == 'actuator_moment':
         moment_rownnz = np.zeros(m.nu, dtype=np.int32)
         moment_rowadr = np.zeros(m.nu, dtype=np.int32)
-        moment_colind = np.zeros(m.nu * m.nv, dtype=np.int32)
-        actuator_moment = np.zeros(m.nu * m.nv)
-        mujoco.mju_dense2sparse(
-            actuator_moment,
-            d.actuator_moment,
-            moment_rownnz,
-            moment_rowadr,
-            moment_colind,
-        )
+        moment_colind = np.zeros(m.nJmom, dtype=np.int32)
+        actuator_moment = np.zeros(m.nJmom)
+        if m.nu:
+          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))
+        result_i.moment_colind[:] = moment_colind
+        result_i.actuator_moment[:] = actuator_moment
         continue
 
       value = getattr(d_i, field.name)
@@ -558,10 +559,10 @@ def put_data(
   moment = np.zeros((m.nu, m.nv))
   mujoco.mju_sparse2dense(
       moment,
-      d.actuator_moment.reshape(-1),
+      d.actuator_moment,
       d.moment_rownnz,
       d.moment_rowadr,
-      d.moment_colind.reshape(-1),
+      d.moment_colind,
   )
   fields['actuator_moment'] = moment
 

mjx/mujoco/mjx/_src/smooth_test.py

@@ -122,10 +122,10 @@ class SmoothTest(absltest.TestCase):
     moment = np.zeros((m.nu, m.nv))
     mujoco.mju_sparse2dense(
         moment,
-        d.actuator_moment.reshape(-1),
+        d.actuator_moment,
         d.moment_rownnz,
         d.moment_rowadr,
-        d.moment_colind.reshape(-1),
+        d.moment_colind,
     )
     _assert_eq(moment, dx.actuator_moment, 'actuator_moment')
 
@@ -193,10 +193,10 @@ class SmoothTest(absltest.TestCase):
     moment = np.zeros((m.nu, m.nv))
     mujoco.mju_sparse2dense(
         moment,
-        d.actuator_moment.reshape(-1),
+        d.actuator_moment,
         d.moment_rownnz,
         d.moment_rowadr,
-        d.moment_colind.reshape(-1),
+        d.moment_colind,
     )
     _assert_eq(moment, dx.actuator_moment, 'actuator_moment')
 

mjx/mujoco/mjx/_src/types.py

@@ -536,6 +536,7 @@ class Model(PyTreeNode):
     nM: number of non-zeros in sparse inertia matrix
     nD: number of non-zeros in sparse dof-dof matrix
     nB: number of non-zeros in sparse body-dof matrix
+    nJmom: number of non-zeros in sparse actuator_moment matrix
     ntree: number of kinematic trees under world body
     ngravcomp: number of bodies with nonzero gravcomp
     nuserdata: size of userdata array
@@ -855,6 +856,7 @@ class Model(PyTreeNode):
   nM: int  # pylint:disable=invalid-name
   nD: int  # pylint:disable=invalid-name
   nB: int  # pylint:disable=invalid-name
+  nJmom: int
   ntree: int = _restricted_to('mujoco')
   ngravcomp: int
   nuserdata: int
@@ -1236,8 +1238,8 @@ class Data(PyTreeNode):
     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)
+    moment_colind: column indices in sparse Jacobian            (nJmom,)
+    actuator_moment: actuator moments                           (nJmom,)
     crb: com-based composite inertia and mass                   (nbody, 10)
     qM: total inertia                                if sparse: (nM,)
                                                      if dense:  (nv, nv)

mjx/mujoco/mjx/integration_test/smooth_test.py

@@ -76,10 +76,10 @@ class TransmissionIntegrationTest(parameterized.TestCase):
     moment = np.zeros((m.nu, m.nv))
     mujoco.mju_sparse2dense(
         moment,
-        d.actuator_moment.reshape(-1),
+        d.actuator_moment,
         d.moment_rownnz,
         d.moment_rowadr,
-        d.moment_colind.reshape(-1),
+        d.moment_colind,
     )
     _assert_eq(
         moment,