PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Quentin Delfosse, Raban Emunds, Paul Seitz, Sebastian Wette, Jannis Blüml, Daniel Kirn, Dominik Mandok, Kristian Kersting

arXiv:2610.23650v1cs.LG

TL;DR

ALE’s CPU-bound emulation and limited semantic access constrain scalable tests of whether agents learn robust causal structure or brittle visual shortcuts. JAXtari reimplements Atari environments in JAX with object-centric state and source-level modifications. It preserves ALE-style learning behavior while delivering higher throughput and making counterfactual evaluation practical at scale.

  • Problem

    ALE lacks native object-centric observations and semantic modification, while CPU-based emulation limits GPU-centered training and tests of robust versus shortcut-based behavior.

    Source: §1 · pp. 1–2

  • Method

    JAXtari reimplements Atari game logic in JAX, exposing structured state and source-level modifications within an accelerator-native framework.

    Source: §1 · p. 2, §2.1 · p. 3

  • Results

    JAXtari-15 reaches at least 1M steps-per-second, versus maxima of 360k for ALE-VectorEnv and 280k for EnvPool, while preserving ALE-style learning behavior.

    Source: §3.1 · p. 6, §6 · p. 10

  • Takeaways & Limitations

    JAXtari makes counterfactual evaluation practical at scale through targeted visual, dynamical, reward-level, and object-level shifts.

    Source: §6 · p. 10

    The current release covers only a subset of the 57 commonly used ALE games and should not yet be treated as a complete replacement or a drop-in substitute for exact emulator-based evaluation.

    Source: §5 · pp. 9–10

Abstract

from arXiv · show

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 evaluating whether agents learn robust causal models or rely on brittle visual heuristics. 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 JAXtari, an open-source, high-performance reimplementation of Atari 2600 environments. By leveraging massive GPU parallelization, JAXtari allows PPO agents to finish training on 100M environment steps on a single A100-GPU in under an hour. Beyond speed, it 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.

1 Introduction

ALE remains a central reasoning benchmark, but CPU-bound emulation, pixel-only observations, and opaque ROMs constrain scalable evaluation and semantic intervention. JAXtari reimplements Atari games in JAX to combine accelerator-native execution with modifiable game logic and object-centric state.

1 Introduction

  • ALE agents can rely on brittle visual shortcuts and fail under simple task or visual changes, motivating controlled counterfactual testing.The cited analyses report shortcut learning and failures under task simplification or visual perturbation.Source: §1 · p. 1
  • Closed-source ROMs obscure game logic and state, while CPU–GPU synchronization penalizes parallel training, population evaluation, and broad hyperparameter sweeps.These constraints make semantic modification and native interpretation difficult in ALE.Source: §1 · p. 2
  • JAXtari reimplements ALE environments in JAX, preserving Atari’s human-designed nature while exposing game logic, object-centric state, and the rendering pipeline.The framework is intended to support high-throughput, counterfactual evaluation.Source: §1 · p. 2
  • A source-level plugin system makes semantic modifications accessible, turning robustness and zero-shot shortcut-learning experiments into few-line changes rather than ROM reverse-engineering.JAXtari also exposes structured object states natively, avoiding post-hoc feature extraction.Source: §1 · p. 2
  • JAXtari targets behavioral and learning-task parity rather than bit-level emulation, and uses standard CleanRL and Gymnasium-compatible baselines to compare learning dynamics and performance rankings.Its pure-JAX, JIT-compiled, vectorizable environments are designed for high-throughput execution.Source: §1 · p. 2

2 JAXtari: The JAX-native Arcade Framework

JAXtari uses functional, JAX-native game reimplementations to support accelerator execution, structured observations, and direct semantic interventions. Its modification and benchmarking interfaces enable zero-shot evaluation under controlled environment shifts.

2 JAXtari: The JAX-native Arcade Framework

  • JAXtari replaces stateful ROM emulation with bottom-up JAX reimplementations of game physics, collisions, object updates, scoring, and termination in a compiled computation graph.The framework is designed around a functional, stateless paradigm compatible with JIT compilation.Source: §2.1 · p. 3

