Add Jdot correction for connect and weld constaints

Measured reduction in constraint violations before/after this change: | Model | Correction ON (Avg Viol) | Correction OFF (Avg Viol) | Reduction | | :--- | :--- | :--- | :--- | | `jdotv_connect_2d.xml` | 3.959e-4 | 1.699e-3 | **76.7%** | | `jdotv_connect_3d.xml` | 1.399e-3 | 5.493e-3 | **74.5%** | | `jdotv_weld_3d.xml` | 9.472e-3 | 1.148e-2 | **17.5%** | PiperOrigin-RevId: 899137525 Change-Id: Ic3e33764ebd64239bab916289c23d80c3da0b51b
2026-04-13 12:47:17 -07:00
parent 0c337799bd
commit 412cee2059
13 changed files with 481 additions and 43 deletions
@@ -29,9 +29,13 @@ General
  may unify the linear and higher-order coefficients into a single array.
 - Added :ref:`midpoint integration<geMidpoint>` for standalone free bodies in ``implicit`` and ``implicitfast``
  :ref:`integrators<geIntegrators>`. This applies the implicit midpoint rule to the rotational dynamics of free bodies
-  with no children, exactly conserving kinetic energy and angular momentum in the absence of external torques. The
+  with no children, conserving kinetic energy to machine precision in the absence of external torques. The
  :ref:`invdiscrete<option-flag-invdiscrete>` flag now also disables midpoint integration, providing an opt-out
  mechanism.
+- Added the centripetal/Coriolis acceleration term :math:`\dot{J}v` to the constraint solver bias for
+  :ref:`connect<equality-connect>` and :ref:`weld<equality-weld>` equality constaints. This significantly improves the
+  stability of constrained mechanisms like four-bar linkages. See :ref:`Dual problem<soDual>` for details.
+
 - Introduced :ref:`mjpEncoder`, the counterpart to :ref:`mjpDecoder` for encoding of :ref:`mjSpec` and :ref:`mjModel`
  into :ref:`mjResource`.

@@ -1296,10 +1296,13 @@ when

 This key identity is essentially Newton's second law projected in constraint space. It is derived by moving the term
 :math:`c` in the equations of motion :eq:`eq:motion` to the right hand side, multiplying by :math:`J M^{-1}` from the
-left, adding :math:`\dot{J} v` to both sides, and substituting the above definitions of :math:`A, \au, \ac`. In terms of
-implementation, we do not actually compute the acceleration term :math:`\dot{J} v`. This is because our optimization
-problems depend on differences of constraint-space accelerations, and so this term would cancel out even if we were to
-compute it.
+left, adding :math:`\dot{J} v` to both sides, and substituting the above definitions of :math:`A, \au, \ac`. Computing
+:math:`\dot{J} v` requires differentiating the constraint Jacobian with respect to time, which is nontrivial.
+Although this term cancels in the identity :eq:`eq:identity` and so does not affect the forward-inverse comparison, its
+omission in the forward dynamics introduces a velocity-dependent bias for any constraint whose Jacobian varies with
+configuration. We compute this term for equality constraints (connect and weld) where Jacobian differentiation
+is tractable. For contacts, the term remains omitted due to the complexity of differentiating the contact frame through
+the collision pipeline.

 Note that the quadratic term in the inverse problem is weighted by :math:`R` instead of :math:`A+R`. This is the key
 structural insight: the :math:`A` matrix cancels entirely, leaving only :math:`R` in the quadratic term. Two
@@ -60,6 +60,9 @@ class ConstraintTest(parameterized.TestCase):
    # sample a mix of active/inactive constraints at different timesteps
    for key in range(3):
      mujoco.mj_resetDataKeyframe(m, d, key)
+      # scale down velocities to minimize Jdotv effect (not in MJX)
+      # TODO(team): remove this change when mjx supports this feature
+      d.qvel[:] *= 1e-2
      if rand_eq_active:
        d.eq_active[:] = np.random.randint(0, 2, size=m.neq)
      mujoco.mj_forward(m, d)
