CBFJAX — Control Barrier Functions in JAX

CBFJAX is a high-performance JAX implementation of Control Barrier Functions (CBFs) for safety-critical control. It provides a clean, functional, JIT-compatible API for building safe controllers — from simple closed-form filters up to Backup-CBFs and NMPC with barrier constraints — and runs efficiently on CPU and GPU.

This project is the JAX successor to the CBFTorch framework.

---

Features

• Pure JAX, end-to-end JIT — barriers, dynamics, and safety filters are all equinox modules with functional semantics; trajectory rollouts use diffrax.
• Higher-Order CBFs (HOCBFs) with automatic differentiation for arbitrary relative degree.
• Predicted-Flow CBFs (P-CBFs) — barriers that are functionals of a predicted flow under a parametrized control plan, certifying safety over the whole prediction horizon in a single QP:
  • FlowBarrier + ParametricFlowSafeControl
  • FlowBarrier2 + ParametricFlowSafeControl2
• A toolbox of controllers and safe-control backends:
  • Closed-form min-intervention safe control (MinIntervCFSafeControl)
  • QP-based safe control with slack variables (MinIntervQPSafeControl)
  • Input-constrained QP (MinIntervInputConstQPSafeControl)
  • Backup-CBF with forward invariance (MinIntervBackupSafeControl)
  • Predicted-flow safe control (ParametricFlowSafeControl, …Control2)
  • MPPI with barrier-aware cost (MPPIControl)
  • NMPC with barrier constraints (acados / do-mpc — optional)
  • Constrained iLQR with barrier-aware cost (trajax — optional)
• Pluggable QP backends behind one interface (params['qp_solver']): qpax (default), jaxopt_osqp (warm-started), mpax, and cvxopt.
• Composable barrier algebra: MultiBarriers, SoftCompositionBarrier, HardCompositionBarrier, BackupBarrier, FlowBarrier, FlowBarrier2.
• Config-driven construction (cbfjax.from_config): a named barrier namespace — define barriers by name, reference them by name, wire one into the filter; unused entries stay available for plotting/analysis.
• Built-in dynamics: unicycle, single/double integrator, kinematic bicycle, inverted pendulum, reduced-order unicycle — plus a generic AffineInControlDynamics base and a CustomDynamics escape hatch.
• 64-bit precision by default for the numerical stability that CBF methods require.

---

Installation

From PyPI

pip install cbfjax

The install pulls in JAX, Equinox, Diffrax, qpax, jaxopt, mpax, matplotlib, NumPy, and SciPy — enough to run every safety filter and every example.

Development install

pip install cbfjax[dev]

Adds the optional solver backends (cvxopt, CasADi + do-mpc for NMPC) and the development tooling (pytest, build, twine, ruff, black, mypy, …).

From source

git clone https://github.com/amirsaeid254/cbfjax.git
cd cbfjax
pip install -e .[dev,examples]

---

Quick Start

Everything is built from one config dict — named barriers, referenced by name, and one controller consuming the barrier it names:

dynamics ──► barriers {name: spec, ...} ──► filter ──► safe action u(x)
                                              ▲
                       desired_control (goal controller / iLQR / MPPI planner)

Minimal example — QP safety filter on a unicycle

import jax.numpy as jnp
import cbfjax

goal = jnp.array([5.0, 5.0])

system = cbfjax.from_config({
    # 1. Dynamics — state [x, y, v, theta], control [a, omega]
    'dynamics': 'unicycle',

    # 2. Barriers — every barrier gets a name
    'barriers': {
        # stay outside the unit disk at the origin (relative degree 2)
        'disk': {'type': 'func',
                 'h': lambda x: jnp.linalg.norm(x[:2]) - 1.0,
                 'rel_deg': 2, 'alphas': (10.0,)},
    },

    # 3. QP-based min-intervention safety filter consuming 'disk'
    'filter': {
        'type': 'min_interv_qp',
        'barrier': 'disk',
        'action_dim': 2,
        'alpha': lambda h: 1.0 * h,
        'params': {'slack_gain': 200.0, 'slacked': True, 'qp_solver': 'qpax'},
        'desired_control': lambda x: 0.5 * jnp.array([goal[0] - x[0],
                                                      goal[1] - x[1]]),
    },
})

# 4. Query the safe action (single state in, single action out)
x0 = jnp.array([-2.0, -2.0, 0.0, 0.0])
u_safe, _ = system.filter.optimal_control(x0, system.filter.get_init_state())
print(u_safe)

Maps and compositions

A map entry turns a geometric world description into member barriers, and compositions reduce them: soft_composition (scalar softmin — what closed-form filters need), hard_composition (scalar exact min), or multi_barrier (one QP constraint per member). Entries nobody consumes are still built — handy for plotting and analysis.

