CRAX

Constrained RL Accelerated with JAX

CRAX: Fast Safe Reinforcement Learning Benchmarking

Tristan Tomilin Mourad Boustani Mickey Beurskens Khoi H. B. Nguyen Thiago D. Simão

Eindhoven University of Technology

Paper Code Play in browser BibTeX

Abstract

Safety is a core concern for deploying reinforcement learning (RL) agents in real-world domains such as robotics and autonomous driving. While benchmarks have been central to progress in RL, existing 3D physics-based safety benchmarks remain computationally slow, limiting large-scale experimentation and rapid prototyping. To address this gap, we propose CRAX (Constrained RL Accelerated with JAX). Built on top of the MuJoCo XLA (MJX) physics engine, CRAX leverages vectorized operations and hardware acceleration, yielding up to 200x faster training over comparable CPU-based safety benchmarks. The benchmark features eight tasks spanning three difficulty levels and multiple agent morphologies. Evaluating seven popular safe RL methods, we find that none dominates across tasks, and that learning safe policies from pixels remains largely unsolved.

Can you beat the agents?

Put your own safe policy to the test. Drive the Point agent through a full episode of a safe navigation task and see where you land on the leaderboard against the eight baselines from the paper. Collect reward, but watch out for the cost counter! You have a safety budget of 25 per episode.

Runs MuJoCo compiled to WebAssembly, right in your browser (~10 MB download).

An episode lasts as long as in training: 1000 steps (80 s), or 2000 steps (160 s) for Push. Use whichever camera you like; the agents never see pixels, only lidar, compass and body sensors. The browser port runs the same MuJoCo models and task logic as python -m crax.play, but on CPU MuJoCo instead of MJX, so trajectories can differ slightly from the benchmark.

Environments

Eight tasks, each with a distinct reward and cost signal. There are five constraint types across the tasks: hazards velocity height contact boundary.

Safe Goal
Point GoalReach the goal without colliding with the hazards
Safe Push
Point PushPush the block to the goal while avoiding obstacles
Safe Circle
Point CircleMove along a circular path while avoiding hazards
Safe Reacher
Reacher ReachReach the target without touching the hazards
Safe Lift
Spider LiftMove forward while lifting designated legs
Safe Pathway
Walker PathwayTraverse the path without stepping on the hazards
Safe Velocity
Ant VelocityMove fast while adhering to a velocity limit
Safe Height
Humanoid HeightMove forward under the height bound

Agents × Tasks

CRAX tasks support different agent morphologies.

3D2DFixed base
TaskPointAntHumanoidSpiderHalfCheetahWalker2dHopperReacher
Goal●●●●
Circle●●●●
Push●●●●
Velocity●●●●●●●
Height●
Lift●
Pathway●●●
Reach●

Highlighted entries are the configurations evaluated in the paper.

Three difficulty levels

Each suite ramps from a lenient setting most methods can solve to one that leaves ample room for improvement: more varied hazards, smaller goals, moving targets, or stricter bounds.

Level 1
Level 1
Level 2
Level 2
Level 3
Level 3

Learning from pixels

Every task exposes an egocentric camera mounted on the agent and an allocentric camera observing it from outside. Frames are rendered on the GPU inside the same compiled loop as the simulation, so pixel observations never leave the device. Policies trained from pixels receive these views at 64×64.

EgocentricAllocentric
Point Goal
EgocentricAllocentric
Walker Pathway
EgocentricAllocentric
Point Circle

Algorithms

All baselines share one fully compiled JAX training loop with hooks for custom losses and constraint handling, so adding a method means writing its loss, not its infrastructure. Select one with --alg.

ppoPPOUnconstrained on-policy reference
ppo_costPPO-CostCost folded into the reward as a penalty
ppo_lagPPO-LagrangianPrimal-dual with a learned multiplier
ppo_pidPPO-PIDPID-controlled Lagrange multiplier
ppo_sautePPO-SautéSafety budget augmented into the state
p3oP3OPenalized proximal policy optimization
focopsFOCOPSFirst-order constrained optimization in policy space
crpoCRPOAlternates reward and cost steps, no multiplier
sacSACUnconstrained off-policy reference
sac_lagSAC-LagrangianOff-policy primal-dual
sac_pidSAC-PIDOff-policy PID Lagrangian

Results

Two orders of magnitude faster

On identical hardware, CRAX runs 25–180× faster than the peak of Safety-Gymnasium. Safety-Gymnasium saturates at 3K–20K steps per second and runs out of CPU memory at a few hundrer parallel environments. CRAX keeps scaling to thousands of parallel environments and peaks at roughly 1M SPS on Point Goal and Ant Velocity and 340K SPS on Humanoid Velocity.

Steps per second of CRAX and Safety-Gymnasium versus number of parallel environments

The full CRAX evaluation suite (algorithms × environments × levels × seeds) is hundreds of runs and trillions of environment steps. CRAX completes it in 2 weeks on a single H100. The same study on Safety-Gymnasium would take close to a year!

Ramp up the seeds

Few seeds give noisy results, and many seeds used to be too expensive. On Humanoid Height, three seeds leave the methods' confidence bands overlapping. CRAX runs twenty seeds in about six hours, enough for a clear ranking.

