🤖 machine learning

CRAX: Fast Safe Reinforcement Learning Benchmarking

CRAX is a fast and safe reinforcement learning benchmark built on the MuJoCo XLA physics engine that leverages hardware acceleration to achieve up to 100x speedups over existing CPU-based safety benchmarks, enabling large-scale evaluation of safe RL methods across diverse environments and revealing key insights into performance-safety trade-offs and curriculum learning benefits.

Original authors: Tristan Tomilin, Mourad Boustani, Mickey Beurskens, Thiago D. Simão

Published 2026-06-19

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Tristan Tomilin, Mourad Boustani, Mickey Beurskens, Thiago D. Simão

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a robot to walk through a minefield. Your goal is to get the robot to the finish line as fast as possible (the Reward), but you must ensure it never steps on a mine (the Safety Constraint). If it steps on a mine, it gets a "cost" penalty. The challenge is finding the perfect balance: going too fast might mean hitting a mine, but moving too slowly means you never finish.

This is the core problem of Safe Reinforcement Learning (Safe RL).

The paper introduces a new tool called CRAX to help researchers solve this problem faster and better. Here is a breakdown of what they did, using simple analogies.

1. The Problem: The "Slow Motion" Bottleneck

Previously, researchers testing these robots had to use standard computer processors (CPUs) to simulate the physics.

The Analogy: Imagine trying to train a marathon runner, but every time they take a step, the simulation freezes for a second to calculate the physics. To run a full marathon, you'd have to wait years.
The Reality: Existing safety benchmarks were accurate but incredibly slow. This made it hard to run thousands of experiments to find the best training methods.

2. The Solution: CRAX (The "Turbo-Charged" Simulator)

The authors built CRAX (Constrained RL Accelerated with JAX).

The Analogy: Instead of training one runner at a time on a slow track, CRAX builds a massive stadium where thousands of robots are running simultaneously on a super-fast track powered by specialized graphics cards (GPUs).
The Result: It is roughly 100 times faster than previous tools. The paper claims that a task that used to take a year to simulate on old computers now takes only two weeks on their new system.

3. The Playground: Six Different "Minefields"

CRAX isn't just one game; it's a collection of six different environments, each with a different type of robot and a different type of danger. Think of it as a video game with different levels:

Safe Goal: A robot must reach a target while avoiding floating hazards.
Safe Push: A robot must push a heavy box to a goal without hitting obstacles.
Safe Spider: A six-legged robot must walk forward while keeping specific legs in the air (like a high-wire act).
Safe Height: A robot must walk forward but stay low to the ground, like ducking under a low ceiling.
Safe Pathway: A robot must hop across a path where stepping on the "wrong" spots incurs a penalty.
Safe Velocity: A robot must run fast but never exceed a specific speed limit.

Each of these games has three difficulty levels:

Easy: Few obstacles, wide margins for error.
Medium: More obstacles, tighter spaces.
Hard: A chaotic minefield where the robot must be extremely precise.

4. The Experiment: Who Wins the Race?

The researchers tested six popular AI training methods (algorithms) in this new, fast playground to see which one was best at balancing speed and safety.

The Findings: There was no single "champion."
- Some methods were very safe but moved very slowly (like a cautious turtle).
- Some moved fast but crashed often (like a reckless racer).
- P3O and FOCOPS were the strongest overall performers, managing to get high scores while staying safe in most scenarios.
- PPOLag was the safest (it almost never crashed), but it often didn't get the highest scores.

5. The "Training Camp" Discovery (Curriculum Learning)

The researchers tested a specific training strategy called Curriculum Learning.

The Analogy: Instead of throwing a student into a final exam immediately, you teach them the basics, then medium-level problems, and finally the hard exam.
The Result: This worked! For many robots, starting on the "Easy" level and gradually moving to "Hard" levels helped them learn better than if they had been thrown straight into the hardest level. It's like learning to ride a bike on a flat driveway before trying to ride down a steep hill.

6. Why This Matters (According to the Paper)

The paper doesn't claim this will immediately fix self-driving cars or save lives in hospitals. Instead, it claims to be a tool for scientists.

Because CRAX is so fast, researchers can now run massive experiments that were previously impossible.
It allows them to test many different ideas quickly to figure out how to make AI agents safer and more efficient in complex 3D worlds.

In summary: CRAX is a super-fast, GPU-powered video game engine designed specifically to teach robots how to be fast without being reckless. It proved that while no single AI method is perfect yet, training robots gradually (from easy to hard) helps them learn better, and having a fast simulator is essential for making progress in this field.

Technical Summary: CRAX (Constrained RL Accelerated with JAX)

Problem Statement

Safety is a critical requirement for deploying reinforcement learning (RL) agents in real-world domains such as robotics and autonomous driving. While benchmarks have historically driven progress in RL (e.g., ImageNet for CNNs, ALE for DQN), existing safety benchmarks with high-fidelity 3D physics remain computationally slow. Most current SafeRL research relies on CPU-based simulations (e.g., Safety Gym, Safety Gymnasium), which limits large-scale experimentation, rapid prototyping, and rigorous statistical analysis. The lack of hardware-accelerated, 3D physics-based safety benchmarks creates a bottleneck for developing and evaluating constrained RL algorithms.

