Add mjModel.ngravcomp, number of bodies with nonzero gravcomp attribute.

PiperOrigin-RevId: 627087335
Change-Id: Ie85ae4150962463b8a8c8435fec068c44877aa16

doc/includes/references.h

@@ -922,6 +922,7 @@ struct mjModel_ {
   int nD;                         // number of non-zeros in sparse dof-dof matrix
   int nB;                         // number of non-zeros in sparse body-dof matrix
   int ntree;                      // number of kinematic trees under world body
+  int ngravcomp;                  // number of bodies with nonzero gravcomp
   int nemax;                      // number of potential equality-constraint rows
   int njmax;                      // number of available rows in constraint Jacobian
   int nconmax;                    // number of potential contacts in contact list

include/mujoco/mjmodel.h

@@ -637,6 +637,7 @@ struct mjModel_ {
   int nD;                         // number of non-zeros in sparse dof-dof matrix
   int nB;                         // number of non-zeros in sparse body-dof matrix
   int ntree;                      // number of kinematic trees under world body
+  int ngravcomp;                  // number of bodies with nonzero gravcomp
   int nemax;                      // number of potential equality-constraint rows
   int njmax;                      // number of available rows in constraint Jacobian
   int nconmax;                    // number of potential contacts in contact list

include/mujoco/mjxmacro.h

@@ -139,6 +139,7 @@
     X   ( njmax )              \
     X   ( nconmax )            \
     XMJV( ntree )              \
+    X   ( ngravcomp )          \
     X   ( nuserdata )          \
     XMJV( nsensordata )        \
     X   ( npluginstate )       \

introspect/structs.py

@@ -1200,6 +1200,11 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='int'),
                  doc='number of kinematic trees under world body',
              ),
+             StructFieldDecl(
+                 name='ngravcomp',
+                 type=ValueType(name='int'),
+                 doc='number of bodies with nonzero gravcomp',
+             ),
              StructFieldDecl(
                  name='nemax',
                  type=ValueType(name='int'),

src/engine/engine_forward.c

@@ -256,6 +256,23 @@ static void clampVec(mjtNum* vec, const mjtNum* range, const mjtByte* limited, i
 
 
 
+// return number of dofs given joint type
+static int jnt_dofnum(mjtJoint type) {
+  switch (type) {
+  case mjJNT_HINGE:
+  case mjJNT_SLIDE:
+    return 1;
+
+  case mjJNT_BALL:
+    return 3;
+
+  case mjJNT_FREE:
+    return 6;
+  }
+}
+
+
+
 // (qpos, qvel, ctrl, act) => (qfrc_actuator, actuator_force, act_dot)
 void mj_fwdActuation(const mjModel* m, mjData* d) {
   TM_START;
@@ -469,7 +486,7 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
   mju_mulMatTVec(d->qfrc_actuator, moment, force, nu, nv);
 
   // actuator-level gravity compensation
-  if (!mjDISABLED(mjDSBL_GRAVITY) && mju_norm3(m->opt.gravity)) {
+  if (m->ngravcomp && !mjDISABLED(mjDSBL_GRAVITY) && mju_norm3(m->opt.gravity)) {
     int njnt = m->njnt;
     for (int i=0; i < njnt; i++) {
       // skip if gravcomp added as passive force
@@ -477,28 +494,10 @@ void mj_fwdActuation(const mjModel* m, mjData* d) {
         continue;
       }
 
-      // get number of dofs for this joint
-      int dofnum;
-      switch (m->jnt_type[i]) {
-      case mjJNT_HINGE:
-      case mjJNT_SLIDE:
-        dofnum = 1;
-        break;
-
-      case mjJNT_BALL:
-        dofnum = 3;
-        break;
-
-      case mjJNT_FREE:
-        dofnum = 6;
-        break;
-      }
-
       // add gravcomp force
+      int dofnum = jnt_dofnum(m->jnt_type[i]);
       int dofadr = m->jnt_dofadr[i];
-      for (int j=0; j < dofnum; j++) {
+      mju_addTo(d->qfrc_actuator + dofadr, d->qfrc_gravcomp + dofadr, dofnum);
-        d->qfrc_actuator[dofadr+j] += d->qfrc_gravcomp[dofadr+j];
-      }
     }
   }
 

src/engine/engine_passive.c

@@ -175,7 +175,7 @@ static void mj_springdamper(const mjModel* m, mjData* d) {
 
 // body-level gravity compensation, return 1 if any, 0 otherwise
 static int mj_gravcomp(const mjModel* m, mjData* d) {
-  if (mjDISABLED(mjDSBL_GRAVITY) || mju_norm3(m->opt.gravity) == 0) {
+  if (!m->ngravcomp || mjDISABLED(mjDSBL_GRAVITY) || mju_norm3(m->opt.gravity) == 0) {
     return 0;
   }
 

src/user/user_model.cc

@@ -2105,6 +2105,13 @@ void mjCModel::CopyTree(mjModel* m) {
     }
   }
 
+  // count bodies with gravity compensation, compute ngravcomp
+  int ngravcomp = 0;
+  for (int i=0; i<nbody; i++) {
+    ngravcomp += (m->body_gravcomp[i] > 0);
+  }
+  m->ngravcomp = ngravcomp;
+
   // compute nM and dof_Madr
   nM = 0;
   for (int i=0; i<nv; i++) {

unity/Runtime/Bindings/MjBindings.cs

@@ -5163,6 +5163,7 @@ public unsafe struct mjModel_ {
   public int nD;
   public int nB;
   public int ntree;
+  public int ngravcomp;
   public int nemax;
   public int njmax;
   public int nconmax;