CRAX: Constrained Reinforcement Learning Accelerated with JAX

🌐 Check out the project page to try out CRAX yourself and claim a spot on the leaderboard!

CRAX is a high-performance benchmark for Constrained Reinforcement Learning (Safe RL) built on top of Brax and MuJoCo XLA (MJX). It provides GPU/TPU-accelerated environments with safety constraints and a suite of state-of-the-art safe RL algorithms.

Features

• Massively parallel simulation: Train agents across thousands of environments simultaneously on GPU/TPU
• Configurable difficulty levels: Environments with varying numbers and types of hazards
• Multiple constraint types: Cylindrical hazards, cubical hazards, boundary constraints, velocity limits, and more
• State-of-the-art algorithms: Implementations of leading constrained RL methods

Safe RL Algorithms

CRAX includes efficient JAX implementations of the following, selected with the --alg flag:

Algorithm	Flag	Description	Reference
PPO	ppo	Vanilla PPO (unconstrained on-policy baseline)	Schulman et al., 2017
PPO-Cost	ppo_cost	PPO with a cost penalty folded into the reward
PPO-Lagrange	ppo_lag	PPO with Lagrangian relaxation for constraints	Ray et al., 2019
PPO-PID	ppo_pid	PPO with PID-controlled Lagrange multiplier	Stooke et al., 2020
PPO-Saute	ppo_saute	State Augmentation for Safe RL	Sootla et al., 2022
FOCOPS	focops	First Order Constrained Optimization in Policy Space	Zhang et al., 2020
P3O	p3o	Penalized Proximal Policy Optimization	Zhang et al., 2022
CRPO	crpo	Constrained Rectified Policy Optimization	Xu et al., 2021
SAC	sac	Soft Actor-Critic (unconstrained off-policy baseline)	Haarnoja et al., 2018
SAC-Lagrange	sac_lag	SAC with Lagrangian relaxation for constraints
SAC-PID	sac_pid	SAC with PID-controlled Lagrange multiplier

All algorithms share a common training infrastructure with hooks for custom loss functions and constraint handling, making it easy to implement new methods. The registry lives in get_algorithm_train_fn in crax/training/run_utils.py; each algorithm adds its own CLI section in crax/training/config.py (--safety_bound, --pid_kp, --nu_lr, --tau, ...), and irrelevant arguments are filtered out per algorithm by filter_kwargs_for_fn.

Environments

CRAX environments follow a task-agent naming convention: safe_<task>_<agent>.

Task	Agents Supported	Description
Goal	Point, Ant, Spider, Humanoid, Swimmer	Navigate to goal positions while avoiding hazards
Circle	Point, Ant, Spider, Humanoid, Swimmer	Orbit a circular path within boundaries
Button	Point, Ant, Spider, Humanoid, Swimmer	Press target buttons while avoiding gremlins and hazards
Push	Point, Ant, Spider, Humanoid, Swimmer	Push a block to a goal while avoiding hazards
Velocity	Ant, Humanoid, HalfCheetah, Hopper, Swimmer, Walker2D	Locomotion under maximum velocity constraints
Lift	Ant, Humanoid, Spider	Locomotion with restricted legs touching the floor
Height	Humanoid, HalfCheetah, Hopper, Walker2D	Locomotion under a height ceiling
Pathway	HalfCheetah, Hopper, Walker2D	Traverse a corridor with hazard gaps
Reacher	Reacher	Robotic arm reaching targets while avoiding obstacles

Difficulty Levels

Every environment offers three difficulty levels that progressively increase constraint difficulty.

Environment	Level 1	Level 2	Level 3
Goal
Circle
Push
Reacher
Pathway
Height
Lift

Installation

CRAX is published on PyPI as crax-bench (the import name is crax):

pip install crax-bench             # environments only
pip install "crax-bench[train]"    # + the RL algorithms in crax.training

To work on CRAX itself, or to use the scripts in results/ and scripts/, install from source:

git clone https://github.com/TTomilin/CRAX.git
cd CRAX

Create and activate a virtual environment using your preferred tool:

# Option A: venv
python3 -m venv .venv && source .venv/bin/activate

# Option B: conda
conda create -n crax python=3.11 && conda activate crax

Then install the package. The base install contains only what is needed to build, step and render environments — no RL dependencies:

pip install -e .

To train agents, install the train extra, which adds the crax.training dependencies (optax, orbax, wandb):

pip install -e ".[train]"

Other extras:

Extra	Adds	For
train	optax, orbax-checkpoint, wandb, imageio, matplotlib	Running the RL algorithms in crax.training
analysis	pandas, matplotlib, wandb	Downloading and plotting results (results/)
gym	gym, dm_env	The Gym/dm_env wrappers in crax.envs.wrappers
torch	torch	The PyTorch wrapper
vision	mujoco-warp, warp-lang	GPU pixel observations (--vision)
play	pygame	Playing environments with the keyboard (crax.play)
cuda	jax[cuda12]	GPU support on CUDA 12
all	train + analysis + gym + vision + play	Everything except CUDA and torch

Extras combine, e.g. a GPU training install with pixel observations:

pip install -e ".[train,vision,cuda]"

See SETUP.md for pinned installs, GPU/headless configuration, wandb and troubleshooting, and docs/VISION_TRAINING.md for training from pixel observations.

Start Here

