Compress dense rows in mjData.actuator_moment

PiperOrigin-RevId: 698516133 Change-Id: Idd861c2b63d77748a6282d72f8241eaeeb88006c
google-deepmind · Nov 20, 2024 · cedd602 · cedd602
1 parent 1767c11
commit cedd602
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 25 deletions.
diff --git 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
+        mju_scl(moment+adr, d->ten_J + id*nv, gear[0], nv);
+
+        // sparsity (compress)
+        nnz = 0;
         for (int j = 0; j < nv; j++) {
-          colind[adr+j] = j;
+          if (moment[adr+j]) {
+            moment[adr+nnz] = moment[adr+j];
+            colind[adr+nnz] = j;
+            nnz++;
+          }
         }
-        rownnz[i] = nv;
-
-        mju_scl(moment+adr, d->ten_J + id*nv, gear[0], nv);
+        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) {
         }
       }
 
-      break;
-
-    case mjTRN_BODY:                  // body (adhesive contacts)
-      // sparsity
+      // sparsity (compress)
+      nnz = 0;
       for (int j = 0; j < nv; j++) {
-        colind[adr+j] = j;
+        if (moment[adr+j]) {
+          moment[adr+nnz] = moment[adr+j];
+          colind[adr+nnz] = j;
+          nnz++;
+        }
       }
-      rownnz[i] = nv;
+      rownnz[i] = nnz;
 
+      break;
+
+    case mjTRN_BODY:                  // body (adhesive contacts)
       // 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
@@ -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);