3 seeds~1 h
Humanoid Height reward curves of all methods averaged over 3 seeds, with wide overlapping confidence bands
Very noisy results
+17× seeds
20 seeds~6 h
Humanoid Height reward curves of all methods averaged over 20 seeds, with tight separated confidence bands
A clear ranking of methods

How many seeds are enough?

CRAX's throughput lets us achieve reliable results. We run 20 seeds on every task and measure how the 95% confidence interval narrows as seeds are added. Relative to the mean, it shrinks from about 90% at two seeds to 40% at three and 22% at five, after which gains are small.

The curves also show that some tasks are far more volatile than others. Point Goal and Point Circle are tight from the start, while Spider Lift and Humanoid Velocity stay wide well past three seeds. A seed count that suffices for one task can mislead on another. 10 seeds keep the interval within 30% of the mean on every task.

Confidence interval width versus number of seeds

Evaluation on Level 1

Seven SafeRL baselines plus unconstrained PPO, trained for 500M steps with 10 seeds on every Level 1 task. The dashed line marks the cost threshold d = 25. Methods strike different trade-offs between performance and safety, and none is best everywhere.

PPO sets an example of obtainable reward when safety is ignored. The Lagrangian methods keep cost close to the threshold: PPOPID stays within it on every task, PPOLag on all but one. PPOCost settles for a compromise with no guarantee of satisfying the constraint, while PPOSauté is safe on Reacher, Height and Pathway but otherwise behaves much like PPO. FOCOPS is strongest on navigation (Goal, Push) and Reacher, yet weakest on forward locomotion (Lift, Height, Pathway).

Reward and cost of baselines on Level 1 tasks
Summary across the 8 environments. Wins: highest reward while safe (cost < 25). Safe%: share of environments where the method is safe.
AlgorithmLevel 1Level 2Level 3Total
Wins ↑Safe% ↑Wins ↑Safe% ↑Wins ↑Safe% ↑Wins ↑Safe% ↑
PPO12511200212
PPOCost088038125150
PPOLag28801000100296
PPOPID0100050188179
PPOSaute138025012125
P3O262375375871
FOCOPS175288388683
CRPO175262038358

Vision makes learning a safe policy harder

Every method earns less reward from egocentric frames than from state vectors, and less still from allocentric frames. No method learns to push the block from allocentric pixels, and on the remaining five environments none learns a meaningful policy from pixels at all.

Reward and cost from vector, egocentric and allocentric observations

Varying safety requirements

Real applications differ in how much risk they tolerate, so CRAX allows the cost bound to be adjustable. We explore the effects of this by training PPOLag on Level 1 with bounds of 15, 25 (the default) and 35.

Tight bound: 15

Point Goal - Level 1The agent steers around the hazards.

Loose bound: 35

Point Goal - Level 1The agent often cuts straight through them.

One full Level 1 episode (1000 steps) of PPOLag per bound. Both collect almost the same reward. Only the risk they take differs.

PPOLag tracks each bound closely. Its cost settles on the dotted line every time. The trade-off is not symmetric: tightening the bound from 25 to 15 costs more reward than loosening it to 35 gains, most visibly on Point Push and Reacher Reach.

Reward and cost of PPOLag on Point Goal, Reacher Reach, Point Push and Point Circle under cost bounds of 15, 25 and 35

Curriculum learning & safety transfer

Can easier levels or an unconstrained warm start help on the hardest tasks? We compare three regimes under the same budget of 500M steps: Normal trains on Level 3 from scratch for all 500M steps. Curriculum steps through Levels 1, 2 and 3 with 167M steps each, carrying parameters over. Transfer trains an unconstrained PPO policy on Level 3 for 250M steps, then fine-tunes it with a safe RL method for the remaining 250M.

Normal: Level 3 directly

Level 3 500M stepsToo careful!

Curriculum: Level 1 → 2 → 3

Level 1 167M stepsLearns to run with few hazards.
Level 2 167M stepsMaintains a good pace while safe.
Level 3 167M stepsEarns twice the reward!

P3O on Walker Pathway. Overlays show the running episode reward and cost; the cost bound is 25.

Neither curriculum learning nor safety transfer helps everywhere. On Walker Pathway, both raise the reward of PPOLag, PPOPID and P3O while staying under the bound, since direct training leaves much of the safety budget unused. On Reacher Reach, curriculum lifts PPOLag and P3O, but transfer breaks safety for every method except FOCOPS. On Point Goal, direct training remains best.

Reward and cost on Level 3 for Normal, Curriculum and Transfer training of PPOLag, PPOPID, P3O and FOCOPS

Get started

Install

git clone https://github.com/TTomilin/CRAX.git
cd CRAX
pip install -e ".[train,cuda]"

Train

python -m crax.training.train_env \
  --env_name safe_goal_point --difficulty 1 \
  --alg ppo_lag

Play with the keyboard

pip install pygame
python -m crax.play --env_name safe_push_point --level 3

Step an environment

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, key = jax.random.split(rng)
    action = jax.random.uniform(key, (env.action_size,),
                                minval=-1.0, maxval=1.0)
    state = step_fn(state, action)
    reward = float(state.reward)
    cost = float(state.info['cost'])  # safety cost per step

Citation

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