CRAX: Fast Safe Reinforcement Learning Benchmarking

Poster A: Monday -- 11:00 - 12:30

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

Keywords: reinforcement learning, safe RL, safety, benchmark, JAX, 3D, physics, simulation, MuJoCo, XLA

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 safety benchmarks with high-fidelity 3D physics remain computationally slow, limiting large-scale experimentation and rapid prototyping. To address this gap, we propose **CRAX** (**C**onstrained **R**L **A**ccelerated with JA**X**). Built on top of the MuJoCo XLA (MJX) physics engine with realistic 3D dynamics, \acronym~leverages vectorized operations and hardware acceleration, yielding up to ~100x speedups over comparable CPU-based safety benchmarks. The benchmark features six environment suites and three agent-specific tasks, each spanning three difficulty levels. Evaluating six popular safe RL methods shows that no single approach dominates across all tasks, and reveals the trade-offs between performance and safety. We find that curriculum learning across difficulty levels and safety transfer can improve performance over direct training in harder settings.