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

Fix flexcomp instability when parent body has non-identity quaternion

Browse Source 
Fast path in mj_flexPassiveInterp, mj_flexPassiveBendInterp, and
mj_flexPassiveStretch assumed body slide joints are world-aligned
(J = I). When parent body has a non-identity quaternion, joint axes
are rotated (J = R_body), causing wrong force mapping and instability
(NaN/Inf in QACC).

Fix: project world-frame forces onto body local frame via
mju_mulMatTVec3(R_body^T, force) before adding to qfrc_spring/damper.
Also fix the derivative paths:
- mjd_flexInterp_kernel fast path: R^T * K_rot * R * vec
- mjd_flexStiff_assemble (CSR): R_bi^T * K_rot_block * R_bj

The CSR-assembled stiffness matrix is the actual path used by the
implicit CG solver (mj_flexCG gate). The test uses solver="CG" to
activate this path; without it, flex stiffness is integrated
explicitly and no derivative fix can help.

Ported from GitHub PR https://github.com/google-deepmind/mujoco/pull/3379
Original author: Devansh (https://github.com/devansh0703)
Fixes https://github.com/google-deepmind/mujoco/issues/3364

PiperOrigin-RevId: 952789284
Change-Id: If924f7160dd16c0cc88170d80e6e605da0fe2e04
This commit is contained in:
Alessio Quaglino
2026-07-23 09:01:51 -07:00
committed by Copybara-Service
parent b04a18a00a
commit fe9dc58477
3 changed files with 146 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_derivative.c
+35 -8
View File
@@ -1108,6 +1108,9 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d,
int* chain_colind = mjSTACKALLOC(d, nv, int);
mjtNum* blk_jac = mjSTACKALLOC(d, 3*nv, mjtNum);
// temp buffer for fast path element result
mjtNum* fast_tmp = mjSTACKALLOC(d, max_dim_c, mjtNum);
// loop over flexes
for (int f=0; f < m->nflex; f++) {
// only process flex_interp
@@ -1167,7 +1170,7 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d,
mju_flexGatherState(m, d, f, xpos, NULL);
// check if centered fast path applies: centered, all nodes on simple slider
// bodies (body_simple == 2 means diag M with sliders only, J = I_3 per node)
// bodies (body_simple == 2 means diag M with sliders only, J = R_body per node)
int use_fast_path = m->flex_centered[f];
if (use_fast_path) {
int nodenum_f = m->flex_nodenum[f];
@@ -1272,22 +1275,32 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d,
}
// fast path: centered flex with 3 translational DOFs per body
// J is identity, so J'*K*J*vec = K*vec — scatter directly
// J = R_body (not I when body is rotated), so J'*K*J*vec = R^T*K*(R*vec)
if (use_fast_path) {
// apply R_body to vec for each node
for (int n = 0; n < npe; n++) {
int dof = m->body_dofadr[bodyid[gindices[n]]];
mji_mulMatVec3(fast_tmp + 3*n, d->xmat + 9*bodyid[gindices[n]], vec + dof);
}
// compute R^T * K_rot * (R*vec) and add to res
for (int a = 0; a < npe; a++) {
int dof_a = m->body_dofadr[bodyid[gindices[a]]];
mjtNum row[3] = {0, 0, 0};
for (int b = 0; b < npe; b++) {
int dof_b = m->body_dofadr[bodyid[gindices[b]]];
// K_rot_cell[3*a, 3*b] is the 3x3 block
int adr = (3*a)*dim_e + 3*b;
for (int r = 0; r < 3; r++) {
mjtNum val = 0;
for (int c = 0; c < 3; c++) {
val += K_rot_cell[adr + r*dim_e + c] * vec[dof_b + c];
row[r] += K_rot_cell[adr + r*dim_e + c] * fast_tmp[3*b + c];
}
res[dof_a + r] += val;
}
}
// apply R^T and add to res
mjtNum qfrc_loc[3];
mji_mulMatTVec3(qfrc_loc, d->xmat + 9*bodyid[gindices[a]], row);
for (int r = 0; r < 3; r++) {
res[dof_a + r] += qfrc_loc[r];
}
}
} else {
// general path: construct sparse Jacobian for this element's nodes
@@ -1963,14 +1976,28 @@ int mjd_flexStiff_assemble(const mjModel* m, mjData* d, int* rownnz, int* rowadr
for (int i = 0; i < npe; i++) {
int si = nslot[m->flex_nodeadr[f] + gindices[i]];
if (si < 0) continue;
int bi = m->flex_nodebodyid[m->flex_nodeadr[f] + gindices[i]];
for (int j = 0; j < npe; j++) {
int sj = nslot[m->flex_nodeadr[f] + gindices[j]];
if (sj < 0) continue;
int bj = m->flex_nodebodyid[m->flex_nodeadr[f] + gindices[j]];
// extract K_rot 3x3 block for (i,j)
mjtNum kb_ij[9];
for (int r = 0; r < 3; r++) {
for (int c = 0; c < 3; c++) {
kb_ij[3*r+c] = iscale*kb[(3*i + r)*dim_e + 3*j + c];
}
}
// apply R_bi^T * kb_ij * R_bj to transform to DOF space
mjtNum tmp[9];
mjtNum blk[9];
mji_mulMatMat3(tmp, kb_ij, d->xmat + 9*bj); // tmp = kb * R_bj
mji_mulMatTMat3(blk, d->xmat + 9*bi, tmp); // blk = R_bi^T * tmp
int pos;
FLEXSTIFF_BLOCK(si, sj, pos);
for (int r = 0; r < 3; r++) {
for (int c = 0; c < 3; c++) {
val[rowadr[vdof[si] + r] + 3*pos + c] += iscale*kb[(3*i + r)*dim_e + 3*j + c];
val[rowadr[vdof[si] + r] + 3*pos + c] += blk[3*r+c];
}
}
}
src/engine/engine_passive.c
+23 -5
View File
@@ -194,8 +194,16 @@ static void mj_flexPassiveInterp(const mjModel* m, mjData* d, int f,
(m->flex_node[3*nidx+0] == 0 &&
m->flex_node[3*nidx+1] == 0 &&
m->flex_node[3*nidx+2] == 0))) {
if (enbl_spring) mji_addTo3(d->qfrc_spring + m->body_dofadr[bid], frc_g+3*i);
if (enbl_damper) mji_addTo3(d->qfrc_damper + m->body_dofadr[bid], dmp_g+3*i);
if (enbl_spring) {
mjtNum qfrc_loc[3];
mju_mulMatTVec3(qfrc_loc, d->xmat+9*bid, frc_g+3*i);
mji_addTo3(d->qfrc_spring + m->body_dofadr[bid], qfrc_loc);
}
if (enbl_damper) {
mjtNum qdmp_loc[3];
mju_mulMatTVec3(qdmp_loc, d->xmat+9*bid, dmp_g+3*i);
mji_addTo3(d->qfrc_damper + m->body_dofadr[bid], qdmp_loc);
}
} else {
if (enbl_spring) mj_applyFT(m, d, frc_g+3*i, 0, xpos_g+3*i, bid, d->qfrc_spring);
if (enbl_damper) mj_applyFT(m, d, dmp_g+3*i, 0, xpos_g+3*i, bid, d->qfrc_damper);
@@ -428,8 +436,16 @@ static void mj_flexPassiveBendInterp(const mjModel* m, mjData* d, int f,
(m->flex_node[3*nidx+0] == 0 &&
m->flex_node[3*nidx+1] == 0 &&
m->flex_node[3*nidx+2] == 0))) {
if (enbl_spring) mji_addTo3(d->qfrc_spring + m->body_dofadr[bid], frc_g+3*i);
if (enbl_damper) mji_addTo3(d->qfrc_damper + m->body_dofadr[bid], dmp_g+3*i);
if (enbl_spring) {
mjtNum qfrc_loc[3];
mju_mulMatTVec3(qfrc_loc, d->xmat+9*bid, frc_g+3*i);
mji_addTo3(d->qfrc_spring + m->body_dofadr[bid], qfrc_loc);
}
if (enbl_damper) {
mjtNum qdmp_loc[3];
mju_mulMatTVec3(qdmp_loc, d->xmat+9*bid, dmp_g+3*i);
mji_addTo3(d->qfrc_damper + m->body_dofadr[bid], qdmp_loc);
}
} else {
if (enbl_spring) mj_applyFT(m, d, frc_g+3*i, 0, xpos_g+3*i, bid, d->qfrc_spring);
if (enbl_damper) mj_applyFT(m, d, dmp_g+3*i, 0, xpos_g+3*i, bid, d->qfrc_damper);
@@ -618,8 +634,10 @@ static void mj_flexPassiveStretch(const mjModel* m, mjData* d, int f,
} else {
int body_dofnum = m->body_dofnum[bid];
int body_dofadr = m->body_dofadr[bid];
mjtNum qfrc_loc[3];
mju_mulMatTVec3(qfrc_loc, d->xmat+9*bid, qfrc+3*v);
for (int x = 0; x < body_dofnum; x++) {
d->qfrc_spring[body_dofadr+x] += qfrc[3*v+x];
d->qfrc_spring[body_dofadr+x] += qfrc_loc[x];
}
}
}