Methodology

To address this gap, the authors propose CRAX (Constrained RL Accelerated with JAX), a novel benchmark built on top of the MuJoCo XLA (MJX) physics engine. MJX provides realistic 3D dynamics while leveraging JAX for vectorized operations and hardware acceleration (GPUs).

Benchmark Design

CRAX consists of six environment suites and three agent-specific tasks, each spanning three difficulty levels:

Safe Navigation: Includes Goal, Button, Circle, and Push tasks where agents must reach targets or manipulate objects while avoiding hazards.
Safe Velocity: Locomotion tasks with strict speed limits.
Safe Pathway: Agents traverse paths with hazard zones, penalizing foot placement errors.
Safe Reach: A planar robotic arm reaching targets while avoiding obstacles.
Safe Height: A humanoid agent walking while maintaining a low center of mass.
Safe Spider: A hexapod agent walking while keeping specific legs airborne.

The benchmark supports various agent morphologies (Point, Ant, Humanoid, Spider, HalfCheetah, Walker2d, Hopper, Reacher) and defines safety via Constrained Markov Decision Processes (CMDPs). The objective is to maximize expected cumulative reward $V^\pi(s)$ while ensuring the expected cumulative discounted cost $C^\pi(s)$ remains below a threshold $d$ .

Experimental Setup

The authors reimplemented seven popular SafeRL algorithms in JAX to evaluate their performance:

PPO: Unconstrained reference.
PPOCost: Treats costs as negative rewards.
PPOLag: Primal-dual approach with a learned Lagrange multiplier.
PPOPID: Refines PPOLag using a PID controller for the multiplier.
PPOSauté: Augments state with a safety budget.
P3O: Progressively increases cost penalties.
FOCOPS: Uses a trust-region formulation to bound expected cost.

Experiments were conducted on a single NVIDIA H100 GPU, running 500 million environment steps per configuration across 5 seeds.

Key Contributions

Hardware-Accelerated SafeRL Benchmark: CRAX is the first benchmark to combine high-fidelity 3D physics, explicit safety constraints, and GPU acceleration. It achieves up to ~100x speedups over comparable CPU-based benchmarks and scales to ~300,000 steps per second (SPS) with thousands of parallel environments.
Comprehensive Algorithm Evaluation: The paper provides a systematic evaluation of six SafeRL methods against an unconstrained baseline across diverse agent morphologies and difficulty levels, identifying specific strengths and limitations of each approach.
Analysis of Training Strategies: The study investigates the efficacy of curriculum learning (training across difficulty levels) and safety transfer (initializing safe policies with unconstrained ones), finding that these strategies can improve performance in challenging settings but are highly environment-dependent.

Results

Algorithm Performance

No Dominant Method: No single algorithm dominated across all tasks. P3O and FOCOPS emerged as the strongest baselines overall.
Safety vs. Performance Trade-offs: PPOLag achieved the highest safety percentage (satisfying all cost bounds in Levels 2 and 3) but often failed to achieve high rewards. PPOPID offered a less conservative trade-off, sacrificing some strict safety adherence for higher performance.
Task Specificity: Methods showed distinct affinities for constraint types. For instance, FOCOPS performed well on navigation and reaching tasks but struggled with forward-locomotion tasks.

Computational Efficiency

Throughput: CRAX achieved ~300K SPS at 8,192 parallel environments, whereas Safety-Gymnasium saturated at ~3K SPS and failed to scale beyond 256 environments due to CPU and memory bottlenecks.
Time Savings: A full evaluation suite (algorithms × environments × difficulty levels × seeds) that would take nearly a year on Safety-Gymnasium was completed in two weeks on a single H100 GPU using CRAX.

Curriculum and Transfer

Curriculum Learning: Improved performance in Safe Pathway and Safe Height but showed mixed results in Safe Goal and Safe Reacher.
Safety Transfer: Initializing safe algorithms with unconstrained PPO policies was ineffective in Safe Reacher (leading to constraint violations) but beneficial in Safe Pathway, where it helped agents utilize the safety budget more effectively than direct training.

Significance and Claims

The paper claims that CRAX serves as a vital tool for the SafeRL community by enabling large-scale experimentation that is currently infeasible with CPU-bound benchmarks. By providing a platform with realistic 3D dynamics and hardware acceleration, CRAX facilitates:

More rigorous testing of algorithms through extensive hyperparameter tuning and statistical significance analysis.
Faster iteration cycles for algorithm development.
The exploration of safety-performance trade-offs across diverse agent morphologies.

The authors acknowledge limitations, noting that MJX currently lacks some features of CPU-based MuJoCo (e.g., certain rigid-body collision types) and that the evaluation is restricted to on-policy methods and state-based observations. They identify off-policy methods, model-based approaches, and multi-agent scenarios as important directions for future work.