The quickest way to see CRAX working is examples/01_run_episode.py.

python examples/01_run_episode.py --env_name safe_goal_point --level 1

The core loop is just a few lines:

import jax
from crax import envs

env = envs.get_environment('safe_goal_point', level=1)
reset_fn, step_fn = jax.jit(env.reset), jax.jit(env.step)

rng = jax.random.PRNGKey(0)
state = reset_fn(rng)
for _ in range(300):
    rng, action_rng = jax.random.split(rng)
    action = jax.random.uniform(action_rng, (env.action_size,), minval=-1.0, maxval=1.0)
    state = step_fn(state, action)
    reward, cost = float(state.reward), float(state.info['cost'])  # safety cost per step

See examples/README.md for the full list of examples.

Playing CRAX

Point-agent environments can be played with a keyboard (requires the play extra):

python -m crax.play --env_name safe_goal_point --level 2

W/S thrust, A/D turn, C cycles the camera, R resets, Space pauses, Esc quits.

Training from the command line

The training entry points live in crax.training and are run as modules:

# Single environment
python -m crax.training.train_env --env_name safe_goal_point --alg ppo_lag --difficulty 1

# Curriculum across difficulty levels
python -m crax.training.train_curriculum --env_name safe_goal_point --alg ppo_lag

# Safety transfer: unsafe PPO pre-training, then safe fine-tuning
python -m crax.training.train_transfer --env_name safe_velocity_ant --alg ppo_lag

Installing the package (with the train extra) also puts them on your PATH as crax-train, crax-train-curriculum and crax-train-transfer. All three share the argument set defined in crax/training/config.py; run with --help for the full list.

Project Structure

crax/ holds the simulation and the environments; its crax/training/ subpackage holds the RL algorithms. Nothing outside crax.training imports it, so environments can be used on their own with a much smaller dependency set (see Installation).

CRAX/
├── crax/                        # Environments, simulation and (in training/) RL algorithms
│   ├── base.py, math.py, ...    # Physics core (Brax/MJX pipelines)
│   ├── generalized/, positional/, spring/, mjx/
│   ├── io/                      # MJCF loading, image + HTML export
│   ├── envs/                    # Environment definitions
│   │   ├── builder.py           # Modular XML scene builder
│   │   ├── difficulty.py        # Difficulty level configurations
│   │   ├── hazards.py           # Hazard generation utilities
│   │   ├── goals.py             # Goal sampling utilities
│   │   └── wrappers/            # Env wrappers (incl. pixel observations)
│   └── training/                # RL algorithms and training infrastructure
│       ├── agents/              # Algorithm implementations (PPO, SAC, etc.)
│       ├── networks.py          # Policy/value network factories
│       ├── acting.py            # Rollout and evaluation loops
│       ├── curriculum.py        # Curriculum training loop
│       ├── transfer.py          # Transfer learning loop
│       ├── checkpoint.py        # Checkpoint save/load
│       ├── config.py            # Shared CLI argument definitions
│       ├── run_utils.py         # Shared training helpers
│       ├── train_env.py         # Entry point: single environment training
│       ├── train_curriculum.py  # Entry point: progressive difficulty training
│       └── train_transfer.py    # Entry point: safety transfer learning
├── results/                     # Result download + plotting pipeline
├── examples/                    # Small self-contained usage examples
├── scripts/                     # Utility & visualization scripts
└── tests/                       # Pytest suite

Architecture

CRAX uses a modular architecture where constrained RL algorithms are wrappers around a base RL algorithm, e.g.:

ppo/train.py          # Base trainer with hooks (loss_fn, post_step_fn, init_aux_state_fn)
    │
    ├── ppo_lag/      # Lagrange multiplier update + lagrange loss
    ├── ppo_pid/      # PID controller + lagrange loss
    ├── focops/       # FOCOPS loss + nu update
    ├── p3o/          # P3O loss + kappa adaptation
    ├── ppo_saute/    # Environment wrapper approach
    └── crpo/         # Rectified reward/cost objective switching

This design minimizes code duplication and makes it easy to add new algorithms.

Contributing

Planned work is tracked in ROADMAP.md. Contributions are welcome. The roadmap flags which items are self-contained enough to pick up directly. For anything substantial, please open an issue first so we can avoid duplicated work.

Research Powered by CRAX

Papers that have used CRAX:

• Lindsay Spoor, Aske Plaat, Thomas Moerland. Evaluation Metrics for Safe Reinforcement Learning, 2026. arXiv:2609.15315

Using CRAX in your work? Feel free to open a pull request to add your paper to this list.

Acknowledgements

CRAX builds on Brax, Google's JAX-based physics and RL library, and on MuJoCo XLA (MJX) and MuJoCo Warp, which power the accelerated simulation and the GPU pixel renderer. We thank the Brax and MuJoCo teams for these foundations. PROVENANCE.md outlines which sources come from Brax. Our navigation and velocity suite tasks are close reimplementations from Safety Gym and Safety-Gymnasium.

Citation

If you use CRAX in your research, please cite:

@article{tomilin2026crax,
  title         = {CRAX: Fast Safe Reinforcement Learning Benchmarking},
  author        = {Tristan Tomilin and Mourad Boustani and Mickey Beurskens and Thiago D. Sim{\~a}o},
  journal       = {arXiv preprint arXiv:2606.20376},
  year          = {2026}
}