From 828052e6f490a8541652ae0cbadc2829286b1062 Mon Sep 17 00:00:00 2001 From: Yuval Tassa Date: Mon, 8 Jun 2026 08:53:35 -0700 Subject: [PATCH] CG solver: Replace PRP+ with Hager-Zhang update MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- doc/changelog.rst | 2 ++ doc/computation/index.rst | 6 +++--- src/engine/engine_solver.c | 42 +++++++++++++++++++++++++++++++++++--- 3 files changed, 44 insertions(+), 6 deletions(-) diff --git a/doc/changelog.rst b/doc/changelog.rst index abf1bd26..4f00fe86 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -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 ` pull request `1374 `__. +- The :ref:`CG solver` 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` element for programmatically creating flex objects with auto-generated bodies, joints, and equality constraints. Exposed as ``body.make_flex()`` in Python. diff --git a/doc/computation/index.rst b/doc/computation/index.rst index 72c71410..a0d890f1 100644 --- a/doc/computation/index.rst +++ b/doc/computation/index.rst @@ -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 diff --git a/src/engine/engine_solver.c b/src/engine/engine_solver.c index 8e1dd40f..8f7e4676 100644 --- a/src/engine/engine_solver.c +++ b/src/engine/engine_solver.c @@ -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++) {