2.1 Architecture

  • The functional step function carries immutable state, observation, reward, and termination data through JAX transformations, and its PRNG key propagates with the game state.Structured state permits interventions between steps; object-centric and rendered-pixel observations are both supported, while jax.vmap enables thousands of parallel environments.Source: §2.1 · pp. 3–4

2.2 The Modification System

  • JAXtari’s internal modifications alter rules or assets during environment construction, while post-step modifications directly intervene on state after each step.The systems can be combined into reusable collections and support changes such as object relocation, variable hiding, or forced termination.Source: §2.2 · p. 4

2.3 Using JAXtari for Benchmarking

  • The benchmark protocol trains agents on unmodified games, then measures zero-shot transfer under controlled visual, object, dynamics, or reward changes that preserve core task structure.Performance differences across changes serve as a proxy for whether agents rely on robust task-relevant structure or configuration-specific heuristics.Source: §2.3 · p. 4

3 Experimental Evaluation

JAXtari is evaluated for learning-task parity with ALE, GPU throughput, and controlled robustness testing. The experiments report faster training, comparable pixel-based learning behavior, and performance drops under many environment modifications.

Experimental Setup

  • The evaluation uses CleanRL PPO conventions, three-seed results, and 15 games spanning six dynamic-control categories, with turn-based and static games omitted.Each selected environment is validated against its corresponding ALE version and supports pre-implemented modifications.Source: Experimental Setup · p. 5

3.1 Experimental Results

  • Pixel-based PPO learning curves track ALE comparably across 15 games, while object-centric curves differ more; both suites retain similar task challenges.The parity criterion is equivalent learning dynamics and exploration challenges, rather than cycle-level emulation.Source: §3.1 · pp. 5–6, Figure 3 · p. 6
  • 100M-step PPO training takes under one hour with pixels and under seventeen minutes with object-centric observations on one A100 GPU.The experiments used 8192 object-centric or 2048 RGB parallel environments.Source: §3.1 · pp. 6–7
  • 5× average pixel-run throughput versus ALE-VectorEnv and EnvPool; JAXtari scales to 1M parallel Pong environments before running out of memory.Across the suite, JAXtari reaches at least 1M steps per second, compared with maxima of 360k for ALE-VectorEnv and 280k for EnvPool.Source: §3.1 · p. 6, Figure 4 · p. 7
  • Zero-shot evaluation finds that agents lose performance on most modified games, including some simplifications and gameplay-preserving sprite swaps.Signed performance change is negative across most game-modification pairs, while absolute change remains well above zero.Source: §3.1 · p. 7, Figure 5 · p. 7
  • Exposed game parameters and visual assets support domain randomization, while native structured states provide a substrate for neuro-symbolic methods.The paper also proposes using structured states to connect multimodal foundation models and code-as-policy systems with environmental perturbation tests.Source: §3.1 · p. 8

4 Related Work

Prior Atari tools, accelerator-native environments, and robustness benchmarks address capabilities that JAXtari combines for Atari-style games. JAXtari reconstructs game logic in JAX to unite GPU execution, semantic modification, pixel compatibility, and native object-centric states.

Related Work

  • Existing Atari extensions typically improve one axis at a time, whereas JAXtari combines accelerator-native execution, semantic editing, pixel compatibility, and native object-centric observations.Examples include HackAtari for semantic interventions, CuLE and EnvPool for acceleration, and AtariARI and OCAtari for structured state.Source: §4 · p. 8, Figure 6 · p. 8
  • Accelerator-native frameworks target domains such as continuous control, classic control, structured decisions, board games, and multi-agent RL rather than ALE-style Atari gameplay.Examples include Brax, Gymnax, Jumanji, Pgx, and JAXMARL.Source: §4 · p. 8
  • JAXtari moves controlled semantic interventions from opaque RAM manipulation to source-level JAX edits that are compositional and scalable across experimental sweeps.Prior work found that agents can fail even when task variants become easier for humans.Source: §4 · p. 9
  • Native object states support relational and neuro-symbolic approaches, including differentiable logic rules, hybrid neural-symbolic policies, and editable knowledge bases.The passage also connects structured representations to multimodal foundation models and code-as-policy systems.Source: §4 · p. 9
  • JAXtari unifies Atari-like task performance evaluation with tests of the causal structures or brittle heuristics agents rely on when succeeding.Its modular API supports direct modifications to dynamics, rules, and visual properties without opaque RAM manipulation.Source: §4 · p. 9

