Add nJmom, the number of non-zeros in sparse actuator_moment matr…

…ix, to `mjModel`. PiperOrigin-RevId: 698343986 Change-Id: I3d4a4dea1ec095d4e7b231c2fffc0e6241cee39f
google-deepmind · Nov 20, 2024 · aae5fd6 · aae5fd6
1 parent cd84011
commit aae5fd6
Show file tree

Hide file tree

Showing 13 changed files with 78 additions and 33 deletions.
diff --git a/doc/includes/references.h b/doc/includes/references.h
@@ -263,8 +263,8 @@ struct mjData_ {
   mjtNum* actuator_length;   // actuator lengths                                 (nu x 1)
   int*    moment_rownnz;     // number of non-zeros in actuator_moment row       (nu x 1)
   int*    moment_rowadr;     // row start address in colind array                (nu x 1)
-  int*    moment_colind;     // column indices in sparse Jacobian                (nu x nv)
-  mjtNum* actuator_moment;   // actuator moments                                 (nu x nv)
+  int*    moment_colind;     // column indices in sparse Jacobian                (nJmom x 1)
+  mjtNum* actuator_moment;   // actuator moments                                 (nJmom x 1)
 
   // computed by mj_fwdPosition/mj_crb
   mjtNum* crb;               // com-based composite inertia and mass             (nbody x 10)
@@ -956,6 +956,7 @@ struct mjModel_ {
   int nB;                         // number of non-zeros in sparse body-dof matrix
   int nC;                         // number of non-zeros in sparse reduced dof-dof matrix
   int nD;                         // number of non-zeros in sparse dof-dof matrix
+  int nJmom;                      // number of non-zeros in sparse actuator_moment matrix
   int ntree;                      // number of kinematic trees under world body
   int ngravcomp;                  // number of bodies with nonzero gravcomp
   int nemax;                      // number of potential equality-constraint rows

diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h
@@ -291,8 +291,8 @@ struct mjData_ {
   mjtNum* actuator_length;   // actuator lengths                                 (nu x 1)
   int*    moment_rownnz;     // number of non-zeros in actuator_moment row       (nu x 1)
   int*    moment_rowadr;     // row start address in colind array                (nu x 1)
-  int*    moment_colind;     // column indices in sparse Jacobian                (nu x nv)
-  mjtNum* actuator_moment;   // actuator moments                                 (nu x nv)
+  int*    moment_colind;     // column indices in sparse Jacobian                (nJmom x 1)
+  mjtNum* actuator_moment;   // actuator moments                                 (nJmom x 1)
 
   // computed by mj_fwdPosition/mj_crb
   mjtNum* crb;               // com-based composite inertia and mass             (nbody x 10)

diff --git a/include/mujoco/mjmodel.h b/include/mujoco/mjmodel.h
@@ -663,6 +663,7 @@ struct mjModel_ {
   int nB;                         // number of non-zeros in sparse body-dof matrix
   int nC;                         // number of non-zeros in sparse reduced dof-dof matrix
   int nD;                         // number of non-zeros in sparse dof-dof matrix
+  int nJmom;                      // number of non-zeros in sparse actuator_moment matrix
   int ntree;                      // number of kinematic trees under world body
   int ngravcomp;                  // number of bodies with nonzero gravcomp
   int nemax;                      // number of potential equality-constraint rows

diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h
@@ -137,6 +137,7 @@
     X   ( nB )                 \
     X   ( nC )                 \
     X   ( nD )                 \
+    X   ( nJmom )              \
     XMJV( ntree )              \
     X   ( ngravcomp )          \
     X   ( nemax )              \
@@ -625,8 +626,8 @@
     X   ( mjtNum,    actuator_length,   nu,          1           ) \
     X   ( int,       moment_rownnz,     nu,          1           ) \
     X   ( int,       moment_rowadr,     nu,          1           ) \
-    X   ( int,       moment_colind,     nu,          MJ_M(nv)    ) \
-    X   ( mjtNum,    actuator_moment,   nu,          MJ_M(nv)    ) \
+    X   ( int,       moment_colind,     nJmom,       1           ) \
+    X   ( mjtNum,    actuator_moment,   nJmom,       1           ) \
     X   ( mjtNum,    crb,               nbody,       10          ) \
     X   ( mjtNum,    qM,                nM,          1           ) \
     X   ( mjtNum,    qLD,               nM,          1           ) \

diff --git a/introspect/structs.py b/introspect/structs.py
@@ -1178,6 +1178,11 @@
                  type=ValueType(name='int'),
                  doc='number of non-zeros in sparse dof-dof matrix',
              ),
+             StructFieldDecl(
+                 name='nJmom',
+                 type=ValueType(name='int'),
+                 doc='number of non-zeros in sparse actuator_moment matrix',
+             ),
              StructFieldDecl(
                  name='ntree',
                  type=ValueType(name='int'),
@@ -5166,15 +5171,15 @@
                      inner_type=ValueType(name='int'),
                  ),
                  doc='column indices in sparse Jacobian',
-                 array_extent=('nu', 'nv'),
+                 array_extent=('nJmom',),
              ),
              StructFieldDecl(
                  name='actuator_moment',
                  type=PointerType(
                      inner_type=ValueType(name='mjtNum'),
                  ),
                  doc='actuator moments',
-                 array_extent=('nu', 'nv'),
+                 array_extent=('nJmom',),
              ),
              StructFieldDecl(
                  name='crb',

diff --git a/mjx/mujoco/mjx/_src/io.py b/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
 

diff --git a/mjx/mujoco/mjx/_src/smooth_test.py b/mjx/mujoco/mjx/_src/smooth_test.py
@@ -122,10 +122,10 @@ def test_smooth(self):
     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 @@ def test_site_transmission(self):
     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')
 

diff --git a/mjx/mujoco/mjx/_src/types.py b/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)

diff --git a/mjx/mujoco/mjx/integration_test/smooth_test.py b/mjx/mujoco/mjx/integration_test/smooth_test.py
@@ -76,10 +76,10 @@ def test_transmission(self, seed):
     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,

diff --git a/src/user/user_model.cc b/src/user/user_model.cc
@@ -777,6 +777,7 @@ void mjCModel::Clear() {
   nM = 0;
   nD = 0;
   nB = 0;
+  nJmom = 0;
   njmax = -1;
   nconmax = -1;
   nmocap = 0;
@@ -2490,7 +2491,7 @@ void mjCModel::CopyTree(mjModel* m) {
       }
     }
   }
-   m->nC = nC = 2 * nOD + nv;
+  m->nC = nC = 2 * nOD + nv;
 }
 
 // copy plugin data
