From cd6db9ebe24d448e47018d357d7d131b42715221 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Sun, 31 May 2026 03:59:31 -0700 Subject: [PATCH] Improve CG solver precision and stability under float32 via line search refactor Fixes #2313 This CL combines two complementary improvements to the Conjugate Gradient (CG) solver, significantly improving numerical precision, stability, and efficiency, particularly in single precision (float32). 1. Line Search Cost Evaluation Refactor: Previously, solver improvement was calculated by subtracting absolute costs: cost(alpha) - cost(0). In highly converged states or single precision, this is susceptible to catastrophic cancellation. We refactor PrimalSearch to compute the cost delta directly, dramatically improving precision. 2. Improved Solver Termination Condition: Near the float32 precision limit, line search deltas can occasionally be slightly negative due to numerical noise. Previously, any value below m->opt.tolerance (including negative values) triggered termination, halting the solver and locking in destabilizing steps. We update the termination condition to require positive improvement (0 < improvement < m->opt.tolerance), allowing the solver to continue iterating and recover stability. Together, these changes yield substantial improvements, see reduced tolerances herein. PiperOrigin-RevId: 924229669 Change-Id: Ic0bbefaed090f3a8b1e79ab8d45422c3e86fb56c --- doc/changelog.rst | 5 +- src/engine/engine_solver.c | 225 +++++++++++++++++++++---- test/engine/engine_core_smooth_test.cc | 2 +- test/engine/engine_inverse_test.cc | 2 +- test/engine/engine_sensor_test.cc | 4 +- test/engine/engine_solver_test.cc | 12 +- 6 files changed, 208 insertions(+), 42 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index 5862cf42..0da93da0 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -7,6 +7,9 @@ Upcoming version (not yet released) General ^^^^^^^ +- Improved primal solver convergence under float32. Improvements initially proposed by :github:user:`n3b` in + :issue:`2313` and :github:user:`denzeler-nvidia` in :doc:`MJWarp ` pull request + `1374 `__. - Added :ref:`mjs_makeFlex`, a new C API function equivalent to the :ref:`flexcomp` element for programmatically creating flex objects with auto-generated bodies, joints, and equality constraints. Exposed as ``body.make_flex()`` in Python. @@ -17,7 +20,7 @@ General .. admonition:: Breaking API changes :class: attention - - The header file ``mjthread.h`` was removed along with the engine threading API. + - The header file ``mjthread.h`` was removed along with the old engine threading API. **Migration:** Use :ref:`mju_threadpool` to set number of worker threads for the engine. diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 5c7fa0ea..8e1dd40f 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -1367,7 +1367,145 @@ struct _mjPrimalPnt { typedef struct _mjPrimalPnt mjPrimalPnt; -// evaluate linesearch cost, return first and second derivatives +// Huber cost of a single friction constraint at a given point x +static mjtNum frictionCost(mjtNum x, mjtNum f, mjtNum Rf, mjtNum D) { + // -bound < x < bound : quadratic + if (-Rf < x && x < Rf) return 0.5*D*x*x; + + // x < -bound : linear negative + else if (x <= -Rf) return f*(-0.5*Rf - x); + + // bound < x : linear positive + else return f*(-0.5*Rf + x); +} + + +// Huber cost difference of a single friction constraint: cost(x) - cost(start) +static mjtNum frictionCostDif(mjtNum start, mjtNum x, mjtNum f, mjtNum Rf, mjtNum D) { + int state_start = (-Rf < start && start < Rf) ? 0 : (start <= -Rf ? -1 : 1); + int state_x = (-Rf < x && x < Rf) ? 0 : (x <= -Rf ? -1 : 1); + + // both quadratic + if (state_start == 0 && state_x == 0) { + return 0.5*D*(x - start)*(x + start); + } + + // both linear negative + if (state_start == -1 && state_x == -1) { + return f*(start - x); + } + + // both linear positive + if (state_start == 1 && state_x == 1) { + return f*(x - start); + } + + // otherwise different zones: compute absolute costs and subtract + return frictionCost(x, f, Rf, D) - frictionCost(start, f, Rf, D); +} + + +// compute cost of an elliptic cone at a given alpha +static mjtNum ellipticCost(const mjtNum* quad, mjtNum alpha, mjtNum mu, mjtNum Dm) { + mjtNum U0 = quad[3], V0 = quad[4], UU = quad[5]; + mjtNum UV = quad[6], VV = quad[7]; + mjtNum N = U0 + alpha*V0; + mjtNum Tsqr = UU + alpha*(2*UV + alpha*VV); + + // no tangential force : top or bottom zone + if (Tsqr <= 0) { + // bottom zone: quadratic cost + if (N < 0) return alpha*alpha*quad[2] + alpha*quad[1] + quad[0]; + + // top zone: nothing to do + } + + // otherwise regular processing + else { + mjtNum T = mju_sqrt(Tsqr); + // N>=mu*T : top zone + if (N >= mu*T) { + // nothing to do + } + + // mu*N+T<=0 : bottom zone + else if (mu*N+T <= 0) { + return alpha*alpha*quad[2] + alpha*quad[1] + quad[0]; + } + + // otherwise middle zone + else { + return 0.5*Dm*(N-mu*T)*(N-mu*T); + } + } + + return 0; +} + + +// compute cost difference of an elliptic cone at a given alpha: cost(alpha) - cost(0) +static mjtNum ellipticCostDif(const mjtNum* quad, mjtNum alpha, mjtNum mu, mjtNum Dm) { + mjtNum U0 = quad[3], V0 = quad[4], UU = quad[5]; + mjtNum UV = quad[6], VV = quad[7]; + + // determine zone and cost at alpha=0 + int zone0 = 0; + mjtNum T0 = 0; + if (UU <= 0) { + zone0 = (U0 < 0) ? 2 : 1; + } else { + T0 = mju_sqrt(UU); + if (U0 >= mu*T0) { + zone0 = 1; // top zone + } else if (mu*U0 + T0 <= 0) { + zone0 = 2; // bottom zone + } else { + zone0 = 3; // middle zone + } + } + + // determine zone and cost at alpha + mjtNum N = U0 + alpha*V0; + mjtNum Tsqr = UU + alpha*(2*UV + alpha*VV); + int zone_alpha = 0; + mjtNum T = 0; + + if (Tsqr <= 0) { + zone_alpha = (N < 0) ? 2 : 1; // bottom or top zone + } else { + T = mju_sqrt(Tsqr); + if (N >= mu*T) { + zone_alpha = 1; // top zone + } else if (mu*N + T <= 0) { + zone_alpha = 2; // bottom zone + } else { + zone_alpha = 3; // middle zone + } + } + + // both top zone + if (zone0 == 1 && zone_alpha == 1) { + return 0; + } + + // both bottom zone + if (zone0 == 2 && zone_alpha == 2) { + return alpha*alpha*quad[2] + alpha*quad[1]; + } + + // both middle zone + if (zone0 == 3 && zone_alpha == 3) { + mjtNum diff_alpha = N - mu*T; + mjtNum diff0 = U0 - mu*T0; + return 0.5*Dm*(diff_alpha - diff0)*(diff_alpha + diff0); + } + + // otherwise different zones: compute absolute costs and subtract + return ellipticCost(quad, alpha, mu, Dm) - ellipticCost(quad, 0, mu, Dm); +} + + +// evaluate shifted linesearch cost: cost(alpha) - cost(0), and derivatives static void PrimalEval(mjPrimalContext* ctx, mjPrimalPnt* p) { int ne = ctx->ne, nf = ctx->nf, nefc = ctx->nefc; @@ -1375,41 +1513,45 @@ static void PrimalEval(mjPrimalContext* ctx, mjPrimalPnt* p) { mjtNum cost = 0, alpha = p->alpha; mjtNum deriv[2] = {0, 0}; - // init quad with Gauss - mjtNum quadTotal[3]; - mju_copy3(quadTotal, ctx->quadGauss); + // init quad with Gauss, shifted: drop quadGauss[0] + mjtNum quadTotal[3] = {0, ctx->quadGauss[1], ctx->quadGauss[2]}; // process constraints for (int i=0; i < nefc; i++) { - // equality + // equality: shifted quad (skip quad[0]) if (i < ne) { - mju_addTo3(quadTotal, ctx->quad+3*i); + quadTotal[1] += ctx->quad[3*i+1]; + quadTotal[2] += ctx->quad[3*i+2]; continue; } - // friction + // friction: compute cost(alpha) - cost(0) directly if (i < ne + nf) { // search point, friction loss, bound (Rf) - mjtNum start = ctx->Jaref[i], dir = ctx->Jv[i]; + mjtNum start = ctx->Jaref[i]; + mjtNum dir = ctx->Jv[i]; mjtNum x = start + alpha*dir; mjtNum f = ctx->efc_frictionloss[i]; + mjtNum D = ctx->efc_D[i]; mjtNum Rf = ctx->efc_R[i]*f; + // cost delta + cost += frictionCostDif(start, x, f, Rf, D); + // -bound < x < bound : quadratic if (-Rf < x && x < Rf) { - mju_addTo3(quadTotal, ctx->quad+3*i); + deriv[0] += D*x*dir; + deriv[1] += D*dir*dir; } // x < -bound : linear negative else if (x <= -Rf) { - mjtNum qf[3] = {f*(-0.5*Rf-start), -f*dir, 0}; - mju_addTo3(quadTotal, qf); + deriv[0] += -f*dir; } // bound < x : linear positive else { - mjtNum qf[3] = {f*(-0.5*Rf+start), f*dir, 0}; - mju_addTo3(quadTotal, qf); + deriv[0] += f*dir; } continue; } @@ -1430,15 +1572,19 @@ static void PrimalEval(mjPrimalContext* ctx, mjPrimalPnt* p) { mjtNum VV = quad[7]; mjtNum Dm = quad[8]; - // compute N, Tsqr + // shifted cost + cost += ellipticCostDif(quad, alpha, mu, Dm); + + // compute N, Tsqr for derivatives mjtNum N = U0 + alpha*V0; mjtNum Tsqr = UU + alpha*(2*UV + alpha*VV); // no tangential force : top or bottom zone if (Tsqr <= 0) { - // bottom zone: quadratic cost + // bottom zone: quadratic derivatives if (N < 0) { - mju_addTo3(quadTotal, quad); + deriv[0] += 2*alpha*quad[2] + quad[1]; + deriv[1] += 2*quad[2]; } // top zone: nothing to do @@ -1456,7 +1602,8 @@ static void PrimalEval(mjPrimalContext* ctx, mjPrimalPnt* p) { // mu*N+T<=0 : bottom zone else if (mu*N+T <= 0) { - mju_addTo3(quadTotal, quad); + deriv[0] += 2*alpha*quad[2] + quad[1]; + deriv[1] += 2*quad[2]; } // otherwise middle zone @@ -1465,9 +1612,6 @@ static void PrimalEval(mjPrimalContext* ctx, mjPrimalPnt* p) { mjtNum N1 = V0; mjtNum T1 = (UV + alpha*VV)/T; mjtNum T2 = VV/T - (UV + alpha*VV)*T1/(T*T); - - // add to cost - cost += 0.5*Dm*(N-mu*T)*(N-mu*T); deriv[0] += Dm*(N-mu*T)*(N1-mu*T1); deriv[1] += Dm*((N1-mu*T1)*(N1-mu*T1) + (N-mu*T)*(-mu*T2)); } @@ -1477,16 +1621,23 @@ static void PrimalEval(mjPrimalContext* ctx, mjPrimalPnt* p) { i += (dim-1); } else { // inequality // search point - mjtNum x = ctx->Jaref[i] + alpha*ctx->Jv[i]; + mjtNum start = ctx->Jaref[i]; + mjtNum x = start + alpha*ctx->Jv[i]; + mjtNum cost0 = start < 0 ? ctx->quad[3*i] : 0; // active if (x < 0) { - mju_addTo3(quadTotal, ctx->quad+3*i); + // shifted quad: add quad[1], quad[2] and (quad[0] - cost0) + quadTotal[0] += ctx->quad[3*i] - cost0; + quadTotal[1] += ctx->quad[3*i+1]; + quadTotal[2] += ctx->quad[3*i+2]; + } else { + cost -= cost0; } } } - // add total quadratic + // add total quadratic (quadTotal[0] contains only shifted residuals) cost += alpha*alpha*quadTotal[2] + alpha*quadTotal[1] + quadTotal[0]; deriv[0] += 2*alpha*quadTotal[2] + quadTotal[1]; deriv[1] += 2*quadTotal[2]; @@ -1536,7 +1687,8 @@ static int updateBracket(mjPrimalContext* ctx, // line search -static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_iterations) { +static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_iterations, + mjtNum* improvement) { int nv = ctx->nv, nefc = ctx->nefc; mjPrimalPnt p0, p1, p2, pmid, p1next, p2next; @@ -1544,6 +1696,7 @@ static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_ite ctx->LSiter = 0; ctx->LSresult = 0; ctx->LSslope = 1; // means not computed + *improvement = 0; // save search vector length, check mjtNum snorm = mju_norm(ctx->search, nv); @@ -1587,6 +1740,7 @@ static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_ite ctx->LSresult = 0; // SUCCESS } ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; + *improvement = -p1.cost; return p1.alpha; } @@ -1642,6 +1796,7 @@ static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_ite // check for convergence if (mju_abs(p1.deriv[0]) < gtol) { ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; + *improvement = -p1.cost; return p1.alpha; // SUCCESS } } @@ -1650,6 +1805,7 @@ static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_ite if (ctx->LSiter >= ls_iterations) { ctx->LSresult = 3; // could not bracket ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; + *improvement = -p1.cost; return p1.alpha; } @@ -1657,6 +1813,7 @@ static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_ite if (!p2update) { ctx->LSresult = 6; // no p2 update ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; + *improvement = -p1.cost; return p1.alpha; } @@ -1686,6 +1843,7 @@ static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_ite } if (bestind >= 0) { ctx->LSslope = mju_abs(candidates[bestind].deriv[0])*slopescl; + *improvement = -candidates[bestind].cost; return candidates[bestind].alpha; // SUCCESS } @@ -1695,25 +1853,28 @@ static mjtNum PrimalSearch(mjPrimalContext* ctx, mjtNum tolerance, mjtNum ls_ite // no update possible: numerical accuracy reached, use midpoint if (!b1 && !b2) { - if (pmid.cost < p0.cost) { + if (pmid.cost < 0) { ctx->LSresult = 0; // SUCCESS } else { ctx->LSresult = 7; // no improvement, could not bracket } ctx->LSslope = mju_abs(pmid.deriv[0])*slopescl; + *improvement = -pmid.cost; return pmid.alpha; } } // choose bracket with best cost - if (p1.cost <= p2.cost && p1.cost < p0.cost) { + if (p1.cost <= p2.cost && p1.cost < 0) { ctx->LSresult = 4; // improvement but no convergence ctx->LSslope = mju_abs(p1.deriv[0])*slopescl; + *improvement = -p1.cost; return p1.alpha; - } else if (p2.cost <= p1.cost && p2.cost < p0.cost) { + } else if (p2.cost <= p1.cost && p2.cost < 0) { ctx->LSresult = 4; // improvement but no convergence ctx->LSslope = mju_abs(p2.deriv[0])*slopescl; + *improvement = -p2.cost; return p2.alpha; } else { ctx->LSresult = 5; // no improvement @@ -2069,7 +2230,9 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i // main loop while (iter < maxiter) { // perform linesearch - alpha = PrimalSearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations); + mjtNum ls_improvement; + alpha = PrimalSearch(&ctx, m->opt.tolerance * m->opt.ls_tolerance, m->opt.ls_iterations, + &ls_improvement); // no improvement: done if (alpha == 0) { @@ -2087,7 +2250,6 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i mju_copy(ctx.Mgradold, ctx.Mgrad, nv); } mju_copyInt(oldstate, ctx.efc_state, nefc); - mjtNum oldcost = ctx.cost; // update PrimalUpdateConstraint(&ctx, flg_Newton & (m->opt.cone == mjCONE_ELLIPTIC)); @@ -2103,7 +2265,7 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i } // scale improvement, gradient, save stats - mjtNum improvement = scale * (oldcost - ctx.cost); + mjtNum improvement = scale * ls_improvement; mjtNum gradient = scale * mju_norm(ctx.grad, nv); saveStats(m, d, island, iter, improvement, gradient, ctx.LSslope, ctx.nactive, nchange, ctx.LSiter, ctx.nupdate); @@ -2112,7 +2274,8 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i iter++; // termination - if (improvement < m->opt.tolerance || gradient < m->opt.tolerance) { + if ((improvement > 0 && improvement < m->opt.tolerance) || + gradient < m->opt.tolerance) { break; } diff --git a/test/engine/engine_core_smooth_test.cc b/test/engine/engine_core_smooth_test.cc index 3120b328..df97cd02 100644 --- a/test/engine/engine_core_smooth_test.cc +++ b/test/engine/engine_core_smooth_test.cc @@ -408,7 +408,7 @@ void TestConnect(const char* const filepath) { } for (int i=0; i < 3; i++) { EXPECT_THAT(data->sensordata[i] - model->sensor_user[i], - MjNear(0, 1e-6, 1e-4)); + MjNear(0, 1e-6, 2e-4)); } mj_deleteData(data); mj_deleteModel(model); diff --git a/test/engine/engine_inverse_test.cc b/test/engine/engine_inverse_test.cc index b698a00b..428ea2ce 100644 --- a/test/engine/engine_inverse_test.cc +++ b/test/engine/engine_inverse_test.cc @@ -72,7 +72,7 @@ TEST_F(InverseTest, ForwardInverseMatch) { mjtNum epsilon; switch (solver) { case mjSOL_PGS: epsilon = MjTol(1e-6, 1e-2); break; - case mjSOL_CG: epsilon = MjTol(1e-3, 1e0); break; + case mjSOL_CG: epsilon = MjTol(1e-9, 1e-1); break; case mjSOL_NEWTON: epsilon = MjTol(1e-10, 1e-2); break; } EXPECT_LT(data->solver_fwdinv[0], epsilon) diff --git a/test/engine/engine_sensor_test.cc b/test/engine/engine_sensor_test.cc index a71c5061..1ab1dc77 100644 --- a/test/engine/engine_sensor_test.cc +++ b/test/engine/engine_sensor_test.cc @@ -973,7 +973,7 @@ TEST_F(SensorTest, ContactNet) { mj_applyFT(model, data, force, torque, point, b1, qfrc.data()); // compare - EXPECT_THAT(qfrc, Pointwise(MjNear(1e-6, 1e-4), qfrc_expected)); + EXPECT_THAT(qfrc, Pointwise(MjNear(1e-6, 2e-4), qfrc_expected)); // check net force, sensor returns body2 -> body1 vector net21 = GetSensor(model, data, "net21"); @@ -992,7 +992,7 @@ TEST_F(SensorTest, ContactNet) { mj_applyFT(model, data, force, torque, point, b2, qfrc.data()); // compare - EXPECT_THAT(qfrc, Pointwise(MjNear(1e-6, 1e-4), qfrc_expected)); + EXPECT_THAT(qfrc, Pointwise(MjNear(1e-6, 2e-4), qfrc_expected)); nconmax = std::max(nconmax, data->ncon); } diff --git a/test/engine/engine_solver_test.cc b/test/engine/engine_solver_test.cc index 1eb99420..064524b1 100644 --- a/test/engine/engine_solver_test.cc +++ b/test/engine/engine_solver_test.cc @@ -64,9 +64,9 @@ TEST_F(SolverTest, IslandsEquivalent) { // quite a lot, even at high iteration count and zero {ls_}tolerance. // Increasing the iteration count higher than 60 does not improve convergence. mjtNum rtol[kNumTol] = { - MjTol(1e-1, 2), - MjTol(3e-2, 1), - MjTol(1.5e-4, 3.6) + MjTol(1e-1, 1.2e-1), + MjTol(3e-2, 2e-2), + MjTol(1.3e-5, 2.8e-3) }; for (int i = 0; i < kNumTol; ++i) { @@ -188,7 +188,7 @@ TEST_F(SolverTest, IslandsEquivalentForward) { int worst_idx = -1; mjtNum scale = 0.5 * (mju_norm(data_noisland->qacc, nv) + mju_norm(data_island->qacc, nv)); - mjtNum rtol = solver == mjSOL_CG ? MjTol(1e-6, 1e-2) + mjtNum rtol = solver == mjSOL_CG ? MjTol(1e-8, 1e-4) : MjTol(1e-13, 1e-3); mjtNum worst_allowed = scale * rtol; @@ -240,7 +240,7 @@ TEST_F(SolverTest, SolversEquivalent) { .tolerances = { .newton = MjTol(1e-13, 1e-5), - .cg = MjTol(1e-5, 1e-2), + .cg = MjTol(1e-13, 1e-5), .pgs_pyramidal = MjTol(1e-12, 1e-5), .pgs_elliptic = MjTol(1e-3, 1e-2), }}, @@ -248,7 +248,7 @@ TEST_F(SolverTest, SolversEquivalent) { .tolerances = { .newton = MjTol(1e-13, 1e-5), - .cg = MjTol(1e-5, 1e-2), + .cg = MjTol(1e-12, 1e-5), .pgs_pyramidal = MjTol(1e-5, 1e-5), .pgs_elliptic = MjTol(1e-8, 1e-4), }},