Apply unsymmetrized fluid derivatives to standalone free bodies in implicitfast.

PiperOrigin-RevId: 948899583
Change-Id: Icfb5a713f89a94e597c7607e9aa10a9e151dc2aa

src/engine/engine_derivative.c

@@ -818,31 +818,35 @@ void mjd_freeBias_vel(const mjModel* m, const mjData* d, int jnt, mjtNum B[36])
 }
 
 
+// return 1 if body is a standalone free body (single free joint, no children)
+mjtBool mj_isFreeBody(const mjModel* m, int body) {
+  // must have exactly one joint, of free type
+  if (m->body_jntnum[body] != 1 || m->jnt_type[m->body_jntadr[body]] != mjJNT_FREE) {
+    return false;
+  }
+
+  int adr = m->jnt_dofadr[m->body_jntadr[body]];
+
+  // must be a standalone 6-DOF tree with no children
+  if (m->tree_dofnum[m->dof_treeid[adr]] != 6 ||
+      m->body_subtreemass[body] != m->body_mass[body]) {
+    return false;
+  }
+
+  return true;
+}
+
+
 // 6x6 block A = M - h * (d qfrc_smooth / d qvel) for the free joint of a standalone body
 //   returns 1 and writes A if jnt is the free joint of a standalone awake body, 0 otherwise
 //   requires valid d->qDeriv rows for the block, computed with flg_bias = 0; the bias
