CRAX: Fast Safe Reinforcement Learning Benchmarking

2026-06-18

AI Summary

Existing high-fidelity 3D safety reinforcement learning (RL) benchmarks suffer from poor computational efficiency, hindering large-scale experimentation and rapid prototyping. This work proposes CRAX, the first benchmark that integrates JAX-based acceleration with high-fidelity 3D safe RL by leveraging the MuJoCo XLA (MJX) physics engine. Through vectorized parallel simulation and hardware acceleration, CRAX achieves approximately 100× speedup over prior benchmarks. It encompasses six environment categories, three task types, and three difficulty levels, enabling systematic evaluation of six state-of-the-art safe RL algorithms. The study reveals fundamental trade-offs between performance and safety and demonstrates that curriculum learning and safe transfer strategies substantially improve agent performance on challenging tasks.

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 safety benchmarks with high-fidelity 3D physics 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 with realistic 3D dynamics, CRAX 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.

Problem

• safe reinforcement learning
• benchmarking
• computational efficiency
• 3D physics simulation
• real-world deployment

Innovation

• safe reinforcement learning
• JAX acceleration
• vectorized simulation
• high-fidelity physics
• benchmarking