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
2026-05-31 03:59:31 -07:00
parent 4358a102cd
commit cd6db9ebe2
6 changed files with 208 additions and 42 deletions
@@ -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 <mjwarp/index>` pull request
+  `1374 <https://github.com/google-deepmind/mujoco_warp/pull/1374>`__.
 - Added :ref:`mjs_makeFlex`, a new C API function equivalent to the :ref:`flexcomp<body-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.

@@ -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;
    }

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