From 6d348751d24846dc188b9fbd990a1dc0241bf207 Mon Sep 17 00:00:00 2001 From: Bardh Hoxha Date: Mon, 10 Aug 2026 12:58:18 -0400 Subject: [PATCH] fix(qp): align fast/cvxopt objective convention, reject dropped equalities, tune PDIPM The registry documents min x'Hx + f'x, but the fast (PDIPM) and cvxopt backends passed H straight into their native 0.5 x'Px + q'x form, so both tracked 2*u_nom instead of u_nom (hidden in safety tests because the active barrier pins the solution). Convert P = 2H at the registry boundary for both; all four backends now agree on the audit probe and a new cross-solver parity suite (interior, single-active, corner, random, and end-to-end CBF-filter cases). The fast backend also silently ignored a_mat/b_vec equality constraints, so MPC through get_solver('fast') returned dynamics-violating trajectories with success status; it now raises NotImplementedError pointing at jaxopt/casadi. PDIPM: factor the reduced Newton matrix once per Mehrotra iteration (predictor and corrector differ only in RHS), compute residuals once for the convergence check and the step, and lower DEFAULT_MAX_ITER 25 -> 16. 16 was calibrated on 250 benign + 250 ill-conditioned near-boundary QPs: deviation vs max_iter=60 is exactly 0.0 on both families (12 fails on 30/250 hard QPs precisely when the barrier is nearly violated; table in-file). Measured 1.64x per solve. fori_loop retained: reverse-mode AD has no while_loop rule and the differentiable-CBF-QP work grads through this solver. The max_iter=25 pins in solve_qp_fast and fast_solver now defer to the solver default. test_fast_qp.py:112 expected [2,4], the registry-convention bug's answer; the correct minimizer of x'x - 2x0 - 4x1 is [1,2] (2Hx + f = 0). --- .../quadratic_program/qp_solver_cvxopt.py | 6 +- .../quadratic_program/qp_solver_fast.py | 9 +- .../quadratic_program/qp_solver_pdipm.py | 101 +++++--- .../quadratic_program/solver_registry.py | 44 +++- tests/test_optimization/test_fast_qp.py | 12 +- .../test_fast_solver_equality_guard.py | 81 +++++++ tests/test_optimization/test_pdipm_qp.py | 33 ++- .../test_solver_convention_parity.py | 228 ++++++++++++++++++ 8 files changed, 462 insertions(+), 52 deletions(-) create mode 100644 tests/test_optimization/test_fast_solver_equality_guard.py create mode 100644 tests/test_optimization/test_solver_convention_parity.py diff --git a/src/cbfkit/optimization/quadratic_program/qp_solver_cvxopt.py b/src/cbfkit/optimization/quadratic_program/qp_solver_cvxopt.py index e8d11b5a..3e3a4427 100644 --- a/src/cbfkit/optimization/quadratic_program/qp_solver_cvxopt.py +++ b/src/cbfkit/optimization/quadratic_program/qp_solver_cvxopt.py @@ -104,8 +104,12 @@ def solve_with_details( ``init_params`` is accepted for interface compatibility but ignored (CVXOPT does not support warm-starting). + + Note: the registry convention is ``min x'Hx + f'x`` while CVXOPT's native + form is ``min 1/2 x'Px + q'x``; the ``P = 2H`` conversion happens here so + all ``get_solver()`` backends agree (see ``solver_registry`` docstring). """ from cbfkit.optimization.quadratic_program.solver_registry import QpSolution - primal, success = solve(h_mat, f_vec, g_mat, h_vec, a_mat, b_vec) + primal, success = solve(2.0 * h_mat, f_vec, g_mat, h_vec, a_mat, b_vec) return QpSolution(primal=primal, status=1 if success else 0, params=None) diff --git a/src/cbfkit/optimization/quadratic_program/qp_solver_fast.py b/src/cbfkit/optimization/quadratic_program/qp_solver_fast.py index 85600b8f..2e613480 100644 --- a/src/cbfkit/optimization/quadratic_program/qp_solver_fast.py +++ b/src/cbfkit/optimization/quadratic_program/qp_solver_fast.py @@ -13,6 +13,7 @@ from jax import Array from cbfkit.optimization.quadratic_program.qp_solver_pdipm import ( + DEFAULT_MAX_ITER, PdipmState, solve_qp_pdipm, ) @@ -24,7 +25,7 @@ def solve_qp_fast( G: Array, h: Array, warm_start: Optional[Array] = None, - max_iter: int = 25, + max_iter: Optional[int] = None, tol: float = 1e-6, ) -> Tuple[Array, int, Array]: """Backward-compatible alias for solve_qp_pdipm. @@ -32,6 +33,10 @@ def solve_qp_fast( Returns ``(x, status, dual)``. For full state including primal slacks (needed for richer warm-starting), call ``solve_qp_pdipm`` directly. + ``max_iter=None`` defers to ``qp_solver_pdipm.DEFAULT_MAX_ITER`` rather + than pinning its own copy of the budget, so tuning the solver default is + not silently masked here. + Note: the legacy interface accepted an Array as ``warm_start`` (the dual only). To preserve that signature, we wrap it in a minimal ``PdipmState`` with slacks defaulting to ones, which the PDIPM init will then clamp @@ -55,7 +60,7 @@ def solve_qp_fast( G, h, warm_start=warm, - max_iter=max_iter, + max_iter=DEFAULT_MAX_ITER if max_iter is None else max_iter, tol=tol, ) return x, status, state.dual diff --git a/src/cbfkit/optimization/quadratic_program/qp_solver_pdipm.py b/src/cbfkit/optimization/quadratic_program/qp_solver_pdipm.py index a3eae3fb..c8bb7c95 100644 --- a/src/cbfkit/optimization/quadratic_program/qp_solver_pdipm.py +++ b/src/cbfkit/optimization/quadratic_program/qp_solver_pdipm.py @@ -46,25 +46,45 @@ def _step_to_boundary(s: Array, ds: Array) -> Array: _EPS_PD = 1e-10 # PD regularization on H before Cholesky -def _solve_newton_reduced(P: Array, G: Array, s: Array, lam: Array, rhs: Array) -> Array: - """Solve (P + G^T diag(lam/s) G + eps*I) dx = rhs via Cholesky. +def _factor_newton_reduced(P: Array, G: Array, s: Array, lam: Array) -> Array: + """Cholesky factor of (P + G^T diag(lam/s) G + eps*I). The eps*I term guards against rounding-induced non-PD near the optimum. + + The predictor and corrector of a Mehrotra iteration share this matrix and + differ only in their right-hand side, so it is factored once per iteration + and reused via ``_solve_with_factor``. """ n = P.shape[0] D = lam / s H = P + G.T @ (D[:, None] * G) + _EPS_PD * jnp.eye(n) - L = jnp.linalg.cholesky(H) - # dx = H^-1 rhs via two triangular solves + return jnp.linalg.cholesky(H) + + +def _solve_with_factor(L: Array, rhs: Array) -> Array: + """Solve H dx = rhs given the Cholesky factor L of H, via two triangular solves.""" y = jax.scipy.linalg.solve_triangular(L, rhs, lower=True) - dx = jax.scipy.linalg.solve_triangular(L.T, y, lower=False) - return dx + return jax.scipy.linalg.solve_triangular(L.T, y, lower=False) _STEP_SAFETY = 0.995 # fraction of step-to-boundary; keeps strict interior - -def _pdipm_iteration( +# Default iteration budget. Because the fori_loop always runs to completion, +# this is the per-solve cost, so it is chosen as the smallest value that is +# still exact rather than a loose upper bound. +# +# Calibrated on two 250-QP suites (n=5, m=24), comparing against max_iter=60: +# - benign QPs (well-conditioned P, slack interior): converged by iteration 8. +# - hard QPs (cond(P) ~ 1e6 from slack penalty weighting, optimum on the +# barrier boundary with many near-active constraints) needed far more: +# max_iter=10 -> 123/250 QPs off by >1e-8; max_iter=12 -> 30/250; +# max_iter=14 -> 1/250; max_iter=16 -> 0/250, all status==1. +# 16 is the first budget that is exact on both suites. Benign problems are +# unaffected in accuracy since they freeze on convergence. +DEFAULT_MAX_ITER = 16 + + +def _residuals( P: Array, q: Array, G: Array, @@ -72,16 +92,33 @@ def _pdipm_iteration( x: Array, s: Array, lam: Array, +) -> Tuple[Array, Array]: + """Dual and primal residuals (r_d, r_p). + + Computed once per iteration and shared by the freeze-on-converge test and + the Newton step, which would otherwise each rebuild them. + """ + r_d = P @ x + G.T @ lam + q + r_p = G @ x + s - h + return r_d, r_p + + +def _pdipm_iteration( + P: Array, + G: Array, + x: Array, + s: Array, + lam: Array, + r_d: Array, + r_p: Array, ) -> Tuple[Array, Array, Array]: - """One Mehrotra predictor-corrector PDIPM step. + """One Mehrotra predictor-corrector PDIPM step, given precomputed residuals. Returns updated (x, s, lam) all strictly positive in s, lam. """ m = G.shape[0] - # --- Residuals (predictor: r_c with mu = 0) --- - r_d = P @ x + G.T @ lam + q - r_p = G @ x + s - h + # Predictor complementarity residual (r_c with mu = 0). r_c_aff = s * lam # --- Predictor (affine) step --- @@ -90,7 +127,9 @@ def _pdipm_iteration( # = -(r_d + G^T ((lam*r_p - s*lam) / s)) # which simplifies via r_c_aff = s*lam to -(r_d + G^T ((lam*r_p - r_c_aff)/s)). rhs_aff = -(r_d + G.T @ ((lam * r_p - r_c_aff) / s)) - dx_aff = _solve_newton_reduced(P, G, s, lam, rhs_aff) + # Predictor and corrector share this factorization; see _factor_newton_reduced. + L = _factor_newton_reduced(P, G, s, lam) + dx_aff = _solve_with_factor(L, rhs_aff) ds_aff = -r_p - G @ dx_aff dlam_aff = -(r_c_aff + lam * ds_aff) / s @@ -110,7 +149,7 @@ def _pdipm_iteration( # --- Corrector --- r_c = s * lam + ds_aff * dlam_aff - sigma * mu rhs = -(r_d + G.T @ ((lam * r_p - r_c) / s)) - dx = _solve_newton_reduced(P, G, s, lam, rhs) + dx = _solve_with_factor(L, rhs) ds = -r_p - G @ dx dlam = -(r_c + lam * ds) / s @@ -125,21 +164,11 @@ def _pdipm_iteration( return x_new, s_new, lam_new -def _combined_residual( - P: Array, - q: Array, - G: Array, - h: Array, - x: Array, - s: Array, - lam: Array, -) -> Array: +def _combined_residual(r_d: Array, r_p: Array, s: Array, lam: Array) -> Array: """Scalar residual used both for in-loop freeze-on-converge and for the post-loop status check: ||r_d||_inf + ||r_p||_inf + mu. """ - r_d = P @ x + G.T @ lam + q - r_p = G @ x + s - h - mu = jnp.sum(s * lam) / G.shape[0] + mu = jnp.sum(s * lam) / s.shape[0] return jnp.max(jnp.abs(r_d)) + jnp.max(jnp.abs(r_p)) + mu @@ -150,7 +179,7 @@ def solve_qp_pdipm( G: Array, h: Array, warm_start: Optional[PdipmState] = None, - max_iter: int = 25, + max_iter: int = DEFAULT_MAX_ITER, tol: float = 1e-6, ) -> Tuple[Array, Array, PdipmState]: """Mehrotra predictor-corrector PDIPM for min 0.5 x^T P x + q^T x s.t. G x <= h. @@ -161,7 +190,14 @@ def solve_qp_pdipm( fixed for the remaining iterations rather than stepping further. This prevents a degenerate (s ≈ 0, lam ≈ 0) post-convergence state from propagating NaN through ``lam/s`` divisions in subsequent iterations. - The total flop count is still ``max_iter`` Newton solves per call. + The total flop count is still ``max_iter`` Newton solves per call, which + is why ``max_iter`` defaults to the tight ``DEFAULT_MAX_ITER`` rather than + a loose upper bound; raise it for QPs harder than the CBF-QP shapes it was + calibrated on (see that constant for the measurements). + + ``fori_loop`` is deliberate and load-bearing: ``lax.while_loop`` has no + reverse-mode differentiation rule, and downstream differentiable-CBF-QP + work takes ``jax.grad`` through this solver. Returns: (x, status, state) where status is: @@ -193,8 +229,10 @@ def solve_qp_pdipm( # --- Outer loop (fixed iterations, freeze-on-converge) --- def body(_, carry): x_, s_, lam_ = carry - already_converged = _combined_residual(P, q, G, h, x_, s_, lam_) < tol - x_new, s_new, lam_new = _pdipm_iteration(P, q, G, h, x_, s_, lam_) + # One residual evaluation feeds both the convergence test and the step. + r_d, r_p = _residuals(P, q, G, h, x_, s_, lam_) + already_converged = _combined_residual(r_d, r_p, s_, lam_) < tol + x_new, s_new, lam_new = _pdipm_iteration(P, G, x_, s_, lam_, r_d, r_p) # If already converged, keep current (good) state rather than stepping. # Both branches of jnp.where are traced; the step is cheap relative to # the cost of NaN propagation if we let lam/s explode after convergence. @@ -206,7 +244,8 @@ def body(_, carry): x_final, s_final, lam_final = lax.fori_loop(0, max_iter, body, (x0, s0, lam0)) # --- Status: solved if final residual norm below tol --- - res_norm = _combined_residual(P, q, G, h, x_final, s_final, lam_final) + r_d_final, r_p_final = _residuals(P, q, G, h, x_final, s_final, lam_final) + res_norm = _combined_residual(r_d_final, r_p_final, s_final, lam_final) status = jnp.where(res_norm < tol, jnp.int32(1), jnp.int32(2)) state = PdipmState(x=x_final, s=s_final, dual=lam_final, iter_num=max_iter) diff --git a/src/cbfkit/optimization/quadratic_program/solver_registry.py b/src/cbfkit/optimization/quadratic_program/solver_registry.py index 694760b1..2359c7f4 100644 --- a/src/cbfkit/optimization/quadratic_program/solver_registry.py +++ b/src/cbfkit/optimization/quadratic_program/solver_registry.py @@ -3,6 +3,19 @@ Provides a common ``QpSolution`` return type and factory functions that wrap each backend (jaxopt, cvxopt, casadi) behind a single callable signature. +**Objective convention.** Every solver returned by :func:`get_solver` solves + +.. math:: + + \\min_x \\; x^T H x + f^T x \\quad \\text{s.t.} \\quad G x \\le h \\;(, A x = b) + +This matches the CBF-CLF-QP generator, which passes ``f = -2 H u_nom`` so the +unconstrained minimizer is ``u_nom``. Backends whose native form is +``min 1/2 x^T P x + q^T x`` (OSQP, CVXOPT, PDIPM) are adapted at this boundary +(``P = 2 H`` or equivalently ``q = f/2``); the raw modules +(``qp_solver_pdipm.solve_qp_pdipm``, ``qp_solver_cvxopt.solve``) keep their +native halved convention. + Usage:: from cbfkit.optimization.quadratic_program import get_solver @@ -168,7 +181,7 @@ def casadi_solver() -> QpSolverCallable: # --------------------------------------------------------------------------- -def fast_solver(max_iter: int = 25, tol: float = 1e-6) -> QpSolverCallable: +def fast_solver(max_iter: Optional[int] = None, tol: float = 1e-6) -> QpSolverCallable: """Fast PDIPM solver for small CBF-CLF problems. Mehrotra predictor-corrector primal-dual interior-point method. Designed @@ -181,15 +194,24 @@ def fast_solver(max_iter: int = 25, tol: float = 1e-6) -> QpSolverCallable: (see ``benchmarks/qp_solver_comparison.py``). JIT-compatible and warm-startable across consecutive control steps. + Inequality constraints only: passing ``a_mat``/``b_vec`` raises + ``NotImplementedError`` rather than dropping them. + Args: - max_iter: Maximum PDIPM iterations (default 25; ~10-15 typically suffice). + max_iter: Maximum PDIPM iterations. ``None`` (default) defers to + ``qp_solver_pdipm.DEFAULT_MAX_ITER`` so the solver's calibrated + budget is not silently pinned here. tol: Combined primal/dual/complementarity residual tolerance. """ from cbfkit.optimization.quadratic_program.qp_solver_pdipm import ( + DEFAULT_MAX_ITER, PdipmState, solve_qp_pdipm, ) + if max_iter is None: + max_iter = DEFAULT_MAX_ITER + def solve_with_details( h_mat: Array, f_vec: Array, @@ -199,8 +221,19 @@ def solve_with_details( b_vec: Optional[Array] = None, init_params: Any = None, ) -> QpSolution: + if a_mat is not None or b_vec is not None: + raise NotImplementedError( + "The 'fast' (PDIPM) backend solves inequality-constrained QPs only " + "(min x'Hx + f'x s.t. Gx <= h), but was given equality constraints " + "via a_mat/b_vec. Silently dropping them would return a solution " + "that violates Ax = b — e.g. an MPC trajectory ignoring its own " + "dynamics. Use get_solver('jaxopt') or get_solver('casadi'), which " + "support equality constraints." + ) + if g_mat is None or h_vec is None: - x = jnp.linalg.solve(h_mat, -f_vec) + # Registry convention min x'Hx + f'x => 2H x* = -f. + x = jnp.linalg.solve(2.0 * h_mat, -f_vec) return QpSolution(primal=x, status=1, params=None) # Extract warm-start state from previous QpSolution.params @@ -213,8 +246,11 @@ def solve_with_details( elif isinstance(init_params, PdipmState): warm = init_params + # Convention adapter: solve_qp_pdipm natively solves + # min 1/2 x'Px + q'x; the registry convention is min x'Hx + f'x, + # so pass P = 2H (the jaxopt wrapper adapts via q = f/2 instead). sol, status, state = solve_qp_pdipm( - h_mat, + 2.0 * h_mat, f_vec, g_mat, h_vec, diff --git a/tests/test_optimization/test_fast_qp.py b/tests/test_optimization/test_fast_qp.py index 59df2354..2412bd31 100644 --- a/tests/test_optimization/test_fast_qp.py +++ b/tests/test_optimization/test_fast_qp.py @@ -1,7 +1,5 @@ """Tests for the fast small-QP solver and its integration with CBF-CLF-QP.""" -import pytest -import jax import jax.numpy as jnp from jax import random @@ -110,15 +108,21 @@ def test_get_fast_solver(self): assert solver.solver_name == "fast" def test_fast_solver_via_registry(self): + """Registry convention: min x'Hx + f'x (matches jaxopt/casadi backends). + + With H=I, f=-2*[1,2] the minimizer is [1,2] — NOT [2,4], which is + what the pre-fix wrapper returned by feeding (H, f) straight into + the native ``min 1/2 x'Px + q'x`` PDIPM form. + """ solver = get_solver("fast") P = jnp.eye(2) q = jnp.array([-2.0, -4.0]) G = jnp.array([[1, 0], [-1, 0], [0, 1], [0, -1.0]]) h = jnp.array( [3.0, 3.0, 5.0, 5.0] - ) # box [-3,3] x [-5,5] — unconstrained opt [2,4] is feasible + ) # box [-3,3] x [-5,5] — unconstrained opt [1,2] is feasible sol = solver(P, q, G, h) - assert jnp.allclose(sol.primal, jnp.array([2.0, 4.0]), atol=1e-4) + assert jnp.allclose(sol.primal, jnp.array([1.0, 2.0]), atol=1e-4) assert int(sol.status) == 1 def test_fast_solver_warm_start_via_registry(self): diff --git a/tests/test_optimization/test_fast_solver_equality_guard.py b/tests/test_optimization/test_fast_solver_equality_guard.py new file mode 100644 index 00000000..e9ef84da --- /dev/null +++ b/tests/test_optimization/test_fast_solver_equality_guard.py @@ -0,0 +1,81 @@ +"""The 'fast' (PDIPM) backend must reject equality constraints, not drop them. + +``fast_solver`` accepts ``a_mat``/``b_vec`` for interface compatibility with the +other registry backends, but its underlying PDIPM solves ``Gx <= h`` only. It +used to ignore those arguments and return ``status == 1``, so +``generate_mpc_solver_quadratic_cost_linear_dynamics`` — which encodes the plant +dynamics as equality constraints — got a "successful" trajectory that satisfied +no dynamics at all. A loud failure is the only safe behaviour here. +""" + +import jax.numpy as jnp +import pytest + +from cbfkit.optimization.quadratic_program import get_solver + + +def _problem(): + """min ||x - [1,2]||^2 s.t. loose box, plus the equality x0 + x1 = 1.""" + H = jnp.eye(2) + f = -2.0 * jnp.array([1.0, 2.0]) + G = jnp.vstack([jnp.eye(2), -jnp.eye(2)]) + h = 10.0 * jnp.ones(4) + a_mat = jnp.array([[1.0, 1.0]]) + b_vec = jnp.array([1.0]) + return H, f, G, h, a_mat, b_vec + + +class TestFastRejectsEqualityConstraints: + def test_raises_with_inequalities_present(self): + H, f, G, h, a_mat, b_vec = _problem() + with pytest.raises(NotImplementedError, match="equality constraints"): + get_solver("fast")(H, f, G, h, a_mat, b_vec) + + def test_raises_on_the_no_inequality_path(self): + """The unconstrained escape branch drops a_mat/b_vec just as silently.""" + H, f, _, _, a_mat, b_vec = _problem() + with pytest.raises(NotImplementedError, match="equality constraints"): + get_solver("fast")(H, f, None, None, a_mat, b_vec) + + @pytest.mark.parametrize("which", ["a_mat", "b_vec"]) + def test_raises_when_only_one_half_is_given(self, which): + """Half a specification is a caller bug; the other backends ignore it.""" + H, f, G, h, a_mat, b_vec = _problem() + kwargs = {"a_mat": a_mat} if which == "a_mat" else {"b_vec": b_vec} + with pytest.raises(NotImplementedError): + get_solver("fast")(H, f, G, h, **kwargs) + + def test_message_names_backend_and_alternatives(self): + H, f, G, h, a_mat, b_vec = _problem() + with pytest.raises(NotImplementedError) as exc: + get_solver("fast")(H, f, G, h, a_mat, b_vec) + message = str(exc.value) + assert "fast" in message + assert "jaxopt" in message or "casadi" in message + + def test_inequality_only_call_still_works(self): + """The guard must not fire on the ordinary CBF-QP path.""" + H, f, G, h, _, _ = _problem() + sol = get_solver("fast")(H, f, G, h) + assert int(sol.status) == 1 + assert jnp.allclose(sol.primal, jnp.array([1.0, 2.0]), atol=1e-4) + + +class TestEqualityCapableBackendsUnaffected: + def test_jaxopt_still_enforces_equality(self): + H, f, G, h, a_mat, b_vec = _problem() + sol = get_solver("jaxopt")(H, f, G, h, a_mat, b_vec) + assert int(sol.status) == 1 + assert jnp.allclose(sol.primal @ jnp.array([1.0, 1.0]), 1.0, atol=1e-3) + + def test_cvxopt_still_enforces_equality(self): + H, f, G, h, a_mat, b_vec = _problem() + try: + solver = get_solver("cvxopt") + sol = solver(H, f, G, h, a_mat, b_vec) + except ImportError: + pytest.skip("cvxopt/kvxopt not installed") + assert int(sol.status) == 1 + assert jnp.allclose(sol.primal @ jnp.array([1.0, 1.0]), 1.0, atol=1e-4) + # Equality-constrained minimizer of ||x-[1,2]||^2 on x0+x1=1 is [0,1]. + assert jnp.allclose(sol.primal, jnp.array([0.0, 1.0]), atol=1e-4) diff --git a/tests/test_optimization/test_pdipm_qp.py b/tests/test_optimization/test_pdipm_qp.py index 7c2d6f6b..9e2a7c8e 100644 --- a/tests/test_optimization/test_pdipm_qp.py +++ b/tests/test_optimization/test_pdipm_qp.py @@ -41,7 +41,7 @@ def test_jit_compatibility(self): class TestNewtonReduced: def test_H_is_positive_definite(self): """H = P + G^T diag(lam/s) G must be PD when P is PD and s, lam > 0.""" - from cbfkit.optimization.quadratic_program.qp_solver_pdipm import _solve_newton_reduced + from cbfkit.optimization.quadratic_program.qp_solver_pdipm import _factor_newton_reduced from jax import random key = random.PRNGKey(0) @@ -49,23 +49,29 @@ def test_H_is_positive_definite(self): G = random.normal(key, (5, 3)) s = jnp.array([0.5, 1.0, 1.5, 2.0, 0.1]) lam = jnp.array([0.2, 0.8, 1.0, 0.05, 2.0]) - rhs = jnp.ones(3) # Build H exactly as the helper would D = lam / s H = P + G.T @ (D[:, None] * G) + 1e-10 * jnp.eye(3) eigs = jnp.linalg.eigvalsh((H + H.T) / 2) assert float(eigs[0]) > 0, f"H not PD, min eig = {float(eigs[0])}" + # The Cholesky factor returned by the helper must reproduce H + L = _factor_newton_reduced(P, G, s, lam) + assert float(jnp.max(jnp.abs(L @ L.T - H))) < 1e-8 def test_solves_system(self): - """_solve_newton_reduced returns dx such that H dx = rhs.""" - from cbfkit.optimization.quadratic_program.qp_solver_pdipm import _solve_newton_reduced + """factor + solve_with_factor returns dx such that H dx = rhs.""" + from cbfkit.optimization.quadratic_program.qp_solver_pdipm import ( + _factor_newton_reduced, + _solve_with_factor, + ) P = jnp.diag(jnp.array([1.0, 1.0])) G = jnp.array([[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]]) s = jnp.array([1.0, 1.0, 1.0]) lam = jnp.array([0.5, 0.5, 0.5]) rhs = jnp.array([1.0, 2.0]) - dx = _solve_newton_reduced(P, G, s, lam, rhs) + L = _factor_newton_reduced(P, G, s, lam) + dx = _solve_with_factor(L, rhs) D = lam / s H = P + G.T @ (D[:, None] * G) + 1e-10 * jnp.eye(2) residual = H @ dx - rhs @@ -82,18 +88,25 @@ def _make_simple_qp(self): return P, q, G, h def test_iteration_keeps_strict_interior(self): - from cbfkit.optimization.quadratic_program.qp_solver_pdipm import _pdipm_iteration + from cbfkit.optimization.quadratic_program.qp_solver_pdipm import ( + _pdipm_iteration, + _residuals, + ) P, q, G, h = self._make_simple_qp() x = jnp.zeros(2) s = jnp.ones(4) lam = jnp.ones(4) - x_new, s_new, lam_new = _pdipm_iteration(P, q, G, h, x, s, lam) + r_d, r_p = _residuals(P, q, G, h, x, s, lam) + x_new, s_new, lam_new = _pdipm_iteration(P, G, x, s, lam, r_d, r_p) assert float(jnp.min(s_new)) > 0.0, f"slack went non-positive: {s_new}" assert float(jnp.min(lam_new)) > 0.0, f"dual went non-positive: {lam_new}" def test_iteration_reduces_residual(self): - from cbfkit.optimization.quadratic_program.qp_solver_pdipm import _pdipm_iteration + from cbfkit.optimization.quadratic_program.qp_solver_pdipm import ( + _pdipm_iteration, + _residuals, + ) P, q, G, h = self._make_simple_qp() x = jnp.zeros(2) @@ -107,7 +120,8 @@ def residual_norm(x_, s_, lam_): return float(jnp.linalg.norm(r_d)) + float(jnp.linalg.norm(r_p)) + r_c r0 = residual_norm(x, s, lam) - x1, s1, lam1 = _pdipm_iteration(P, q, G, h, x, s, lam) + r_d, r_p = _residuals(P, q, G, h, x, s, lam) + x1, s1, lam1 = _pdipm_iteration(P, G, x, s, lam, r_d, r_p) r1 = residual_norm(x1, s1, lam1) assert r1 < r0, f"residual did not decrease: {r0} -> {r1}" @@ -211,7 +225,6 @@ class TestWarmStart: def test_warm_start_reduces_iters_to_convergence(self): """Warm-started solve from prior optimum converges with extremely loose iter budget.""" from cbfkit.optimization.quadratic_program.qp_solver_pdipm import ( - PdipmState, solve_qp_pdipm, ) diff --git a/tests/test_optimization/test_solver_convention_parity.py b/tests/test_optimization/test_solver_convention_parity.py new file mode 100644 index 00000000..071e4827 --- /dev/null +++ b/tests/test_optimization/test_solver_convention_parity.py @@ -0,0 +1,228 @@ +"""Cross-solver objective-convention parity tests. + +Registry convention (see ``solver_registry`` module docstring): every solver +returned by ``get_solver()`` solves + + min_x x^T H x + f^T x s.t. G x <= h + +Historically the jaxopt wrapper adapted ``(H, 0.5 f)`` into OSQP's +``min 1/2 x^T Q x + c^T x`` form, while the fast (PDIPM) and cvxopt registry +wrappers passed ``(H, f)`` through unchanged — silently solving +``min 1/2 x^T H x + f^T x`` instead. For the CBF-CLF-QP generator (which +passes ``f = -2 H u_nom``) that meant tracking ``2 u_nom`` instead of +``u_nom`` whenever constraints were inactive. These tests pin the convention +for every registered backend and add a controller-level regression with +geometry chosen so the active set does NOT mask the objective-center error. +""" + +import jax.numpy as jnp +import pytest +from jax import random + +from cbfkit.optimization.quadratic_program import get_solver + +JIT_SOLVERS = ["jaxopt", "fast"] +ALL_SOLVERS = ["jaxopt", "fast", "cvxopt", "casadi"] + + +def _get_solver_or_skip(name: str): + if name in ("cvxopt", "casadi"): + try: + solver = get_solver(name) + # Import errors surface on first call for lazily-imported backends + solver(jnp.eye(1), jnp.zeros(1), jnp.ones((1, 1)), jnp.ones(1)) + except ImportError: + pytest.skip(f"{name} not installed") + return solver + return get_solver(name) + + +# -- Analytic problems in the registry convention ------------------------- + + +@pytest.mark.parametrize("name", ALL_SOLVERS) +def test_interior_optimum(name): + """min ||x - [1,2]||^2 with loose box: H=I, f=-2*[1,2] => x* = [1,2].""" + solver = _get_solver_or_skip(name) + H = jnp.eye(2) + f = -2.0 * jnp.array([1.0, 2.0]) + G = jnp.vstack([jnp.eye(2), -jnp.eye(2)]) + h = 10.0 * jnp.ones(4) + sol = solver(H, f, G, h) + assert int(sol.status) == 1 + assert jnp.allclose( + sol.primal, jnp.array([1.0, 2.0]), atol=2e-3 + ), f"{name}: expected registry-convention minimizer [1,2], got {sol.primal}" + + +@pytest.mark.parametrize("name", ALL_SOLVERS) +def test_single_active_constraint(name): + """min ||x - [2,0]||^2 s.t. x0 <= 1 => x* = [1,0] (projection onto halfspace).""" + solver = _get_solver_or_skip(name) + H = jnp.eye(2) + f = -2.0 * jnp.array([2.0, 0.0]) + G = jnp.array([[1.0, 0.0]]) + h = jnp.array([1.0]) + sol = solver(H, f, G, h) + assert int(sol.status) == 1 + assert jnp.allclose( + sol.primal, jnp.array([1.0, 0.0]), atol=2e-3 + ), f"{name}: expected [1,0], got {sol.primal}" + + +@pytest.mark.parametrize("name", ALL_SOLVERS) +def test_two_active_constraints_corner(name): + """min ||x - [2,3]||^2 s.t. x <= [1,1] => x* = [1,1] (corner).""" + solver = _get_solver_or_skip(name) + H = jnp.eye(2) + f = -2.0 * jnp.array([2.0, 3.0]) + G = jnp.eye(2) + h = jnp.ones(2) + sol = solver(H, f, G, h) + assert int(sol.status) == 1 + assert jnp.allclose( + sol.primal, jnp.array([1.0, 1.0]), atol=2e-3 + ), f"{name}: expected [1,1], got {sol.primal}" + + +@pytest.mark.parametrize("name", ALL_SOLVERS) +def test_audit_probe_matches_u_nom(name): + """The exact probe from the solver audit: H=I, f=-2*u_nom, loose box. + + u_nom = [1.0, -0.5]; pre-fix the fast and cvxopt backends returned + [2.0, -1.0] with a "solved" status. + """ + solver = _get_solver_or_skip(name) + u_nom = jnp.array([1.0, -0.5]) + H = jnp.eye(2) + f = -2.0 * H @ u_nom + G = jnp.vstack([jnp.eye(2), -jnp.eye(2)]) + h = 10.0 * jnp.ones(4) + sol = solver(H, f, G, h) + assert int(sol.status) == 1 + assert jnp.allclose(sol.primal, u_nom, atol=1e-4), f"{name}: got {sol.primal}" + + +# -- Random cross-solver agreement (JIT backends) -------------------------- + + +@pytest.mark.parametrize("seed", range(5)) +def test_fast_matches_jaxopt_random(seed): + """fast and jaxopt agree on the primal for random well-conditioned QPs.""" + key = random.PRNGKey(seed) + k1, k2, k3, k4 = random.split(key, 4) + n, m = 4, 8 + H = jnp.diag(jnp.abs(random.normal(k1, (n,))) + 0.5) + f = random.normal(k2, (n,)) + G = random.normal(k3, (m, n)) + h = jnp.abs(random.normal(k4, (m,))) + 0.5 + + sols = {} + for name in JIT_SOLVERS: + sol = get_solver(name)(H, f, G, h) + assert int(sol.status) == 1, f"{name} failed on seed {seed}" + sols[name] = sol.primal + + assert jnp.allclose( + sols["fast"], sols["jaxopt"], atol=2e-3 + ), f"seed {seed}: fast={sols['fast']} vs jaxopt={sols['jaxopt']}" + + +# -- Controller-level regression ------------------------------------------- + + +def _make_cbf_controller(solver_name: str, alpha: float = 1.0): + from cbfkit.certificates import generate_certificate + from cbfkit.certificates.conditions.barrier_conditions.zeroing_barriers import ( + linear_class_k, + ) + from cbfkit.controllers.cbf_clf import vanilla_cbf_clf_qp_controller + from cbfkit.systems.single_integrator.dynamics import two_dimensional_single_integrator + + dynamics = two_dimensional_single_integrator() + + def h(x): + return (x[0] - 2.0) ** 2 + x[1] ** 2 - 0.25 + + barriers = generate_certificate(h, linear_class_k(alpha), input_style="state") + return vanilla_cbf_clf_qp_controller( + control_limits=jnp.array([1.0, 1.0]), + dynamics_func=dynamics, + barriers=barriers, + solver=get_solver(solver_name), + ) + + +@pytest.mark.parametrize("name", JIT_SOLVERS) +def test_cbf_filter_inactive_returns_u_nom(name): + """With the CBF inactive, the filter must return u_nom — NOT 2*u_nom. + + At x=[0,0]: grad h = [-4, 0], constraint 4*u0 <= h(x) = 3.75, so + u_nom = [0.3, -0.2] is strictly feasible and must pass through unchanged. + The pre-fix fast wrapper returned clip(2*u_nom) here. + """ + from cbfkit.utils.user_types import ControllerData + + controller = _make_cbf_controller(name) + u, data = controller( + 0.0, jnp.array([0.0, 0.0]), jnp.array([0.3, -0.2]), random.PRNGKey(0), ControllerData() + ) + assert not bool(data.error) + assert jnp.allclose(u, jnp.array([0.3, -0.2]), atol=2e-3), f"{name}: u={u}" + + +@pytest.mark.parametrize("name", JIT_SOLVERS) +def test_cbf_filter_active_non_masking_geometry(name): + """CBF active with an off-axis nominal so the projection exposes any + objective-center error. + + At x=[1.2, 0]: h = 0.39, grad h = [-1.6, 0] => constraint u0 <= 0.24375. + Projecting u_nom=[0.8, 0.6] gives u* = [0.24375, 0.6]. Projecting the + doubled center [1.6, 1.2] (pre-fix fast) would give u1 = 1.0 (clipped), + so the u1 component discriminates the conventions. + """ + from cbfkit.utils.user_types import ControllerData + + expected = jnp.array([0.39 / 1.6, 0.6]) + controller = _make_cbf_controller(name) + u, data = controller( + 0.0, jnp.array([1.2, 0.0]), jnp.array([0.8, 0.6]), random.PRNGKey(0), ControllerData() + ) + assert not bool(data.error) + assert jnp.allclose(u, expected, atol=5e-3), f"{name}: u={u}, expected {expected}" + + +def test_cbf_filter_audit_case_fast_vs_jaxopt(): + """Audit case u_nom=[1.0, 0.2] at x=[0,0]: both backends must give u[1]=0.2. + + grad h = [-4, 0] involves only u0, so u1 is untouched by the CBF row and + reports the objective center directly: 0.2 after the fix, 0.4 before it. + """ + from cbfkit.utils.user_types import ControllerData + + u_nom = jnp.array([1.0, 0.2]) + controls = {} + for name in JIT_SOLVERS: + u, data = _make_cbf_controller(name)( + 0.0, jnp.array([0.0, 0.0]), u_nom, random.PRNGKey(0), ControllerData() + ) + assert not bool(data.error) + controls[name] = u + + for name, u in controls.items(): + assert jnp.allclose(u[1], 0.2, atol=2e-3), f"{name}: u={u}, expected u[1]=0.2" + assert jnp.allclose( + controls["fast"], controls["jaxopt"], atol=2e-3 + ), f"fast={controls['fast']} vs jaxopt={controls['jaxopt']}" + + +@pytest.mark.parametrize("name", JIT_SOLVERS) +def test_unconstrained_escape_convention(name): + """No inequality constraints: x* = -(2H)^-1 f in the registry convention.""" + if name == "jaxopt": + pytest.skip("jaxopt path requires params_ineq or params_eq; escape is fast-only") + solver = get_solver(name) + H = jnp.diag(jnp.array([1.0, 4.0])) + f = jnp.array([-2.0, -8.0]) + sol = solver(H, f) + assert jnp.allclose(sol.primal, jnp.array([1.0, 1.0]), atol=1e-6), sol.primal