Add zero-iteration early exit to the primal solvers, certified by the duality gap.

The primal cost has curvature of at least M in every zone, making it strongly
convex in the M-norm and bounding the suboptimality of any point by the
Fenchel duality gap at its constraint forces:

  cost(qacc) - cost* <= 0.5*grad'*M^-1*grad

Since M's factorization always exists, this certificate is evaluable before
the solver does any work: one triangular solve and one dot product. When the
warmstarted solution is already certified to satisfy the tolerance, CG and
Newton now return with zero iterations; for Newton this skips building and
factorizing the Hessian. If the certificate declines, Newton gets a second
exit after factorization: the Newton decrement, checked before the first
line search.

Because the gap bounds cost suboptimality, stiff constraints can convert it
into force errors of order sqrt(2*gap*stiffness). Newton solutions are
characteristically force-accurate, so Newton zero-iteration exits also
require the gradient criterion, preserving constraint-force accuracy at
rest; CG solutions are characteristically cost-accurate and exit on the gap
alone.

On a settling pile of 50 boxes (300 dofs, ~200 contacts), end-to-end time
per step drops 13% over a settle-then-rest run and 27% in the quiescent
limit, with Newton iterations falling from 0.98 to 0.40 per step.

Tests: WarmstartZeroIterations sweeps solver/cone/jacobian on a settled box,
asserting zero iterations, forward/inverse consistency, and agreement with a
tolerance=0 control solve from the same state. WarmstartZeroIterationsIslands
checks per-island exits with a kicked box next to a settled one.
RefsiteConservesMomentum now requests an exact solve (tolerance=0), since it
asserts momentum conservation tighter than the solver tolerance contract.
PiperOrigin-RevId: 947993735
Change-Id: I2fd855774bff619709b2c386f1ba2714286e0821

doc/XMLreference.rst

@@ -418,9 +418,10 @@ adjust it properly through the XML.
 
 :at:`iterations`: :at-val:`int, "100"`
    Maximum number of iterations of the constraint solver. When the warmstart attribute of :ref:`flag <option-flag>` is
-   enabled (which is the default), accurate results are obtained with fewer iterations. Larger and more complex systems
-   with many interacting constraints require more iterations. Note that mjData.solver contains statistics about solver
-   convergence, also shown in the profiler.
+   enabled (which is the default), accurate results are obtained with fewer iterations; if the warmstarted solution
+   already satisfies the tolerance, the CG and Newton solvers terminate with zero iterations. Larger and more complex
+   systems with many interacting constraints require more iterations. Note that mjData.solver contains statistics about
+   solver convergence, also shown in the profiler.
 
 .. _option-tolerance:
 
@@ -428,7 +429,9 @@ adjust it properly through the XML.
    Tolerance threshold used for early termination of the iterative solver. For PGS, the threshold is applied to the cost
    improvement between two iterations. For CG and Newton, it is applied to the smaller of the cost improvement and the
    gradient norm. For Newton, it is additionally applied to the Newton decrement :math:`\tfrac{1}{2} g^T H^{-1} g`, the
-   predicted cost improvement of the next iteration. Set the tolerance to 0 to disable early termination.
+   predicted cost improvement of the next iteration. Before the first iteration, CG and Newton also apply it to a
+   :ref:`convergence certificate<soAlgorithms>` of the warmstarted solution, possibly terminating with zero iterations.
+   Set the tolerance to 0 to disable early termination.
 
 .. _option-ls_iterations:
 

doc/changelog.rst

@@ -13,6 +13,11 @@ General
   early-termination criterion of the :ref:`Newton solver<soAlgorithms>`, alongside cost improvement and gradient norm.
   This reduces iteration counts at no accuracy cost. Proposed by :github:user:`adenzler-nvidia` in
   :doc:`MJWarp <mjwarp/index>` pull request `1520 <https://github.com/google-deepmind/mujoco_warp/pull/1520>`__.
+- The CG and Newton solvers now terminate with zero iterations when a duality-gap certificate proves that the
+  warmstarted solution already satisfies the tolerance. The certificate requires only the existing mass-matrix
+  factorization, so quiescent scenes skip Hessian construction, factorization and the line search entirely. Newton
+  zero-iteration exits additionally require the gradient criterion, preserving Newton's characteristic force-level
+  accuracy. See :ref:`Warmstart<soAlgorithms>` in the Computation chapter for details.
 - :ref:`mj_encode` now supports encoding of MJB and TXT files.
 - :ref:`mj_setConst` now recomputes the ``mjModel.{body,geom,site}_sameframe`` flags, to account for changes in
   body/geom/site frames after compilation.

