diff --git a/doc/changelog.rst b/doc/changelog.rst index 5b9df28c..bb6af630 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -2,6 +2,14 @@ Changelog ========= +Upcoming version (not yet released) +----------------------------------- + +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. + Version 3.10.0 (June 22, 2026) ------------------------------ diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 8f7e4676..e3f9f174 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -26,6 +26,7 @@ #include "engine/engine_core_smooth.h" #include "engine/engine_core_util.h" #include "engine/engine_memory.h" +#include "engine/engine_macro.h" #include "engine/engine_util_blas.h" #include "engine/engine_util_errmem.h" #include "engine/engine_util_misc.h" @@ -375,6 +376,27 @@ static int dualStateChange(const mjData* d, int* state, int* oldstate, } +// project onto friction ellipsoid, write to force[i+1..i+dim-1] +// project tangential force onto friction ellipsoid: sum(f_t[j]^2/mu[j]^2) <= f_n^2 +// if feasible is true, only scale down if outside the ellipsoid +// if feasible is false, always scale to the boundary +static void projectEllipsoid(mjtNum* friction, mjtNum normal, const mjtNum* mu, + int dim, int feasible) { + mjtNum s = 0; + for (int j=0; j < dim-1; j++) { + s += friction[j]*friction[j] / (mu[j]*mu[j]); + } + + mjtNum normal2 = normal*normal; + if (!feasible || s > normal2) { + mjtNum scl = mju_sqrt(normal2 / mju_max(mjMINVAL, s)); + for (int j=0; j < dim-1; j++) { + friction[j] *= scl; + } + } +} + + // solve QCQP and project onto friction ellipsoid, write to force[i+1..i+dim-1] static void solveQCQP(mjtNum* force, int i, int dim, mjtNum* Ac, mjtNum* bc, const mjtNum* mu) { @@ -392,14 +414,7 @@ static void solveQCQP(mjtNum* force, int i, int dim, // on constraint: put v on ellipsoid, in case QCQP is approximate if (flg_active) { - mjtNum s = 0; - for (int j=0; j < dim-1; j++) { - s += v[j]*v[j] / (mu[j]*mu[j]); - } - s = mju_sqrt(force[i]*force[i] / mju_max(mjMINVAL, s)); - for (int j=0; j < dim-1; j++) { - v[j] *= s; - } + projectEllipsoid(v, force[i], mu, dim, /*feasible=*/0); } // assign @@ -407,6 +422,31 @@ static void solveQCQP(mjtNum* force, int i, int dim, } +// project contact force block onto friction cone (pyramidal or elliptic) +static void projectCone(mjtNum* force, const mjtNum* mu, int dim, int type) { + // elliptic cone: project onto friction ellipsoid + if (type == mjCNSTR_CONTACT_ELLIPTIC) { + // clamp normal force + if (force[0] < 0) { + mju_zero(force, dim); + } else { + projectEllipsoid(force+1, force[0], mu, dim, /*feasible=*/1); + } + } + + // pyramidal or scalar: clamp to non-negative + else { + if (force[0] < 0) { + force[0] = 0; + } + } +} + + +// global variable to toggle Nesterov momentum (for benchmarks/tests) +mjTHREADLOCAL int mj_nesterov_momentum = 1; + + //---------------------------- PGS solver ---------------------------------------------------------- // core PGS solver: iterates over constraints specified by efclist @@ -423,6 +463,16 @@ static void solPGS(const mjModel* m, mjData* d, int island, int* oldstate = mjSTACKALLOC(d, 2*nefc, int); int* blockstart = oldstate + nefc; + // Nesterov momentum + mjtBool nesterov = (mj_nesterov_momentum != 0); + mjtNum* force_prev = NULL; + mjtNum* force_momentum = NULL; + if (nesterov) { + force_prev = mjSTACKALLOC(d, nefc, mjtNum); + force_momentum = mjSTACKALLOC(d, nefc, mjtNum); + mju_gather(force_prev, force, efclist, nefc); + } + int island_stat = mjMAX(0, island); // island index for diagnostic stats mjtNum scale = 1 / (m->stat.meaninertia * mjMAX(1, m->nv)); @@ -452,7 +502,56 @@ static void solPGS(const mjModel* m, mjData* d, int island, // main iteration int iter = 0; + int nesterov_k = 0; // Nesterov counter (resets on adaptive restart) while (iter < maxiter) { + // Nesterov momentum extrapolation + if (nesterov) { + mjtNum beta = 0; + if (iter > 0) { + beta = (mjtNum)(nesterov_k - 1) / (mjtNum)(nesterov_k + 2); + } + + // update with momentum, save pre-extrapolation value + if (beta > 0) { + for (int c=0; c < nefc; c++) { + int i = efclist ? efclist[c] : c; + mjtNum f_save = force[i]; + force[i] += beta*(force[i] - force_prev[c]); + force_prev[c] = f_save; + } + + // friction loss: project onto bounds + for (int c=ne; c < ne+nf; c++) { + int i = efclist ? efclist[c] : c; + force[i] = mju_clip(force[i], -floss[i], floss[i]); + } + + // contact force: project onto friction cone + for (int c=ne+nf; c < nefc; ) { + int i = efclist ? efclist[c] : c; + int dim = 1; + int type = d->efc_type[i]; + const mjtNum* mu = NULL; + + if (type == mjCNSTR_CONTACT_ELLIPTIC) { + dim = d->contact[d->efc_id[i]].dim; + mu = d->contact[d->efc_id[i]].friction; + } + + projectCone(force+i, mu, dim, type); + c += dim; + } + } + + // iter == 0 or beta <= 0 (nesterov_k <= 1): just save current force + else { + mju_gather(force_prev, force, efclist, nefc); + } + + // save extrapolated point for gradient restart check + mju_gather(force_momentum, force, efclist, nefc); + } + // clear improvement mjtNum improvement = 0; @@ -591,6 +690,27 @@ static void solPGS(const mjModel* m, mjData* d, int island, improvement *= scale; saveStats(m, d, island_stat, iter, improvement, 0, 0, nactive, nchange, 0, 0); + // Nesterov gradient restart (O'Donoghue-Candès): reset when correction opposes extrapolation + if (nesterov) { + mjtBool restart = false; + if (iter > 0) { + mjtNum dot_corr_extr = 0; + for (int c=0; c < nefc; c++) { + int i = efclist ? efclist[c] : c; + mjtNum correction = force[i] - force_momentum[c]; + mjtNum extrapolation = force_momentum[c] - force_prev[c]; + dot_corr_extr += correction * extrapolation; + } + restart = (dot_corr_extr < 0); + } + + if (restart) { + nesterov_k = 0; + } else { + nesterov_k++; + } + } + // increment iteration count iter++; diff --git a/src/engine/engine_util_misc.c b/src/engine/engine_util_misc.c index 4a4e1f9a..3c438ba7 100644 --- a/src/engine/engine_util_misc.c +++ b/src/engine/engine_util_misc.c @@ -2129,8 +2129,13 @@ void mju_n2d(double* res, const mjtNum* vec, int n) { } -// gather +// gather: res[i] = vec[ind[i]], or copy if ind is NULL void mju_gather(mjtNum* restrict res, const mjtNum* restrict vec, const int* restrict ind, int n) { + if (!ind) { + mju_copy(res, vec, n); + return; + } + for (int i=0; i < n; i++) { res[i] = vec[ind[i]]; } @@ -2146,8 +2151,13 @@ void mju_gatherMasked(mjtNum* restrict res, const mjtNum* restrict vec, } -// scatter +// scatter: res[ind[i]] = vec[i], or copy if ind is NULL void mju_scatter(mjtNum* restrict res, const mjtNum* restrict vec, const int* restrict ind, int n) { + if (!ind) { + mju_copy(res, vec, n); + return; + } + for (int i=0; i < n; i++) { res[ind[i]] = vec[i]; } diff --git a/src/engine/engine_util_misc.h b/src/engine/engine_util_misc.h index 8617b922..80dbcd1f 100644 --- a/src/engine/engine_util_misc.h +++ b/src/engine/engine_util_misc.h @@ -281,13 +281,13 @@ MJAPI void mju_d2n(mjtNum* res, const double* vec, int n); // convert from mjtNum to double MJAPI void mju_n2d(double* res, const mjtNum* vec, int n); -// gather mjtNums +// gather mjtNums: res[i] = vec[ind[i]], or copy if ind is NULL MJAPI void mju_gather(mjtNum* res, const mjtNum* vec, const int* ind, int n); // gather mjtNums, set to 0 at negative indices MJAPI void mju_gatherMasked(mjtNum* res, const mjtNum* vec, const int* ind, int n); -// scatter mjtNums +// scatter mjtNums: res[ind[i]] = vec[i], or copy if ind is NULL MJAPI void mju_scatter(mjtNum* res, const mjtNum* vec, const int* ind, int n); // gather integers diff --git a/test/engine/engine_pgs_convergence_test.cc b/test/engine/engine_pgs_convergence_test.cc new file mode 100644 index 00000000..996b14cd --- /dev/null +++ b/test/engine/engine_pgs_convergence_test.cc @@ -0,0 +1,274 @@ +// Copyright 2026 DeepMind Technologies Limited +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// PGS with/without Nesterov momentum: compare convergence and timing. +// +// Rolls out Newton ground truth on 2humanoid100.xml, then evaluates PGS with +// and without Nesterov momentum at various iteration budgets. + + +#include // NOLINT +#include +#include // NOLINT +#include +#include + +#include +#include +#include +#include "test/fixture.h" + +extern "C" thread_local int mj_nesterov_momentum; + +namespace mujoco { +namespace { + +using ::testing::NotNull; + +mjtNum gettm(void) { + using Clock = std::chrono::steady_clock; + using Microseconds = std::chrono::duration; + static const Clock::time_point tm_start = Clock::now(); + return Microseconds(Clock::now() - tm_start).count(); +} + +using PgsConvergenceTest = MujocoTest; + +// get total solver iterations across all islands +int get_total_solver_iters(const mjData* d) { + int nisland = mjMAX(1, mjMIN(d->nisland, mjNISLAND)); + int total = 0; + for (int i = 0; i < nisland; i++) { + total += d->solver_niter[i]; + } + return total; +} + +// run solver benchmark at various iteration counts, print table +void run_benchmark(mjModel* model, mjData* data, + const std::vector& all_qpos, + const std::vector& all_qvel, + const std::vector& all_warmstart, + const std::vector& all_qacc, + int nq, int nv, int kStride, int kNumEval, + const int* kIterCounts, int kNumIter, + const char* label) { + std::printf("\n %s:\n", label); + std::printf(" %6s | %11s | %11s | %10s | %11s\n", + "Iters", "Mean Err", "Max Err", "Mean Iters", "Solver us"); + std::printf(" %s\n", + "-------+-------------+-------------+------------+-----------"); + + for (int c = 0; c < kNumIter; c++) { + model->opt.iterations = kIterCounts[c]; + + for (int i = 0; i < mjNTIMER; i++) { + data->timer[i].duration = 0; + data->timer[i].number = 0; + } + + int total_iters = 0; + mjtNum sum_rel_err = 0; + mjtNum max_rel_err = 0; + + for (int e = 0; e < kNumEval; e++) { + int idx = e * kStride; + + mju_copy(data->qpos, all_qpos.data() + idx*nq, nq); + mju_copy(data->qvel, all_qvel.data() + idx*nv, nv); + mju_copy(data->qacc_warmstart, all_warmstart.data() + idx*nv, nv); + + mj_forward(model, data); + total_iters += get_total_solver_iters(data); + + mjtNum newton_norm = mju_norm(all_qacc.data() + idx*nv, nv); + mjtNum err = 0; + for (int j = 0; j < nv; j++) { + mjtNum diff = data->qacc[j] - all_qacc[idx*nv + j]; + err += diff * diff; + } + mjtNum rel_err = mju_sqrt(err) / mju_max(newton_norm, 1e-10); + sum_rel_err += rel_err; + if (rel_err > max_rel_err) { + max_rel_err = rel_err; + } + } + + mjtNum solver_time = data->timer[mjTIMER_CONSTRAINT].duration; + mjtNum mean_iters = static_cast(total_iters) / kNumEval; + std::printf(" %6d | %11.4e | %11.4e | %10.2f | %11.2f\n", + kIterCounts[c], sum_rel_err / kNumEval, max_rel_err, + mean_iters, solver_time); + } + std::printf(" %s\n", + "-------+-------------+-------------+------------+-----------"); +} + +// run pipeline mode: consecutive steps with tolerance, print summary +void run_pipeline(mjModel* model, mjData* data, int kNumSteps, + const char* label) { + std::printf("\n %s Pipeline mode (mj_step, tolerance = 1e-8):\n", label); + + mj_resetData(model, data); + for (int i = 0; i < mjNTIMER; i++) { + data->timer[i].duration = 0; + data->timer[i].number = 0; + } + + int pipe_total_iters = 0; + for (int i = 0; i < kNumSteps; i++) { + mj_step(model, data); + pipe_total_iters += get_total_solver_iters(data); + } + + mjtNum pipe_constraint = data->timer[mjTIMER_CONSTRAINT].duration; + mjtNum pipe_step = data->timer[mjTIMER_STEP].duration; + int pipe_step_count = data->timer[mjTIMER_STEP].number; + + std::printf(" Steps/s : %.0f\n", + pipe_step_count > 0 ? 1e6 * pipe_step_count / pipe_step : 0.0); + std::printf(" us/step (total) : %.1f\n", + pipe_step_count > 0 ? pipe_step / pipe_step_count : 0.0); + std::printf(" us/step (constr) : %.1f\n", + pipe_step_count > 0 ? pipe_constraint / pipe_step_count : 0.0); + std::printf(" Iters/step : %.2f\n", + pipe_step_count > 0 + ? static_cast(pipe_total_iters) / pipe_step_count + : 0.0); +} + +TEST_F(PgsConvergenceTest, PGSConvergence) { + static const char* const kPath = + "engine/testdata/forward/perf/2humanoid100_PGS.xml"; + const std::string xml_path = GetTestDataFilePath(kPath); + char error[1024]; + mjModel* model = + mj_loadXML(xml_path.c_str(), nullptr, error, sizeof(error)); + ASSERT_THAT(model, NotNull()) << error; + + // disable islands: monolithic solver makes statistics simpler + model->opt.disableflags |= mjDSBL_ISLAND; + + int nq = model->nq; + int nv = model->nv; + + // stronger sideways gravity + model->opt.gravity[0] = -2; + model->opt.gravity[1] = -2; + model->opt.gravity[2] = -10; + + // configure Newton ground truth + model->opt.solver = mjSOL_NEWTON; + model->opt.jacobian = mjJAC_SPARSE; + model->opt.tolerance = 0; + model->opt.iterations = 10; + model->opt.disableflags &= ~mjDSBL_WARMSTART; + + mjData* data = mj_makeData(model); + mjcb_time = gettm; + + // roll out Newton for kNumSteps, save pre-step state and post-step qacc + constexpr int kNumSteps = 1000; + std::vector all_qpos(kNumSteps * nq); + std::vector all_qvel(kNumSteps * nv); + std::vector all_warmstart(kNumSteps * nv); + std::vector all_qacc(kNumSteps * nv); + + std::printf("Rolling out ground truth Newton steps...\n"); + for (int i = 0; i < kNumSteps; i++) { + mju_copy(all_qpos.data() + i*nq, data->qpos, nq); + mju_copy(all_qvel.data() + i*nv, data->qvel, nv); + mju_copy(all_warmstart.data() + i*nv, data->qacc_warmstart, nv); + + mj_step(model, data); + + mju_copy(all_qacc.data() + i*nv, data->qacc, nv); + } + + // evaluation points + constexpr int kNumEval = 100; + constexpr int kStride = kNumSteps / kNumEval; + + // iteration counts to test + constexpr int kIterCounts[] = {5, 10, 20, 40, 80, 160, 320}; + constexpr int kNumIter = sizeof(kIterCounts) / sizeof(kIterCounts[0]); + + std::printf("\nPGS vs Nesterov PGS: 2humanoid100_PGS.xml\n"); + std::printf(" %d Newton steps, %d evaluation points\n", + kNumSteps, kNumEval); + std::printf(" nv = %d, nq = %d, nefc = %d\n", nv, nq, data->nefc); + + // switch to PGS + model->opt.solver = mjSOL_PGS; + model->opt.tolerance = 0; + model->opt.disableflags &= ~mjDSBL_WARMSTART; + + // 1. Row PGS (default) + mj_nesterov_momentum = 0; + run_benchmark(model, data, all_qpos, all_qvel, all_warmstart, all_qacc, + nq, nv, kStride, kNumEval, kIterCounts, kNumIter, + "Row PGS (Warmstart, tolerance = 0)"); + + // 2. Nesterov PGS (via REFSAFE hack) + mj_nesterov_momentum = 1; + run_benchmark(model, data, all_qpos, all_qvel, all_warmstart, all_qacc, + nq, nv, kStride, kNumEval, kIterCounts, kNumIter, + "Nesterov PGS (Warmstart, tolerance = 0)"); + + // 3. Pipeline: Row PGS + mj_nesterov_momentum = 0; + model->opt.tolerance = 1e-8; + model->opt.iterations = 100; + run_pipeline(model, data, kNumSteps, "Row PGS"); + + // 4. Pipeline: Nesterov PGS + mj_nesterov_momentum = 1; + model->opt.tolerance = 1e-8; + model->opt.iterations = 100; + run_pipeline(model, data, kNumSteps, "Nesterov PGS"); + + std::printf("\n"); + + // ========== PASS 2: ISLANDS ENABLED ========== + model->opt.disableflags &= ~mjDSBL_ISLAND; + + // do one forward pass to get nisland + mj_resetData(model, data); + mj_forward(model, data); + std::printf("\n============================================\n"); + std::printf("PASS 2: ISLANDS ENABLED\n"); + std::printf("============================================\n"); + + // Pipeline: Row PGS with islands + mj_nesterov_momentum = 0; + model->opt.tolerance = 1e-8; + model->opt.iterations = 100; + run_pipeline(model, data, kNumSteps, "Row PGS (islands)"); + + // Pipeline: Nesterov PGS with islands + mj_nesterov_momentum = 1; + model->opt.tolerance = 1e-8; + model->opt.iterations = 100; + run_pipeline(model, data, kNumSteps, "Nesterov PGS (islands)"); + + // Reset to default + mj_nesterov_momentum = 1; + + std::printf("\n"); + mj_deleteData(data); + mj_deleteModel(model); +} + +} // namespace +} // namespace mujoco diff --git a/test/engine/engine_solver_test.cc b/test/engine/engine_solver_test.cc index 064524b1..0684aade 100644 --- a/test/engine/engine_solver_test.cc +++ b/test/engine/engine_solver_test.cc @@ -241,15 +241,15 @@ TEST_F(SolverTest, SolversEquivalent) { { .newton = MjTol(1e-13, 1e-5), .cg = MjTol(1e-13, 1e-5), - .pgs_pyramidal = MjTol(1e-12, 1e-5), - .pgs_elliptic = MjTol(1e-3, 1e-2), + .pgs_pyramidal = MjTol(1e-13, 1e-5), + .pgs_elliptic = MjTol(1e-3, 1e-3), }}, {.path = kHumanoidPath, .tolerances = { .newton = MjTol(1e-13, 1e-5), .cg = MjTol(1e-12, 1e-5), - .pgs_pyramidal = MjTol(1e-5, 1e-5), + .pgs_pyramidal = MjTol(1e-12, 1e-5), .pgs_elliptic = MjTol(1e-8, 1e-4), }}, }; diff --git a/test/engine/testdata/forward/perf/2humanoid100_PGS.xml b/test/engine/testdata/forward/perf/2humanoid100_PGS.xml new file mode 100644 index 00000000..5b549c4f --- /dev/null +++ b/test/engine/testdata/forward/perf/2humanoid100_PGS.xml @@ -0,0 +1,119 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/test/engine/testdata/forward/perf/humanoid.xml b/test/engine/testdata/forward/perf/humanoid.xml new file mode 100644 index 00000000..59885477 --- /dev/null +++ b/test/engine/testdata/forward/perf/humanoid.xml @@ -0,0 +1,252 @@ + +