Merge pull request #3437 from smallquail:flex-rotation-fix

PiperOrigin-RevId: 955127380
Change-Id: Ie31ab56f7d64dd1a2266828c0fcd3683b55768a5

src/engine/engine_derivative.c

@@ -245,7 +245,7 @@ void mjd_subQuat(const mjtNum qa[4], const mjtNum qb[4], mjtNum Da[9], mjtNum Db
 
   // add term linear in K * K
   mjtNum KK[9];
-  mju_mulMatMat3(KK, K, K);
+  mji_mulMatMat3(KK, K, K);
   mjtNum coef = 1.0 - (half_angle < 6e-8 ? 1.0 : half_angle / mju_tan(half_angle));
   mju_addToScl(Da_tmp, KK, coef, 9);
 
@@ -313,7 +313,7 @@ void mjd_quatIntegrate(const mjtNum vel[3], mjtNum scale,
     if (Dvel || Dscale) Dvel_[i] = b*eye[i] + c*cross[i] + d*outer[i];
   }
   if (Dvel) mju_copy9(Dvel, Dvel_);
-  if (Dscale) mju_mulMatVec3(Dscale, Dvel_, vel);
+  if (Dscale) mji_mulMatVec3(Dscale, Dvel_, vel);
 }
 
 
@@ -1247,9 +1247,9 @@ static void mjd_flexInterp_kernel(const mjModel* m, mjData* d,
             }
 
             // tmp = K * R
-            mju_mulMatMat3(tmp, blk, R);
+            mji_mulMatMat3(tmp, blk, R);
             // blk = RT * tmp = RT * K * R
-            mju_mulMatMat3(blk, RT, tmp);
+            mji_mulMatMat3(blk, RT, tmp);
 
             // store in K_rot_cell at (a, b)
             int adr_out = (3*a)*dim_e + 3*b;
@@ -1455,12 +1455,21 @@ void mjd_flexStretch_mul(const mjModel* m, mjData* d, mjtNum* res, const mjtNum*
         int v0 = vert[edge[e][0]], v1 = vert[edge[e][1]];
         int b0 = bodyid[v0],       b1 = bodyid[v1];
         g[e] = 0;
+
+        // the vertex bodies' slide dofs are expressed in their own (possibly rotated) frame while
+        // the stiffness is built from world-space edge vectors, so the operator must be sandwiched
+        // with R (dof -> world) and R^T (world -> dof); mj_flexPassiveStretch applies the same R^T
+        // to its world-space force. R = I for the common case of an unrotated parent body.
+        mjtNum w0[3] = {0}, w1[3] = {0};
+        if (m->body_dofnum[b0]) {
+          mji_mulMatVec3(w0, d->xmat + 9*b0, vec + m->body_dofadr[b0]);
+        }
+        if (m->body_dofnum[b1]) {
+          mji_mulMatVec3(w1, d->xmat + 9*b1, vec + m->body_dofadr[b1]);
+        }
         for (int x = 0; x < 3; x++) {
           dvec[e][x] = xpos[3*v0+x] - xpos[3*v1+x];
-          mjtNum dv = 0;
+          g[e] += dvec[e][x]*(w0[x] - w1[x]);
-          if (m->body_dofnum[b0]) dv += vec[m->body_dofadr[b0]+x];
-          if (m->body_dofnum[b1]) dv -= vec[m->body_dofadr[b1]+x];
-          g[e] += dvec[e][x]*dv;
         }
       }
 
@@ -1483,9 +1492,21 @@ void mjd_flexStretch_mul(const mjModel* m, mjData* d, mjtNum* res, const mjtNum*
         }
         coef *= 2*scale;
         int b0 = bodyid[vert[edge[e][0]]], b1 = bodyid[vert[edge[e][1]]];
+        mjtNum rw[3], rl[3];
         for (int x = 0; x < 3; x++) {
-          if (m->body_dofnum[b0]) res[m->body_dofadr[b0]+x] += coef*dvec[e][x];
+          rw[x] = coef*dvec[e][x];
-          if (m->body_dofnum[b1]) res[m->body_dofadr[b1]+x] -= coef*dvec[e][x];
+        }
+        if (m->body_dofnum[b0]) {   // world -> dof frame
+          mji_mulMatTVec3(rl, d->xmat + 9*b0, rw);
+          for (int x = 0; x < 3; x++) {
+            res[m->body_dofadr[b0]+x] += rl[x];
+          }
+        }
+        if (m->body_dofnum[b1]) {
+          mji_mulMatTVec3(rl, d->xmat + 9*b1, rw);
+          for (int x = 0; x < 3; x++) {
+            res[m->body_dofadr[b1]+x] -= rl[x];
+          }
         }
       }
     }