doc/computation/index.rst

@@ -1445,6 +1445,14 @@ representations of the constraint Jacobian and related matrices.
    bootstraps the solver when constraints persist across time steps, but avoids carrying over stale forces from
    constraints that have disappeared.
 
+   Because every zone of the piecewise-quadratic cost has curvature of at least :math:`M`, the cost is strongly convex
+   in the :math:`M`-norm, which bounds the suboptimality of any point by the duality gap at its constraint forces:
+   :math:`\text{cost}(a) - \text{cost}^* \le \tfrac{1}{2} g^T M^{-1} g`. Before starting iterations, the CG and Newton
+   solvers evaluate this certificate at the warmstarted point, using the already-computed factorization of
+   :math:`M`. If it is below tolerance, convergence is proven and the solver returns immediately with zero iterations;
+   in the Newton case this skips constructing and factorizing the Hessian. In a quiescent, well-warmstarted scene this
+   eliminates nearly the entire cost of the constraint solver.
+
 .. _soIsland:
 
 Constraint islands

doc/programming/simulation.rst

@@ -1105,9 +1105,10 @@ does not need timing, and in that case there is no reason to call timing functio
 One part of the simulation pipeline that needs to be monitored closely is the iterative constraint solver. The
 simplest diagnostic here is ``mjData.solver_niter`` which shows how many iterations the solver took on the last call to
 mj_step or ``mj_forward``. Note that the solver has tolerance parameters for early termination, so this number is
-usually smaller than the maximum number of iterations allowed. The array ``mjData.solver`` contains one
-:ref:`mjSolverStat` data structure per iteration of the constraint solver, with information about the constraint state
-and line search.
+usually smaller than the maximum number of iterations allowed; it can be 0 when a warmstarted solution is already
+certified as converged, in which case no iterations are performed and no statistics are written. The array
+``mjData.solver`` contains one :ref:`mjSolverStat` data structure per iteration of the constraint solver, with
+information about the constraint state and line search.
 
 When the option :at:`fwdinv` is enabled in ``mjModel.opt.enableflags``, the field ``mjData.fwdinv`` is also populated.
 It contains the difference between the forward and inverse dynamics, in terms of generalized forces and constraint

src/engine/engine_solver.c

@@ -1376,14 +1376,18 @@ static void PrimalUpdateConstraint(mjPrimalContext* ctx, int flg_HessianCone) {
 }
 
 
