Constraint solvers
Compose CBF action filters, CFS trajectory projections and adaptive augmented-Lagrangian scheduling with the numerical interfaces below.
Constraints participate inside the sampling loop. A candidate batch has
shape (N, H, action_dim): filters correct candidate actions, rollout residuals
change their weights, and manifold geometry shapes the next generative update.
MDOC, MD-COAS, 2GO and MGA compose these operations differently. See
Batched constraints and generative samples
for the data flow, tensor shapes, trained-diffusion DPCC/SafeDiffuser paths and
an executable JAX example.
| Solver | What it provides | Implementation | Upstream reference |
|---|---|---|---|
| Closed form | Euclidean halfspace projection, box clipping and unconstrained solves for special cases. | ClosedFormSolver | Local implementation |
| JAXopt OSQP | A general convex-QP interface built on jaxopt.OSQP. |
JAXOPTOsqpSolver | JAXopt documentation |
| OSQP | Sparse CPU convex-QP solves through the OSQP Python interface. | OSQPSolver | OSQP documentation |
| CVXOPT | CPU convex-QP solves, including trajectory smoothness and slack formulations. | CVXOPTSolver | CVXOPT documentation |
Install optional numerical packages with:
python -m pip install -e ".[optimization]"
Compose constraint handling
| Component | Use it for | Implementation |
|---|---|---|
| CBF filters | Correcting actions along rollouts, including Cartesian-to-joint Jacobian lifting. | Closed-form · QP filter · Joint lift |
| CFS projections | Building local obstacle halfspaces and correcting per-step or full-horizon actions. | Per step · Full horizon · Joint lift |
| Adaptive ALM scheduling | Updating feasibility penalties and projection effort from constraint feedback. | ALMAdaptiveConstraintScheduler |
CBF and CFS identify constraint-handling methods; the numerical solver is a separate choice. See the original CBF-QP paper and CFS paper for the underlying formulations, then inspect MDOC and MD-COAS recipes for their integration.
Choose the right solve path
- Batching follows the integration. JAX filters map over candidate
trajectories with
vmap; a state-dependent rollout usesscanover time. Full-horizon CFS couples the actions within each trajectory. These are different from solving one independent QP for every(sample, time)pair. - Closed form is specialized. Its halfspace and box paths perform Euclidean projection and clipping. They are not general QP solves for arbitrary Hessians.
- Fast JAX helpers are approximate projections. Functions such as
solve_slack_qp_jax()use finite, projection-style updates. They are distinct from theJAXOPTOsqpSolver.solve_qp()interface that invokesjaxopt.OSQP. - Host-facing QP wrappers are separate.
OSQPSolverandCVXOPTSolveruse CPU Python/NumPy interfaces. The currentJAXOPTOsqpSolver.solve_qp()wrapper also returns NumPy arrays. Do not put these wrappers directly inside ajax.jit/vmapsampling loop; use the existing pure-JAX filter path or an explicit host-side integration. - QPAX is an optional operator fallback. The per-step operator can call QPAX; this repository does not provide a separate QPAX solver class.
Use the compatibility matrix for backend and device status. Constraint margins, slack, iteration limits and the selected task determine the behavior of an integration; a solver name alone does not establish task safety or numerical convergence.