Search papers, labs, and topics across Lattice.
This paper introduces CRAX, a fast benchmarking tool for safe reinforcement learning (RL) that leverages the MuJoCo XLA physics engine for significant computational speedups. By utilizing vectorized operations and hardware acceleration, CRAX achieves up to ~100x faster performance compared to existing CPU-based safety benchmarks, facilitating large-scale experimentation. The evaluation of six popular safe RL methods within CRAX reveals no single approach excels universally, highlighting the nuanced trade-offs between performance and safety across various tasks and difficulty levels.
CRAX accelerates safe RL benchmarking by up to 100x, revealing critical trade-offs in performance and safety that challenge conventional wisdom in the field.
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.