-// update grad, Mgrad
-static void PrimalUpdateGradient(mjPrimalContext* ctx, int flg_Newton) {
+// update grad = M*qacc - qfrc_smooth - qfrc_constraint
+static void PrimalUpdateGrad(mjPrimalContext* ctx) {
   int nv = ctx->nv;
-
-  // grad = M*qacc - qfrc_smooth - qfrc_constraint
   for (int i=0; i < nv; i++) {
     ctx->grad[i] = ctx->Ma[i] - ctx->qfrc_smooth[i] - ctx->qfrc_constraint[i];
   }
+}
+
+
+// update Mgrad;  Newton: Mgrad = H \ grad,  CG: Mgrad = M \ grad
+static void PrimalUpdateMgrad(mjPrimalContext* ctx, int flg_Newton) {
+  int nv = ctx->nv;
 
   // Newton: Mgrad = H \ grad
   if (flg_Newton) {
@@ -1404,6 +1408,13 @@ static void PrimalUpdateGradient(mjPrimalContext* ctx, int flg_Newton) {
 }
 
 
+// update grad, Mgrad
+static void PrimalUpdateGradient(mjPrimalContext* ctx, int flg_Newton) {
+  PrimalUpdateGrad(ctx);
+  PrimalUpdateMgrad(ctx, flg_Newton);
+}
+
+
 // prepare quadratic polynomials and contact cone quantities
 static void PrimalPrepare(mjPrimalContext* ctx) {
   int nv = ctx->nv, nefc = ctx->nefc;
@@ -2326,15 +2337,7 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i
 
   // first update
   PrimalUpdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC));
-  if (flg_Newton) {
-    // compute and factorize Hessian
-    MakeHessian(d, &ctx);
-    FactorizeHessian(d, &ctx, /*flg_recompute=*/0);
-  }
-  PrimalUpdateGradient(&ctx, flg_Newton);
-
-  // start both with preconditioned gradient
-  mju_scl(ctx.search, ctx.Mgrad, -1, nv);
+  PrimalUpdateGrad(&ctx);
 
   // compute and save scaling factor
   mjtNum scale;
@@ -2350,8 +2353,42 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i
   }
   ctx.scale = scale;
 
+  // Mgrad = M \ grad: the CG preconditioned gradient, also the convergence certificate
+  PrimalUpdateMgrad(&ctx, /*flg_Newton=*/0);
+
+  // convergence certificate: the cost is strongly convex in the M-norm, bounding the
+  // suboptimality by the duality gap at the current constraint forces:
+  //   cost(qacc) - cost* <= 0.5 * grad'*M^-1*grad
+  // if already below tolerance (e.g. good warmstart), skip the Hessian and the main loop
+  int flg_gap = mju_max(0, 0.5*scale*mju_dot(ctx.grad, ctx.Mgrad, nv)) < m->opt.tolerance;
+
+  // the gap bounds the *cost* suboptimality; on stiff constraints this permits force
+  // errors of order sqrt(2*gap*stiffness). Newton solutions are characteristically
+  // force-accurate, so Newton zero-iteration exits also require the gradient criterion;
+  // CG solutions are characteristically cost-accurate and exit on the gap alone
+  int flg_gradient = scale*mju_norm(ctx.grad, nv) < m->opt.tolerance;
+  int flg_certificate = flg_gap && (!flg_Newton || flg_gradient);
+  int flg_done = flg_certificate;
+
+  // Newton: compute and factorize Hessian, Mgrad = H \ grad
+  if (!flg_done && flg_Newton) {
+    MakeHessian(d, &ctx);
+    FactorizeHessian(d, &ctx, /*flg_recompute=*/0);
+    PrimalUpdateMgrad(&ctx, /*flg_Newton=*/1);
+
+    // Newton decrement already below tolerance: converged, skip the first line search
+    // (gradient-gated like the certificate: H^-1 suppresses stiff-direction force errors)
+    flg_done = flg_gradient &&
+               mju_max(0, 0.5*scale*mju_dot(ctx.grad, ctx.Mgrad, nv)) < m->opt.tolerance;
+  }
+
+  // start both with preconditioned gradient
+  if (!flg_done) {
+    mju_scl(ctx.search, ctx.Mgrad, -1, nv);
+  }
+
   // main loop
-  while (iter < maxiter) {
+  while (!flg_done && iter < maxiter) {
     // perform linesearch
     mjtNum ls_improvement;
     alpha = PrimalSearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations,
@@ -2470,8 +2507,8 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i
     // update solver iterations
     d->solver_niter[island_stat] += iter;
 
-    // set solver_nnz
-    if (flg_Newton) {
+    // set solver_nnz; if the certificate fired, no Hessian was built: report Jacobian nnz
+    if (flg_Newton && !flg_certificate) {
       if (mj_isSparse(m)) {
         // two L factors if Lcone is present
         int num_factors = 1 + (ctx.Lcone != NULL);

test/engine/engine_solver_test.cc

@@ -17,6 +17,7 @@
 #include <algorithm>
 #include <cstdlib>
 #include <string>
+#include <vector>
 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
@@ -463,6 +464,139 @@ TEST_F(SolverTest, NewtonDecrementTermination) {
   mj_deleteModel(model);
 }
 
+// a settled, warmstarted scene certifies convergence and solves in zero iterations
+TEST_F(SolverTest, WarmstartZeroIterations) {
+  std::string xml = R"(
+  <mujoco>
+    <worldbody>
+      <geom type="plane" size="1 1 .1"/>
+      <body pos="0 0 0.1">
+        <freejoint/>
+        <geom type="box" size="0.1 0.1 0.1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+
+  char error[1024];
+  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  MjDataPtr data = MakeData(model);
+  model->opt.disableflags |= mjDSBL_ISLAND;  // monolithic solve: stats in slot 0
+  model->opt.enableflags |= mjENBL_FWDINV;
+
+  int nv = model->nv;
+  int state_size = mj_stateSize(model.get(), mjSTATE_FULLPHYSICS);
+  std::vector<mjtNum> state(state_size);
+  std::vector<mjtNum> qacc(nv), qfrc(nv);
+
+  // float32 cannot resolve the default tolerance: use a resolvable one
+  const mjtNum tolerance = MjTol(1e-8, 1e-6);
+
+  for (mjtSolver solver : {mjSOL_CG, mjSOL_NEWTON}) {
+    for (mjtCone cone : {mjCONE_PYRAMIDAL, mjCONE_ELLIPTIC}) {
+      for (mjtJacobian jacobian : {mjJAC_DENSE, mjJAC_SPARSE}) {
+        std::string config = std::string(solver == mjSOL_CG ? "CG" : "Newton") +
+                             (cone == mjCONE_ELLIPTIC ? "/elliptic" : "/pyramidal") +
+                             (jacobian == mjJAC_SPARSE ? "/sparse" : "/dense");
+        model->opt.solver = solver;
+        model->opt.cone = cone;
+        model->opt.jacobian = jacobian;
+        model->opt.tolerance = tolerance;
+
+        // settle the box on the plane
+        mj_resetData(model.get(), data.get());
+        for (int i=0; i < 500; i++) {
+          mj_step(model.get(), data.get());
+        }
+        mj_getState(model.get(), data.get(), state.data(), mjSTATE_FULLPHYSICS);
+
+        // solve once more: certificate fires, forward/inverse stay consistent
+        mj_forward(model.get(), data.get());
+        EXPECT_EQ(data->solver_niter[0], 0) << config;
+
+        // thresholds here and below are ~10x above measured, per precision
+        EXPECT_LT(data->solver_fwdinv[0], MjTol(1e-12, 1e-4)) << config;
+        EXPECT_LT(data->solver_fwdinv[1], MjTol(1e-2, 2e-1)) << config;
+        mju_copy(qacc.data(), data->qacc, nv);
+        mju_copy(qfrc.data(), data->qfrc_constraint, nv);
+
+        // control arm: tolerance = 0 disables the certificate, full solve from
+        // the same state must agree with the skipped solve
+        model->opt.tolerance = 0;
+        mj_setState(model.get(), data.get(), state.data(), mjSTATE_FULLPHYSICS);
+        mj_forward(model.get(), data.get());
+        mjtNum dqacc = 0, dqfrc = 0;
+        for (int j=0; j < nv; j++) {
+          dqacc = max(dqacc, std::abs(qacc[j] - data->qacc[j]));
+          dqfrc = max(dqfrc, std::abs(qfrc[j] - data->qfrc_constraint[j]));
+        }
+        EXPECT_LT(dqacc, MjTol(2e-4, 1.5e-3)) << config;
+        EXPECT_LT(dqfrc, MjTol(2e-2, 4e-1)) << config;
+
+        // guard: a perturbed scene does not certify
+        model->opt.tolerance = tolerance;
+        mj_setState(model.get(), data.get(), state.data(), mjSTATE_FULLPHYSICS);
+        data->qfrc_applied[0] = 5;
+        mj_forward(model.get(), data.get());
+        EXPECT_GT(data->solver_niter[0], 0) << config;
+        data->qfrc_applied[0] = 0;
+      }
+    }
+  }
+}
+
+// per-island certificates: settled islands solve in zero iterations while
+// islands with new loads solve normally
+TEST_F(SolverTest, WarmstartZeroIterationsIslands) {
+  std::string xml = R"(
+  <mujoco>
+    <worldbody>
+      <geom type="plane" size="2 2 .1"/>
+      <body pos="-0.5 0 0.1">
+        <freejoint/>
+        <geom type="box" size="0.1 0.1 0.1"/>
+      </body>
+      <body pos="0.5 0 0.1">
+        <freejoint/>
+        <geom type="box" size="0.1 0.1 0.1"/>
+      </body>
+    </worldbody>
+  </mujoco>
+  )";
+
+  char error[1024];
+  MjModelPtr model = LoadModelFromString(xml, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  MjDataPtr data = MakeData(model);
+
+  // float32 cannot resolve the default tolerance: use a resolvable one
+  model->opt.tolerance = MjTol(1e-8, 1e-6);
+
+  // settle both boxes, islands enabled (default)
+  for (int i=0; i < 500; i++) {
+    mj_step(model.get(), data.get());
+  }
+
+  // both islands certify: zero iterations everywhere
+  mj_forward(model.get(), data.get());
+  ASSERT_EQ(data->nisland, 2);
+  EXPECT_EQ(data->solver_niter[0], 0);
+  EXPECT_EQ(data->solver_niter[1], 0);
+
+  // kick the second box: its island solves, the settled island still certifies
+  data->qfrc_applied[6] = 5;
+  mj_forward(model.get(), data.get());
+  ASSERT_EQ(data->nisland, 2);
+  int island1 = data->dof_island[0];
+  int island2 = data->dof_island[6];
+  ASSERT_GE(island1, 0);
+  ASSERT_GE(island2, 0);
+  ASSERT_NE(island1, island2);
+  EXPECT_EQ(data->solver_niter[island1], 0);
+  EXPECT_GT(data->solver_niter[island2], 0);
+}
+
 // tolerance == 0 disables early termination, including the Newton decrement
 TEST_F(SolverTest, ZeroToleranceDisablesTermination) {
   const std::string xml_path = GetTestDataFilePath(kHumanoidPath);