@@ -49,6 +49,9 @@ class ForwardTest(absltest.TestCase):
    d.xfrc_applied[0, 2] = 0.1  # torque
    d.xfrc_applied[1, 4] = 0.3  # linear force
    mujoco.mj_step(m, d, 20)  # get some dynamics going
+    # scale down velocities to minimize Jdotv effect (not in MJX)
+    # TODO(team): remove this change when mjx supports this feature
+    d.qvel[:] *= 1e-2
    mujoco.mj_forward(m, d)

    mx = mjx.put_model(m)
@@ -92,6 +95,9 @@ class ForwardTest(absltest.TestCase):
    d.xfrc_applied[0, 2] = 0.1  # torque
    d.xfrc_applied[1, 4] = 0.3  # linear force
    mujoco.mj_step(m, d, 20)  # get some dynamics going
+    # scale down velocities to minimize Jdotv effect (not in MJX)
+    # TODO(team): remove this change when mjx supports this feature
+    d.qvel[:] *= 1e-2

    dx = jax.jit(mjx.step)(mjx.put_model(m), mjx.put_data(m, d))
    mujoco.mj_step(m, d)
@@ -612,6 +612,39 @@ void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum*
 }


+// compute global anchor points for connect/weld equality constraints
+static void mj_equalityAnchors(const mjModel* m, const mjData* d, int eq_id,
+                               mjtNum pos1[3], mjtNum pos2[3],
+                               int* body1, int* body2) {
+  mjtEq type = (mjtEq) m->eq_type[eq_id];
+  int obj1 = m->eq_obj1id[eq_id];
+  int obj2 = m->eq_obj2id[eq_id];
+
+  if (m->eq_objtype[eq_id] == mjOBJ_BODY) {
+    const mjtNum* data = m->eq_data + mjNEQDATA*eq_id;
+    if (type == mjEQ_CONNECT) {
+      mju_mulMatVec3(pos1, d->xmat + 9*obj1, data);
+      mju_addTo3(pos1, d->xpos + 3*obj1);
+      mju_mulMatVec3(pos2, d->xmat + 9*obj2, data + 3);
+      mju_addTo3(pos2, d->xpos + 3*obj2);
+    } else {
+      // weld uses data+3*(1-j) for anchor
+      mju_mulMatVec3(pos1, d->xmat + 9*obj1, data + 3);
+      mju_addTo3(pos1, d->xpos + 3*obj1);
+      mju_mulMatVec3(pos2, d->xmat + 9*obj2, data);
+      mju_addTo3(pos2, d->xpos + 3*obj2);
+    }
+    *body1 = obj1;
+    *body2 = obj2;
+  } else {
+    mju_copy3(pos1, d->site_xpos + 3*obj1);
+    mju_copy3(pos2, d->site_xpos + 3*obj2);
+    *body1 = m->site_bodyid[obj1];
+    *body2 = m->site_bodyid[obj2];
+  }
+}
+
+
 //--------------------- instantiate constraints by type --------------------------------------------

 // equality constraints