@@ -4156,6 +4157,13 @@ void mjCModel::TryCompile(mjModel*& m, mjData*& d, const mjVFS* vfs) {
   // copy objects outsite kinematic tree (including keyframes)
   CopyObjects(m);
 
+  // compute nJmom
+  for (int i = 0; i < nu; i++) {
+    // dense rows
+    nJmom += nv;
+  }
+  m->nJmom = nJmom;
+
   // scale mass
   if (compiler.settotalmass>0) {
     mj_setTotalmass(m, compiler.settotalmass);
@@ -4292,7 +4300,7 @@ bool mjCModel::CopyBack(const mjModel* m) {
       neq!=m->neq || ntendon!=m->ntendon || nwrap!=m->nwrap || nsensor!=m->nsensor ||
       nnumeric!=m->nnumeric || nnumericdata!=m->nnumericdata || ntext!=m->ntext ||
       ntextdata!=m->ntextdata || nnames!=m->nnames || nM!=m->nM || nD!=m->nD || nC!=m->nC ||
-      nB!=m->nB || nemax!=m->nemax || nconmax!=m->nconmax || njmax!=m->njmax ||
+      nB!=m->nB || nJmom!=m->nJmom ||nemax!=m->nemax || nconmax!=m->nconmax || njmax!=m->njmax ||
       npaths!=m->npaths) {
     errInfo = mjCError(0, "incompatible models in CopyBack");
     return false;

diff --git a/src/user/user_model.h b/src/user/user_model.h
@@ -118,6 +118,7 @@ class mjCModel_ : public mjsElement {
   int nB;              // number of non-zeros in sparse body-dof matrix
   int nC;              // number of non-zeros in reduced sparse dof-dof matrix
   int nD;              // number of non-zeros in sparse dof-dof matrix
+  int nJmom;           // number of non-zeros in sparse actuator_moment matrix
 
   // statistics, as computed by mj_setConst
   double meaninertia_auto;  // mean diagonal inertia, as computed by mj_setConst

diff --git a/test/user/user_model_test.cc b/test/user/user_model_test.cc
@@ -116,6 +116,30 @@ TEST_F(UserCModelTest, SameFrame) {
   mj_deleteModel(model);
 }
 
+TEST_F(UserCModelTest, ActuatorSparsity) {
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <worldbody>
+      <body>
+        <geom size="1"/>
+        <joint name="a"/>
+        <body>
+          <geom size="1"/>
+          <joint name="b"/>
+        </body>
+      </body>
+    </worldbody>
+    <actuator>
+      <motor joint="a"/>
+      <motor joint="b"/>
+    </actuator>
+  </mujoco>
+  )";
+  mjModel* m = LoadModelFromString(xml);
+  ASSERT_EQ(m->nJmom, 4);
+  mj_deleteModel(m);
+}
+
 
 // ------------- test automatic inference of nuser_xxx -------------------------
 

diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs
@@ -5244,6 +5244,7 @@ public unsafe struct mjModel_ {
   public int nB;
   public int nC;
   public int nD;
+  public int nJmom;
   public int ntree;
   public int ngravcomp;
   public int nemax;