From cedd602bc130327bf96d39e717125b2a51f3c3d7 Mon Sep 17 00:00:00 2001
From: Yuval Tassa <tassa@google.com>
Date: Wed, 20 Nov 2024 14:08:58 -0800
Subject: [PATCH] Compress dense rows in `mjData.actuator_moment`

PiperOrigin-RevId: 698516133
Change-Id: Idd861c2b63d77748a6282d72f8241eaeeb88006c
---
 src/engine/engine_core_smooth.c | 70 +++++++++++++++++++++------------
 src/engine/engine_print.c       |  2 +-
 2 files changed, 46 insertions(+), 26 deletions(-)

diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c
index d14015aa..6cc5a21f 100644
--- a/src/engine/engine_core_smooth.c
+++ b/src/engine/engine_core_smooth.c
@@ -879,7 +879,7 @@ void mj_transmission(const mjModel* m, mjData* d) {
   // compute lengths and moments
   for (int i=0; i < nu; i++) {
     rowadr[i] = i == 0 ? 0 : rowadr[i-1] + rownnz[i-1];
-    int adr = rowadr[i];
+    int nnz, adr = rowadr[i];
 
     // extract info
     int id = m->actuator_trnid[2*i];
@@ -1008,12 +1008,6 @@ void mj_transmission(const mjModel* m, mjData* d) {
         mj_jacSite(m, d, jac, 0, id);
         mju_subFrom(jac, jacS, 3*nv);
 
-        // sparsity
-        for (int j = 0; j < nv; j++) {
-          colind[adr+j] = j;
-        }
-        rownnz[i] = nv;
-
         // clear moment
         mju_zero(moment + adr, nv);
 
@@ -1029,6 +1023,17 @@ void mj_transmission(const mjModel* m, mjData* d) {
         for (int j = 0; j < nv; j++) {
           moment[adr+j] *= gear[0];
         }
+
+        // sparsity (compress)
+        nnz = 0;
+        for (int j = 0; j < nv; j++) {
+          if (moment[adr+j]) {
+            moment[adr+nnz] = moment[adr+j];
+            colind[adr+nnz] = j;
+            nnz++;
+          }
+        }
+        rownnz[i] = nnz;
       }
       break;
 
@@ -1045,23 +1050,22 @@ void mj_transmission(const mjModel* m, mjData* d) {
 
         mju_scl(moment + adr, d->ten_J + ten_J_rowadr, gear[0], ten_J_rownnz);
       } else {
-        // sparsity
-        for (int j = 0; j < nv; j++) {
-          colind[adr+j] = j;
-        }
-        rownnz[i] = nv;
-
         mju_scl(moment+adr, d->ten_J + id*nv, gear[0], nv);
+
+        // sparsity (compress)
+        nnz = 0;
+        for (int j = 0; j < nv; j++) {
+          if (moment[adr+j]) {
+            moment[adr+nnz] = moment[adr+j];
+            colind[adr+nnz] = j;
+            nnz++;
+          }
+        }
+        rownnz[i] = nnz;
       }
       break;
 
     case mjTRN_SITE:                    // site
-      // sparsity
-      for (int j = 0; j < nv; j++) {
-        colind[adr+j] = j;
-      }
-      rownnz[i] = nv;
-
       // get site translation (jac) and rotation (jacS) Jacobians in global frame
       mj_jacSite(m, d, jac, jacS, id);
 
@@ -1192,15 +1196,20 @@ void mj_transmission(const mjModel* m, mjData* d) {
         }
       }
 
+      // sparsity (compress)
+      nnz = 0;
+      for (int j = 0; j < nv; j++) {
+        if (moment[adr+j]) {
+          moment[adr+nnz] = moment[adr+j];
+          colind[adr+nnz] = j;
+          nnz++;
+        }
+      }
+      rownnz[i] = nnz;
+
       break;
 
     case mjTRN_BODY:                  // body (adhesive contacts)
-      // sparsity
-      for (int j = 0; j < nv; j++) {
-        colind[adr+j] = j;
-      }
-      rownnz[i] = nv;
-
       // cannot compute meaningful length, set to 0
       length[i] = 0;
 
@@ -1299,6 +1308,17 @@ void mj_transmission(const mjModel* m, mjData* d) {
         }
       }
 
+      // sparsity (compress)
+      nnz = 0;
+      for (int j = 0; j < nv; j++) {
+        if (moment[adr+j]) {
+          moment[adr+nnz] = moment[adr+j];
+          colind[adr+nnz] = j;
+          nnz++;
+        }
+      }
+      rownnz[i] = nnz;
+
       break;
 
     default:
diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c
index f7132bc4..68e48571 100644
--- a/src/engine/engine_print.c
+++ b/src/engine/engine_print.c
@@ -1084,7 +1084,7 @@ void mj_printFormattedData(const mjModel* m, mjData* d, const char* filename,
   }
 
   printArray("ACTUATOR_LENGTH", m->nu, 1, d->actuator_length, fp, float_format);
-  printSparsity("actuator_moments", m->nu, m->nv,
+  printSparsity("actuator_moment", m->nu, m->nv,
                 d->moment_rowadr, d->moment_rownnz, d->moment_colind, fp);
   printSparse("ACTUATOR_MOMENT", d->actuator_moment, m->nu, d->moment_rownnz,
               d->moment_rowadr, d->moment_colind, fp, float_format);