5 Limitations and Future Work

JAXtari’s current game coverage and parity are bounded: it is not a complete ALE replacement or a bit-level ROM emulator. Its modification flexibility also requires care in designing comparable benchmarks, while future work targets broader game and modification coverage.

5 Limitations and Future Work

  • JAXtari currently covers a subset of ALE’s 57 games and omits the original ROM mode and difficulty switches, so it is not yet a complete ALE replacement.The project plans to expand game coverage.Source: §5 · p. 9
  • JAXtari cannot guarantee bit-level equivalence with Atari ROMs; timing, collision handling, object behavior, and rare edge cases may diverge, partly because of different random-number generation.The authors position it as a high-fidelity reimplementation for scalable intervention, not an exact emulator substitute.Source: §5 · pp. 9–10
  • Some modifications can change task structure, objective, difficulty, or optimal policy, making comparisons ambiguous unless results distinguish the semantic category of each change.The authors plan a larger curated modification library and validated semantic variations.Source: §5 · p. 10

6 Conclusion

JAXtari reimplements Atari-style environments in JAX to support scalable, editable RL, preserving ALE-style learning behavior while increasing throughput and enabling counterfactual evaluation. The paper presents it as an open-source benchmark substrate and describes its broader research impact and supporting appendices.

6 Conclusion

  • JAXtari preserves ALE-style learning behavior while providing substantially higher throughput and practical counterfactual evaluation at scale.The framework exposes object-centric state alongside pixels and supports targeted visual, dynamical, reward-level, and object-level shifts.Source: §6 · p. 10
  • JAXtari’s lower compute requirements are intended to make large-scale RL experimentation and systematic robustness and shortcut-learning evaluation more accessible.The paper states that the release includes no pretrained agents, collects no user data, and involves no human subjects.Source: §6 · p. 10

Appendix Overview

  • The appendices cover artifact access, validation, metrics, training setup, extended results, and game and modification documentation.Appendix B includes a Frostbite verification checklist and documented hardware divergences.Source: Appendix Overview · p. 14

A Code and Data Availability

The artifact provides JAX-native Atari environments, observation interfaces, wrappers, modification tools, and experiment scripts. The accompanying material documents licensing, reproducibility, usage, and a multi-stage process for checking behavioral fidelity against ALE.

A.1 Artifact Overview

  • The released artifact includes JAX-native environments, pixel and object-centric observations, Gymnasium-compatible wrappers, modification plugins, and reproduction scripts.The supported games and modifications are listed in Appendix F.Source: §A.1 · p. 15

A.2 License and Intellectual Property

  • JAXtari is a bottom-up reimplementation that does not distribute or execute proprietary Atari ROMs, and its code, wrappers, configurations, and plugins use the MIT license.The authors disclaim affiliation with Atari and place responsibility for third-party ROM and asset compliance on users.Source: §A.2 · p. 15

A.3 Reproducibility Statement

  • Reported experiments use fixed seeding, average three seeds, and include exact configuration files and scripts to support reproduction.The scripts log environment, observation mode, modification, agent, seed, returns, episode length, runtime, software version, and hardware.Source: §A.3 · p. 15

A.4 Usage

  • The usage appendix covers installation and a minimal environment example, and notes that modifications can be supplied when creating an environment.The passage also states that the same interface supports batched execution with jax.vmap and compiled rollouts with jax.jit.Source: §A.4 · p. 16

B Human + Agentic Validation Pipeline and Implementation Verification

  • The validation protocol starts from documented game specifications, then audits source code for omitted rules and undocumented behavior.Specifications cover actions, mechanics, rewards, and termination; the audits assess recall and precision.Source: B Human + Agentic Validation Pipeline… · p. 17, §B.1 · p. 17

