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), }},