@@ -670,21 +703,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
    switch ((mjtEq) m->eq_type[i]) {
    case mjEQ_CONNECT:              // connect bodies with ball joint
      // find global points, body semantic
-      if (m->eq_objtype[i] == mjOBJ_BODY) {
-        for (int j=0; j < 2; j++) {
-          mju_mulMatVec3(pos[j], d->xmat + 9*id[j], data + 3*j);
-          mju_addTo3(pos[j], d->xpos + 3*id[j]);
-          body_id[j] = id[j];
-        }
-      }
-
-      // find global points, site semantic
-      else {
-        for (int j=0; j < 2; j++) {
-          mju_copy3(pos[j], d->site_xpos + 3*id[j]);
-          body_id[j] = m->site_bodyid[id[j]];
-        }
-      }
+      mj_equalityAnchors(m, d, i, pos[0], pos[1], body_id, body_id + 1);

      // compute position error
      mju_sub3(cpos, pos[0], pos[1]);
@@ -702,22 +721,7 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {

    case mjEQ_WELD:                 // fix relative position and orientation
      // find global points, body semantic
-      if (m->eq_objtype[i] == mjOBJ_BODY) {
-        for (int j=0; j < 2; j++) {
-          mjtNum* anchor = data + 3*(1-j);
-          mju_mulMatVec3(pos[j], d->xmat + 9*id[j], anchor);
-          mju_addTo3(pos[j], d->xpos + 3*id[j]);
-          body_id[j] = id[j];
-        }
-      }
-
-      // find global points, site semantic
-      else {
-        for (int j=0; j < 2; j++) {
-          mju_copy3(pos[j], d->site_xpos + 3*id[j]);
-          body_id[j] = m->site_bodyid[id[j]];
-        }
-      }
+      mj_equalityAnchors(m, d, i, pos[0], pos[1], body_id, body_id + 1);

      // compute position error
      mju_sub3(cpos, pos[0], pos[1]);
@@ -1133,6 +1137,208 @@ void mj_instantiateEquality(const mjModel* m, mjData* d) {
  mj_freeStack(d);
 }

+// subtract Jdot*v correction from result vector for equality constraints
+void mj_Jdotv(const mjModel* m, mjData* d, mjtNum* result) {
+  int nv = m->nv, ne = d->ne;
+
+  // nothing to do
+  if (!ne || !nv) {
+    return;
+  }
+
+  int issparse = mj_isSparse(m);
+
+  mj_markStack(d);
+
+  // allocate scratch for jacDot matrices (translational and rotational)
+  int* chain = issparse ? mjSTACKALLOC(d, nv, int) : NULL;
+  mjtNum* jacdot1 = NULL;
+  mjtNum* jacdot2 = NULL;
+  mjtNum* jacrdot1 = NULL;
+  mjtNum* jacrdot2 = NULL;
+
+  // iterate over equality constraint efc rows
+  int row = 0;
+  while (row < ne) {
+    int eq_id = d->efc_id[row];
+    mjtEq type = (mjtEq) m->eq_type[eq_id];
+
+    // connect or weld: compute Jdot*v for translational part
+    if (type == mjEQ_CONNECT || type == mjEQ_WELD) {
+      mjtNum* data = m->eq_data + mjNEQDATA*eq_id;
+
+      // allocate translational scratch on first connect or weld
+      if (!jacdot1) {
+        jacdot1 = mjSTACKALLOC(d, 3*nv, mjtNum);
+        jacdot2 = mjSTACKALLOC(d, 3*nv, mjtNum);
+      }
+
+      // allocate rotational scratch on first weld
+      if (type == mjEQ_WELD && !jacrdot1) {
+        jacrdot1 = mjSTACKALLOC(d, 3*nv, mjtNum);
+        jacrdot2 = mjSTACKALLOC(d, 3*nv, mjtNum);
+      }
+
+      // compute global anchor points and body ids
+      int obj1 = m->eq_obj1id[eq_id];
+      int obj2 = m->eq_obj2id[eq_id];
+      mjtNum pos1[3], pos2[3];
+      int body1, body2;
+      mj_equalityAnchors(m, d, eq_id, pos1, pos2, &body1, &body2);
+
+      // compute jacDot*v for each body point
+      mjtNum jdv1[3], jdv2[3];
+      mjtNum jrdv1[3] = {0}, jrdv2[3] = {0};
+      if (issparse) {
+        // get merged chain for the two bodies
+        int NV = mj_mergeChain(m, chain, body1, body2, /*flg_skipcommon=*/0);
+
+        if (NV) {
+          // sparse: translational and rotational
+          mjtNum* jacr1 = (type == mjEQ_WELD) ? jacrdot1 : NULL;
+          mjtNum* jacr2 = (type == mjEQ_WELD) ? jacrdot2 : NULL;
+          mj_jacDotSparse(m, d, jacdot1, jacr1, pos1, body1, NV, chain);
+          mj_jacDotSparse(m, d, jacdot2, jacr2, pos2, body2, NV, chain);
+
+          // translational jdv = jacDot * qvel
+          mju_dotSparseX3(jdv1, jdv1+1, jdv1+2, jacdot1, jacdot1+NV, jacdot1+2*NV,
+                          d->qvel, NV, chain);
+          mju_dotSparseX3(jdv2, jdv2+1, jdv2+2, jacdot2, jacdot2+NV, jacdot2+2*NV,
+                          d->qvel, NV, chain);
+
+          // rotational jdv for welds
+          if (type == mjEQ_WELD) {
+            mju_dotSparseX3(jrdv1, jrdv1+1, jrdv1+2, jacrdot1, jacrdot1+NV, jacrdot1+2*NV,
+                            d->qvel, NV, chain);
+            mju_dotSparseX3(jrdv2, jrdv2+1, jrdv2+2, jacrdot2, jacrdot2+NV, jacrdot2+2*NV,
+                            d->qvel, NV, chain);
+          }
+        } else {
+          mju_zero3(jdv1);
+          mju_zero3(jdv2);
+        }
+      } else {
+        // dense: translational and rotational
+        mjtNum* jacr1 = (type == mjEQ_WELD) ? jacrdot1 : NULL;
+        mjtNum* jacr2 = (type == mjEQ_WELD) ? jacrdot2 : NULL;
+        mj_jacDot(m, d, jacdot1, jacr1, pos1, body1);
+        mj_jacDot(m, d, jacdot2, jacr2, pos2, body2);
+
+        // translational jdv = jacDot * qvel
+        mju_mulMatVec(jdv1, jacdot1, d->qvel, 3, nv);
+        mju_mulMatVec(jdv2, jacdot2, d->qvel, 3, nv);
+
+        // rotational jdv for welds
+        if (type == mjEQ_WELD) {
+          mju_mulMatVec(jrdv1, jacrdot1, d->qvel, 3, nv);
+          mju_mulMatVec(jrdv2, jacrdot2, d->qvel, 3, nv);
+        }
+      }
+
+      // subtract translational Jdot*v
+      result[row+0] -= jdv1[0] - jdv2[0];
+      result[row+1] -= jdv1[1] - jdv2[1];
+      result[row+2] -= jdv1[2] - jdv2[2];
+
+      // advance past translational rows
+      row += 3;
+
+      // weld: compute rotational Jdot*v
+      if (type == mjEQ_WELD) {
+        mjtNum torquescale = data[10];
+
+        // get body quaternions and relpose, following mj_instantiateEquality
+        mjtNum q0r[4], negq1[4];   // q0r = q0*relpose, negq1 = neg(q1)
+        if (m->eq_objtype[eq_id] == mjOBJ_BODY) {
+          mjtNum* relpose = data+6;
+          mju_mulQuat(q0r, d->xquat+4*body1, relpose);
+          mju_negQuat(negq1, d->xquat+4*body2);
+        } else {
+          mju_mulQuat(q0r, d->xquat+4*body1, m->site_quat+4*obj1);
+          mjtNum qsite1[4];
+          mju_mulQuat(qsite1, d->xquat+4*body2, m->site_quat+4*obj2);
+          mju_negQuat(negq1, qsite1);
+        }
+
+        // angular velocities from cvel (first 3 components are angular)
+        const mjtNum* omega1 = d->cvel+6*body1;
+        const mjtNum* omega2 = d->cvel+6*body2;
+
+        // relative angular velocity: domega = omega1 - omega2
+        mjtNum domega[3];
+        mju_sub3(domega, omega1, omega2);
+
+        // quaternion derivatives: qdot = 0.5 * q * (0, omega)
+        mjtNum qdot0[4];
+        if (m->eq_objtype[eq_id] == mjOBJ_BODY) {
+          mju_derivQuat(qdot0, d->xquat+4*body1, omega1);
+        } else {
+          mjtNum qfull0[4];
+          mju_mulQuat(qfull0, d->xquat+4*body1, m->site_quat+4*obj1);
+          mju_derivQuat(qdot0, qfull0, omega1);
+        }
+        mjtNum qdot0r[4];  // d/dt(q0 * relpose) = qdot0 * relpose
+        if (m->eq_objtype[eq_id] == mjOBJ_BODY) {
+          mju_mulQuat(qdot0r, qdot0, data+6);
+        } else {
+          mju_copy4(qdot0r, qdot0);
+        }
+
+        // neg(qdot1): d/dt(neg(q1)) = neg(qdot1)
+        mjtNum negqdot1[4];
+        if (m->eq_objtype[eq_id] == mjOBJ_BODY) {
+          mjtNum qdot1[4];
+          mju_derivQuat(qdot1, d->xquat+4*body2, omega2);
+          mju_negQuat(negqdot1, qdot1);
+        } else {
+          mjtNum qfull1[4], qdot1[4];
+          mju_mulQuat(qfull1, d->xquat+4*body2, m->site_quat+4*obj2);
+          mju_derivQuat(qdot1, qfull1, omega2);
+          mju_negQuat(negqdot1, qdot1);
+        }
+
+        // Jdot_rot * v differentiates: 0.5 * neg(q1) * (J0-J1)*v * q0*relpose
+        // three terms from product rule:
+
+        // djrdv = Jrdot0*v - Jrdot1*v (rotational jacDot difference * v)
+        mjtNum djrdv[3];
+        mju_sub3(djrdv, jrdv1, jrdv2);
+
+        // term1: neg(qdot1) * domega * q0r
+        mjtNum t1a[4], t1[4];
+        mju_mulQuatAxis(t1a, negqdot1, domega);
+        mju_mulQuat(t1, t1a, q0r);
+
+        // term2: neg(q1) * djrdv * q0r
+        mjtNum t2a[4], t2[4];
+        mju_mulQuatAxis(t2a, negq1, djrdv);
+        mju_mulQuat(t2, t2a, q0r);
+
+        // term3: neg(q1) * domega * qdot0r
+        mjtNum t3a[4], t3[4];
+        mju_mulQuatAxis(t3a, negq1, domega);
+        mju_mulQuat(t3, t3a, qdot0r);
+
+        // combine: 0.5 * (term1 + term2 + term3), take vector part, scale
+        result[row+0] -= 0.5 * (t1[1] + t2[1] + t3[1]) * torquescale;
+        result[row+1] -= 0.5 * (t1[2] + t2[2] + t3[2]) * torquescale;
+        result[row+2] -= 0.5 * (t1[3] + t2[3] + t3[3]) * torquescale;
+
+        row += 3;
+      }
+    }
+
+    // other types: advance past all rows with this efc_id
+    else {
+      while (row < ne && d->efc_id[row] == eq_id) {
+        row++;
+      }
+    }
+  }
+
+  mj_freeStack(d);
+}
+

 // return number of constraint non-zeros, handle dense and dof-less cases
 static inline int mj_addConstraintCount(const mjModel* m, int size, int NV) {
@@ -2838,6 +3044,11 @@ void mj_referenceConstraint(const mjModel* m, mjData* d) {
    d->efc_aref[i] = -KBIP[4*i+1]*d->efc_vel[i]
                     -KBIP[4*i]*KBIP[4*i+2]*(d->efc_pos[i]-d->efc_margin[i]);
  }
+
+  // subtract Jdot*v correction for connect/weld equality constraints
+  if (d->ne > 0) {
+    mj_Jdotv(m, d, d->efc_aref);
+  }
 }


@@ -18,6 +18,7 @@
 #include <mujoco/mjdata.h>
 #include <mujoco/mjexport.h>
 #include <mujoco/mjmodel.h>
+#include <mujoco/mjtnum.h>

 #ifdef __cplusplus
 extern "C" {
@@ -35,6 +36,9 @@ MJAPI void mj_mulJacVec(const mjModel* m, const mjData* d, mjtNum* res, const mj
 // multiply JacobianT by vector
 MJAPI void mj_mulJacTVec(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);

+// subtract Jdot*v correction from result vector
+MJAPI void mj_Jdotv(const mjModel* m, mjData* d, mjtNum* result);
+

 //-------------------------- utility functions -----------------------------------------------------

@@ -32,6 +32,11 @@ extern "C" {
 MJAPI mjtNum mju_dotSparse2(const mjtNum* vec1, const int* ind1, int nnz1,
                            const mjtNum* vec2, const int* ind2, int nnz2);

+// dot-productX3, first vector is sparse; supernode of size 3
+void mju_dotSparseX3(mjtNum* res0, mjtNum* res1, mjtNum* res2,
+                     const mjtNum* vec10, const mjtNum* vec11, const mjtNum* vec12,
+                     const mjtNum* vec2, int nnz1, const int* ind1);
+
 // convert matrix from dense to sparse
 //  nnz is size of res and colind, return 1 if too small, 0 otherwise
 MJAPI int mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
@@ -787,6 +787,116 @@ TEST_F(CoreConstraintTest, ContactSharedDofJacobian) {
  mj_deleteModel(model);
 }

+static const char* const kJdotvConnect2dPath =
+    "engine/testdata/core_constraint/jdotv_connect_2d.xml";
+static const char* const kJdotvConnect3dPath =
+    "engine/testdata/core_constraint/jdotv_connect_3d.xml";
+static const char* const kJdotvWeld3dPath =
+    "engine/testdata/core_constraint/jdotv_weld_3d.xml";
+
+// validate mj_Jdotv against finite-differenced constraint Jacobian
+TEST_F(CoreConstraintTest, JdotvFiniteDifference) {
+
+  for (const char* path : {kJdotvConnect2dPath,
+                           kJdotvConnect3dPath,
+                           kJdotvWeld3dPath}) {
+    const std::string xml_path = GetTestDataFilePath(path);
+    char err[1024];
+    mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, err, sizeof(err));
+    ASSERT_THAT(m, NotNull()) << err << " for " << path;
+    int nv = m->nv;
+    mjData* d = mj_makeData(m);
+
+    // simulate for 1 second to accumulate velocity
+    while (d->time < 1.0) {
+      mj_step(m, d);
+    }
+
+    // forward to populate constraints
+    mj_forward(m, d);
+    ASSERT_GT(d->ne, 0) << "no equality constraints for " << path;
+    int ne = d->ne;
+
+    // get dense J_0 (ne x nv)
+    std::vector<mjtNum> J0(ne * nv);
+    if (mj_isSparse(m)) {
+      mju_sparse2dense(J0.data(), d->efc_J, ne, nv,
+                       d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
+    } else {
+      mju_copy(J0.data(), d->efc_J, ne * nv);
+    }
+
+    // compute mj_Jdotv at current state
+    std::vector<mjtNum> jdv(ne, 0);
+    mj_Jdotv(m, d, jdv.data());
+
+    // save qpos and qvel
+    std::vector<mjtNum> qpos0(m->nq), qvel0(nv);
+    mju_copy(qpos0.data(), d->qpos, m->nq);
+    mju_copy(qvel0.data(), d->qvel, nv);
+
+    // integrate qpos forward by h using qvel
+    const mjtNum h = MjTol(1e-7, 5e-4);
+    mj_integratePos(m, d->qpos, d->qvel, h);
+    mj_forward(m, d);
+
+    // get dense J_h (ne x nv)
+    ASSERT_EQ(d->ne, ne) << "constraint count changed after integration";
+    std::vector<mjtNum> Jh(ne * nv);
+    if (mj_isSparse(m)) {
+      mju_sparse2dense(Jh.data(), d->efc_J, ne, nv,
+                       d->efc_J_rownnz, d->efc_J_rowadr, d->efc_J_colind);
+    } else {
+      mju_copy(Jh.data(), d->efc_J, ne * nv);
+    }
+
+    // FD: Jdotv_fd[i] = -sum_j (Jh[i,j] - J0[i,j]) / h * qvel[j]
+    // (negated because mj_Jdotv subtracts)
+    std::vector<mjtNum> jdv_fd(ne, 0);
+    for (int i = 0; i < ne; i++) {
+      for (int j = 0; j < nv; j++) {
+        jdv_fd[i] -= (Jh[i*nv+j] - J0[i*nv+j]) / h * qvel0[j];
+      }
+    }
+
+    // compare
+    EXPECT_THAT(AsVector(jdv.data(), ne),
+                Pointwise(MjNear(1e-4, 1e-2), AsVector(jdv_fd.data(), ne)))
+        << "Jdotv FD mismatch for " << path;
+
+    mj_deleteData(d);
+    mj_deleteModel(m);
+  }
+}
+
+// Test 2: forward-inverse identity preserved with Jdot*v correction
+TEST_F(CoreConstraintTest, JdotvFwdInvIdentity) {
+  for (const char* path : {kJdotvConnect2dPath,
+                           kJdotvConnect3dPath,
+                           kJdotvWeld3dPath}) {
+    const std::string xml_path = GetTestDataFilePath(path);
+    char err[1024];
+    mjModel* m = mj_loadXML(xml_path.c_str(), nullptr, err, sizeof(err));
+    ASSERT_THAT(m, NotNull()) << err;
+    mjData* d = mj_makeData(m);
+
+    // give initial velocity
+    for (int i = 0; i < m->nv; i++) d->qvel[i] = 0.5 * (i + 1);
+
+    // forward (with correction ON by default)
+    mj_forward(m, d);
+    mj_compareFwdInv(m, d);
+    mjtNum fwdinv = d->solver_fwdinv[0];
+
+    mjtNum epsilon = MjTol(1e-10, 1e-2);
+    EXPECT_LT(fwdinv, epsilon)
+        << "fwdinv broken for " << path
+        << " (fwdinv=" << fwdinv << ")";
+
+    mj_deleteData(d);
+    mj_deleteModel(m);
+  }
+}

 }  // namespace
 }  // namespace mujoco
@@ -55,16 +55,21 @@ TEST_F(SolverTest, IslandsEquivalent) {
  mjData* data_island = mj_makeData(model);
  mjData* data_noisland = mj_makeData(model);

-  // Below are 3 tolerances associated with 3 different iteration counts,
-  // they are only moderately tight, 12x higher than x86-64 failure on Linux,
-  // i.e. in that case the test fails with rtol smaller than {6e-3, 6e-4, 6e-5}.
+  constexpr int kNumTol = 3;
+  mjtNum maxiter[kNumTol] = {30,   40,   60};
+  // Below are 3 tolerances associated with 3 different iteration counts.
+  // Tolerances are set to be ~12x higher than failure thresholds.
+  // For float32, failure thresholds are ~6000x larger than for float64.
+  // Line 99 adds a 500x factor for float32, so we need another ~12x in rtol.
  // The point of this test is to show that CG convergence is actually not very
  // precise, simply changing whether islands are used changes the solution by
  // quite a lot, even at high iteration count and zero {ls_}tolerance.
  // Increasing the iteration count higher than 60 does not improve convergence.
-  constexpr int kNumTol = 3;
-  mjtNum maxiter[kNumTol] = {30,   40,   60};
-  mjtNum rtol[kNumTol] =    {6e-2, 6e-3, 6e-4};
+  mjtNum rtol[kNumTol] = {
+      MjTol(6e-2, 7.2e-1),
+      MjTol(6e-3, 7.2e-2),
+      MjTol(6e-4, 7.2e-3)
+  };

  for (int i = 0; i < kNumTol; ++i) {
    model->opt.iterations = maxiter[i];
@@ -0,0 +1,29 @@
+<mujoco>
+  <option>
+   <flag contact="disable"/>
+  </option>
+
+  <worldbody>
+    <light pos="0 -1 1"/>
+
+    <!-- body 1: pendulum on world hinge, horizontal pointing right -->
+    <body name="upper" pos="0 0 1">
+      <joint name="hinge1" type="hinge" axis="0 1 0"/>
+      <geom type="capsule" size="0.03" fromto="0 0 0 0.5 0 0"/>
+      <site name="tip1" pos="0.5 0 0"/>
+    </body>
+
+    <!-- body 2: planar free body (slide x, slide z, hinge y) -->
+    <body name="lower" pos="0.5 0 1" euler="0 -90 0">
+      <joint type="slide" axis="1 0 0"/>
+      <joint type="slide" axis="0 0 1"/>
+      <joint type="hinge" axis="0 1 0"/>
+      <geom type="capsule" size="0.04" fromto="0 0 0 0.4 0 0"/>
+      <site name="tip2" pos="0 0 0"/>
+    </body>
+  </worldbody>
+
+  <equality>
+    <connect body1="upper" body2="lower" anchor="0.5 0 0"/>
+  </equality>
+</mujoco>
@@ -0,0 +1,29 @@
+<mujoco>
+  <option integrator="implicit">
+   <flag contact="disable"/>
+  </option>
+
+  <worldbody>
+    <light pos="0 -1 1"/>
+
+    <!-- body 1: pendulum on world ball joint, tilted -->
+    <body name="upper">
+      <joint type="ball"/>
+      <geom type="box" size="0.03" fromto="0 0 0 0.5 0 0"/>
+      <geom type="box" size="0.03" fromto="0.5 0 0 0.5 0.3 0"/>
+      <site name="tip1" pos="0.5 0.3 0"/>
+    </body>
+
+    <!-- body 2: free body, positioned at tip of upper -->
+    <body name="lower" pos="0.5 0.3 0">
+      <freejoint/>
+      <geom type="box" size="0.04" fromto="0 0 0 0.5 0 0"/>
+      <geom type="box" size="0.04" fromto="0.5 0 0 0.5 0.3 0"/>
+      <site name="tip2" pos="0 0 0"/>
+    </body>
+  </worldbody>
+
+  <equality>
+    <connect site1="tip1" site2="tip2"/>
+  </equality>
+</mujoco>
@@ -0,0 +1,29 @@
+<mujoco>
+  <option integrator="implicit">
+   <flag contact="disable"/>
+  </option>
+
+  <worldbody>
+    <light pos="0 -1 1"/>
+
+    <!-- body 1: pendulum on world ball joint, tilted -->
+    <body name="upper">
+      <joint type="ball"/>
+      <geom type="box" size="0.03" fromto="0 0 0 0.5 0 0"/>
+      <geom type="box" size="0.03" fromto="0.5 0 0 0.5 0.3 0"/>
+      <site name="tip1" pos="0.5 0.3 0"/>
+    </body>
+
+    <!-- body 2: free body, positioned at tip of upper -->
+    <body name="lower" pos="0.5 0.3 0">
+      <freejoint/>
+      <geom type="box" size="0.04" fromto="0 0 0 0.5 0 0"/>
+      <geom type="box" size="0.04" fromto="0.5 0 0 0.5 0.3 0"/>
+      <site name="tip2" pos="0 0 0"/>
+    </body>
+  </worldbody>
+
+  <equality>
+    <weld site1="tip1" site2="tip2" solimp="0 0.96 0.01" torquescale="0.1"/>
+  </equality>
+</mujoco>
@@ -365,7 +365,7 @@ describe('MuJoCo WASM Bindings', () => {
      mujoco.mj_constraintUpdate(
          model!, data!, res.GetView(), cost, /*flg_coneHessian=*/ 1);

-      expect(cost.GetView()[0]).toBeCloseTo(3355.837);
+      expect(cost.GetView()[0]).toBeCloseTo(3357.584);

      res.delete();
      cost.delete();