B.1 The Four-Stage Validation Protocol

  • Human experts verify mapped dynamics and investigate flagged behavior by comparing the original ALE ROM with JAXtari.This human review follows the autonomous specification and code-audit stages.Source: §B.1 · p. 17
  • CleanRL PPO agents are trained on JAXtari and ALE with identical hyperparameters, and their learning curves, asymptotic returns, and exploration rates are compared.This final stage checks that JAXtari presents an equivalent learning challenge without unintended shortcuts.Source: §B.1 · p. 17

B.2 Case Study: Validation Checklist for Frostbite

  • The Frostbite case study presents tables for checking mapped action-space and core-dynamics specifications and for checking the reward system.The action-space and dynamics checklist is described as verified against ALE.Source: §B.2 · p. 17

B.3 Documented Implementation Differences Between JAXtari and ALE

JAXtari’s accelerator-native design introduces implementation differences from ALE, including explicit functional randomness and typed state representations. These choices support deterministic parallel execution and avoid decoding object coordinates from opaque RAM.

B.3 Documented Implementation Differences Between JAXtari and ALE

  • JAXtari uses functional JAX PRNG splitting rather than ALE’s implicit hardware-related randomness, providing deterministic behavior across parallel threads.ALE randomness is described as arising from uninitialized memory and Stella hardware counter side-effects.Source: §B.3 · p. 17
  • JAXtari stores state in typed PyTrees of JAX arrays, whereas ALE object coordinates must be inferred by decoding opaque RAM bytes.The passage says this representation provides exact floating- or fixed-point entity bounds.Source: §B.3 · p. 18
  • To avoid GPU branching bottlenecks, JAXtari vectorizes continuous entity velocities across arrays instead of using discrete sub-pixel scanline counters.Source: §B.3 · p. 18

C Evaluation Metrics and Computational Benchmarking

The section defines robust summaries for final performance and learning dynamics, a metric for changes under environmental perturbations, and hardware-efficiency measures.

Reliable Performance Aggregation

  • IQM summarizes final performance over the middle 50% of outcomes, with 95% stratified bootstrap confidence intervals; learning curves report mean rewards and standard deviation across three seeds.The IQM uses the 25th–75th percentiles to reduce sensitivity to outliers and high variance.Source: C Evaluation Metrics and Computational… · p. 19

Performance Change (PC)

  • Performance Change (PC) measures the relative shift between a base-trained agent’s scores on modified and original environments, adjusted by their random baselines.PC = 0 indicates perfect performance invariance, while PC < 0 signifies a collapse in the perturbation.Source: C Evaluation Metrics and Computational… · p. 19

Throughput and Hardware Efficiency

  • Throughput is frames processed per second across parallel environments, while total GPU hours track computational cost.The experiments use a single A100 GPU, although the framework can scale to multiple GPUs.Source: C Evaluation Metrics and Computational… · p. 19

D Agent Architectures and Training Setup

The setup compares JAXtari PPO agents with established Atari environments using matched training configurations, pixel and object-centric inputs, and multi-seed evaluation.

D Agent Architectures and Training Setup

  • PPO training uses JAX-based implementations where possible and compares JAXtari with ALE, EnvPool, and ALE-VectorEnv using matched baseline architectures and hyperparameters.The implementations also support non-JAX frameworks through DLPack and provide a Gym wrapper.Source: D Agent Architectures and Training Setup · p. 20, §D.1 · p. 20

D.2 Preprocessing

  • Both agents share preprocessing parameters matching standard ALE preprocessing used in the CleanRL PPO baseline6: up to 30 random no-ops at reset, episodic-life termination, four-frame skipping and stacking, and max pooling over the last two frames.Source: §D.2 · pp. 20–21
  • Object-centric inputs use four-frame skipping and stacking, with features normalized to [0, 1] using feature-specific bounds.Source: §D.2 · p. 20

