chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add mjData flags marking lazily evaluated pipeline stages

Browse Source 
PiperOrigin-RevId: 861763477
Change-Id: Ib613c56949264573a16306800658f868d2aa7665
This commit is contained in:
Yuval Tassa
2026-01-27 09:44:16 -08:00
committed by Copybara-Service
parent ba32568138
commit 5fb47df197
12 changed files with 124 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
doc/includes/references.h
+6
View File
@@ -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
+6
View File
@@ -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
+4
View File
@@ -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
+20
View File
@@ -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
+6
View File
@@ -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
+15 -38
View File
@@ -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
+6 -2
View File
@@ -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
+6
View File
@@ -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
+1
View File
@@ -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
+22 -33
View File
@@ -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;
}
unity/Runtime/Bindings/MjBindings.cs
+4
View File
@@ -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
+28
View File
@@ -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)