Compress actuator_moment memory using nJmom <= nu x nv.

PiperOrigin-RevId: 698446692
Change-Id: I49c9633e12129a1e690724db82d1f11204e41d9c

src/engine/engine_core_smooth.c

@@ -1087,7 +1087,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
         int refid = m->actuator_trnid[2*i+1];
         if (!jacref) jacref = mj_stackAllocNum(d, 3*nv);
 
-        // intialize last dof address for each body
+        // initialize last dof address for each body
         int b0 = m->body_weldid[m->site_bodyid[id]];
         int b1 = m->body_weldid[m->site_bodyid[refid]];
         int dofadr0 = m->body_dofadr[b0] + m->body_dofnum[b0] - 1;

src/user/user_model.cc

@@ -2579,6 +2579,65 @@ void mjCModel::CopyPlugins(mjModel* m) {
   }
 }
 
+
+
+// compute non-zeros in actuator_moment matrix
+int mjCModel::CountNJmom(const mjModel* m) {
+  int nu = m->nu;
+  int nv = m->nv;
+
+  int count = 0;
+  for (int i = 0; i < nu; i++) {
+    // extract info
+    int id = m->actuator_trnid[2 * i];
+
+    // process according to transmission type
+    switch ((mjtTrn)m->actuator_trntype[i]) {
+      case mjTRN_JOINT:
+      case mjTRN_JOINTINPARENT:
+        switch ((mjtJoint)m->jnt_type[id]) {
+          case mjJNT_SLIDE:
+          case mjJNT_HINGE:
+            count += 1;
+            break;
+
+          case mjJNT_BALL:
+            count += 3;
+            break;
+
+          case mjJNT_FREE:
+            count += 6;
+            break;
+        }
+        break;
+      // TODO(taylorhowell): improve upper bounds
+      case mjTRN_SLIDERCRANK:
+        count += nv;
+        break;
+
+      case mjTRN_TENDON:
+        count += nv;
+        break;
+
+      case mjTRN_SITE:
+        count += nv;
+        break;
+
+      case mjTRN_BODY:
+        count += nv;
+        break;
+
+      default:
+        // SHOULD NOT OCCUR
+        throw mjCError(0, "unknown transmission type");
+        break;
+    }
+  }
+  return count;
+}
+
+
+
 // copy objects outside kinematic tree
 void mjCModel::CopyObjects(mjModel* m) {
   int adr, bone_adr, vert_adr, normal_adr, face_adr, texcoord_adr;
@@ -4157,12 +4216,8 @@ 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;
+  // compute non-zeros in actuator_moment
+  m->nJmom = nJmom = CountNJmom(m);
 
   // scale mass
   if (compiler.settotalmass>0) {

src/user/user_model.h

@@ -324,6 +324,7 @@ class mjCModel : public mjCModel_, private mjSpec {
   void CopyObjects(mjModel*);           // copy objects outside kinematic tree
   void CopyTree(mjModel*);              // copy objects inside kinematic tree
   void CopyPlugins(mjModel*);           // copy plugin data
+  int CountNJmom(const mjModel* m);     // compute number of non-zeros in actuator_moment matrix
 
   // objects created here
   std::vector<mjCFlex*>     flexes_;      // list of flexes

test/user/user_model_test.cc

@@ -136,7 +136,7 @@ TEST_F(UserCModelTest, ActuatorSparsity) {
   </mujoco>
   )";
   mjModel* m = LoadModelFromString(xml);
-  ASSERT_EQ(m->nJmom, 4);
+  ASSERT_EQ(m->nJmom, 2);
   mj_deleteModel(m);
 }