diff --git a/doc/includes/references.h b/doc/includes/references.h index 2ead7899..181959ca 100644 --- a/doc/includes/references.h +++ b/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 diff --git a/include/mujoco/mjdata.h b/include/mujoco/mjdata.h index c96cd6cc..01c1736b 100644 --- a/include/mujoco/mjdata.h +++ b/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 diff --git a/include/mujoco/mjxmacro.h b/include/mujoco/mjxmacro.h index 2252656d..6b876a94 100644 --- a/include/mujoco/mjxmacro.h +++ b/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 ) diff --git a/python/mujoco/introspect/structs.py b/python/mujoco/introspect/structs.py index 67b9f2b5..f0e4a848 100644 --- a/python/mujoco/introspect/structs.py +++ b/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'), diff --git a/src/engine/engine_core_smooth.c b/src/engine/engine_core_smooth.c index 287b83ba..ba5c7254 100644 --- a/src/engine/engine_core_smooth.c +++ b/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; } diff --git a/src/engine/engine_forward.c b/src/engine/engine_forward.c index 03e86c61..eae2a198 100644 --- a/src/engine/engine_forward.c +++ b/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); } diff --git a/src/engine/engine_inverse.c b/src/engine/engine_inverse.c index ae4e1191..982228e2 100644 --- a/src/engine/engine_inverse.c +++ b/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); } diff --git a/src/engine/engine_io.c b/src/engine/engine_io.c index 01bc2fda..51e4d833 100644 --- a/src/engine/engine_io.c +++ b/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) diff --git a/src/engine/engine_print.c b/src/engine/engine_print.c index 08e2af3e..66291819 100644 --- a/src/engine/engine_print.c +++ b/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, " "); \ diff --git a/src/engine/engine_sensor.c b/src/engine/engine_sensor.c index 99d5d006..faeaa3b9 100644 --- a/src/engine/engine_sensor.c +++ b/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 - if (subtreeVel == 0 && + // call mj_subtreeVel for sensors that need it (unless already computed) + 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 - // TODO(b/247107630): add a flag to allow plugin to specify whether it actually needs this + // compute subtree_linvel, subtree_angmom (unless already computed) + // 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 - if (rnePost == 0 && (type == mjSENS_ACCELEROMETER || - type == mjSENS_FORCE || - type == mjSENS_TORQUE || - type == mjSENS_FRAMELINACC || - type == mjSENS_FRAMEANGACC || - type == mjSENS_USER)) { - // compute cacc, cfrc_int, cfrc_ext + // call mj_rnePostConstraint for sensors that need it (unless already computed) + if (!d->flg_rnepost && + (type == mjSENS_ACCELEROMETER || + type == mjSENS_FORCE || + type == mjSENS_TORQUE || + type == mjSENS_FRAMELINACC || + type == mjSENS_FRAMEANGACC || + 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 - // TODO(b/247107630): add a flag to allow plugin to specify whether it actually needs this - mj_rnePostConstraint(m, d); - - // mark computed - rnePost = 1; - } + // compute cacc, cfrc_int, cfrc_ext (function handles early return) + // TODO(b/247107630): add a flag to allow plugin to specify whether it actually needs this + mj_rnePostConstraint(m, d); 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; } diff --git a/unity/Runtime/Bindings/MjBindings.cs b/unity/Runtime/Bindings/MjBindings.cs index 48bb77b1..732a7d31 100644 --- a/unity/Runtime/Bindings/MjBindings.cs +++ b/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; diff --git a/wasm/codegen/generated/bindings.cc b/wasm/codegen/generated/bindings.cc index 7a964c34..65e08abc 100644 --- a/wasm/codegen/generated/bindings.cc +++ b/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)