Add the Newton decrement as a termination criterion of the Newton solver.

After an accepted line-search step, the solver has already rebuilt the gradient
and Hessian and solved for the next search direction, so the Newton decrement
0.5*g'*H^-1*g -- the quadratic model's predicted cost improvement of the next
iteration -- costs one dot product. Terminating when it falls below tolerance
avoids running one more iteration only to observe a correspondingly small
actual improvement.

This is a C port of Alain's proposal in MJWarp:
https://github.com/google-deepmind/mujoco_warp/pull/1520

PiperOrigin-RevId: 947768034
Change-Id: I94e5c71a4e2b4a7775611edd1dad254bba2633b4

doc/XMLreference.rst

@@ -427,7 +427,8 @@ adjust it properly through the XML.
 :at:`tolerance`: :at-val:`real, "1e-8"`
    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. Set the tolerance to 0 to disable early termination.
+   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.
 
 .. _option-ls_iterations:
 

doc/changelog.rst

@@ -9,6 +9,10 @@ General
 ^^^^^^^
 - Added Nesterov momentum extrapolation with adaptive gradient restart (O'Donoghue-Candès) to the PGS solver,
   significantly improving convergence. Overall PGS now requires ~2x fewer iterations.
+- Added the Newton decrement -- the quadratic model's predicted cost improvement of the next iteration -- as a third
+  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>`__.
 - :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

@@ -1392,7 +1392,10 @@ representations of the constraint Jacobian and related matrices.
    This algorithm implements the exact Newton method, with analytical second-order derivatives and Cholesky
    factorization of the Hessian. The line-search is the same as in the CG method. When constraint states change between
    iterations (e.g., a constraint transitions from quadratic to linear), the Hessian factorization is updated
-   incrementally via rank-1 Cholesky updates, avoiding full refactorization. It is the default solver.
+   incrementally via rank-1 Cholesky updates, avoiding full refactorization. Early termination is triggered when any of
+   three quantities falls below :ref:`tolerance<option-tolerance>`: the cost improvement of the last iteration, the
+   gradient norm, and the Newton decrement :math:`\tfrac{1}{2} g^T H^{-1} g` -- the predicted cost improvement of the
+   next iteration. It is the default solver.
 
 **PGS** : Projected Gauss-Seidel method
    This is the most common algorithm used in physics simulators, and used to be the default in MuJoCo, until we

src/engine/engine_solver.c

@@ -2393,12 +2393,17 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i
     saveStats(m, d, island, iter, improvement, gradient, ctx.LSslope,
               ctx.nactive, nchange, ctx.LSiter, ctx.nupdate);
 
+    // Newton decrement: 0.5*grad'*H^-1*grad, the model's predicted improvement of the
+    // next step; clamp to 0 so that tolerance == 0 keeps early termination disabled
+    mjtNum decrement = flg_Newton ? mju_max(0, 0.5*scale*mju_dot(ctx.grad, ctx.Mgrad, nv)) : 0;
+
     // increment iteration count
     iter++;
 
     // termination
     if ((improvement > 0 && improvement < m->opt.tolerance) ||
-         gradient < m->opt.tolerance) {
+         gradient < m->opt.tolerance ||
+         (flg_Newton && decrement < m->opt.tolerance)) {
       break;
     }
 

test/engine/engine_core_smooth_test.cc

@@ -595,6 +595,9 @@ TEST_F(CoreSmoothTest, RefsiteConservesMomentum) {
   ASSERT_THAT(model, NotNull());
   mjData* data = mj_makeData(model);
 
+  // this test asserts tight momentum conservation: solve exactly, no early termination
+  model->opt.tolerance = 0;
+
   data->ctrl[0] = 1;
   data->ctrl[1] = -1;
 

test/engine/engine_solver_test.cc

@@ -391,5 +391,101 @@ TEST_F(SolverTest, EllipticLineSearchPrecisionDiagnostics) {
   }
 }
 
+// Newton terminates early when the decrement predicts sub-tolerance improvement
+TEST_F(SolverTest, NewtonDecrementTermination) {
+  const std::string xml_path = GetTestDataFilePath(kHumanoidPath);
+  char error[1024];
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  model->opt.solver = mjSOL_NEWTON;
+  model->opt.disableflags |= mjDSBL_ISLAND;  // monolithic solve
+  model->opt.iterations = 100;
+  const mjtNum tolerance = MjTol(1e-6, 1e-4);
+
+  int state_size = mj_stateSize(model, mjSTATE_FULLPHYSICS);
+  mjtNum* state = (mjtNum*) mju_malloc(sizeof(mjtNum)*state_size);
+
+  mjData* data = mj_makeData(model);
+  mjData* data_test = mj_makeData(model);
+  mjData* data_deep = mj_makeData(model);
+  mj_resetDataKeyframe(model, data, 0);
+
+  int nfired = 0;
+  mjtNum max_leftover = 0;
+  for (int step = 0; step < 200; step++) {
+    model->opt.tolerance = tolerance;
+    mj_step(model, data);
+    mj_getState(model, data, state, mjSTATE_FULLPHYSICS);
+
+    // solve with the test tolerance
+    mj_setState(model, data_test, state, mjSTATE_FULLPHYSICS);
+    mj_forward(model, data_test);
+
+    // reference: tolerance 0 runs until the line search finds no improvement
+    model->opt.tolerance = 0;
+    mj_setState(model, data_deep, state, mjSTATE_FULLPHYSICS);
+    mj_forward(model, data_deep);
+
+    // accuracy: both runs produce identical iterates up to the test run's
+    // stopping point, so the cost improvement forgone by early termination is
+    // the sum of the deep run's remaining (scaled) improvements
+    int niter = data_test->solver_niter[0];
+    int niter_deep = std::min(data_deep->solver_niter[0], mjNSOLVER);
+    mjtNum leftover = 0;
+    for (int i = niter; i < niter_deep; i++) {
+      leftover += max(static_cast<mjtNum>(0), data_deep->solver[i].improvement);
+    }
+    max_leftover = max(max_leftover, leftover);
+
+    // count decrement terminations: the test run stopped while both existing
+    // criteria were above tolerance, and the deep run shows that the next
+    // iteration would have improved the cost by less than tolerance
+    if (niter > 0 && niter < std::min(model->opt.iterations, mjNSOLVER) &&
+        data_deep->solver_niter[0] > niter) {
+      const mjSolverStat& last = data_test->solver[niter - 1];
+      const mjSolverStat& next = data_deep->solver[niter];
+      if (last.improvement >= tolerance && last.gradient >= tolerance &&
+          next.improvement < tolerance) {
+        nfired++;
+      }
+    }
+  }
+
+  EXPECT_LT(max_leftover, 10*tolerance)
+      << "early termination forgoes more than a small multiple of tolerance";
+  EXPECT_GT(nfired, 0)
+      << "no state exercised the Newton decrement termination criterion";
+
+  mj_deleteData(data_deep);
+  mj_deleteData(data_test);
+  mj_deleteData(data);
+  mju_free(state);
+  mj_deleteModel(model);
+}
+
+// tolerance == 0 disables early termination, including the Newton decrement
+TEST_F(SolverTest, ZeroToleranceDisablesTermination) {
+  const std::string xml_path = GetTestDataFilePath(kHumanoidPath);
+  char error[1024];
+  mjModel* model = mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error));
+  ASSERT_THAT(model, NotNull()) << error;
+  model->opt.solver = mjSOL_NEWTON;
+  model->opt.disableflags |= mjDSBL_ISLAND | mjDSBL_WARMSTART;
+  model->opt.tolerance = 0;
+  model->opt.iterations = 3;
+
+  mjData* data = mj_makeData(model);
+  for (mjtCone cone : {mjCONE_PYRAMIDAL, mjCONE_ELLIPTIC}) {
+    model->opt.cone = cone;
+    mj_resetDataKeyframe(model, data, 0);
+    mj_forward(model, data);
+    EXPECT_EQ(data->solver_niter[0], 3)
+        << "cone: " << (cone == mjCONE_ELLIPTIC ? "elliptic" : "pyramidal");
+  }
+
+  mj_deleteData(data);
+  mj_deleteModel(model);
+}
+
 }  // namespace
 }  // namespace mujoco