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

Simplify mj_comPos, avoid stack allocation

Browse Source 
PiperOrigin-RevId: 809762290
Change-Id: I87824279ff5f7049b1f06b900d4d3e33f5fafe0f
This commit is contained in:
Yuval Tassa
2025-09-21 14:05:15 -07:00
committed by Copybara-Service
parent 0141c6b8b7
commit d55da23235
1 changed files with 15 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_core_smooth.c
+15 -24
View File
@@ -181,33 +181,24 @@ void mj_kinematics(const mjModel* m, mjData* d) {
// map inertias and motion dofs to global frame centered at subtree-CoM // map inertias and motion dofs to global frame centered at subtree-CoM
void mj_comPos(const mjModel* m, mjData* d) { void mj_comPos(const mjModel* m, mjData* d) {
int nbody = m->nbody, njnt = m->njnt; int nbody = m->nbody, njnt = m->njnt;
mjtNum offset[3], axis[3];
mj_markStack(d);
mjtNum* mass_subtree = mjSTACKALLOC(d, m->nbody, mjtNum);
// clear subtree // subtree_com: initialize with body moment
mju_zero(mass_subtree, m->nbody); for (int i=0; i < nbody; i++) {
mju_zero(d->subtree_com, m->nbody*3); mju_scl3(d->subtree_com+3*i, d->xipos+3*i, m->body_mass[i]);
}
// backwards pass over bodies: compute subtree_com and mass_subtree // subtree_com: accumulate to parent in backward pass
for (int i=nbody-1; i >= 0; i--) { for (int i=nbody-1; i > 0; i--) {
// add local info int j = m->body_parentid[i];
mju_addToScl3(d->subtree_com+3*i, d->xipos+3*i, m->body_mass[i]); mju_addTo3(d->subtree_com+3*j, d->subtree_com+3*i);
mass_subtree[i] += m->body_mass[i]; }
// add to parent, except for world // subtree_com: normalize
if (i) { for (int i=0; i < nbody; i++) {
int j = m->body_parentid[i]; if (m->body_subtreemass[i] < mjMINVAL) {
mju_addTo3(d->subtree_com+3*j, d->subtree_com+3*i);
mass_subtree[j] += mass_subtree[i];
}
// compute local com
if (mass_subtree[i] < mjMINVAL) {
mju_copy3(d->subtree_com+3*i, d->xipos+3*i); mju_copy3(d->subtree_com+3*i, d->xipos+3*i);
} else { } else {
mju_scl3(d->subtree_com+3*i, d->subtree_com+3*i, mju_scl3(d->subtree_com+3*i, d->subtree_com+3*i, 1.0/m->body_subtreemass[i]);
1.0/mjMAX(mjMINVAL, mass_subtree[i]));
} }
} }
@@ -216,6 +207,7 @@ void mj_comPos(const mjModel* m, mjData* d) {
// map inertias to frame centered at subtree_com // map inertias to frame centered at subtree_com
for (int i=1; i < nbody; i++) { for (int i=1; i < nbody; i++) {
mjtNum offset[3];
mju_sub3(offset, d->xipos+3*i, d->subtree_com+3*m->body_rootid[i]); mju_sub3(offset, d->xipos+3*i, d->subtree_com+3*m->body_rootid[i]);
mju_inertCom(d->cinert+10*i, m->body_inertia+3*i, d->ximat+9*i, mju_inertCom(d->cinert+10*i, m->body_inertia+3*i, d->ximat+9*i,
offset, m->body_mass[i]); offset, m->body_mass[i]);
@@ -228,6 +220,7 @@ void mj_comPos(const mjModel* m, mjData* d) {
int bi = m->jnt_bodyid[j]; int bi = m->jnt_bodyid[j];
// compute com-anchor vector // compute com-anchor vector
mjtNum offset[3], axis[3];
mju_sub3(offset, d->subtree_com+3*m->body_rootid[bi], d->xanchor+3*j); mju_sub3(offset, d->subtree_com+3*m->body_rootid[bi], d->xanchor+3*j);
// create motion dof // create motion dof
@@ -264,8 +257,6 @@ void mj_comPos(const mjModel* m, mjData* d) {
break; break;
} }
} }
mj_freeStack(d);
} }