From fe9dc584772bf618144542682b10aa2c8d4a0120 Mon Sep 17 00:00:00 2001 From: Alessio Quaglino Date: Thu, 23 Jul 2026 09:01:51 -0700 Subject: [PATCH] Fix flexcomp instability when parent body has non-identity quaternion 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 --- src/engine/engine_derivative.c | 43 ++++++++++++--- src/engine/engine_passive.c | 28 ++++++++-- test/engine/engine_passive_test.cc | 88 ++++++++++++++++++++++++++++++ 3 files changed, 146 insertions(+), 13 deletions(-) diff --git a/src/engine/engine_derivative.c b/src/engine/engine_derivative.c index d68318f7..f7c3f413 100644 --- a/src/engine/engine_derivative.c +++ b/src/engine/engine_derivative.c @@ -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]; } } } diff --git a/src/engine/engine_passive.c b/src/engine/engine_passive.c index ad4916fb..f977f814 100644 --- a/src/engine/engine_passive.c +++ b/src/engine/engine_passive.c @@ -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]; } } } diff --git a/test/engine/engine_passive_test.cc b/test/engine/engine_passive_test.cc index ab7de3be..223d7b59 100644 --- a/test/engine/engine_passive_test.cc +++ b/test/engine/engine_passive_test.cc @@ -14,6 +14,7 @@ // Tests for engine/engine_core_smooth.c. +#include #include #include @@ -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"( + + + )"; + static constexpr char nonrotated_xml[] = R"( + + + )"; + + 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 +