CG solver: Replace PRP+ with Hager-Zhang update

Replace the Polak-Ribière-Plus (PRP+) conjugate direction update with the Hager-Zhang formula in `mj_solPrimal`. While this change has negligible effect under float64, it leads to a significant 17.5% throughput speedup over PRP+ under float32 (measured via `engine_cg_convergence_test`). This performance gain is driven by: * A 9.4% reduction in CG iterations per step. * A 12.1% reduction in line search evaluations per step. The full output of the comparison is ``` ================================================================ 1/4: HZ + float64 ================================================================ CG Convergence: 2humanoid100.xml 1000 Newton steps, 100 evaluation points nv = 654, nq = 756 metric: ||qacc_cg - qacc_newton|| / ||qacc_newton|| Warmstart (tolerance = 0): Iters | Mean Err | Max Err | Mean Iters | LS evals -------+-------------+-------------+------------+--------- 5 | 3.8532e-01 | 1.4285e+00 | 5.00 | 22368 10 | 2.0809e-01 | 9.9315e-01 | 10.00 | 44591 20 | 6.4977e-02 | 2.3410e-01 | 20.00 | 89417 40 | 5.2839e-03 | 2.1916e-02 | 40.00 | 180984 80 | 3.9709e-05 | 3.0010e-04 | 80.00 | 364241 160 | 3.2498e-09 | 3.5480e-08 | 160.00 | 730131 -------+-------------+-------------+------------+--------- No warmstart (tolerance = 0): Iters | Mean Err | Max Err | Mean Iters | LS evals -------+-------------+-------------+------------+--------- 5 | 9.0740e-01 | 5.0045e+00 | 5.00 | 22123 10 | 3.9506e-01 | 2.1927e+00 | 10.00 | 44879 20 | 1.2496e-01 | 6.8911e-01 | 20.00 | 90341 40 | 1.2846e-02 | 6.9533e-02 | 40.00 | 182869 80 | 1.0288e-04 | 8.6457e-04 | 80.00 | 368312 160 | 5.9820e-09 | 5.0903e-08 | 160.00 | 738479 -------+-------------+-------------+------------+--------- Tolerance sweep (iterations = 100, warmstart): Tol | Mean Err | Max Err | Mean Iters | Max Iters | Solver us | LS evals -----------+-------------+-------------+------------+------------+-------------+--------- 1e-04 | 1.1297e-02 | 8.9331e-02 | 34.88 | 51 | 131002.94 | 12052 1e-06 | 1.0907e-03 | 7.7468e-03 | 50.87 | 72 | 184172.44 | 16700 1e-08 | 1.1485e-04 | 1.0085e-03 | 66.82 | 91 | 232189.34 | 20194 1e-10 | 1.1357e-05 | 8.1547e-05 | 81.97 | 100 | 276114.42 | 23264 1e-12 | 3.7695e-06 | 2.8592e-05 | 90.47 | 100 | 299526.58 | 24964 0 | 3.5019e-06 | 2.8592e-05 | 100.00 | 100 | 2062208.36 | 456276 -----------+-------------+-------------+------------+------------+-------------+--------- Total solver time: 3185214.08 us, avg time per iter: 74.9445 us Pipeline mode (consecutive mj_step, tolerance = 1e-8): 1000 steps, nv = 654 Steps/s : 380 us/step (total) : 2630.0 us/step (constr) : 2103.7 (80.0%) CG iters/step : 63.28 LS evals/step : 190.75 us/iter : 33.24 ================================================================ 2/4: PRP+ + float64 ================================================================ CG Convergence: 2humanoid100.xml 1000 Newton steps, 100 evaluation points nv = 654, nq = 756 metric: ||qacc_cg - qacc_newton|| / ||qacc_newton|| Warmstart (tolerance = 0): Iters | Mean Err | Max Err | Mean Iters | LS evals -------+-------------+-------------+------------+--------- 5 | 3.8533e-01 | 1.4285e+00 | 5.00 | 22228 10 | 2.0808e-01 | 9.9315e-01 | 10.00 | 44873 20 | 6.4978e-02 | 2.3410e-01 | 20.00 | 89349 40 | 5.2895e-03 | 2.1916e-02 | 40.00 | 179827 80 | 4.0188e-05 | 3.0010e-04 | 80.00 | 363740 160 | 3.2891e-09 | 3.5480e-08 | 160.00 | 733905 -------+-------------+-------------+------------+--------- No warmstart (tolerance = 0): Iters | Mean Err | Max Err | Mean Iters | LS evals -------+-------------+-------------+------------+--------- 5 | 9.0740e-01 | 5.0045e+00 | 5.00 | 22590 10 | 3.9506e-01 | 2.1927e+00 | 10.00 | 45093 20 | 1.2496e-01 | 6.8911e-01 | 20.00 | 90393 40 | 1.2846e-02 | 6.9533e-02 | 40.00 | 182422 80 | 1.0288e-04 | 8.6457e-04 | 80.00 | 367264 160 | 5.9811e-09 | 5.0903e-08 | 160.00 | 739556 -------+-------------+-------------+------------+--------- Tolerance sweep (iterations = 100, warmstart): Tol | Mean Err | Max Err | Mean Iters | Max Iters | Solver us | LS evals -----------+-------------+-------------+------------+------------+-------------+--------- 1e-04 | 1.1325e-02 | 8.9941e-02 | 34.95 | 52 | 129456.81 | 12060 1e-06 | 1.0910e-03 | 7.7469e-03 | 50.85 | 72 | 181004.51 | 16687 1e-08 | 1.1488e-04 | 1.0085e-03 | 66.83 | 91 | 228183.13 | 20189 1e-10 | 1.1401e-05 | 8.1547e-05 | 81.98 | 100 | 269635.40 | 23259 1e-12 | 3.8040e-06 | 2.8592e-05 | 90.52 | 100 | 293811.25 | 24967 0 | 3.5543e-06 | 2.8592e-05 | 100.00 | 100 | 2051927.47 | 456000 -----------+-------------+-------------+------------+------------+-------------+--------- Total solver time: 3154018.57 us, avg time per iter: 74.1895 us Pipeline mode (consecutive mj_step, tolerance = 1e-8): 1000 steps, nv = 654 Steps/s : 382 us/step (total) : 2616.6 us/step (constr) : 2091.8 (79.9%) CG iters/step : 63.57 LS evals/step : 193.35 us/iter : 32.91 ================================================================ 3/4: HZ + float32 ================================================================ CG Convergence: 2humanoid100.xml 1000 Newton steps, 100 evaluation points nv = 654, nq = 756 metric: ||qacc_cg - qacc_newton|| / ||qacc_newton|| Warmstart (tolerance = 0): Iters | Mean Err | Max Err | Mean Iters | LS evals -------+-------------+-------------+------------+--------- 5 | 3.6748e-01 | 1.2281e+00 | 5.00 | 19239 10 | 2.0034e-01 | 6.8689e-01 | 10.00 | 39001 20 | 6.1972e-02 | 1.9859e-01 | 20.00 | 80112 40 | 4.0704e-03 | 1.5797e-02 | 40.00 | 168211 80 | 2.4380e-05 | 1.5803e-04 | 79.45 | 347735 160 | 7.7413e-07 | 4.4732e-06 | 157.27 | 704210 -------+-------------+-------------+------------+--------- No warmstart (tolerance = 0): Iters | Mean Err | Max Err | Mean Iters | LS evals -------+-------------+-------------+------------+--------- 5 | 9.3385e-01 | 4.9648e+00 | 5.00 | 19272 10 | 3.6117e-01 | 1.9618e+00 | 10.00 | 37982 20 | 9.7490e-02 | 5.1451e-01 | 20.00 | 76502 40 | 9.5924e-03 | 5.5088e-02 | 40.00 | 162782 80 | 6.7689e-05 | 5.1113e-04 | 79.91 | 344281 160 | 1.5541e-06 | 7.6955e-06 | 157.30 | 697076 -------+-------------+-------------+------------+--------- Tolerance sweep (iterations = 100, warmstart): Tol | Mean Err | Max Err | Mean Iters | Max Iters | Solver us | LS evals -----------+-------------+-------------+------------+------------+-------------+--------- 1e-04 | 1.5891e-02 | 1.2929e-01 | 32.56 | 50 | 136254.00 | 10980 1e-06 | 1.5719e-03 | 1.0990e-02 | 47.42 | 68 | 193566.00 | 15418 1e-08 | 1.6260e-04 | 1.0680e-03 | 62.70 | 86 | 252661.00 | 20162 1e-10 | 1.5664e-05 | 1.0999e-04 | 78.10 | 100 | 311535.00 | 24509 1e-12 | 2.8152e-06 | 1.5236e-05 | 90.39 | 100 | 356814.00 | 27902 0 | 2.4989e-06 | 1.5025e-05 | 99.05 | 100 | 1980580.00 | 436825 -----------+-------------+-------------+------------+------------+-------------+--------- Total solver time: 3231410.00 us, avg time per iter: 78.7726 us Pipeline mode (consecutive mj_step, tolerance = 1e-8): 1000 steps, nv = 654 Steps/s : 349 us/step (total) : 2862.7 us/step (constr) : 2379.3 (83.1%) CG iters/step : 61.97 LS evals/step : 199.01 us/iter : 38.39 ================================================================ 4/4: PRP+ + float32 ================================================================ CG Convergence: 2humanoid100.xml 1000 Newton steps, 100 evaluation points nv = 654, nq = 756 metric: ||qacc_cg - qacc_newton|| / ||qacc_newton|| Warmstart (tolerance = 0): Iters | Mean Err | Max Err | Mean Iters | LS evals -------+-------------+-------------+------------+--------- 5 | 3.6753e-01 | 1.2281e+00 | 5.00 | 19052 10 | 2.0028e-01 | 6.8689e-01 | 10.00 | 38761 20 | 6.1737e-02 | 1.9863e-01 | 20.00 | 79937 40 | 4.0915e-03 | 1.5821e-02 | 39.99 | 168159 80 | 2.4360e-05 | 1.5808e-04 | 79.59 | 348385 160 | 7.8367e-07 | 4.1549e-06 | 157.04 | 700760 -------+-------------+-------------+------------+--------- No warmstart (tolerance = 0): Iters | Mean Err | Max Err | Mean Iters | LS evals -------+-------------+-------------+------------+--------- 5 | 9.3385e-01 | 4.9648e+00 | 5.00 | 19494 10 | 3.6117e-01 | 1.9618e+00 | 10.00 | 38480 20 | 9.7485e-02 | 5.1451e-01 | 20.00 | 77213 40 | 9.5904e-03 | 5.5078e-02 | 40.00 | 162472 80 | 6.7647e-05 | 5.1087e-04 | 79.87 | 342990 160 | 1.4979e-06 | 7.2291e-06 | 157.92 | 699052 -------+-------------+-------------+------------+--------- Tolerance sweep (iterations = 100, warmstart): Tol | Mean Err | Max Err | Mean Iters | Max Iters | Solver us | LS evals -----------+-------------+-------------+------------+------------+-------------+--------- 1e-04 | 1.5998e-02 | 1.3172e-01 | 32.49 | 50 | 135111.00 | 10953 1e-06 | 1.5878e-03 | 1.1052e-02 | 47.33 | 67 | 190914.00 | 15418 1e-08 | 1.5901e-04 | 1.0680e-03 | 62.65 | 84 | 250609.00 | 20215 1e-10 | 1.5953e-05 | 1.1116e-04 | 78.01 | 100 | 307067.00 | 24562 1e-12 | 2.8520e-06 | 1.5174e-05 | 90.27 | 100 | 353010.00 | 27931 0 | 2.5161e-06 | 1.5129e-05 | 99.39 | 100 | 1999478.00 | 438790 -----------+-------------+-------------+------------+------------+-------------+--------- Total solver time: 3236189.00 us, avg time per iter: 78.9045 us Pipeline mode (consecutive mj_step, tolerance = 1e-8): 1000 steps, nv = 654 Steps/s : 298 us/step (total) : 3352.8 us/step (constr) : 2832.0 (84.5%) CG iters/step : 68.43 LS evals/step : 226.33 us/iter : 41.39 ``` PiperOrigin-RevId: 928590104 Change-Id: I1a96730f50f444d6141d8978feb3519009daf320
2026-06-08 08:53:35 -07:00
parent 7e790879ef
commit 828052e6f4
3 changed files with 44 additions and 6 deletions
@@ -13,6 +13,8 @@ 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>`__.
+- The :ref:`CG solver<soAlgorithms>` now uses the Hager-Zhang conjugate direction update instead of the
+  Polak-Ribiere-Plus formula. This improves convergence and leads to a significant speedup under float32.
 - 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.
@@ -1384,9 +1384,9 @@ Each solver algorithm can be used with both pyramidal and elliptic friction cone
 representations of the constraint Jacobian and related matrices.

 **CG** : conjugate gradient method
-   This algorithm uses the non-linear conjugate gradient method with the Polak-Ribiere-Plus formula (non-negative
-   :math:`\beta`). Line-search is exact, using Newton's method in one dimension with analytical second derivatives on
-   the piecewise-quadratic cost. CG has no setup cost.
+   This algorithm uses the non-linear conjugate gradient method with the Hager-Zhang formula. Line-search is
+   exact, using Newton's method in one dimension with analytical second derivatives on the piecewise-quadratic
+   cost. CG has no setup cost.

 **Newton** : Newton's method
   This algorithm implements the exact Newton method, with analytical second-order derivatives and Cholesky
@@ -920,7 +920,8 @@ typedef struct {
  // CG arrays (PrimalAllocate, CG only)
  mjtNum* gradold;        // previous gradient                            (nv x 1)
  mjtNum* Mgradold;       // previous preconditioned gradient             (nv x 1)
-  mjtNum* Mgraddif;       // gradient difference                          (nv x 1)
+  mjtNum* graddif;        // grad - gradold                               (nv x 1)
+  mjtNum* Mgraddif;       // M\(grad - gradold)                           (nv x 1)

  // Newton arrays, known-size (PrimalAllocate)
  mjtNum* D;              // constraint inertia                           (nefc x 1)
@@ -1091,7 +1092,7 @@ static void PrimalAllocate(const mjModel* m, mjData* d, mjPrimalContext* ctx, in
      if (is_elliptic) nNum += nv*nv;  // Lcone (dense)
    }
  } else {
-    nNum += 3*nv;                      // CG arrays
+    nNum += 4*nv;                      // CG arrays
  }

  // add island matrix sizes
@@ -1185,6 +1186,7 @@ static void PrimalAllocate(const mjModel* m, mjData* d, mjPrimalContext* ctx, in
  } else {
    ctx->gradold  = numblock;  numblock += nv;
    ctx->Mgradold = numblock;  numblock += nv;
+    ctx->graddif  = numblock;  numblock += nv;
    ctx->Mgraddif = numblock;  numblock += nv;
  }

@@ -2283,7 +2285,40 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i
    if (flg_Newton) {
      mju_scl(ctx.search, ctx.Mgrad, -1, nv);
    } else {
-      // Polak-Ribiere
+#ifndef mjCG_PRP
+      // Hager-Zhang conjugate direction update
+      mjtNum d_dot_y, y_dot_My, y_dot_Mgrad, d_dot_grad;
+      mjtNum beta_hz;
+      mjtNum d_norm, grad_norm, eta_k;
+      const mjtNum eta = 0.01;
+
+      // graddif = grad - gradold, Mgraddif = Mgrad - Mgradold
+      mju_sub(ctx.graddif, ctx.grad, ctx.gradold, nv);
+      mju_sub(ctx.Mgraddif, ctx.Mgrad, ctx.Mgradold, nv);
+
+      // compute d'*y; restart to steepest descent if conjugacy is lost
+      d_dot_y = mju_dot(ctx.search, ctx.graddif, nv);
+      if (d_dot_y < mjMINVAL) {
+        beta = 0;
+      } else {
+        // compute remaining inner products for the HZ formula
+        y_dot_My = mju_dot(ctx.graddif, ctx.Mgraddif, nv);
+        y_dot_Mgrad = mju_dot(ctx.graddif, ctx.Mgrad, nv);
+        d_dot_grad = mju_dot(ctx.search, ctx.grad, nv);
+
+        // primary Hager-Zhang beta coefficient
+        beta_hz = (y_dot_Mgrad - 2*(y_dot_My/d_dot_y)*d_dot_grad) / d_dot_y;
+
+        // dynamic truncation threshold to ensure d is not orthogonal to grad
+        d_norm = mju_norm(ctx.search, nv);
+        grad_norm = mju_norm(ctx.grad, nv);
+        eta_k = -1.0 / mju_max(mjMINVAL, d_norm * mju_min(eta, grad_norm));
+
+        // apply lower bound
+        beta = mju_max(eta_k, beta_hz);
+      }
+#else
+      // Polak-Ribiere-Plus conjugate direction update
      mju_sub(ctx.Mgraddif, ctx.Mgrad, ctx.Mgradold, nv);
      beta = mju_dot(ctx.grad, ctx.Mgraddif, nv) /
             mju_max(mjMINVAL, mju_dot(ctx.gradold, ctx.Mgradold, nv));
@@ -2292,6 +2327,7 @@ static void mj_solPrimal(const mjModel* m, mjData* d, int island, int maxiter, i
      if (beta < 0) {
        beta = 0;
      }
+#endif

      // update
      for (int i=0; i < nv; i++) {