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];
}
}
}
test/engine/engine_passive_test.cc
+88
View File
@@ -14,6 +14,7 @@
// Tests for engine/engine_core_smooth.c.
#include <cmath>
#include <limits>
#include <string>
@@ -1006,5 +1007,92 @@ TEST_F(ElasticityTest, PinnedVertexBendingForce) {
EXPECT_TRUE(has_nonzero) << "bending should produce nonzero spring forces";
}
// verify that a flexcomp with dof="trilinear" inside a parent body
// with non-identity quaternion produces:
// 1. restoring forces (not expanding) for small perturbations
// 2. identical qfrc_passive between rotated and non-rotated models
// 3. stable simulation with implicit integration (no NaN, bounded qacc)
TEST_F(ElasticityTest, TrilinearParentBodyRotation) {
static constexpr char rotated_xml[] = R"(
<mujoco>
<option gravity="0 0 0" integrator="implicitfast"
timestep="0.0005" solver="CG"/>
<worldbody>
<body name="base" pos="0 0 0" quat="0.7071 0 0.7071 0">
<flexcomp type="grid" count="3 3 3" spacing=".05 .05 .05"
dim="3" radius=".001" mass=".005" name="soft" dof="trilinear">
<elasticity young="1e4" poisson="0.1" damping="0.1"/>
<contact selfcollide="none" internal="false"/>
</flexcomp>
</body>
</worldbody>
</mujoco>
)";
static constexpr char nonrotated_xml[] = R"(
<mujoco>
<option gravity="0 0 0" integrator="implicitfast"
timestep="0.0005" solver="CG"/>
<worldbody>
<body name="base" pos="0 0 0">
<flexcomp type="grid" count="3 3 3" spacing=".05 .05 .05"
dim="3" radius=".001" mass=".005" name="soft" dof="trilinear">
<elasticity young="1e4" poisson="0.1" damping="0.1"/>
<contact selfcollide="none" internal="false"/>
</flexcomp>
</body>
</worldbody>
</mujoco>
)";
char error[1024] = {0};
// --- test 1: force sign and rotational invariance ---
MjModelPtr m_rot = LoadModelFromString(rotated_xml, error, sizeof(error));
ASSERT_THAT(m_rot.get(), NotNull()) << error;
MjDataPtr d_rot = MakeData(m_rot);
MjModelPtr m_non = LoadModelFromString(nonrotated_xml, error, sizeof(error));
ASSERT_THAT(m_non.get(), NotNull()) << error;
MjDataPtr d_non = MakeData(m_non);
// apply small +X perturbation to DOF 0 (node 0, X slide)
d_rot->qpos[0] = 1e-6;
d_non->qpos[0] = 1e-6;
mj_forward(m_rot.get(), d_rot.get());
mj_forward(m_non.get(), d_non.get());
// check force sign on the displaced DOF: positive qpos -> negative qfrc (restoring)
EXPECT_LT(d_rot->qfrc_passive[0], 0)
<< "rotated model: expected restoring force on displaced DOF 0";
EXPECT_LT(d_non->qfrc_passive[0], 0)
<< "non-rotated model: expected restoring force on displaced DOF 0";
// check rotational invariance: forces match within numerical precision
// (tol ~ 1e-13 due to floating-point differences in polar decomposition path)
const mjtNum tol = MjTol(1e-12, 1e-5);
for (int i = 0; i < m_rot->nv; i++) {
EXPECT_NEAR(d_rot->qfrc_passive[i], d_non->qfrc_passive[i], tol)
<< "rotated/non-rotated qfrc mismatch at DOF " << i;
}
// --- test 2: implicit integration stability (200 steps, no NaN) ---
mjtNum max_qacc = 0;
for (int step = 0; step < 200; step++) {
mj_step(m_rot.get(), d_rot.get());
for (int i = 0; i < m_rot->nv; i++) {
ASSERT_TRUE(std::isfinite(d_rot->qacc[i]))
<< "NaN/Inf in qacc at DOF " << i << " at step " << step;
mjtNum a = mju_abs(d_rot->qacc[i]);
if (a > max_qacc) max_qacc = a;
}
if (HasFatalFailure()) return;
}
EXPECT_LT(max_qacc, 1e4);
}
} // namespace
} // namespace mujoco