From a1b18e707a10b9d2be1bda6b86eec684835d4490 Mon Sep 17 00:00:00 2001
From: Taylor Howell <taylorhowell@google.com>
Date: Wed, 20 Nov 2024 10:52:20 -0800
Subject: [PATCH] Compress `actuator_moment` memory using `nJmom` <= `nu` x
 `nv`.

PiperOrigin-RevId: 698446692
Change-Id: I49c9633e12129a1e690724db82d1f11204e41d9c
---
 src/engine/engine_core_smooth.c |  2 +-
 src/user/user_model.cc          | 67 ++++++++++++++++++++++++++++++---
 src/user/user_model.h           |  1 +
 test/user/user_model_test.cc    |  2 +-
 4 files changed, 64 insertions(+), 8 deletions(-)

diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
index e9bf33cc..6e9e7b2e 100644
--- a/src/engine/engine_core_smooth.c
+++ b/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;
diff --git a/src/user/user_model.cc b/src/user/user_model.cc
index 927a8b0b..fb21b1b1 100644
--- a/src/user/user_model.cc
+++ b/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) {
diff --git a/src/user/user_model.h b/src/user/user_model.h
index d4d7dffa..1b248ea2 100644
--- a/src/user/user_model.h
+++ b/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
diff --git a/test/user/user_model_test.cc b/test/user/user_model_test.cc
index 53c6277b..811004a7 100644
--- a/test/user/user_model_test.cc
+++ b/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);
 }