Restrict midpoint integration to unconstrained free bodies in implicitfast

PiperOrigin-RevId: 908750768
Change-Id: I9a45a160ac757cc82bfe54871609956769988369

doc/changelog.rst

@@ -8,6 +8,11 @@ Upcoming version (not yet released)
 - Added island support for the :ref:`PGS solver<soAlgorithms>`.
 - Added support for :ref:`elastic2d<body-flexcomp-elastic2d>` for trilinear and quadratic flex
   :ref:`dofs<body-flexcomp-dof>`.
+- :ref:`Midpoint integration<geMidpoint>` is now restricted to the ``implicitfast``
+  :ref:`integrator<geIntegrators>` and is disabled when fluid forces are active
+  (nonzero :ref:`density<option-density>` or :ref:`viscosity<option-viscosity>`).
+  Midpoint integration treats external forces as zero-order-hold constants, which causes
+  energy gain in the presence of contacts and in fluid media.
 
 Python
 ^^^^^^

doc/computation/index.rst

@@ -575,7 +575,7 @@ Solving for :math:`v_{t+h}`, we obtain the implicit-in-velocity update
 
 .. _geMidpoint:
 
-Midpoint integration for free bodies
+Midpoint integration for free bodies in vacuum
    The implicit-in-velocity update :eq:`eq_implicit_update` treats the acceleration as a function of velocity and
    linearizes. While effective for damping-like forces, it is sub-optimal for rotational dynamics, where
    Coriolis and gyroscopic forces are *quadratic* in angular velocity. For this case, a better approach is to directly
@@ -606,7 +606,9 @@ Midpoint integration for free bodies
    Casimir function of the `Lie-Poisson <https://en.wikipedia.org/wiki/Poisson_bracket>`__ structure, the midpoint
    method is a symmetric (time-reversible) and second-order accurate *Poisson integrator*.
 
-   **Eligibility.** Midpoint integration is only applied to free bodies with no child bodies.
+   **Eligibility.** Midpoint integration is only applied when using the ``implicitfast`` integrator, to
+   free bodies with no child bodies, and only when the medium has zero :ref:`density<option-density>` and
+   :ref:`viscosity<option-viscosity>`.
 
    **Performance.** While the midpoint method carries computational overhead, we've found it to be
    negligible compared to the rest of the pipeline, on the order of 1% in the worst case.
@@ -652,9 +654,8 @@ Fast implicit-in-velocity (``implicitfast``)
    scenarios which are not common and already well-handled by the Runge-Kutta integrator (see below). Because the RNE
    derivatives are also the main source of asymmetry of :math:`D`, by dropping them and symmetrizing, we can use the
    faster :math:`L^TL` rather than :math:`LU` decomposition.
-
-Both ``implicit`` and ``implicitfast`` apply :ref:`midpoint integration<geMidpoint>` to eligible free bodies,
-providing exact energy conservation for spinning objects at negligible additional cost.
+   The ``implicitfast`` integrator applies :ref:`midpoint integration<geMidpoint>` to eligible free bodies in vacuum,
+   providing exact energy conservation for spinning objects at negligible additional cost.
 
 4th-order Runge-Kutta (``RK4``)
    One advantage of our continuous-time formulation is that we can use higher order integrators such as Runge-Kutta or
@@ -688,10 +689,11 @@ providing exact energy conservation for spinning objects at negligible additiona
      increased stability, and is therefore a strict improvement. It is the recommended integrator for most models.
     **implicit**:
      The benefit over ``implicitfast`` is the implicit integration of Coriolis and centripetal forces for *coupled*
-     rotational systems such as multi-link pendula. Both ``implicitfast`` and ``implicit`` apply :ref:`midpoint
-     integration<geMidpoint>` to eligible free bodies with no children, for example
-     `gyroscopic.xml <../_static/gyroscopic.xml>`__ shows an ellipsoid rolling on an
-     inclined plane; both ``implicitfast`` and ``implicit`` handle this case well, while ``Euler`` quickly diverges.
+     rotational systems such as multi-link pendula. Note that ``implicit`` does not apply :ref:`midpoint
+     integration<geMidpoint>` (only ``implicitfast`` does), but its RNE derivatives provide comparable stability
+     for free-body rotation. For example, `gyroscopic.xml <../_static/gyroscopic.xml>`__ shows an ellipsoid rolling
+     on an inclined plane; both ``implicitfast`` and ``implicit`` handle this case well, while ``Euler`` quickly
+     diverges.
     **RK4**:
      This integrator is best for systems which are energy conserving, or almost energy-conserving. `pendulum.xml
      <../_static/pendulum.xml>`__ shows a complicated pendulum mechanism which diverges quickly using ``Euler`` or

