diff --git a/doc/XMLreference.rst b/doc/XMLreference.rst index 801cee92..2211d359 100644 --- a/doc/XMLreference.rst +++ b/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: diff --git a/doc/changelog.rst b/doc/changelog.rst index e4a3e361..61307d3d 100644 --- a/doc/changelog.rst +++ b/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`, alongside cost improvement and gradient norm. + This reduces iteration counts at no accuracy cost. Proposed by :github:user:`adenzler-nvidia` in + :doc:`MJWarp ` pull request `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. diff --git a/doc/computation/index.rst b/doc/computation/index.rst index a0d890f1..e1170505 100644 --- a/doc/computation/index.rst +++ b/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`: 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 diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 0a45fdf3..7dbd1b87 100644 --- a/src/engine/engine_solver.c +++ b/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; } diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 55a9cb83..80542a8a 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/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; diff --git a/test/engine/engine_solver_test.cc b/test/engine/engine_solver_test.cc index 73bf15cf..55766a7c 100644 --- a/test/engine/engine_solver_test.cc +++ b/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(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