D.2.1 Native Downscaling

  • Native downscaling scales sprites at initialization and renders directly at target resolution instead of constructing full 210×160 images, reducing memory use and computation.The rendering step is separate from game logic, so this preprocessing can occur inside the renderer.Source: §D.2.1 · p. 21

D.3.1 Pixel-Based Architectures

  • Pixel-based PPO uses a DQN-like convolutional network, while object-centric PPO uses a fully connected network with 461 and 512 nodes in its first two layers.The first object-centric layer is sized to match the convolutional layer’s parameter count for a fair input comparison.Source: §D.3.1 · p. 21, §D.3.2 · p. 21

D.4.1 PPO

  • The PPO hyperparameter table specifies identical settings for pixel-based and object-centric training runs.Source: Table 3 · p. 22

D.5 Training Protocol

  • Agents are trained from scratch across three independent seeds and evaluated with greedy policies; experiments use at most one 40GB A100 GPU each.The reported main PPO learning-curve computation used 180 GPU hours.Source: §D.5 · p. 22, §D.6 · p. 22

E Extended Results

  • The extended-results section adds throughput scaling, wall-clock time to 100M steps, and agent returns to relate compute speed to learning outcomes.Source: E Extended Results · p. 23

E.1 Throughput and Scaling Analysis

JAXtari scales parallel Atari simulation on a single GPU and makes 100M-step PPO runs practical; longer budgets particularly benefit difficult exploration games.

Throughput and scaling analysis

  • The analysis combines scaling curves, time to 100M steps, and resulting returns to relate throughput to learning outcomes rather than treating speed alone as the outcome.The reported runs use a single GPU and JAX vmap and JIT compilation with rollouts kept on-device.Source: §E.1 · pp. 23–24

Throughput scaling

  • JAXtari’s throughput is not limited by CPU-thread scaling, unlike EnvPool and ALE-VectorEnv, whose throughput curves scale with CPU threads in the reported experiments.The comparison reports agents through out-of-memory errors or 65K environments; the EnvPool and ALE-VectorEnv curves are averaged over five games.Source: Figure 7 · p. 23

Wall-Clock Throughput to 100M Steps

  • 8 minutes on average for object-centric PPO and 32 minutes on average for pixel-based PPO to reach 100M steps on a single GPU.The runs use 8192 parallel environments for object-centric inputs and 2048 for pixel inputs; Figure 9 reports an average maximum of 12 minutes for object-centric runs, while all pixel runs finish under an hour.Source: Figure 8 · p. 24, §E.1 · p. 24, Figure 9 · p. 25

Does More Training Time Yield Better Performance?

  • Longer training budgets yield higher final scores and richer behavior in difficult games such as Seaquest and Montezuma’s Revenge, while many other games saturate earlier.The text attributes the difference to sparse, delayed, or hierarchical rewards and reports high scores after 100M steps in 43 minutes; Figure 10 reports over 150k Seaquest points before truncation at 108000 frames.Source: §E.1 · p. 26, Figure 10 · p. 26

Future investigation

  • JAXtari’s single-GPU throughput makes future evaluations at 200M, 500M, or even 1B steps feasible without large-scale distributed infrastructure.Source: §E.1 · p. 26

E.2 Can JAXtari variations be used to evaluate robustness?

JAXtari’s modification system supports zero-shot robustness evaluation across game variants, while open-source sprite replacements provide a visual alternative to licensed Atari sprites.

E.2 Can JAXtari variations be used to evaluate robustness?

  • The evaluation reports performance change per game and per modification, with Figure 11 presenting per-game changes and improvements or drops across ten runs per seed.The figure reports IQM with 95% confidence intervals.Source: §E.2 · p. 26, Figure 11 · p. 27

F.1 JAXtari with open-source sprite-replacements

  • The open-source sprite package replaces original sprites with single-color rectangles for users without Atari ROM licenses, while preserving game mechanics.Source: §F.1 · p. 28, Figure 12 · p. 28

F.2 Creating new Modifications

This section describes how JAXtari’s modular system supports direct edits to game semantics, with examples spanning gameplay mechanics, assets, and constants across its environment catalog. It also shows that modification support varies by game, from visual examples and multiple gameplay changes to games with no implemented modifications.

