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
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;
}