JAXtari: High-Throughput and Easy-to-Modify Arcade Learning Environment

The Arcade Learning Environment (ALE) has been a cornerstone of Deep Reinforcement Learning research, providing a standardized benchmark for evaluating general agents. However, its lack of native support for object-extraction and environment modification restricts research into agent generalization and reasoning.Furthermore, the ALE's dependence on CPU-based emulation creates a significant bottleneck in modern GPU-accelerated training pipelines, limiting both experimental scale and iteration frequency. In this work, we introduce JAXatari, an open-source, high-performance reimplementation of Atari 2600 environments. By leveraging massive parallelization on GPUs, PPO agents finish training with 100M frames on a single A100-GPU in under an hour. Beyond speed, JAXatari provides object-centric observations and a flexible modification system, allowing researchers to, e.g. develop neuro-symbolic methods or discover novel shortcuts and subsequently train and adapt their agents rapidly, sometimes in only minutes.