src/engine/engine_forward.c

@@ -1597,16 +1597,75 @@ static void flexInterp_solve(const mjModel* m, mjData* d, const FlexInterpContex
 
 
 // return 1 if free joint is eligible for midpoint quaternion integration:
-//   standalone 6-DOF tree with no children
-static int midpoint_eligible(const mjModel* m, int jnt) {
-  if (m->jnt_type[jnt] == mjJNT_FREE) {
-    int body = m->jnt_bodyid[jnt];
-    int treeid = m->dof_treeid[m->jnt_dofadr[jnt]];
-    return m->tree_dofnum[treeid] == 6 &&
-           m->body_subtreemass[body] == m->body_mass[body];
+//   standalone 6-DOF tree with no children, awake, and unconstrained
+static int midpoint_eligible(const mjModel* m, const mjData* d, int jnt) {
+  if (m->jnt_type[jnt] != mjJNT_FREE) {
+    return 0;
   }
 
-  return 0;
+  int body = m->jnt_bodyid[jnt];
+  int adr = m->jnt_dofadr[jnt];
+  int tree = m->dof_treeid[adr];
+
+  // must be standalone 6-DOF tree with no children
+  if (m->tree_dofnum[tree] != 6 ||
+      m->body_subtreemass[body] != m->body_mass[body]) {
+    return 0;
+  }
+
+  // must be awake
+  if (!d->tree_awake[tree]) {
+    return 0;
+  }
+
+  // must be unconstrained
+  if (d->nefc) {
+    // islands enabled: O(1) lookup
+    if (!mjDISABLED(mjDSBL_ISLAND)) {
+      if (d->dof_island[adr] >= 0) {
+        return 0;
+      }
+    }
+
+    // islands disabled: check if any constraint involves this tree
+    else {
+      for (int c=0; c < d->nefc; c++) {
+        int type = d->efc_type[c];
+        int id = d->efc_id[c];
+
+        // contact: check if either geom belongs to this body
+        if (type == mjCNSTR_CONTACT_FRICTIONLESS ||
+            type == mjCNSTR_CONTACT_PYRAMIDAL ||
+            type == mjCNSTR_CONTACT_ELLIPTIC) {
+          int g1 = d->contact[id].geom[0];
+          int g2 = d->contact[id].geom[1];
+          if (g1 >= 0 && m->geom_bodyid[g1] == body) return 0;
+          if (g2 >= 0 && m->geom_bodyid[g2] == body) return 0;
+        }
+
+        // connect or weld: check if either body is this body
+        else if (type == mjCNSTR_EQUALITY &&
+                 (m->eq_type[id] == mjEQ_CONNECT || m->eq_type[id] == mjEQ_WELD)) {
+          int b1 = m->eq_obj1id[id];
+          int b2 = m->eq_obj2id[id];
+          if (m->eq_objtype[id] == mjOBJ_SITE) {
+            b1 = m->site_bodyid[b1];
+            b2 = m->site_bodyid[b2];
+          }
+          if (b1 == body || b2 == body) return 0;
+        }
+
+        // tendon limit or friction: check first two trees
+        else if (type == mjCNSTR_LIMIT_TENDON || type == mjCNSTR_FRICTION_TENDON) {
+          if (m->tendon_treeid[2*id] == tree ||
+              m->tendon_treeid[2*id+1] == tree) return 0;
+        }
+      }
+    }
+  }
+
+  // otherwise eligible
+  return 1;
 }
 
 
@@ -1980,11 +2039,12 @@ void mj_implicitSkip(const mjModel* m, mjData* d, int skipfactor) {
   // count and list joints of free bodies eligible for midpoint integration
   int nfree = 0;
   int* free_jntid = NULL;
-  if (!mjENABLED(mjENBL_INVDISCRETE)) {
+  if (!mjENABLED(mjENBL_INVDISCRETE) &&
+      m->opt.integrator == mjINT_IMPLICITFAST &&
+      m->opt.density == 0 && m->opt.viscosity == 0) {
     free_jntid = mjSTACKALLOC(d, m->njnt, int);
     for (int j=0; j < m->njnt; j++) {
-      // add to list if eligible and awake
-      if (midpoint_eligible(m, j) && d->tree_awake[m->dof_treeid[m->jnt_dofadr[j]]]) {
+      if (midpoint_eligible(m, d, j)) {
         free_jntid[nfree++] = j;
       }
     }

test/engine/engine_derivative_test.cc

@@ -1745,79 +1745,5 @@ TEST_F(DerivativeTest, FlexInterpDerivativesDeformed) {
   mj_deleteModel(model);
 }
 
-TEST_F(DerivativeTest, MidpointFluidAccuracy) {
-  const std::string xml_path =
-      GetTestDataFilePath(kTumblingThinObjectEllipsoidPath);
-  char error[1024];
-  mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
-  ASSERT_THAT(m, NotNull()) << error;
-
-  mjtNum dt_small = 1e-4;
-  mjtNum dt_large = m->opt.timestep;  // 2e-3, the default
-  mjtNum duration = 0.5;
-
-  mjData* d_ref = mj_makeData(m);
-  mjData* d_midpoint = mj_makeData(m);
-  mjData* d_nomidpoint = mj_makeData(m);
-
-  // give initial angular velocity for tumbling
-  mj_resetData(m, d_ref);
-  mj_resetData(m, d_midpoint);
-  mj_resetData(m, d_nomidpoint);
-  d_ref->qvel[3] = 5;
-  d_ref->qvel[4] = 3;
-  d_ref->qvel[5] = 1;
-  d_midpoint->qvel[3] = 5;
-  d_midpoint->qvel[4] = 3;
-  d_midpoint->qvel[5] = 1;
-  d_nomidpoint->qvel[3] = 5;
-  d_nomidpoint->qvel[4] = 3;
-  d_nomidpoint->qvel[5] = 1;
-
-  int nsteps_large = static_cast<int>(duration / dt_large);
-  int substeps = static_cast<int>(dt_large / dt_small);
-
-  mjtNum error_midpoint = 0;
-  mjtNum error_nomidpoint = 0;
-
-  for (int i = 0; i < nsteps_large; i++) {
-    // reference: RK4 at small timestep
-    m->opt.integrator = mjINT_RK4;
-    m->opt.timestep = dt_small;
-    m->opt.enableflags &= ~mjENBL_INVDISCRETE;
-    for (int j = 0; j < substeps; j++) {
-      mj_step(m, d_ref);
-    }
-
-    // implicit with midpoint (default)
-    m->opt.integrator = mjINT_IMPLICIT;
-    m->opt.timestep = dt_large;
-    m->opt.enableflags &= ~mjENBL_INVDISCRETE;
-    mj_step(m, d_midpoint);
-
-    // implicit without midpoint
-    m->opt.enableflags |= mjENBL_INVDISCRETE;
-    mj_step(m, d_nomidpoint);
-
-    // accumulate position errors
-    for (int k = 0; k < 7; k++) {
-      mjtNum diff_mid = d_ref->qpos[k] - d_midpoint->qpos[k];
-      mjtNum diff_nomid = d_ref->qpos[k] - d_nomidpoint->qpos[k];
-      error_midpoint += diff_mid * diff_mid;
-      error_nomidpoint += diff_nomid * diff_nomid;
-    }
-  }
-
-  // expect midpoint to be more accurate
-  EXPECT_LT(error_midpoint, error_nomidpoint)
-      << "implicit midpoint should be more accurate than implicit without "
-      << "midpoint for a free body with fluid forces";
-
-  mj_deleteData(d_nomidpoint);
-  mj_deleteData(d_midpoint);
-  mj_deleteData(d_ref);
-  mj_deleteModel(m);
-}
-
 }  // namespace
 }  // namespace mujoco

test/engine/engine_forward_test.cc

@@ -783,6 +783,140 @@ TEST_F(ImplicitIntegratorTest, MidpointFullNewtonConvergence) {
   EXPECT_LT((mjtNum)total_iter / ncases, 3.0);
 }
 
+// verify midpoint eligibility: compare with/without invdiscrete
+//   if trajectories differ, midpoint was applied
+//   if trajectories match, midpoint was skipped
+TEST_F(ImplicitIntegratorTest, MidpointEligibility) {
+  // free body with asymmetric inertia, optionally near a plane
+  static constexpr char xml[] = R"(
+  <mujoco>
+    <option integrator="implicitfast" timestep="0.01">
+      <flag energy="enable"/>
+    </option>
+    <worldbody>
+      <geom type="plane" size="5 5 0.1"/>
+      <body name="free" pos="0 0 2">
+        <freejoint/>
+        <geom type="ellipsoid" size="0.3 0.2 0.1" mass="1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+
+  char error[1024];
+  mjModel* m = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(m, NotNull()) << error;
+  mjData* d1 = mj_makeData(m);
+  mjData* d2 = mj_makeData(m);
+  int nsteps = 50;
+
+  auto spin_and_compare = [&](const char* label,
+                              bool expect_midpoint) {
+    mj_resetData(m, d1);
+    mj_resetData(m, d2);
+    d1->qvel[3] = d2->qvel[3] = 5;
+    d1->qvel[4] = d2->qvel[4] = 3;
+    d1->qvel[5] = d2->qvel[5] = 1;
+
+    // d1: midpoint enabled (default)
+    m->opt.enableflags &= ~mjENBL_INVDISCRETE;
+    for (int i = 0; i < nsteps; i++) mj_step(m, d1);
+
+    // d2: midpoint disabled
+    m->opt.enableflags |= mjENBL_INVDISCRETE;
+    mj_resetData(m, d2);
+    d2->qvel[3] = 5; d2->qvel[4] = 3; d2->qvel[5] = 1;
+    for (int i = 0; i < nsteps; i++) mj_step(m, d2);
+    m->opt.enableflags &= ~mjENBL_INVDISCRETE;
+
+    // compare angular velocities
+    mjtNum diff = 0;
+    for (int k = 3; k < 6; k++) {
+      mjtNum d = d1->qvel[k] - d2->qvel[k];
+      diff += d * d;
+    }
+    if (expect_midpoint) {
+      EXPECT_GT(diff, 1e-6)
+          << label << ": expected midpoint to be applied";
+    } else {
+      EXPECT_LT(diff, 1e-20)
+          << label << ": expected midpoint to be skipped";
+    }
+  };
+
+  // case 1: free body in vacuum, implicitfast -> midpoint applied
+  m->opt.integrator = mjINT_IMPLICITFAST;
+  m->opt.density = 0;
+  m->opt.viscosity = 0;
+  spin_and_compare("vacuum+implicitfast", true);
+
+  // case 2: implicit integrator -> midpoint NOT applied
+  m->opt.integrator = mjINT_IMPLICIT;
+  spin_and_compare("vacuum+implicit", false);
+
+  // case 3: fluid (nonzero density) -> midpoint NOT applied
+  m->opt.integrator = mjINT_IMPLICITFAST;
+  m->opt.density = 1.2;
+  spin_and_compare("fluid+implicitfast", false);
+  m->opt.density = 0;
+
+  // case 4: fluid (nonzero viscosity) -> midpoint NOT applied
+  m->opt.viscosity = 0.001;
+  spin_and_compare("viscosity+implicitfast", false);
+  m->opt.viscosity = 0;
+
+  // case 5: body with active contacts -> midpoint NOT applied
+  // test both island-enabled and island-disabled branches
+  for (int disable_island = 0; disable_island < 2; disable_island++) {
+    m->opt.integrator = mjINT_IMPLICITFAST;
+    if (disable_island) {
+      m->opt.disableflags |= mjDSBL_ISLAND;
+    } else {
+      m->opt.disableflags &= ~mjDSBL_ISLAND;
+    }
+
+    mj_resetData(m, d1);
+    mj_resetData(m, d2);
+    d1->qpos[2] = d2->qpos[2] = 0.05;
+    d1->qvel[3] = d2->qvel[3] = 5;
+    d1->qvel[4] = d2->qvel[4] = 3;
+    d1->qvel[5] = d2->qvel[5] = 1;
+
+    // verify contacts are active
+    mj_forward(m, d1);
+    ASSERT_GT(d1->ncon, 0) << "body should be in contact with the plane";
+
+    // single step with midpoint enabled
+    mj_resetData(m, d1);
+    d1->qpos[2] = 0.05;
+    d1->qvel[3] = 5; d1->qvel[4] = 3; d1->qvel[5] = 1;
+    m->opt.enableflags &= ~mjENBL_INVDISCRETE;
+    mj_step(m, d1);
+
+    // single step with midpoint disabled
+    mj_resetData(m, d2);
+    d2->qpos[2] = 0.05;
+    d2->qvel[3] = 5; d2->qvel[4] = 3; d2->qvel[5] = 1;
+    m->opt.enableflags |= mjENBL_INVDISCRETE;
+    mj_step(m, d2);
+    m->opt.enableflags &= ~mjENBL_INVDISCRETE;
+
+    mjtNum diff = 0;
+    for (int k = 0; k < m->nv; k++) {
+      mjtNum d = d1->qvel[k] - d2->qvel[k];
+      diff += d * d;
+    }
+    EXPECT_LT(diff, 1e-20)
+        << "contact (island " << (disable_island ? "disabled" : "enabled")
+        << "): expected midpoint to be skipped";
+  }
+  m->opt.disableflags &= ~mjDSBL_ISLAND;
+
+  mj_deleteData(d2);
+  mj_deleteData(d1);
+  mj_deleteModel(m);
+}
+
 TEST_F(ForwardTest, ControlClamping) {
   static constexpr char xml[] = R"(
   <mujoco>