CRAX: Fast Safe Reinforcement Learning Benchmarking

Safe reinforcement learning (SafeRL) is critical for deploying autonomous agents in real-world environments like robotics. However, existing safety benchmarks often rely on CPU-based physics simulations, which are computationally expensive and create a bottleneck for large-scale experimentation, hyperparameter tuning, and ablation studies. The authors introduce CRAX (Constrained RL Accelerated with JAX) to address this by providing a suite of 3D physics-based environments that leverage GPU acceleration via the MuJoCo XLA (MJX) backend.

CRAX provides six environment suites and three agent-specific tasks, each designed with three levels of difficulty. The benchmark explicitly models safety through cost signals, requiring agents to navigate the trade-off between reward maximization and constraint satisfaction. The authors evaluate six popular SafeRL baseline algorithms—including PPO-Lagrangian, PPO-PID, and FOCOPS—across these environments. They also assess the effectiveness of curriculum learning and safety transfer techniques to improve performance in high-difficulty settings.

The empirical evaluation demonstrates that CRAX achieves significant speedups, reaching up to 300,000 steps per second and outperforming existing CPU-based benchmarks by approximately two orders of magnitude. The results reveal that no single algorithm dominates across all tasks, highlighting distinct trade-offs between safety adherence and reward performance. While PPO-Lagrangian is the most consistent in satisfying safety constraints, it often sacrifices significant reward. Furthermore, the study finds that curriculum learning and safety transfer can enhance performance on difficult tasks, though the effectiveness of these strategies is highly dependent on the specific environment and algorithm.