autotuner

SCP weight autotuning strategies.

AcceptanceRatioAutotuner

Bases: AutotuningBase

Four-bucket acceptance-ratio update on lam_prox.

The acceptance ratio :math:\rho (actual over predicted nonlinear reduction) selects one of four buckets — REJECT / ACCEPT_HIGHER / ACCEPT_CONSTANT / ACCEPT_LOWER — via a jnp.where cascade so the whole update traces under jax.jit. Subclasses supply the multiplier rule (how lam_vc / lam_vb_* respond on acceptance) through :meth:_update_multipliers, and must assign an :class:AcceptanceRatioHyper (or subclass) instance to self.hyper in __init__ — the shared body reads its knobs off state.hyper.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/acceptance_ratio.py

class AcceptanceRatioAutotuner(AutotuningBase):
    """Four-bucket acceptance-ratio update on ``lam_prox``.

    The acceptance ratio :math:`\\rho` (actual over predicted nonlinear
    reduction) selects one of four buckets — ``REJECT`` / ``ACCEPT_HIGHER`` /
    ``ACCEPT_CONSTANT`` / ``ACCEPT_LOWER`` — via a ``jnp.where`` cascade so the
    whole update traces under ``jax.jit``. Subclasses supply the multiplier
    rule (how ``lam_vc`` / ``lam_vb_*`` respond on acceptance) through
    :meth:`_update_multipliers`, and must assign an
    :class:`AcceptanceRatioHyper` (or subclass) instance to ``self.hyper`` in
    ``__init__`` — the shared body reads its knobs off ``state.hyper``.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    COLUMNS: List[Column] = [
        Column("J_nonlin", "J_nonlin", 8, "{: .1e}", None, Verbosity.STANDARD),
        Column("J_lin", "J_lin", 8, "{: .1e}", None, Verbosity.STANDARD),
        Column("pred_reduction", "pred_red", 9, "{: .1e}", min_verbosity=Verbosity.FULL),
        Column("actual_reduction", "act_red", 9, "{: .1e}", min_verbosity=Verbosity.FULL),
        Column(
            "acceptance_ratio",
            "acc_ratio",
            9,
            "{: .2e}",
            color_acceptance_ratio,
            Verbosity.STANDARD,
        ),
        Column("lam_prox", "lam_prox", 8, "{: .1e}", min_verbosity=Verbosity.FULL),
        Column("adaptive_state", "Adaptive", 16, "{}", color_adaptive_state, Verbosity.FULL),
    ]

    def citation(self) -> List[str]:
        """Return BibTeX citations for the SCvx acceptance-ratio update rule."""
        return [
            r"""@inproceedings{mao2016scvx,
  title={Successive convexification of non-convex optimal control problems and its
    convergence properties},
  author={Mao, Yuanqi and Szmuk, Michael and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  booktitle={2016 IEEE 55th Conference on Decision and Control (CDC)},
  pages={3636--3641},
  year={2016},
  publisher={IEEE}
}""",
            r"""@misc{mao2019scvx,
  title={Successive Convexification: A Superlinearly Convergent Algorithm for
    Non-convex Optimal Control Problems},
  author={Mao, Yuanqi and Szmuk, Michael and Xu, Xiangru and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  year={2019},
  eprint={1804.06539},
  archivePrefix={arXiv},
  primaryClass={math.OC}
}""",
        ]

    def _update_multipliers(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        params: dict,
        settings: Config,
        new_lam_prox: jnp.ndarray,
    ) -> dict:
        """Return the multiplier-weight updates to apply on an accepted iterate.

        Each returned entry maps an :class:`AlgorithmState` field name
        (``lam_vc`` / ``lam_vb_nodal`` / ``lam_vb_cross``) to its updated
        value, computed against the *candidate* trajectory; the shared body
        gates each by acceptance (a rejected iterate keeps the previous value).
        The default holds every multiplier constant — return ``{}``.
        """
        return {}

    def update_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        settings: Config,
        params: dict,
    ) -> "AlgorithmState":
        """Return the next-iterate state per the acceptance-ratio rules.

        Pure functional update — see class docstring.
        """
        candidate_x_prop = candidate.x_prop_plus[1:]
        nonlin_cost, nonlin_pen, nodal_pen = calculate_nonlinear_penalty(
            candidate_x_prop,
            candidate.x,
            candidate.u,
            state.lam_vc,
            state.lam_vb_nodal,
            state.lam_vb_cross,
            state.lam_cost,
            nodal_constraints,
            params,
            settings,
        )
        J_nonlin = nonlin_cost + nonlin_pen + nodal_pen

        lam_cost_next = self._relaxed_lam_cost(state)

        def first_iter(state):
            # Iter 1: accept unconditionally, leave weights at their init values,
            # only refresh trajectory + propagation fields.
            return state.replace(
                x=candidate.x,
                u=candidate.u,
                x_prop=candidate.x_prop,
                x_prop_plus=candidate.x_prop_plus,
                lam_cost=lam_cost_next,
                J_nonlin=J_nonlin,
                adaptive_state_code=jnp.asarray(int(AdaptiveStateCode.INITIAL), dtype=jnp.int32),
            )

        def later_iter(state):
            # Recompute the previous iterate's J_nonlin from the pytree fields
            # (state.x/state.u were the *previous* accepted iterate).
            prev_x_prop = state.x_prop_plus[1:]
            prev_cost, prev_pen, prev_nodal_pen = calculate_nonlinear_penalty(
                prev_x_prop,
                state.x,
                state.u,
                state.lam_vc,
                state.lam_vb_nodal,
                state.lam_vb_cross,
                state.lam_cost,
                nodal_constraints,
                params,
                settings,
            )
            prev_J_nonlin = prev_cost + prev_pen + prev_nodal_pen

            actual = prev_J_nonlin - J_nonlin
            predicted = prev_J_nonlin - candidate.J_lin
            # If predicted reduction is exactly zero, force the reject bucket
            # (rho = -inf) deterministically instead of raising.
            safe_pred = jnp.where(predicted == 0.0, 1.0, predicted)
            rho = jnp.where(predicted == 0.0, -jnp.inf, actual / safe_pred)

            is_reject = rho < state.hyper.eta_0
            is_accept_higher = (rho >= state.hyper.eta_0) & (rho < state.hyper.eta_1)
            is_accept_constant = (rho >= state.hyper.eta_1) & (rho < state.hyper.eta_2)
            # is_accept_lower implicit (else)
            accepted = ~is_reject

            # Compute both lam_prox candidates and gate.
            lp_higher = jnp.minimum(state.hyper.lam_prox_max, state.hyper.gamma_1 * state.lam_prox)
            lp_lower = jnp.maximum(state.hyper.lam_prox_min, state.hyper.gamma_2 * state.lam_prox)
            new_lam_prox = jnp.where(
                is_reject | is_accept_higher,
                lp_higher,
                jnp.where(is_accept_constant, state.lam_prox, lp_lower),
            )

            # Multiplier (virtual-control / virtual-buffer) updates: computed
            # against the *candidate* trajectory and gated by `accepted`. Reject
            # keeps the previous values (so the next subproblem doesn't see a
            # bumped penalty that wasn't earned). The default hook returns {} —
            # multipliers held constant.
            multiplier_updates = self._update_multipliers(
                state, candidate, nodal_constraints, params, settings, new_lam_prox
            )
            gated_multipliers = {
                name: jnp.where(accepted, value, getattr(state, name))
                for name, value in multiplier_updates.items()
            }

            code = jnp.where(
                is_reject,
                jnp.int32(AdaptiveStateCode.REJECT),
                jnp.where(
                    is_accept_higher,
                    jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
                    jnp.where(
                        is_accept_constant,
                        jnp.int32(AdaptiveStateCode.ACCEPT_CONSTANT),
                        jnp.int32(AdaptiveStateCode.ACCEPT_LOWER),
                    ),
                ),
            )

            return state.replace(
                x=jnp.where(accepted, candidate.x, state.x),
                u=jnp.where(accepted, candidate.u, state.u),
                x_prop=jnp.where(accepted, candidate.x_prop, state.x_prop),
                x_prop_plus=jnp.where(accepted, candidate.x_prop_plus, state.x_prop_plus),
                lam_prox=new_lam_prox,
                lam_cost=lam_cost_next,
                J_nonlin=jnp.where(accepted, J_nonlin, state.J_nonlin),
                predicted_reduction=predicted,
                actual_reduction=actual,
                acceptance_ratio=rho,
                adaptive_state_code=code,
                **gated_multipliers,
            )

        return jax.lax.cond(state.k == 1, first_iter, later_iter, state)

citation() -> List[str]

Return BibTeX citations for the SCvx acceptance-ratio update rule.

Source code in openscvx/algorithms/autotuner/acceptance_ratio.py

    def citation(self) -> List[str]:
        """Return BibTeX citations for the SCvx acceptance-ratio update rule."""
        return [
            r"""@inproceedings{mao2016scvx,
  title={Successive convexification of non-convex optimal control problems and its
    convergence properties},
  author={Mao, Yuanqi and Szmuk, Michael and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  booktitle={2016 IEEE 55th Conference on Decision and Control (CDC)},
  pages={3636--3641},
  year={2016},
  publisher={IEEE}
}""",
            r"""@misc{mao2019scvx,
  title={Successive Convexification: A Superlinearly Convergent Algorithm for
    Non-convex Optimal Control Problems},
  author={Mao, Yuanqi and Szmuk, Michael and Xu, Xiangru and A{\c{c}}{\i}kme{\c{s}}e, Beh{\c{c}}et},
  year={2019},
  eprint={1804.06539},
  archivePrefix={arXiv},
  primaryClass={math.OC}
}""",
        ]

update_weights(state: AlgorithmState, candidate: CandidateIterate, nodal_constraints: LoweredJaxConstraints, settings: Config, params: dict) -> AlgorithmState

Return the next-iterate state per the acceptance-ratio rules.

Pure functional update — see class docstring.

Source code in openscvx/algorithms/autotuner/acceptance_ratio.py

def update_weights(
    self,
    state: "AlgorithmState",
    candidate: "CandidateIterate",
    nodal_constraints: "LoweredJaxConstraints",
    settings: Config,
    params: dict,
) -> "AlgorithmState":
    """Return the next-iterate state per the acceptance-ratio rules.

    Pure functional update — see class docstring.
    """
    candidate_x_prop = candidate.x_prop_plus[1:]
    nonlin_cost, nonlin_pen, nodal_pen = calculate_nonlinear_penalty(
        candidate_x_prop,
        candidate.x,
        candidate.u,
        state.lam_vc,
        state.lam_vb_nodal,
        state.lam_vb_cross,
        state.lam_cost,
        nodal_constraints,
        params,
        settings,
    )
    J_nonlin = nonlin_cost + nonlin_pen + nodal_pen

    lam_cost_next = self._relaxed_lam_cost(state)

    def first_iter(state):
        # Iter 1: accept unconditionally, leave weights at their init values,
        # only refresh trajectory + propagation fields.
        return state.replace(
            x=candidate.x,
            u=candidate.u,
            x_prop=candidate.x_prop,
            x_prop_plus=candidate.x_prop_plus,
            lam_cost=lam_cost_next,
            J_nonlin=J_nonlin,
            adaptive_state_code=jnp.asarray(int(AdaptiveStateCode.INITIAL), dtype=jnp.int32),
        )

    def later_iter(state):
        # Recompute the previous iterate's J_nonlin from the pytree fields
        # (state.x/state.u were the *previous* accepted iterate).
        prev_x_prop = state.x_prop_plus[1:]
        prev_cost, prev_pen, prev_nodal_pen = calculate_nonlinear_penalty(
            prev_x_prop,
            state.x,
            state.u,
            state.lam_vc,
            state.lam_vb_nodal,
            state.lam_vb_cross,
            state.lam_cost,
            nodal_constraints,
            params,
            settings,
        )
        prev_J_nonlin = prev_cost + prev_pen + prev_nodal_pen

        actual = prev_J_nonlin - J_nonlin
        predicted = prev_J_nonlin - candidate.J_lin
        # If predicted reduction is exactly zero, force the reject bucket
        # (rho = -inf) deterministically instead of raising.
        safe_pred = jnp.where(predicted == 0.0, 1.0, predicted)
        rho = jnp.where(predicted == 0.0, -jnp.inf, actual / safe_pred)

        is_reject = rho < state.hyper.eta_0
        is_accept_higher = (rho >= state.hyper.eta_0) & (rho < state.hyper.eta_1)
        is_accept_constant = (rho >= state.hyper.eta_1) & (rho < state.hyper.eta_2)
        # is_accept_lower implicit (else)
        accepted = ~is_reject

        # Compute both lam_prox candidates and gate.
        lp_higher = jnp.minimum(state.hyper.lam_prox_max, state.hyper.gamma_1 * state.lam_prox)
        lp_lower = jnp.maximum(state.hyper.lam_prox_min, state.hyper.gamma_2 * state.lam_prox)
        new_lam_prox = jnp.where(
            is_reject | is_accept_higher,
            lp_higher,
            jnp.where(is_accept_constant, state.lam_prox, lp_lower),
        )

        # Multiplier (virtual-control / virtual-buffer) updates: computed
        # against the *candidate* trajectory and gated by `accepted`. Reject
        # keeps the previous values (so the next subproblem doesn't see a
        # bumped penalty that wasn't earned). The default hook returns {} —
        # multipliers held constant.
        multiplier_updates = self._update_multipliers(
            state, candidate, nodal_constraints, params, settings, new_lam_prox
        )
        gated_multipliers = {
            name: jnp.where(accepted, value, getattr(state, name))
            for name, value in multiplier_updates.items()
        }

        code = jnp.where(
            is_reject,
            jnp.int32(AdaptiveStateCode.REJECT),
            jnp.where(
                is_accept_higher,
                jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
                jnp.where(
                    is_accept_constant,
                    jnp.int32(AdaptiveStateCode.ACCEPT_CONSTANT),
                    jnp.int32(AdaptiveStateCode.ACCEPT_LOWER),
                ),
            ),
        )

        return state.replace(
            x=jnp.where(accepted, candidate.x, state.x),
            u=jnp.where(accepted, candidate.u, state.u),
            x_prop=jnp.where(accepted, candidate.x_prop, state.x_prop),
            x_prop_plus=jnp.where(accepted, candidate.x_prop_plus, state.x_prop_plus),
            lam_prox=new_lam_prox,
            lam_cost=lam_cost_next,
            J_nonlin=jnp.where(accepted, J_nonlin, state.J_nonlin),
            predicted_reduction=predicted,
            actual_reduction=actual,
            acceptance_ratio=rho,
            adaptive_state_code=code,
            **gated_multipliers,
        )

    return jax.lax.cond(state.k == 1, first_iter, later_iter, state)

AcceptanceRatioHyper

Bases: LamCostRelaxHyper

Knobs shared by every acceptance-ratio autotuner.

Extends the shared :class:LamCostRelaxHyper lam_cost relaxation knobs with the acceptance-ratio decision thresholds and the lam_prox clips.

Source code in openscvx/algorithms/autotuner/acceptance_ratio.py

class AcceptanceRatioHyper(LamCostRelaxHyper):
    """Knobs shared by every acceptance-ratio autotuner.

    Extends the shared :class:`LamCostRelaxHyper` ``lam_cost`` relaxation knobs
    with the acceptance-ratio decision thresholds and the ``lam_prox`` clips.
    """

    gamma_1: float = 2.0
    gamma_2: float = 0.5
    eta_0: float = 1e-2
    eta_1: float = 1e-1
    eta_2: float = 0.8
    lam_prox_min: float = 1e-3
    lam_prox_max: float = 1e4

AdaptiveProximalWeight

Bases: AcceptanceRatioAutotuner

PTR-style proximal adaptation with fixed virtual-penalty weights.

Inherits the four-bucket acceptance-ratio update on lam_prox from :class:~openscvx.algorithms.autotuner.acceptance_ratio.AcceptanceRatioAutotuner and keeps the default multiplier hook, so the virtual-control and virtual-buffer weights are carried unchanged.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/adaptive_proximal_weight.py

class AdaptiveProximalWeight(AcceptanceRatioAutotuner):
    """PTR-style proximal adaptation with fixed virtual-penalty weights.

    Inherits the four-bucket acceptance-ratio update on ``lam_prox`` from
    :class:`~openscvx.algorithms.autotuner.acceptance_ratio.AcceptanceRatioAutotuner`
    and keeps the default multiplier hook, so the virtual-control and
    virtual-buffer weights are carried unchanged.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    def __init__(
        self,
        gamma_1: float = 2.0,
        gamma_2: float = 0.5,
        eta_0: float = 1e-2,
        eta_1: float = 1e-1,
        eta_2: float = 0.8,
        lam_prox_min: float = 1e-3,
        lam_prox_max: float = 1e4,
        lam_cost_drop: int = -1,
        lam_cost_relax: float = 1.0,
    ):
        self.hyper = AcceptanceRatioHyper(
            gamma_1=gamma_1,
            gamma_2=gamma_2,
            eta_0=eta_0,
            eta_1=eta_1,
            eta_2=eta_2,
            lam_prox_min=lam_prox_min,
            lam_prox_max=lam_prox_max,
            lam_cost_drop=lam_cost_drop,
            lam_cost_relax=lam_cost_relax,
        )

AdaptiveProximalWeightSpec

Bases: BaseModel

Validates AdaptiveProximalWeight configuration from dict/YAML input.

Source code in openscvx/algorithms/autotuner/adaptive_proximal_weight.py

class AdaptiveProximalWeightSpec(BaseModel):
    """Validates AdaptiveProximalWeight configuration from dict/YAML input."""

    type: Literal["AdaptiveProximalWeight"] = "AdaptiveProximalWeight"
    gamma_1: float = 2.0
    gamma_2: float = 0.5
    eta_0: float = 1e-2
    eta_1: float = 1e-1
    eta_2: float = 0.8
    lam_prox_min: float = 1e-3
    lam_prox_max: float = 1e4
    lam_cost_drop: int = -1
    lam_cost_relax: float = 1.0

    model_config = ConfigDict(extra="forbid")

    def build(self) -> AdaptiveProximalWeight:
        return AdaptiveProximalWeight(**self.model_dump(exclude={"type"}, exclude_unset=True))

AugmentedLagrangian

Bases: AcceptanceRatioAutotuner

Augmented Lagrangian autotuner.

Uses an acceptance-ratio :math:\rho between predicted and actual reduction in the nonlinear objective to drive a four-bucket update of the trust-region weight lam_prox (inherited from :class:~openscvx.algorithms.autotuner.acceptance_ratio.AcceptanceRatioAutotuner) and to update the virtual-control / virtual-buffer weights from constraint violations via the multiplier hook.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/augmented_lagrangian.py

class AugmentedLagrangian(AcceptanceRatioAutotuner):
    """Augmented Lagrangian autotuner.

    Uses an acceptance-ratio :math:`\\rho` between predicted and actual
    reduction in the nonlinear objective to drive a four-bucket update of
    the trust-region weight ``lam_prox`` (inherited from
    :class:`~openscvx.algorithms.autotuner.acceptance_ratio.AcceptanceRatioAutotuner`)
    and to update the virtual-control / virtual-buffer weights from constraint
    violations via the multiplier hook.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    def __init__(
        self,
        rho_init: float = 1.0,
        rho_max: float = 1e2,
        gamma_1: float = 2.0,
        gamma_2: float = 0.5,
        eta_0: float = 1e-2,
        eta_1: float = 1e-1,
        eta_2: float = 0.8,
        ep: float = 0.99,
        eta_lambda: float = 1e1,
        lam_vc_max: float = 1e5,
        lam_prox_min: float = 1e-3,
        lam_prox_max: float = 1e4,
        lam_cost_drop: int = -1,
        lam_cost_relax: float = 1.0,
    ):
        """Initialize Augmented Lagrangian autotuning parameters.

        Every numeric knob is a declared hyperparameter: it lives on the frozen
        :class:`AugmentedLagrangianHyper` container at ``self.hyper`` and is a
        per-solve / batched override target (see :class:`AutotuningBase`).
        Knobs may still be read and written as plain attributes
        (``autotuner.lam_prox_max = 1e6``); the attribute proxy on
        :class:`AutotuningBase` routes the access into ``hyper``.

        Args:
            rho_init: Reserved for a future penalty-growth rule; currently
                unused. Passing a non-default value warns. Defaults to 1.0.
            rho_max: Reserved for a future penalty-growth rule; currently
                unused. Passing a non-default value warns. Defaults to 1e2.
            gamma_1: Factor to increase trust region weight when ratio is low.
                Defaults to 2.0.
            gamma_2: Factor to decrease trust region weight when ratio is high.
                Defaults to 0.5.
            eta_0: Acceptance ratio threshold below which solution is rejected.
                Defaults to 1e-2.
            eta_1: Threshold above which solution is accepted with constant weight.
                Defaults to 1e-1.
            eta_2: Threshold above which solution is accepted with lower weight.
                Defaults to 0.8.
            ep: Threshold for virtual control weight update (nu > ep vs nu <= ep).
                Must lie in (0, 1). Defaults to 0.99.
            eta_lambda: Step size for virtual control weight update. Defaults to 1e1.
            lam_vc_max: Maximum virtual control penalty weight. Defaults to 1e5.
            lam_prox_min: Minimum trust region (proximal) weight. Defaults to 1e-3.
            lam_prox_max: Maximum trust region (proximal) weight. Defaults to 1e4.
            lam_cost_drop: Iteration after which cost relaxation applies
                (``state.k > lam_cost_drop``). ``-1`` relaxes from the first
                iteration; the default ``lam_cost_relax=1.0`` makes that a
                no-op. Defaults to -1.
            lam_cost_relax: Factor applied to lam_cost after lam_cost_drop.
                Defaults to 1.0.
        """
        if rho_init != 1.0 or rho_max != 1e2:
            warnings.warn(
                "AugmentedLagrangian rho_init/rho_max are reserved for a future "
                "penalty-growth rule and are currently unused — setting them has "
                "no effect on the solve.",
                UserWarning,
                stacklevel=2,
            )
        self.hyper = AugmentedLagrangianHyper(
            rho_init=rho_init,
            rho_max=rho_max,
            gamma_1=gamma_1,
            gamma_2=gamma_2,
            eta_0=eta_0,
            eta_1=eta_1,
            eta_2=eta_2,
            ep=ep,
            eta_lambda=eta_lambda,
            lam_vc_max=lam_vc_max,
            lam_prox_min=lam_prox_min,
            lam_prox_max=lam_prox_max,
            lam_cost_drop=lam_cost_drop,
            lam_cost_relax=lam_cost_relax,
        )

    # -----------------------------------------------------------------------
    # Multiplier (virtual-control / virtual-buffer) updates from violation
    # -----------------------------------------------------------------------

    def _update_multipliers(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        params: dict,
        settings: Config,
        new_lam_prox: jnp.ndarray,
    ) -> dict:
        """Update the virtual-control and virtual-buffer weights from violation.

        Each weight follows the same piecewise (linear / quadratic) update keyed
        on the scaled constraint violation; the shared body gates the result by
        iterate acceptance.
        """
        candidate_x_prop = candidate.x_prop_plus[1:]
        return {
            "lam_vc": self._update_virtual_control_weights(
                state, candidate, candidate_x_prop, settings, new_lam_prox
            ),
            "lam_vb_nodal": self._update_virtual_buffer_nodal_weights(
                state, candidate, nodal_constraints, params, new_lam_prox
            ),
            "lam_vb_cross": self._update_virtual_buffer_cross_weights(
                state, candidate, nodal_constraints, params, new_lam_prox
            ),
        }

    def _update_virtual_control_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        candidate_x_prop: jnp.ndarray,
        settings: Config,
        lam_prox: jnp.ndarray,
    ) -> jnp.ndarray:
        """Update virtual control penalty weights from state violation.

        Computes scaled violation ``nu = inv_S_x @ |x[1:] - x_prop|`` and
        applies a piecewise update: linear in ``nu`` when ``nu > ep``,
        quadratic otherwise. Result is clipped to ``lam_vc_max``.
        """
        nu = (settings.sim.inv_S_x @ jnp.abs(candidate.x[1:] - candidate_x_prop).T).T
        # TODO: (haynec) use per-variable lam_prox to scale VC updates proportionally
        lam_prox_scalar = jnp.max(lam_prox)
        scale = state.hyper.eta_lambda * (1.0 / (2.0 * lam_prox_scalar))
        case1 = state.lam_vc + nu * scale
        case2 = state.lam_vc + (nu**2) / state.hyper.ep * scale
        vc_new = jnp.where(nu > state.hyper.ep, case1, case2)
        return jnp.minimum(state.hyper.lam_vc_max, vc_new)

    def _update_virtual_buffer_nodal_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        params: dict,
        lam_prox: jnp.ndarray,
    ) -> jnp.ndarray:
        """Update virtual-buffer weights for nodal constraints.

        Evaluates each nodal constraint, computes the same piecewise update
        as for virtual control, and writes the result back into the
        ``(N, n_nodal)`` weight array via ``.at[...].set(...)``.
        """
        lam_vb_new = state.lam_vb_nodal
        lam_prox_scalar = jnp.max(lam_prox)
        scale = state.hyper.eta_lambda * (1.0 / (2.0 * lam_prox_scalar))

        for idx, constraint in enumerate(nodal_constraints.nodal):
            g = constraint.func(candidate.x, candidate.u, 0, params)
            nu = jnp.abs(g) if constraint.is_equality else jnp.maximum(0.0, g)

            if constraint.nodes is not None:
                nodes_array = jnp.asarray(constraint.nodes)
                nu_slice = nu[nodes_array]
                current = state.lam_vb_nodal[nodes_array, idx]
            else:
                nu_slice = nu
                current = state.lam_vb_nodal[:, idx]

            case1 = current + nu_slice * scale
            case2 = current + (nu_slice**2) / state.hyper.ep * scale
            updated = jnp.where(nu_slice > state.hyper.ep, case1, case2)

            if constraint.nodes is not None:
                lam_vb_new = lam_vb_new.at[nodes_array, idx].set(updated)
            else:
                lam_vb_new = lam_vb_new.at[:, idx].set(updated)

        # Intentional: the virtual-buffer weights share the virtual-control cap
        # ``lam_vc_max`` rather than a dedicated knob — one ceiling governs every
        # penalty weight in this autotuner.
        return jnp.minimum(state.hyper.lam_vc_max, lam_vb_new)

    def _update_virtual_buffer_cross_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        params: dict,
        lam_prox: jnp.ndarray,
    ) -> jnp.ndarray:
        """Update virtual-buffer weights for cross-node constraints."""
        lam_vb_new = state.lam_vb_cross
        lam_prox_scalar = jnp.max(lam_prox)
        scale = state.hyper.eta_lambda * (1.0 / (2.0 * lam_prox_scalar))

        for idx, constraint in enumerate(nodal_constraints.cross_node):
            g = constraint.func(candidate.x, candidate.u, params)
            nu = jnp.sum(jnp.abs(g) if constraint.is_equality else jnp.maximum(0.0, g))
            current = state.lam_vb_cross[idx]
            case1 = current + nu * scale
            case2 = current + (nu**2) / state.hyper.ep * scale
            updated = jnp.where(nu > state.hyper.ep, case1, case2)
            lam_vb_new = lam_vb_new.at[idx].set(updated)

        # Intentional: shares the virtual-control cap ``lam_vc_max`` — see
        # :py:meth:`_update_virtual_buffer_nodal_weights`.
        return jnp.minimum(state.hyper.lam_vc_max, lam_vb_new)

    def citation(self) -> List[str]:
        """Return BibTeX citations for the augmented-Lagrangian multiplier updates.

        Extends the inherited SCvx acceptance-ratio references with Luo's
        dissertation, whose penalty-weight updates this autotuner implements.
        """
        return super().citation() + [
            r"""@phdthesis{luo2025modeling,
  title={Modeling and Algorithms for Nonconvex Trajectory Generation Problems:
    From Constraint Reformulations and First-Order Proximal Methods to
    Structure-Exploiting Convex Solvers},
  author={Luo, Dayou},
  school={University of Washington},
  year={2025},
  url={https://hdl.handle.net/1773/53364}
}"""
        ]

__init__(rho_init: float = 1.0, rho_max: float = 100.0, gamma_1: float = 2.0, gamma_2: float = 0.5, eta_0: float = 0.01, eta_1: float = 0.1, eta_2: float = 0.8, ep: float = 0.99, eta_lambda: float = 10.0, lam_vc_max: float = 100000.0, lam_prox_min: float = 0.001, lam_prox_max: float = 10000.0, lam_cost_drop: int = -1, lam_cost_relax: float = 1.0)

Initialize Augmented Lagrangian autotuning parameters.

Every numeric knob is a declared hyperparameter: it lives on the frozen :class:AugmentedLagrangianHyper container at self.hyper and is a per-solve / batched override target (see :class:AutotuningBase). Knobs may still be read and written as plain attributes (autotuner.lam_prox_max = 1e6); the attribute proxy on :class:AutotuningBase routes the access into hyper.

Parameters:

Name	Type	Description	Default
rho_init	float	Reserved for a future penalty-growth rule; currently unused. Passing a non-default value warns. Defaults to 1.0.	1.0
rho_max	float	Reserved for a future penalty-growth rule; currently unused. Passing a non-default value warns. Defaults to 1e2.	100.0
gamma_1	float	Factor to increase trust region weight when ratio is low. Defaults to 2.0.	2.0
gamma_2	float	Factor to decrease trust region weight when ratio is high. Defaults to 0.5.	0.5
eta_0	float	Acceptance ratio threshold below which solution is rejected. Defaults to 1e-2.	0.01
eta_1	float	Threshold above which solution is accepted with constant weight. Defaults to 1e-1.	0.1
eta_2	float	Threshold above which solution is accepted with lower weight. Defaults to 0.8.	0.8
ep	float	Threshold for virtual control weight update (nu > ep vs nu <= ep). Must lie in (0, 1). Defaults to 0.99.	0.99
eta_lambda	float	Step size for virtual control weight update. Defaults to 1e1.	10.0
lam_vc_max	float	Maximum virtual control penalty weight. Defaults to 1e5.	100000.0
lam_prox_min	float	Minimum trust region (proximal) weight. Defaults to 1e-3.	0.001
lam_prox_max	float	Maximum trust region (proximal) weight. Defaults to 1e4.	10000.0
lam_cost_drop	int	Iteration after which cost relaxation applies (state.k > lam_cost_drop). -1 relaxes from the first iteration; the default lam_cost_relax=1.0 makes that a no-op. Defaults to -1.	-1
lam_cost_relax	float	Factor applied to lam_cost after lam_cost_drop. Defaults to 1.0.	1.0

Source code in openscvx/algorithms/autotuner/augmented_lagrangian.py

def __init__(
    self,
    rho_init: float = 1.0,
    rho_max: float = 1e2,
    gamma_1: float = 2.0,
    gamma_2: float = 0.5,
    eta_0: float = 1e-2,
    eta_1: float = 1e-1,
    eta_2: float = 0.8,
    ep: float = 0.99,
    eta_lambda: float = 1e1,
    lam_vc_max: float = 1e5,
    lam_prox_min: float = 1e-3,
    lam_prox_max: float = 1e4,
    lam_cost_drop: int = -1,
    lam_cost_relax: float = 1.0,
):
    """Initialize Augmented Lagrangian autotuning parameters.

    Every numeric knob is a declared hyperparameter: it lives on the frozen
    :class:`AugmentedLagrangianHyper` container at ``self.hyper`` and is a
    per-solve / batched override target (see :class:`AutotuningBase`).
    Knobs may still be read and written as plain attributes
    (``autotuner.lam_prox_max = 1e6``); the attribute proxy on
    :class:`AutotuningBase` routes the access into ``hyper``.

    Args:
        rho_init: Reserved for a future penalty-growth rule; currently
            unused. Passing a non-default value warns. Defaults to 1.0.
        rho_max: Reserved for a future penalty-growth rule; currently
            unused. Passing a non-default value warns. Defaults to 1e2.
        gamma_1: Factor to increase trust region weight when ratio is low.
            Defaults to 2.0.
        gamma_2: Factor to decrease trust region weight when ratio is high.
            Defaults to 0.5.
        eta_0: Acceptance ratio threshold below which solution is rejected.
            Defaults to 1e-2.
        eta_1: Threshold above which solution is accepted with constant weight.
            Defaults to 1e-1.
        eta_2: Threshold above which solution is accepted with lower weight.
            Defaults to 0.8.
        ep: Threshold for virtual control weight update (nu > ep vs nu <= ep).
            Must lie in (0, 1). Defaults to 0.99.
        eta_lambda: Step size for virtual control weight update. Defaults to 1e1.
        lam_vc_max: Maximum virtual control penalty weight. Defaults to 1e5.
        lam_prox_min: Minimum trust region (proximal) weight. Defaults to 1e-3.
        lam_prox_max: Maximum trust region (proximal) weight. Defaults to 1e4.
        lam_cost_drop: Iteration after which cost relaxation applies
            (``state.k > lam_cost_drop``). ``-1`` relaxes from the first
            iteration; the default ``lam_cost_relax=1.0`` makes that a
            no-op. Defaults to -1.
        lam_cost_relax: Factor applied to lam_cost after lam_cost_drop.
            Defaults to 1.0.
    """
    if rho_init != 1.0 or rho_max != 1e2:
        warnings.warn(
            "AugmentedLagrangian rho_init/rho_max are reserved for a future "
            "penalty-growth rule and are currently unused — setting them has "
            "no effect on the solve.",
            UserWarning,
            stacklevel=2,
        )
    self.hyper = AugmentedLagrangianHyper(
        rho_init=rho_init,
        rho_max=rho_max,
        gamma_1=gamma_1,
        gamma_2=gamma_2,
        eta_0=eta_0,
        eta_1=eta_1,
        eta_2=eta_2,
        ep=ep,
        eta_lambda=eta_lambda,
        lam_vc_max=lam_vc_max,
        lam_prox_min=lam_prox_min,
        lam_prox_max=lam_prox_max,
        lam_cost_drop=lam_cost_drop,
        lam_cost_relax=lam_cost_relax,
    )

citation() -> List[str]

Return BibTeX citations for the augmented-Lagrangian multiplier updates.

Extends the inherited SCvx acceptance-ratio references with Luo's dissertation, whose penalty-weight updates this autotuner implements.

Source code in openscvx/algorithms/autotuner/augmented_lagrangian.py

    def citation(self) -> List[str]:
        """Return BibTeX citations for the augmented-Lagrangian multiplier updates.

        Extends the inherited SCvx acceptance-ratio references with Luo's
        dissertation, whose penalty-weight updates this autotuner implements.
        """
        return super().citation() + [
            r"""@phdthesis{luo2025modeling,
  title={Modeling and Algorithms for Nonconvex Trajectory Generation Problems:
    From Constraint Reformulations and First-Order Proximal Methods to
    Structure-Exploiting Convex Solvers},
  author={Luo, Dayou},
  school={University of Washington},
  year={2025},
  url={https://hdl.handle.net/1773/53364}
}"""
        ]

AugmentedLagrangianSpec

Bases: BaseModel

Validates AugmentedLagrangian configuration from dict/YAML input.

Source code in openscvx/algorithms/autotuner/augmented_lagrangian.py

class AugmentedLagrangianSpec(BaseModel):
    """Validates AugmentedLagrangian configuration from dict/YAML input."""

    type: Literal["AugmentedLagrangian"] = "AugmentedLagrangian"
    rho_init: float = 1.0
    rho_max: float = 1e2
    gamma_1: float = 2.0
    gamma_2: float = 0.5
    eta_0: float = 1e-2
    eta_1: float = 1e-1
    eta_2: float = 0.8
    ep: float = 0.99
    eta_lambda: float = 1e1
    lam_vc_max: float = 1e5
    lam_prox_min: float = 1e-3
    lam_prox_max: float = 1e4
    lam_cost_drop: int = -1
    lam_cost_relax: float = 1.0

    model_config = ConfigDict(extra="forbid")

    def build(self) -> AugmentedLagrangian:
        return AugmentedLagrangian(**self.model_dump(exclude={"type"}, exclude_unset=True))

ConstantProximalWeight

Bases: AutotuningBase

Hold lam_prox constant; relax lam_cost after lam_cost_drop.

Useful when you want a fixed trust-region size without adaptation.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/constant_proximal_weight.py

class ConstantProximalWeight(AutotuningBase):
    """Hold ``lam_prox`` constant; relax ``lam_cost`` after ``lam_cost_drop``.

    Useful when you want a fixed trust-region size without adaptation.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    COMPUTES_ACCEPTANCE_METRICS = False

    def __init__(
        self,
        lam_cost_drop: int = -1,
        lam_cost_relax: float = 1.0,
    ):
        self.hyper = ConstantProximalWeightHyper(
            lam_cost_drop=lam_cost_drop,
            lam_cost_relax=lam_cost_relax,
        )

    def update_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        settings: Config,
        params: dict,
    ) -> "AlgorithmState":
        """Return the next-iterate state.

        Pure functional update — see class docstring.
        """
        lam_cost_next = self._relaxed_lam_cost(state)

        return state.replace(
            x=candidate.x,
            u=candidate.u,
            x_prop=candidate.x_prop,
            x_prop_plus=candidate.x_prop_plus,
            lam_cost=lam_cost_next,
            adaptive_state_code=jnp.asarray(
                int(AdaptiveStateCode.ACCEPT_CONSTANT), dtype=jnp.int32
            ),
        )

update_weights(state: AlgorithmState, candidate: CandidateIterate, nodal_constraints: LoweredJaxConstraints, settings: Config, params: dict) -> AlgorithmState

Return the next-iterate state.

Pure functional update — see class docstring.

Source code in openscvx/algorithms/autotuner/constant_proximal_weight.py

def update_weights(
    self,
    state: "AlgorithmState",
    candidate: "CandidateIterate",
    nodal_constraints: "LoweredJaxConstraints",
    settings: Config,
    params: dict,
) -> "AlgorithmState":
    """Return the next-iterate state.

    Pure functional update — see class docstring.
    """
    lam_cost_next = self._relaxed_lam_cost(state)

    return state.replace(
        x=candidate.x,
        u=candidate.u,
        x_prop=candidate.x_prop,
        x_prop_plus=candidate.x_prop_plus,
        lam_cost=lam_cost_next,
        adaptive_state_code=jnp.asarray(
            int(AdaptiveStateCode.ACCEPT_CONSTANT), dtype=jnp.int32
        ),
    )

ConstantProximalWeightSpec

Bases: BaseModel

Validates ConstantProximalWeight configuration from dict/YAML input.

Source code in openscvx/algorithms/autotuner/constant_proximal_weight.py

class ConstantProximalWeightSpec(BaseModel):
    """Validates ConstantProximalWeight configuration from dict/YAML input."""

    type: Literal["ConstantProximalWeight"] = "ConstantProximalWeight"
    lam_cost_drop: int = -1
    lam_cost_relax: float = 1.0

    model_config = ConfigDict(extra="forbid")

    def build(self) -> ConstantProximalWeight:
        return ConstantProximalWeight(**self.model_dump(exclude={"type"}, exclude_unset=True))

RampProximalWeight

Bases: AutotuningBase

Ramp lam_prox toward lam_prox_max then hold.

update_weights is a pure functional update on the :class:AlgorithmState pytree; see the base-class contract.

Source code in openscvx/algorithms/autotuner/ramp_proximal_weight.py

class RampProximalWeight(AutotuningBase):
    """Ramp ``lam_prox`` toward ``lam_prox_max`` then hold.

    ``update_weights`` is a pure functional update on the
    :class:`AlgorithmState` pytree; see the base-class contract.
    """

    COMPUTES_ACCEPTANCE_METRICS = False

    def __init__(
        self,
        ramp_factor: float = 1.0,
        lam_prox_max: float = 1e3,
        lam_cost_drop: int = -1,
        lam_cost_relax: float = 1.0,
    ):
        self.hyper = RampProximalWeightHyper(
            ramp_factor=ramp_factor,
            lam_prox_max=lam_prox_max,
            lam_cost_drop=lam_cost_drop,
            lam_cost_relax=lam_cost_relax,
        )

    def update_weights(
        self,
        state: "AlgorithmState",
        candidate: "CandidateIterate",
        nodal_constraints: "LoweredJaxConstraints",
        settings: Config,
        params: dict,
    ) -> "AlgorithmState":
        """Return the next-iterate state.

        Pure functional update — see class docstring.
        """
        lam_cost_next = self._relaxed_lam_cost(state)

        was_at_max = jnp.all(state.lam_prox >= state.hyper.lam_prox_max)
        new_lam_prox = jnp.minimum(
            state.lam_prox * state.hyper.ramp_factor, state.hyper.lam_prox_max
        )

        code = jnp.where(
            was_at_max,
            jnp.int32(AdaptiveStateCode.ACCEPT_CONSTANT),
            jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
        )

        return state.replace(
            x=candidate.x,
            u=candidate.u,
            x_prop=candidate.x_prop,
            x_prop_plus=candidate.x_prop_plus,
            lam_prox=new_lam_prox,
            lam_cost=lam_cost_next,
            adaptive_state_code=code,
        )

update_weights(state: AlgorithmState, candidate: CandidateIterate, nodal_constraints: LoweredJaxConstraints, settings: Config, params: dict) -> AlgorithmState

Return the next-iterate state.

Pure functional update — see class docstring.

Source code in openscvx/algorithms/autotuner/ramp_proximal_weight.py

def update_weights(
    self,
    state: "AlgorithmState",
    candidate: "CandidateIterate",
    nodal_constraints: "LoweredJaxConstraints",
    settings: Config,
    params: dict,
) -> "AlgorithmState":
    """Return the next-iterate state.

    Pure functional update — see class docstring.
    """
    lam_cost_next = self._relaxed_lam_cost(state)

    was_at_max = jnp.all(state.lam_prox >= state.hyper.lam_prox_max)
    new_lam_prox = jnp.minimum(
        state.lam_prox * state.hyper.ramp_factor, state.hyper.lam_prox_max
    )

    code = jnp.where(
        was_at_max,
        jnp.int32(AdaptiveStateCode.ACCEPT_CONSTANT),
        jnp.int32(AdaptiveStateCode.ACCEPT_HIGHER),
    )

    return state.replace(
        x=candidate.x,
        u=candidate.u,
        x_prop=candidate.x_prop,
        x_prop_plus=candidate.x_prop_plus,
        lam_prox=new_lam_prox,
        lam_cost=lam_cost_next,
        adaptive_state_code=code,
    )

RampProximalWeightSpec

Bases: BaseModel

Validates RampProximalWeight configuration from dict/YAML input.

Source code in openscvx/algorithms/autotuner/ramp_proximal_weight.py

class RampProximalWeightSpec(BaseModel):
    """Validates RampProximalWeight configuration from dict/YAML input."""

    type: Literal["RampProximalWeight"] = "RampProximalWeight"
    ramp_factor: float = 1.0
    lam_prox_max: float = 1e3
    lam_cost_drop: int = -1
    lam_cost_relax: float = 1.0

    model_config = ConfigDict(extra="forbid")

    def build(self) -> RampProximalWeight:
        return RampProximalWeight(**self.model_dump(exclude={"type"}, exclude_unset=True))