F.2 Creating new Modifications

  • JAXtari provides a modular modification system for editing game semantics directly at the source-code level.Modifications are organized around internal and post-step plugins.Source: §F.2 · p. 28
  • Internal modifications alter constants, collision rules, methods, speeds, spawning, or sprites before JIT compilation, while post-step modifications manipulate state after a transition.Post-step changes can relocate objects, mask state variables, change rewards, or alter episode termination.Source: §F.2 · p. 28
  • A defined modification can be injected during environment initialization through the mods argument to jaxtari.make.The passage gives the form env = jaxtari.make("game_name", mods=["custom_mod"]).Source: §F.2 · p. 28

Example: Modifying Kangaroo

  • In Kangaroo, a modification can test causal mechanics by disabling monkeys, for example by patching their controller or forcing their active status off.The example motivates checking whether an agent learned climbing rather than relying on punching monkeys on the first floor.Source: Example: Modifying Kangaroo · pp. 28–29, Figure 13 · p. 29
  • Kangaroo modifications can also replace coconuts with larger fireballs or swap game assets and constants to examine agent behavior under changed conditions.The example proposes this change after an agent learns to dodge coconuts; Figure 15 describes asset and physical-constant changes for out-of-distribution scenarios.Source: Example: Modifying Kangaroo · p. 29, Figure 15 · p. 30
  • Unlike HackAtari’s emulator-RAM interventions, JAXtari modifies high-level JAX state or overrides Python methods before JIT compilation.Keeping the environment functional and stateless preserves compatibility with jax.vmap and jax.jit, enabling thousands of modified instances on one GPU.Source: Example: Modifying Kangaroo · p. 29

F.3 List of Games and Modifications

  • JAXtari supports 45 environments, with the appendix providing game descriptions and examples of visual modifications for selected games.Figure 16 presents ALE and JAXtari renderings for all 45 supported games.Source: §F.3 · p. 31, Figure 16 · p. 31, Figure 17 · p. 32

F.37.2 Available Modifications for Skiing

  • Skiing asks players to pass gate pairs quickly while avoiding penalties and collisions, and Figure 32 shows a selection of visual modifications.Missing a gate adds a five-second penalty; collisions with gates or trees briefly knock the skier down.Source: §F.37.1 · p. 74, Figure 32 · p. 75

F.38.1 SpaceInvaders details

  • SpaceInvaders has the player shoot descending invaders before they reach Earth; the game ends when lives are lost or invaders reach the ground.Source: §F.38.1 · p. 76

F.39.2 Available Modifications for Spacewar

F.40.2 Available Modifications for Tennis

  • Tennis follows standard rules for an orange player against a computer-controlled blue opponent, and Figure 33 shows a selection of visual modifications.A set requires six games with a two-game margin, and at 6-6 play continues until one player leads by two games.Source: §F.40.1 · p. 78, Figure 33 · p. 79

F.42.2 Available Modifications for Timepilot

F.45.2 Available Modifications for Venture

  • Venture involves clearing dungeon chambers of monsters and collecting treasure while avoiding hallway lurkers, and Figure 34 shows a selection of visual modifications.Monsters become increasingly aggressive if the player lingers in a room.Source: §F.45.1 · p. 84, Figure 34 · p. 85

F.46.2 Available Modifications for VideoCube

  • No modifications are implemented for VideoCube or WordZapper; these games involve swapping adjacent cube tiles to solve colors and shooting letters in a prescribed order.WordZapper’s pace increases over time, making speed and accuracy important.Source: §F.46.2 · p. 86, §F.46.1 · p. 86, §F.48.2 · p. 88

F.47.2 Available Modifications for VideoPinball

  • VideoPinball lists four modifications: constant_ball_dynamics, low_tilt_effect, never_activate_tilt_mode, and no_scoring_cooldown.They respectively constrain trajectory changes, reduce tilt effects, prevent tilt mode, and remove the post-target cooldown.Source: §F.47.2 · p. 87
Loading 2610.23650v1…