Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX

TL;DR

Octax addresses the scalability bottleneck in RL environment research by delivering GPU-accelerated CHIP-8 environments in JAX. It combines end-to-end CHIP-8 emulation with vectorized GPU execution, enabling thousands of parallel game instances while preserving faithful mechanics. The authors demonstrate PPO learning across 16 CHIP-8 games and show substantial speedups and scalability compared to CPU-based implementations; they also introduce an LLM-assisted pipeline for automated environment generation, including Level 1–3 Target Shooter tasks with meaningful difficulty gradients. The work offers a practical platform for large-scale, reproducible RL research on classic arcade dynamics and curriculum learning.

Abstract

Reinforcement learning (RL) research requires diverse, challenging environments that are both tractable and scalable. While modern video games may offer rich dynamics, they are computationally expensive and poorly suited for large-scale experimentation due to their CPU-bound execution. We introduce Octax, a high-performance suite of classic arcade game environments implemented in JAX, based on CHIP-8 emulation, a predecessor to Atari, which is widely adopted as a benchmark in RL research. Octax provides the JAX community with a long-awaited end-to-end GPU alternative to the Atari benchmark, offering image-based environments, spanning puzzle, action, and strategy genres, all executable at massive scale on modern GPUs. Our JAX-based implementation achieves orders-of-magnitude speedups over traditional CPU emulators while maintaining perfect fidelity to the original game mechanics. We demonstrate Octax's capabilities by training RL agents across multiple games, showing significant improvements in training speed and scalability compared to existing solutions. The environment's modular design enables researchers to easily extend the suite with new games or generate novel environments using large language models, making it an ideal platform for large-scale RL experimentation.

Figures (7)

• Figure 1: Overview of CHIP-8 game environments implemented in Octax.
• Figure 2: Octax architecture: ROM loading, CHIP-8 emulation pipeline, and RL environment integration. The system transforms game ROMs through fetch-decode-execute cycles into vectorized JAX operations suitable for GPU acceleration.
• Figure 3: PPO learning curves across 16 games: Interquartile Mean (IQM) returns using 10th-90th percentile ranges over 5M timesteps, with confidence intervals computed across 12 random seeds.
• Figure 4: Performance scaling of Octax and EnvPool across parallelization levels. The solid purple line is the number of steps per second (higher is better), and the dashed green line is the total execution time in seconds (lower is better).
• Figure 5: Rendering of the Target Shooter game, generated by an LLM, showing the player (left, circular object) and target (right, cross-shaped object).