-//   derivative excluded from qDeriv is added here via freeBias_vel_blocks
+//   derivative excluded from qDeriv is added here via mjd_freeBias_vel
 int mjd_freeMhat(const mjModel* m, const mjData* d, int jnt, mjtNum h, mjtNum A[36]) {
-  // must be a free joint
-  if (m->jnt_type[jnt] != mjJNT_FREE) {
-    return 0;
-  }
-
   int body = m->jnt_bodyid[jnt];
   int adr = m->jnt_dofadr[jnt];
-  int tree = m->dof_treeid[adr];
-  mjtNum mass = m->body_mass[body];
 
-  // must be a standalone 6-DOF tree with no children, awake
-  if (m->tree_dofnum[tree] != 6 ||
-      m->body_subtreemass[body] != mass ||
-      !d->tree_awake[tree]) {
-    return 0;
-  }
-
-  // D rows of a standalone free body are exactly the 6x6 block (D sparsity is tree-local);
-  // guard the gathers below against any violation of this invariant
-  if (m->D_rownnz[adr] != 6) {
+  // must be a standalone free body, awake
+  if (!mj_isFreeBody(m, body) || !d->tree_awake[m->dof_treeid[adr]]) {
     return 0;
   }
 
@@ -872,6 +876,7 @@ int mjd_freeMhat(const mjModel* m, const mjData* d, int jnt, mjtNum h, mjtNum A[
   mjtNum s[3];
   mji_sub3(s, d->xipos + 3*body, d->xpos + 3*body);
 
+  mjtNum mass = m->body_mass[body];
   mjtNum lin[9], rot[9];
   freeBias_vel_blocks(mass, d->xmat + 9*body, d->ximat + 9*body,
                       m->body_inertia + 3*body, s, d->qvel + adr + 3, lin, rot);
@@ -2482,8 +2487,8 @@ void mjd_ellipsoidFluid(const mjModel* m, mjData* d, int bodyid) {
 
     mjd_addedMassForces(B, lvel, m->opt.density, virtual_mass, virtual_inertia);
 
-    // make B symmetric if integrator is IMPLICITFAST
-    if (m->opt.integrator == mjINT_IMPLICITFAST) {
+    // make B symmetric if integrator is IMPLICITFAST, except for standalone free bodies
+    if (m->opt.integrator == mjINT_IMPLICITFAST && !mj_isFreeBody(m, bodyid)) {
       mju_symmetrize(B, B, 6);
     }
 

src/engine/engine_derivative.h

@@ -43,9 +43,11 @@ MJAPI void mjd_passive_vel(const mjModel* m, mjData* d);
 // subtract (d qfrc_bias / d qvel) from qDeriv (dense version)
 MJAPI void mjd_rne_vel_dense(const mjModel* m, mjData* d);
 
+// return 1 if body is a standalone free body: a free joint with no children
+mjtBool mj_isFreeBody(const mjModel* m, int body);
+
 // 6x6 block B = d qfrc_bias / d qvel for the free joint of a standalone body
-MJAPI void mjd_freeBias_vel(const mjModel* m, const mjData* d, int jnt,
-                            mjtNum B[36]);
+MJAPI void mjd_freeBias_vel(const mjModel* m, const mjData* d, int jnt, mjtNum B[36]);
 
 // 6x6 block A = M - h * (d qfrc_smooth / d qvel) for the free joint of a standalone body
 //   returns 1 and writes A if jnt is the free joint of a standalone awake body, 0 otherwise

src/engine/engine_forward.c

@@ -1377,7 +1377,7 @@ int mj_flexCG(const mjModel* m) {
 // fully implicit in velocity, possibly skipping factorization
 void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
   TM_START;
-  int nD = m->nD, nC = m->nC;
+  int nD = m->nD, nC = m->nC, njnt = m->njnt;
 
   mj_markStack(d);
   mjtNum* qfrc = mjSTACKALLOC(d, m->nv, mjtNum);
@@ -1427,6 +1427,18 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
 
       // set qH = M - dt*qDeriv
       mju_addScl(d->qH, d->M, d->qH, -m->opt.timestep, nC);
+
+      // standalone free bodies: reset qH block rows to M; their qDeriv rows may be asymmetric and
+      // are handled by the local LU solve; we reset to keep LTL well-defined
+      for (int j=0; j < njnt; j++) {
+        if (m->jnt_type[j] != mjJNT_FREE || !mj_isFreeBody(m, m->jnt_bodyid[j])) {
+          continue;
+        }
+        int adr = m->jnt_dofadr[j];
+        for (int r=0; r < 6; r++) {
+          mju_copy(d->qH + m->M_rowadr[adr+r], d->M + m->M_rowadr[adr+r], m->M_rownnz[adr+r]);
+        }
+      }
     } else {
       mjERROR("integrator must be implicit or implicitfast");
     }

test/engine/engine_forward_test.cc

@@ -16,7 +16,6 @@
 
 #include "src/engine/engine_forward.h"
 
-#include <algorithm>
 #include <array>
 #include <cmath>
 #include <cstdlib>
@@ -68,7 +67,6 @@ using ::testing::Pointwise;
 
 using ::testing::_;
 using ::testing::Gt;
-using ::testing::HasSubstr;
 using ::testing::Ne;
 using ::testing::NotNull;
 
@@ -467,7 +465,8 @@ TEST_F(ImplicitIntegratorTest, EnergyConservation) {
 // free-body local solve: implicitfast matches implicit exactly for a standalone
 // free body
 TEST_F(ImplicitIntegratorTest, FreeBodyMatchesImplicit) {
-  static constexpr char xml[] = R"(
+  // damped free body in vacuum
+  static constexpr char xml1[] = R"(
   <mujoco>
     <option timestep="0.005"/>
     <worldbody>
@@ -479,36 +478,54 @@ TEST_F(ImplicitIntegratorTest, FreeBodyMatchesImplicit) {
   </mujoco>
   )";
 
-  char error[1024];
-  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
-  ASSERT_THAT(model.get(), NotNull()) << error;
-  MjDataPtr d1 = MakeData(model);
-  MjDataPtr d2 = MakeData(model);
-  mjModel* m = model.get();
+  // free body in fluid with wind, ellipsoid fluid model (asymmetric lift
+  // derivatives)
+  static constexpr char xml2[] = R"(
+  <mujoco>
+    <option timestep="0.005" density="1.2" viscosity="0.002" wind="1 2 3"/>
+    <worldbody>
+      <body pos="0.1 -0.2 0.5" euler="20 -30 40">
+        <joint type="free"/>
+        <geom type="ellipsoid" size=".1 .2 .3" mass="2" pos=".04 -.02 .03"
+              fluidshape="ellipsoid"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
 
-  // tumbling initial velocity
-  mj_resetData(m, d1.get());
-  d1->qvel[3] = 5;
-  d1->qvel[4] = -3;
-  d1->qvel[5] = 2;
+  int xml_idx = 1;
+  for (auto xml : {xml1, xml2}) {
+    SCOPED_TRACE(testing::Message() << "XML case " << xml_idx++);
+    char error[1024];
+    MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
+    ASSERT_THAT(model.get(), NotNull()) << error;
+    MjDataPtr d1 = MakeData(model);
+    MjDataPtr d2 = MakeData(model);
+    mjModel* m = model.get();
 
-  // step both integrators from identical states, re-synchronizing each step
-  // to avoid chaotic divergence of tumbling trajectories
-  int nstate = mj_stateSize(m, mjSTATE_INTEGRATION);
-  std::vector<mjtNum> state(nstate);
-  mjtNum tol = MjTol(1e-14, 1e-6);
-  for (int i = 0; i < 50; i++) {
-    mj_getState(m, d1.get(), state.data(), mjSTATE_INTEGRATION);
-    mj_setState(m, d2.get(), state.data(), mjSTATE_INTEGRATION);
+    // tumbling initial velocity
+    mj_resetData(m, d1.get());
+    d1->qvel[3] = 5;
+    d1->qvel[4] = -3;
+    d1->qvel[5] = 2;
 
-    m->opt.integrator = mjINT_IMPLICITFAST;
-    mj_step(m, d1.get());
-    m->opt.integrator = mjINT_IMPLICIT;
-    mj_step(m, d2.get());
+    // step both integrators from identical states, re-synchronizing each step
+    // to avoid chaotic divergence of tumbling trajectories
+    int nstate = mj_stateSize(m, mjSTATE_INTEGRATION);
+    std::vector<mjtNum> state(nstate);
+    for (int i = 0; i < 50; i++) {
+      mj_getState(m, d1.get(), state.data(), mjSTATE_INTEGRATION);
+      mj_setState(m, d2.get(), state.data(), mjSTATE_INTEGRATION);
 
-    for (int k = 0; k < m->nv; k++) {
-      EXPECT_NEAR(d1->qvel[k], d2->qvel[k], tol)
-          << "step " << i << " dof " << k;
+      m->opt.integrator = mjINT_IMPLICITFAST;
+      mj_step(m, d1.get());
+      m->opt.integrator = mjINT_IMPLICIT;
+      mj_step(m, d2.get());
+
+      for (int k = 0; k < m->nv; k++) {
+        EXPECT_NEAR(d1->qvel[k], d2->qvel[k], MjTol(1e-14, 1e-6))
+            << "step " << i << " dof " << k;
+      }
     }
   }
 }