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
@@ -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>