cfg = {
    'dynamics': 'unicycle',
    'barriers': {
        'map': {'type': 'map',
                'geoms': (
                    ('cylinder', {'center': (2.0, 2.0), 'radius': 0.5}),
                    ('norm_box', {'center': (-1.0, 3.0), 'size': (1.0, 1.0)}),
                    ('norm_boundary', {'center': (0.0, 0.0), 'size': (10.0, 10.0)}),
                ),
                'velocity': (2, (-2.0, 2.0)),
                'cfg': {'softmin_rho': 20, 'pos_barrier_rel_deg': 2,
                        'vel_barrier_rel_deg': 1, 'obstacle_alpha': (10.0,),
                        'boundary_alpha': (10.0,), 'velocity_alpha': ()}},
        'rows': {'type': 'multi_barrier', 'barriers': ['map']},
    },
    'filter': {'type': 'min_interv_input_const_qp', 'barrier': 'rows', ...},
}
system = cbfjax.from_config(cfg)
system.barriers['map']     # the Map instance — plotting, member access
system.barriers['rows']    # the MultiBarriers the filter consumes

Closed-loop simulation

trajs = system.filter.get_optimal_trajs(
    x0=x0[None],           # (batch, state_dim)
    sim_time=10.0,
    timestep=0.01,
    method='euler',
)
print(trajs.shape)  # (T, batch, state_dim)

More end-to-end scripts live under examples/unicycle/ (closed-form, QP, input-constrained QP, NMPC, iLQR, hierarchical) and examples/unicycle/backup/ (Backup-CBF).

cd examples/unicycle
python 03_unicycle_qp.py

---

Architecture

cbfjax/
├── barriers/                       # CBF & HOCBF
│   ├── barrier.py                  #  Single barrier
│   ├── multi_barrier.py            #  Multiple barriers
│   ├── composite_barrier.py        #  Soft / hard composition
│   ├── backup_barrier.py           #  Backup-CBF
│   ├── parametric_flow_barrier.py  #  Predicted-flow CBF
│   └── parametric_flow_barrier2.py
├── dynamics/                       # Affine-in-control system dynamics
│   ├── base_dynamic.py
│   ├── unicycle.py
│   ├── unicycle_reduced_order.py
│   ├── double_integrator.py
│   ├── single_integrator.py
│   ├── bicycle.py
│   └── inverted_pendulum.py
├── controls/                       # Nominal/optimal controllers
│   ├── base_control.py
│   ├── mppi_control.py             #  MPPI (GPU-parallel, vmap+scan)
│   ├── parametric_control.py       #  ZOH / FOH control plans u_p(τ; θ)
│   ├── ilqr_control.py             #  (optional: trajax)
│   ├── nmpc_control.py             #  (optional: casadi + acados/do-mpc)
│   └── control_types.py
├── safe_controls/                  # Safety filters
│   ├── base_safe_control.py
│   ├── closed_form_safe_control.py
│   ├── qp_safe_control.py
│   ├── backup_safe_control.py
│   ├── parametric_flow_safe_control.py   #  Predicted-flow safe control
│   ├── parametric_flow_safe_control2.py
│   ├── nmpc_safe_control.py        #  (optional)
│   └── ilqr_safe_control.py        #  (optional)
├── factory.py                      # from_config: named-barrier construction
├── utils/
│   ├── integration.py              #  Diffrax-based ODE rollouts
│   ├── qp.py                       #  QP backends (qpax, OSQP, mpax, cvxopt)
│   ├── make_map.py                 #  Map / barrier factory
│   ├── jax2casadi/                 #  JAX → CasADi conversion (used by NMPC)
│   ├── profile_utils.py
│   └── utils.py
└── config.py                       # JAX configuration helpers

---

Key concepts

Control Barrier Functions

For a control-affine system ẋ = f(x) + g(x) u and a safe set C = {x | h(x) ≥ 0}, a barrier function h ensures forward invariance of C whenever there exists u such that

L_f h(x) + L_g h(x) · u ≥ -α(h(x))

where α is a class-K function.

Higher-Order CBFs

For barriers of relative degree r > 1, CBFJAX automatically constructs the HOCBF series ψ_0, ψ_1, …, ψ_r from a user-provided list of class-K functions (via Barrier(barrier_func=h, rel_deg=r, alphas=[α_1, …, α_r], dynamics=dyn) or a {'type': 'func', 'h': h, 'rel_deg': r, 'alphas': [...]} config entry).

---

Citation

If you use CBFJAX in your research, please cite it as:

@article{safari2026predicted,
  title={Predicted-Flow Control Barrier Functions for Real-Time Safe Optimal Control},
  author={Safari, Amirsaeid and Hoagg, Jesse B},
  journal={arXiv preprint arXiv:2606.00297},
  year={2026}
}

Related work

• CBFTorch — PyTorch implementation of CBFs.