Improve CRB code:

- remove mj_crbSkip since this was only utilized during compile time to compute m->dof_m0, replaced with mj_crbDiag to compute only the diagonals for dof_m0,
- change row backward pass to forward pass to optimize processing diagonal blocks,
- and various code refactoring.

PiperOrigin-RevId: 556799924
Change-Id: I0a18d7a45862860fdddd7150540ddc8d64b01b29

src/engine/engine_core_smooth.c

@@ -957,16 +957,17 @@ void mj_transmission(const mjModel* m, mjData* d) {
 
 //-------------------------- inertia ---------------------------------------------------------------
 
-// composite rigid body inertia algorithm, with skipsimple
-void mj_crbSkip(const mjModel* m, mjData* d, int skipsimple) {
-  mjtNum tmp[6];
+// composite rigid body inertia algorithm
+void mj_crb(const mjModel* m, mjData* d) {
+  mjtNum buf[6];
   mjtNum* crb = d->crb;
+  int last_body = m->nbody - 1, nv = m->nv;
 
   // crb = cinert
   mju_copy(crb, d->cinert, 10*m->nbody);
 
   // backward pass over bodies, accumulate composite inertias
-  for (int i=m->nbody-1; i > 0; i--) {
+  for (int i=last_body; i > 0; i--) {
     if (m->body_parentid[i] > 0) {
       mju_addTo(crb+10*m->body_parentid[i], crb+10*i, 10);
     }
@@ -975,45 +976,38 @@ void mj_crbSkip(const mjModel* m, mjData* d, int skipsimple) {
   // clear qM
   mju_zero(d->qM, m->nM);
 
-  // dense backward pass over dofs
-  for (int i=m->nv-1; i >= 0; i--) {
-    // copy
-    if (skipsimple && m->dof_simplenum[i]) {
-      d->qM[m->dof_Madr[i]] = m->dof_M0[i];
-    }
+  // dense forward pass over dofs
+  for (int i=0; i < nv; i++) {
+    // process block of diagonals (simple bodies)
+    if (m->dof_simplenum[i]) {
+      int n = i + m->dof_simplenum[i];
+      for (; i < n; i++) {
+        d->qM[m->dof_Madr[i]] = m->dof_M0[i];
+      }
 
-    // compute
-    else {
-      // init M(i,i) with armature inertia
-      int Madr_ij = m->dof_Madr[i];
-      d->qM[Madr_ij] = m->dof_armature[i];
-
-      // precompute tmp = crb * cdof
-      mju_mulInertVec(tmp, crb+10*m->dof_bodyid[i], d->cdof+6*i);
-
-      // sparse backward pass over ancestors
-      int j = i;
-      while (j >= 0) {
-        // M(i,j) += cdof_j * crb_body(i) * cdof_i = cdof_j * tmp
-        d->qM[Madr_ij] += mju_dot(d->cdof+6*j, tmp, 6);
-
-        // advance to parent
-        j = m->dof_parentid[j];
-        Madr_ij++;
+      // finish or else fall through with next row
+      if (i == nv) {
+        break;
       }
     }
+
+    // init M(i,i) with armature inertia
+    int Madr_ij = m->dof_Madr[i];
+    d->qM[Madr_ij] = m->dof_armature[i];
+
+    // precompute buf = crb_body_i * cdof_i
+    mju_mulInertVec(buf, crb+10*m->dof_bodyid[i], d->cdof+6*i);
+
+    // sparse backward pass over ancestors
+    for (int j=i; j >= 0; j = m->dof_parentid[j]) {
+      // M(i,j) += cdof_j * (crb_body_i * cdof_i)
+      d->qM[Madr_ij++] += mju_dot(d->cdof+6*j, buf, 6);
+    }
   }
 }
 
 
 
-// composite rigid body inertia algorithm
-void mj_crb(const mjModel* m, mjData* d) {
-  mj_crbSkip(m, d, 1);
-}
-
-
-
 // sparse L'*D*L factorizaton of inertia-like matrix M, assumed spd
 void mj_factorI(const mjModel* m, mjData* d, const mjtNum* M, mjtNum* qLD, mjtNum* qLDiagInv,
                 mjtNum* qLDiagSqrtInv) {

src/engine/engine_core_smooth.h

@@ -42,9 +42,6 @@ MJAPI void mj_transmission(const mjModel* m, mjData* d);
 
 //-------------------------- inertia ---------------------------------------------------------------
 
-// composite rigid body inertia algorithm, with skip
-void mj_crbSkip(const mjModel* m, mjData* d, int skipsimple);
-
 // composite rigid body inertia algorithm
 MJAPI void mj_crb(const mjModel* m, mjData* d);
 

src/engine/engine_setconst.c

@@ -29,6 +29,34 @@
 #include "engine/engine_util_misc.h"
 #include "engine/engine_util_spatial.h"
 
+
+// compute dof_M0 via composite rigid body algorithm
+static void mj_setM0(mjModel* m, mjData* d) {
+  mjtNum buf[6];
+  mjtNum* crb = d->crb;
+  int last_body = m->nbody - 1, nv = m->nv;
+
+  // copy cinert into crb
+  mju_copy(crb, d->cinert, 10*m->nbody);
+
+  // backward pass over bodies, accumulate composite inertias
+  for (int i=last_body; i > 0; i--) {
+    if (m->body_parentid[i] > 0) {
+      mju_addTo(crb+10*m->body_parentid[i], crb+10*i, 10);
+    }
+  }
+
+  for (int i=0; i < nv; i++) {
+    // precomute buf = crb_body_i * cdof_i
+    mju_mulInertVec(buf, crb+10*m->dof_bodyid[i], d->cdof+6*i);
+
+    // dof_M0(i) = armature inertia + cdof_i * (crb_body_i * cdof_i)
+    m->dof_M0[i] = m->dof_armature[i] + mju_dot(d->cdof+6*i, buf, 6);
+  }
+}
+
+
+
 // set quantities that depend on qpos0
 static void set0(mjModel* m, mjData* d) {
   int id, id1, id2, dnum, nv = m->nv;
@@ -60,14 +88,12 @@ static void set0(mjModel* m, mjData* d) {
   mj_kinematics(m, d);
   mj_comPos(m, d);
   mj_camlight(m, d);
-  mj_crbSkip(m, d, 0);
 
-  // save dof_M0
-  for (int i=0; i < nv; i++) {
-    m->dof_M0[i] = d->qM[m->dof_Madr[i]];
-  }
+  // compute dof_M0 for CRB algorithm
+  mj_setM0(m, d);
 
-  // run remaining computations (factorM needs dof_M0)
+  // run remaining computations
+  mj_crb(m, d);
   mj_factorM(m, d);
   mj_tendon(m, d);
   mj_transmission(m, d);

src/user/user_model.cc

@@ -1667,19 +1667,14 @@ void mjCModel::CopyTree(mjModel* m) {
   m->nB = nB;
 
   // set dof_simplenum
-  int scnt = 0;
-  for (int i=nv-1; i>=0; i--) {
-    // dof in simple body
+  int count = 0;
+  for (int i=nv-1; i >= 0; i--) {
     if (m->body_simple[m->dof_bodyid[i]]) {
-      scnt++;
-      m->dof_simplenum[i] = scnt;
-    }
-
-    // dof in regular body
-    else {
-      scnt = 0;
-      m->dof_simplenum[i] = 0;
+      count++;    // increment counter
+    } else {
+      count = 0;  // reset
     }
+    m->dof_simplenum[i] = count;
   }
 }