Accelerate inertia computation for simple bodies.

PiperOrigin-RevId: 571411319
Change-Id: I82732e95795cc0f0645f8273a1b084f2a9187a75

doc/includes/references.h

@@ -909,7 +909,7 @@ struct mjModel_ {
   int*      body_treeid;          // id of body's kinematic tree; -1: static  (nbody x 1)
   int*      body_geomnum;         // number of geoms                          (nbody x 1)
   int*      body_geomadr;         // start addr of geoms; -1: no geoms        (nbody x 1)
-  mjtByte*  body_simple;          // body is simple (has diagonal M)          (nbody x 1)
+  mjtByte*  body_simple;          // 1: diagonal M; 2: diag M, no rotations   (nbody x 1)
   mjtByte*  body_sameframe;       // inertial frame is same as body frame     (nbody x 1)
   mjtNum*   body_pos;             // position offset rel. to parent body      (nbody x 3)
   mjtNum*   body_quat;            // orientation offset rel. to parent body   (nbody x 4)

include/mujoco/mjmodel.h

@@ -638,7 +638,7 @@ struct mjModel_ {
   int*      body_treeid;          // id of body's kinematic tree; -1: static  (nbody x 1)
   int*      body_geomnum;         // number of geoms                          (nbody x 1)
   int*      body_geomadr;         // start addr of geoms; -1: no geoms        (nbody x 1)
-  mjtByte*  body_simple;          // body is simple (has diagonal M)          (nbody x 1)
+  mjtByte*  body_simple;          // 1: diagonal M; 2: diag M, no rotations   (nbody x 1)
   mjtByte*  body_sameframe;       // inertial frame is same as body frame     (nbody x 1)
   mjtNum*   body_pos;             // position offset rel. to parent body      (nbody x 3)
   mjtNum*   body_quat;            // orientation offset rel. to parent body   (nbody x 4)

introspect/structs.py

@@ -1253,7 +1253,7 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=PointerType(
                      inner_type=ValueType(name='mjtByte'),
                  ),
-                 doc='body is simple (has diagonal M)          (nbody x 1)',
+                 doc='1: diagonal M; 2: diag M, no rotations   (nbody x 1)',
              ),
              StructFieldDecl(
                  name='body_sameframe',

src/engine/engine_setconst.c

@@ -59,7 +59,7 @@ static void mj_setM0(mjModel* m, mjData* d) {
 
 // set quantities that depend on qpos0
 static void set0(mjModel* m, mjData* d) {
-  int id, id1, id2, dnum, nv = m->nv;
+  int nv = m->nv;
   mjtNum A[36] = {0}, pos[3], quat[4];
   mj_markStack(d);
   mjtNum* jac = mj_stackAllocNum(d, 6*nv);
@@ -112,53 +112,85 @@ static void set0(mjModel* m, mjData* d) {
 
   // compute body_invweight0
   m->body_invweight0[0] = m->body_invweight0[1] = 0.0;
-  for (int i=1; i < m->nbody; i++) {
+  for (int i=1; i<m->nbody; i++) {
-    if (nv) {
+    // static bodies: zero invweight0
-      // inverse spatial inertia:  A = J*inv(M)*J'
+    if (m->body_weldid[i] == 0) {
-      mj_jacBodyCom(m, d, jac, jac+3*nv, i);
+      m->body_invweight0[2*i] = m->body_invweight0[2*i+1] = 0;
-      mj_solveM(m, d, tmp, jac, 6);
-      mju_mulMatMatT(A, jac, tmp, 6, nv, 6);
     }
 
-    // average diagonal and assign
+    // accelerate simple bodies with no rotations
-    m->body_invweight0[2*i] = (A[0] + A[7] + A[14])/3;
+    else if (m->body_simple[i]==2) {
-    m->body_invweight0[2*i+1] = (A[21] + A[28] + A[35])/3;
+      mjtNum mass = m->body_mass[i];
+      if (!mass) {  // SHOULD NOT OCCUR
+        mjERROR("moving body %d has 0 mass", i);
+      }
+      m->body_invweight0[2*i+0] = 1/mju_max(mjMINVAL, mass);
+      m->body_invweight0[2*i+1] = 0;
+    }
+
+    // general body: full inertia
+    else {
+      if (nv) {
+        // inverse spatial inertia: A = J*inv(M)*J'
+        mj_jacBodyCom(m, d, jac, jac+3*nv, i);
+        mj_solveM(m, d, tmp, jac, 6);
+        mju_mulMatMatT(A, jac, tmp, 6, nv, 6);
+      }
+
+      // average diagonal and assign
+      m->body_invweight0[2*i] = (A[0] + A[7] + A[14])/3;
+      m->body_invweight0[2*i+1] = (A[21] + A[28] + A[35])/3;
+    }
   }
 
   // compute dof_invweight0
-  for (int i=0; i < m->njnt; i++) {
+  for (int i=0; i<m->njnt; i++) {
-    id = m->jnt_dofadr[i];
+    // simple body with no rotations: no off-diagonal inertia
-
+    if (m->body_simple[m->jnt_bodyid[i]] == 2) {
-    // get number of components
+      int id = m->jnt_dofadr[i];
-    if (m->jnt_type[i] == mjJNT_FREE) {
+      int bi = m->jnt_bodyid[i];
-      dnum = 6;
+      mjtNum mass = m->body_mass[bi];
-    } else if (m->jnt_type[i] == mjJNT_BALL) {
+      if (!mass) {  // SHOULD NOT OCCUR
-      dnum = 3;
+        mjERROR("moving body %d has 0 mass", bi);
-    } else {
-      dnum = 1;
-    }
-
-    // inverse joint inertia:  A = J*inv(M)*J'
-    if (nv) {
-      mju_zero(jac, dnum*nv);
-      for (int j=0; j < dnum; j++) {
-        jac[j*(nv+1) + id] = 1;
       }
-      mj_solveM(m, d, tmp, jac, dnum);
+      m->dof_invweight0[id] = 1/mju_max(mjMINVAL, mass);
-      mju_mulMatMatT(A, jac, tmp, dnum, nv, dnum);
     }
 
-    // average diagonal and assign
+    // general joint: full inertia
-    if (dnum == 6) {
-      m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] =
-        (A[0] + A[7] + A[14])/3;
-      m->dof_invweight0[id+3] = m->dof_invweight0[id+4] = m->dof_invweight0[id+5] =
-        (A[21] + A[28] + A[35])/3;
-    } else if (dnum == 3)
-      m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] =
-        (A[0] + A[4] + A[8])/3;
     else {
-      m->dof_invweight0[id] = A[0];
+      int dnum, id = m->jnt_dofadr[i];
+
+      // get number of components
+      if (m->jnt_type[i] == mjJNT_FREE) {
+        dnum = 6;
+      } else if (m->jnt_type[i] == mjJNT_BALL) {
+        dnum = 3;
+      } else {
+        dnum = 1;
+      }
+
+      // inverse joint inertia:  A = J*inv(M)*J'
+      if (nv) {
+        mju_zero(jac, dnum*nv);
+        for (int j=0; j < dnum; j++) {
+          jac[j*(nv+1) + id] = 1;
+        }
+        mj_solveM(m, d, tmp, jac, dnum);
+        mju_mulMatMatT(A, jac, tmp, dnum, nv, dnum);
+      }
+
+      // average diagonal and assign
+      if (dnum == 6) {
+        m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] =
+          (A[0] + A[7] + A[14])/3;
+        m->dof_invweight0[id+3] = m->dof_invweight0[id+4] = m->dof_invweight0[id+5] =
+          (A[21] + A[28] + A[35])/3;
+      } else if (dnum == 3) {
+        m->dof_invweight0[id] = m->dof_invweight0[id+1] = m->dof_invweight0[id+2] =
+          (A[0] + A[4] + A[8])/3;
+      } else {
+        m->dof_invweight0[id] = A[0];
+      }
     }
   }
 
@@ -198,8 +230,8 @@ static void set0(mjModel* m, mjData* d) {
   // compute missing eq_data for body constraints
   for (int i=0; i < m->neq; i++) {
     // get ids
-    id1 = m->eq_obj1id[i];
+    int id1 = m->eq_obj1id[i];
-    id2 = m->eq_obj2id[i];
+    int id2 = m->eq_obj2id[i];
 
     // connect constraint
     if (m->eq_type[i] == mjEQ_CONNECT) {
@@ -239,8 +271,8 @@ static void set0(mjModel* m, mjData* d) {
   // camera compos0, pos0, mat0
   for (int i=0; i < m->ncam; i++) {
     // get body ids
-    id = m->cam_bodyid[i];              // camera body
+    int id = m->cam_bodyid[i];              // camera body
-    id1 = m->cam_targetbodyid[i];       // target body
+    int id1 = m->cam_targetbodyid[i];       // target body
 
     // compute positional offsets
     mju_sub3(m->cam_pos0+3*i, d->cam_xpos+3*i, d->xpos+3*id);
@@ -253,8 +285,8 @@ static void set0(mjModel* m, mjData* d) {
   // light compos0, pos0, dir0
   for (int i=0; i < m->nlight; i++) {
     // get body ids
-    id = m->light_bodyid[i];            // light body
+    int id = m->light_bodyid[i];            // light body
-    id1 = m->light_targetbodyid[i];     // target body
+    int id1 = m->light_targetbodyid[i];     // target body
 
     // compute positional offsets
     mju_sub3(m->light_pos0+3*i, d->light_xpos+3*i, d->xpos+3*id);

src/user/user_model.cc

@@ -1520,6 +1520,17 @@ void mjCModel::CopyTree(mjModel* m) {
       qposadr += nPOS[pj->type];
     }
 
+    // simple body with no rotational dofs: promote to simple level 2
+    if (m->body_simple[i]) {
+      m->body_simple[i] = 2;
+      for (int j=0; j<(int)pb->joints.size(); j++) {
+        if (pb->joints[j]->type!=mjJNT_SLIDE) {
+        m->body_simple[i] = 1;
+        break;
+        }
+      }
+    }
+
     // loop over geoms for this body
     for (int j=0; j<(int)pb->geoms.size(); j++) {
       // get pointer and id

test/engine/engine_solver_test.cc

@@ -107,7 +107,7 @@ TEST_F(SolverTest, OneBigIsland) {
   mjData* data_noisland = mj_makeData(model);
 
   int nv = model->nv;
-  mjtNum tol = 1e-9;
+  mjtNum tol = 1e-8;
 
   // save current (default) iterations
   int iterations_default = model->opt.iterations;