@@ -1945,11 +1966,18 @@ int mjd_flexStiff_assemble(const mjModel* m, mjData* d, int* rownnz, int* rowadr
                 }
               }
             }
+            // blk is world-space but the destination dofs are the vertex bodies' own (possibly
+            // rotated) slide axes: blk_dof = R_bi^T * blk_world * R_bj, matching the force path
+            int bi = m->flex_vertbodyid[m->flex_vertadr[f] + vert[i]];
+            int bj = m->flex_vertbodyid[m->flex_vertadr[f] + vert[j]];
+            mjtNum tmp[9], blkd[9];
+            mji_mulMatMat3(tmp, blk, d->xmat + 9*bj);      // tmp  = blk * R_bj
+            mji_mulMatTMat3(blkd, d->xmat + 9*bi, tmp);    // blkd = R_bi^T * tmp
             int pos;
             FLEXSTIFF_BLOCK(si, sj, pos);
             for (int k = 0; k < 3; k++) {
               for (int c = 0; c < 3; c++) {
-                val[rowadr[vdof[si] + k] + 3*pos + c] += blk[3*k+c];
+                val[rowadr[vdof[si] + k] + 3*pos + c] += blkd[3*k+c];
               }
             }
           }

test/engine/engine_passive_test.cc

@@ -1093,6 +1093,94 @@ TEST_F(ElasticityTest, TrilinearParentBodyRotation) {
   EXPECT_LT(max_qacc, 1e4);
 }
 
+
+// A dim=2 flexcomp with stretch elasticity inside a parent body with a
+// non-identity quaternion. The implicit metric assembles the stretch
+// stiffness from world-space edge vectors, but the vertex bodies' slide
+// dofs live in the (rotated) parent frame. Without the R^T (.) R change
+// of basis the metric stops being the Jacobian of the passive force,
+// which shows up as a loss of rotational invariance and, at stiffnesses
+// the unrotated model handles comfortably, as divergence.
+TEST_F(ElasticityTest, StretchParentBodyRotation) {
+  static constexpr char rotated_xml[] = R"(
+  <mujoco>
+    <option gravity="0 0 0" integrator="implicitfast"
+            timestep="0.001" solver="CG"/>
+    <worldbody>
+      <body name="base" pos="0 0 0" quat="0.7071 0.7071 0 0">
+        <flexcomp type="grid" count="5 5 1" spacing=".05 .05 .05"
+                  dim="2" radius=".001" mass=".01" name="sheet">
+          <elasticity young="1e5" poisson="0" thickness="1e-3"
+                      elastic2d="stretch"/>
+          <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.001" solver="CG"/>
+    <worldbody>
+      <body name="base" pos="0 0 0">
+        <flexcomp type="grid" count="5 5 1" spacing=".05 .05 .05"
+                  dim="2" radius=".001" mass=".01" name="sheet">
+          <elasticity young="1e5" poisson="0" thickness="1e-3"
+                      elastic2d="stretch"/>
+          <contact selfcollide="none" internal="false"/>
+        </flexcomp>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+
+  char error[1024] = {0};
+
+  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);
+
+  // stretch one dof; the response is expressed in the parent frame in
+  // both models, so the passive force and the implicit acceleration must
+  // agree regardless of the parent's orientation
+  d_rot->qpos[0] = 1e-4;
+  d_non->qpos[0] = 1e-4;
+
+  mj_forward(m_rot.get(), d_rot.get());
+  mj_forward(m_non.get(), d_non.get());
+
+  EXPECT_LT(d_rot->qfrc_passive[0], 0)
+      << "expected a restoring force on the stretched dof";
+
+  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;
+  }
+
+  // qacc exercises the metric itself (the force is only its right-hand
+  // side): a metric in the wrong basis breaks this invariance even though
+  // the force above is already correct
+  for (int i = 0; i < m_rot->nv; i++) {
+    EXPECT_NEAR(d_rot->qacc[i], d_non->qacc[i], MjTol(1e-9, 1e-3))
+        << "rotated/non-rotated qacc mismatch at dof " << i;
+  }
+
+  // and the rotated model must integrate stably
+  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;
+    }
+    if (HasFatalFailure()) return;
+  }
+}
+
 }  // namespace
 }  // namespace mujoco
-