Add mjData flags marking lazily evaluated pipeline stages

PiperOrigin-RevId: 861763477
Change-Id: Ib613c56949264573a16306800658f868d2aa7665

doc/includes/references.h

@@ -199,6 +199,12 @@ struct mjData_ {
   int     nparent_awake;     // number of bodies with awake parents
   int     nv_awake;          // number of awake dofs
 
+  // flags marking lazily evaluated stages
+  mjtByte flg_energypos;     // has mj_energyPos been called
+  mjtByte flg_energyvel;     // has mj_energyVel been called
+  mjtByte flg_subtreevel;    // has mj_subtreeVel been called
+  mjtByte flg_rnepost;       // has mj_rnePostConstraint been called
+
   // global properties
   mjtNum  time;              // simulation time
   mjtNum  energy[2];         // potential, kinetic energy

include/mujoco/mjdata.h

@@ -233,6 +233,12 @@ struct mjData_ {
   int     nparent_awake;     // number of bodies with awake parents
   int     nv_awake;          // number of awake dofs
 
+  // flags marking lazily evaluated stages
+  mjtByte flg_energypos;     // has mj_energyPos been called
+  mjtByte flg_energyvel;     // has mj_energyVel been called
+  mjtByte flg_subtreevel;    // has mj_subtreeVel been called
+  mjtByte flg_rnepost;       // has mj_rnePostConstraint been called
+
   // global properties
   mjtNum  time;              // simulation time
   mjtNum  energy[2];         // potential, kinetic energy

include/mujoco/mjxmacro.h

@@ -915,6 +915,10 @@
     X( int,       nbody_awake        ) \
     X( int,       nparent_awake      ) \
     X( int,       nv_awake           ) \
+    X( mjtByte,   flg_energypos      ) \
+    X( mjtByte,   flg_energyvel      ) \
+    X( mjtByte,   flg_subtreevel     ) \
+    X( mjtByte,   flg_rnepost        ) \
     X( mjtNum,    time               ) \
     X( uintptr_t, threadpool         )
 

python/mujoco/introspect/structs.py

@@ -5331,6 +5331,26 @@ STRUCTS: Mapping[str, StructDecl] = dict([
                  type=ValueType(name='int'),
                  doc='number of awake dofs',
              ),
+             StructFieldDecl(
+                 name='flg_energypos',
+                 type=ValueType(name='mjtByte'),
+                 doc='has mj_energyPos been called',
+             ),
+             StructFieldDecl(
+                 name='flg_energyvel',
+                 type=ValueType(name='mjtByte'),
+                 doc='has mj_energyVel been called',
+             ),
+             StructFieldDecl(
+                 name='flg_subtreevel',
+                 type=ValueType(name='mjtByte'),
+                 doc='has mj_subtreeVel been called',
+             ),
+             StructFieldDecl(
+                 name='flg_rnepost',
+                 type=ValueType(name='mjtByte'),
+                 doc='has mj_rnePostConstraint been called',
+             ),
              StructFieldDecl(
                  name='time',
                  type=ValueType(name='mjtNum'),

src/engine/engine_core_smooth.c

@@ -2414,6 +2414,9 @@ void mj_subtreeVel(const mjModel* m, mjData* d) {
   }
 
   mj_freeStack(d);
+
+  // mark as computed
+  d->flg_subtreevel = 1;
 }
 
 
@@ -2674,6 +2677,9 @@ void mj_rnePostConstraint(const mjModel* m, mjData* d) {
   for (int j=nbody-1; j > 0; j--) {
     mju_addTo(d->cfrc_int+6*m->body_parentid[j], d->cfrc_int+6*j, 6);
   }
+
+  // mark as computed
+  d->flg_rnepost = 1;
 }
 
 

src/engine/engine_forward.c

@@ -152,6 +152,9 @@ void mj_fwdKinematics(const mjModel* m, mjData* d) {
 void mj_fwdPosition(const mjModel* m, mjData* d) {
   TM_START1;
 
+  // clear position-dependent flags for lazy evaluation
+  d->flg_energypos = 0;
+
   TM_START;
   mj_fwdKinematics(m, d);
 
@@ -218,6 +221,10 @@ void mj_fwdPosition(const mjModel* m, mjData* d) {
 void mj_fwdVelocity(const mjModel* m, mjData* d) {
   TM_START;
 
+  // clear velocity-dependent flags for lazy evaluation
+  d->flg_subtreevel = 0;
+  d->flg_energyvel = 0;
+
   // flexedge velocity: always sparse
   mju_mulMatVecSparse(d->flexedge_velocity, d->flexedge_J, d->qvel, m->nflexedge,
                       m->flexedge_J_rownnz, m->flexedge_J_rowadr, m->flexedge_J_colind, NULL);
@@ -1173,36 +1180,6 @@ void mj_implicit(const mjModel* m, mjData* d) {
 }
 
 
-// return 1 if potential energy was computed by sensor, 0 otherwise
-static int energyPosSensor(const mjModel* m) {
-  if (mjDISABLED(mjDSBL_SENSOR)) {
-    return 0;
-  }
-
-  for (int i=0; i < m->nsensor; i++) {
-    if (m->sensor_type[i] == mjSENS_E_POTENTIAL) {
-      return 1;
-    }
-  }
-  return 0;
-}
-
-
-// return 1 if kinetic energy was computed by sensor, 0 otherwise
-static int energyVelSensor(const mjModel* m) {
-  if (mjDISABLED(mjDSBL_SENSOR)) {
-    return 0;
-  }
-
-  for (int i=0; i < m->nsensor; i++) {
-    if (m->sensor_type[i] == mjSENS_E_KINETIC) {
-      return 1;
-    }
-  }
-  return 0;
-}
-
-
 //-------------------------- top-level API ---------------------------------------------------------
 
 // forward dynamics with skip; skipstage is mjtStage
@@ -1213,13 +1190,11 @@ void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor)
   if (skipstage < mjSTAGE_POS) {
     mj_fwdPosition(m, d);
 
-    int energyPos = 0;
     if (!skipsensor) {
       mj_sensorPos(m, d);
-      energyPos = energyPosSensor(m);
     }
 
-    if (!energyPos) {
+    if (!d->flg_energypos) {
       if (mjENABLED(mjENBL_ENERGY)) {
         mj_energyPos(m, d);
       } else {
@@ -1232,13 +1207,11 @@ void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor)
   if (skipstage < mjSTAGE_VEL) {
     mj_fwdVelocity(m, d);
 
-    int energyVel = 0;
     if (!skipsensor) {
       mj_sensorVel(m, d);
-      energyVel = energyVelSensor(m);
     }
 
-    if (mjENABLED(mjENBL_ENERGY) && !energyVel) {
+    if (mjENABLED(mjENBL_ENERGY) && !d->flg_energyvel) {
       mj_energyVel(m, d);
     }
   }
@@ -1252,6 +1225,7 @@ void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor)
   mj_fwdAcceleration(m, d);
   mj_fwdConstraint(m, d);
   if (!skipsensor) {
+    d->flg_rnepost = 0;  // clear flag for lazy evaluation
     mj_sensorAcc(m, d);
   }
 
@@ -1310,18 +1284,21 @@ void mj_step1(const mjModel* m, mjData* d) {
   mj_checkVel(m, d);
   mj_fwdPosition(m, d);
   mj_sensorPos(m, d);
-  if (!energyPosSensor(m)) {
+
+  if (!d->flg_energypos) {
     if (mjENABLED(mjENBL_ENERGY)) {
       mj_energyPos(m, d);
     } else {
       d->energy[0] = d->energy[1] = 0;
     }
   }
+
   mj_fwdVelocity(m, d);
   mj_sensorVel(m, d);
-  if (mjENABLED(mjENBL_ENERGY) && !energyVelSensor(m)) {
+  if (mjENABLED(mjENBL_ENERGY) && !d->flg_energyvel) {
     mj_energyVel(m, d);
   }
+
   if (mjcb_control) {
     mjcb_control(m, d);
   }

src/engine/engine_inverse.c

@@ -39,6 +39,9 @@ void mj_invPosition(const mjModel* m, mjData* d) {
   TM_START1;
   TM_START;
 
+  // clear flag for lazy evaluation
+  d->flg_energypos = 0;
+
   mj_kinematics(m, d);
   mj_comPos(m, d);
   mj_camlight(m, d);
@@ -191,7 +194,7 @@ void mj_inverseSkip(const mjModel* m, mjData* d,
     if (!skipsensor) {
       mj_sensorPos(m, d);
     }
-    if (mjENABLED(mjENBL_ENERGY)) {
+    if (mjENABLED(mjENBL_ENERGY) && !d->flg_energypos) {
       mj_energyPos(m, d);
     }
   }
@@ -202,7 +205,7 @@ void mj_inverseSkip(const mjModel* m, mjData* d,
     if (!skipsensor) {
       mj_sensorVel(m, d);
     }
-    if (mjENABLED(mjENBL_ENERGY)) {
+    if (mjENABLED(mjENBL_ENERGY) && !d->flg_energyvel) {
       mj_energyVel(m, d);
     }
   }
@@ -224,6 +227,7 @@ void mj_inverseSkip(const mjModel* m, mjData* d,
   mj_tendonBias(m, d, d->qfrc_inverse);
 
   if (!skipsensor) {
+    d->flg_rnepost = 0;  // clear flag for lazy evaluation
     mj_sensorAcc(m, d);
   }
 

src/engine/engine_io.c

@@ -1327,6 +1327,12 @@ static void _resetData(const mjModel* m, mjData* d, unsigned char debug_value) {
   d->time = 0;
   mju_zero(d->energy, 2);
 
+  // clear lazy evaluation flags
+  d->flg_energypos = 0;
+  d->flg_energyvel = 0;
+  d->flg_subtreevel = 0;
+  d->flg_rnepost = 0;
+
   //------------------------------ clear buffer, set defaults
 
   // fill buffer with debug_value (normally 0)

src/engine/engine_print.c

@@ -1149,6 +1149,7 @@ void mj_printFormattedData(const mjModel* m, const mjData* d, const char* filena
         d->name,                                                              \
         int : INT_FORMAT,                                                     \
         mjtSize : SIZE_FORMAT,                                                \
+        mjtByte : INT_FORMAT,                                                 \
         default : NULL);                                                      \
     if (format) {                                                             \
       fprintf(fp, "  ");                                                      \

src/engine/engine_sensor.c

@@ -857,7 +857,6 @@ void mj_sensorVel(const mjModel* m, mjData* d) {
   int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody;
 
   // process sensors matching stage
-  int subtreeVel = 0;
   for (int i=0; i < m->nsensor; i++) {
     // skip sensor plugins -- these are handled after builtin sensor types
     if (m->sensor_type[i] == mjSENS_PLUGIN) {
@@ -878,16 +877,12 @@ void mj_sensorVel(const mjModel* m, mjData* d) {
       reftype = m->sensor_reftype[i];
       adr = m->sensor_adr[i];
 
-      // call mj_subtreeVel when first relevant sensor is encountered
+      // call mj_subtreeVel for sensors that need it (unless already computed)
-      if (subtreeVel == 0 &&
+      if (!d->flg_subtreevel &&
           (type == mjSENS_SUBTREELINVEL ||
            type == mjSENS_SUBTREEANGMOM ||
            type == mjSENS_USER)) {
-        // compute subtree_linvel, subtree_angmom
         mj_subtreeVel(m, d);
-
-        // mark computed
-        subtreeVel = 1;
       }
 
       // process according to type
@@ -1016,13 +1011,10 @@ void mj_sensorVel(const mjModel* m, mjData* d) {
         if (!plugin->compute) {
           mjERROR("`compute` is null for plugin at slot %d", slot);
         }
-        if (subtreeVel == 0) {
+        // compute subtree_linvel, subtree_angmom (unless already computed)
-          // compute subtree_linvel, subtree_angmom
+        // TODO(b/247107630): add a flag to allow plugin to specify whether it actually needs this
-          // TODO(b/247107630): add a flag to allow plugin to specify whether it actually needs this
+        if (!d->flg_subtreevel) {
           mj_subtreeVel(m, d);
-
-          // mark computed
-          subtreeVel = 1;
         }
         plugin->compute(m, d, i, mjPLUGIN_SENSOR);
       }
@@ -1050,7 +1042,6 @@ void mj_sensorAcc(const mjModel* m, mjData* d) {
   int sleep_filter = mjENABLED(mjENBL_SLEEP) && d->nbody_awake < m->nbody;
 
   // process sensors matching stage
-  int rnePost = 0;
   for (int i=0; i < m->nsensor; i++) {
     // skip sleeping sensor
     if (sleep_filter && mj_sleepState(m, d, mjOBJ_SENSOR, i) == mjS_ASLEEP) {
@@ -1069,18 +1060,15 @@ void mj_sensorAcc(const mjModel* m, mjData* d) {
       objid = m->sensor_objid[i];
       adr = m->sensor_adr[i];
 
-      // call mj_rnePostConstraint when first relevant sensor is encountered
+      // call mj_rnePostConstraint for sensors that need it (unless already computed)
-      if (rnePost == 0  && (type == mjSENS_ACCELEROMETER ||
+      if (!d->flg_rnepost &&
-                            type == mjSENS_FORCE         ||
+          (type == mjSENS_ACCELEROMETER ||
-                            type == mjSENS_TORQUE        ||
+           type == mjSENS_FORCE         ||
-                            type == mjSENS_FRAMELINACC   ||
+           type == mjSENS_TORQUE        ||
-                            type == mjSENS_FRAMEANGACC   ||
+           type == mjSENS_FRAMELINACC   ||
-                            type == mjSENS_USER)) {
+           type == mjSENS_FRAMEANGACC   ||
-        // compute cacc, cfrc_int, cfrc_ext
+           type == mjSENS_USER)) {
         mj_rnePostConstraint(m, d);
-
-        // mark computed
-        rnePost = 1;
       }
 
       // process according to type
@@ -1550,14 +1538,9 @@ void mj_sensorAcc(const mjModel* m, mjData* d) {
         if (!plugin->compute) {
           mjERROR("`compute` is null for plugin at slot %d", slot);
         }
-        if (rnePost == 0) {
+        // compute cacc, cfrc_int, cfrc_ext (function handles early return)
-          // compute cacc, cfrc_int, cfrc_ext
+        // TODO(b/247107630): add a flag to allow plugin to specify whether it actually needs this
-          // TODO(b/247107630): add a flag to allow plugin to specify whether it actually needs this
+        mj_rnePostConstraint(m, d);
-          mj_rnePostConstraint(m, d);
-
-          // mark computed
-          rnePost = 1;
-        }
         plugin->compute(m, d, i, mjPLUGIN_SENSOR);
       }
     }
@@ -1672,6 +1655,9 @@ void mj_energyPos(const mjModel* m, mjData* d) {
       }
     }
   }
+
+  // mark as computed
+  d->flg_energypos = 1;
 }
 
 
@@ -1685,4 +1671,7 @@ void mj_energyVel(const mjModel* m, mjData* d) {
   d->energy[1] = 0.5*mju_dot(vec, d->qvel, m->nv);
 
   mj_freeStack(d);
+
+  // mark as computed
+  d->flg_energyvel = 1;
 }

unity/Runtime/Bindings/MjBindings.cs

@@ -4940,6 +4940,10 @@ public unsafe struct mjData_ {
   public int nbody_awake;
   public int nparent_awake;
   public int nv_awake;
+  public byte flg_energypos;
+  public byte flg_energyvel;
+  public byte flg_subtreevel;
+  public byte flg_rnepost;
   public double time;
   public fixed double energy[2];
   public void* buffer;

wasm/codegen/generated/bindings.cc

@@ -6205,6 +6205,30 @@ struct MjData {
   void set_nv_awake(int value) {
     ptr_->nv_awake = value;
   }
+  mjtByte flg_energypos() const {
+    return ptr_->flg_energypos;
+  }
+  void set_flg_energypos(mjtByte value) {
+    ptr_->flg_energypos = value;
+  }
+  mjtByte flg_energyvel() const {
+    return ptr_->flg_energyvel;
+  }
+  void set_flg_energyvel(mjtByte value) {
+    ptr_->flg_energyvel = value;
+  }
+  mjtByte flg_subtreevel() const {
+    return ptr_->flg_subtreevel;
+  }
+  void set_flg_subtreevel(mjtByte value) {
+    ptr_->flg_subtreevel = value;
+  }
+  mjtByte flg_rnepost() const {
+    return ptr_->flg_rnepost;
+  }
+  void set_flg_rnepost(mjtByte value) {
+    ptr_->flg_rnepost = value;
+  }
   mjtNum time() const {
     return ptr_->time;
   }
@@ -11075,6 +11099,10 @@ EMSCRIPTEN_BINDINGS(mujoco_bindings) {
     .property("flexvert_J", &MjData::flexvert_J)
     .property("flexvert_length", &MjData::flexvert_length)
     .property("flexvert_xpos", &MjData::flexvert_xpos)
+    .property("flg_energypos", &MjData::flg_energypos, &MjData::set_flg_energypos, reference())
+    .property("flg_energyvel", &MjData::flg_energyvel, &MjData::set_flg_energyvel, reference())
+    .property("flg_rnepost", &MjData::flg_rnepost, &MjData::set_flg_rnepost, reference())
+    .property("flg_subtreevel", &MjData::flg_subtreevel, &MjData::set_flg_subtreevel, reference())
     .property("geom_xmat", &MjData::geom_xmat)
     .property("geom_xpos", &MjData::geom_xpos)
     .property("iLD", &MjData::iLD)