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.