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

JAXtari: High-Throughput and Easy-to-Modify Arcade Learning Environment Quentin Delfosse1,2,*,† , Raban Emunds1,3,* , Paul Seitz1,* , Sebastian Wette1,* , Jannis Blüml1,3,* , Daniel Kirn1 , Dominik Mandok1 , and Kristian Kersting1,3,4,5

arXiv:2610.23650v1 [cs.LG] 16 Oct 2026

1

Department of Computer Science, Technical University Darmstadt, Germany 2 Google Intrinsic AI Research 3 Hessian Center for Artificial Intelligence (hessian.AI) 4 Centre for Cognitive Science, Darmstadt 5 German Research Center for Artificial Intelligence (DFKI)

Abstract 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 objectcentric 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

The Arcade Learning Environment (ALE) (Bellemare et al., 2013) has shaped deep reinforcement learning (RL) for over a decade (Mnih et al., 2015; Machado et al., 2018; Badia et al., 2020). As a suite of games originally designed by professional developers for human players, Atari possesses a unique advantage: it is inherently free of the "experimenter bias" that often plagues custombuilt RL environments. Consequently, ALE remains the primary RL benchmark (cf. Figure 1), particularly for studying tasks that require abstract, relational concepts and long-horizon planning, standing in contrast to continuous control benchmarks (e.g. MuJoCo (Todorov et al., 2012), DM Control (Tunyasuvunakool et al., 2020), D4RL (Fu et al., 2020)), that require low-level motor actuation control. However, despite superhuman aggregate scores achieved on the diverse set of reasoning tasks, ALE is far from being solved. Recent analyses reveal that high performance often stems from shortcut learning rather than robust human-like cognition; agents overfit to highly specific visual cues and frequently fail completely when a task is trivially simplified or visually perturbed (Farebrother et al., 2018; Taiga et al., 2023; Delfosse et al., 2025). Diagnosing and mitigating this brittleness requires * These authors contributed equally to this work.

†Work done while at TU Darmstadt; now at Google Intrinsic AI Research

Preprint.

ALE

100M

JAXtari

68

↑

OPEN SOURCE & TRANSPARENT

EnvPool

COMPILED & CLOSED SOURCE

ACTIONS + STATE NEXT STATE REWARDS

ACTIONS ↑ ↑

1K

OBSERVATIONS, REWARDS

1

1M

RENDERING ball:{x,y, } player:{x,y, } enemy:{x,y, } score1:{x,y, } score2:{x,y, }

MODIFYING

OC-STATE

Figure 1: JAXtari modernizes ALE as a benchmark for scalable and robust evaluation. ALE is the most-used RL benchmark (left), but its setup is constrained by CPU-bound emulation (center), pixel-only observations, and ROM-level barriers to semantic modification (right). JAXtari preserves ALE’s visual fidelity and dynamics, while reimplementing game logic in JAX, enabling acceleratornative throughput, direct source-level modifications, and native access to object-centric and symbolic state. controlled, counterfactual testing—such as systematically altering object locations, simplifying mechanics, or shifting domains, as introduced by HackAtari (Delfosse et al., 2024a, 2025). Unfortunately ALE prevents it, as it emulates closed-source binary ROMs that are semantically opaque. Both the games logic and their states are not natively interpretable, preventing straightforward implementations of neurosymbolic RL methods and of controlled semantic modifications. Furthermore, the CPU-bound nature of the emulation is completely out of step with modern accelerator-centric training pipelines. The continuous host–device synchronization required to transfer data between the CPU and GPU heavily penalizes wall-clock time, acting as a bottleneck for highly parallel algorithms, population-based training, and broad hyperparameter sweeps. To resolve this dichotomy between the necessity of Atari as a reasoning benchmark and the limitations of its legacy architecture, we introduce and open source JAXtari1 , a framework of faithful, fully JAX-based reimplementations of the ALE environments. By reconstructing the game logic from the ground up rather than wrapping legacy ROMs, we preserve the human-designed, bias-free nature of Atari while making the environments accelerator-native, semantically editable, and inherently object-centric. JAXtari can expose the underlying game logic, rich object-centric (OC) states, and rendering pipeline directly to the user, bridging the gap between high-throughput scalability and the deep, counterfactual evaluation required to diagnose and mitigate goal misgeneralization. Our main contributions are: (i) Flexible Modification API: We introduce a modular plugin system for editing game semantics directly at the source-code level, turning robustness evaluations and zero-shot shortcut-learning experiments into simple, few-line changes rather than complex ROM reverse-engineering projects. (ii) Object-Centric by Design: Unlike the constraints of ALE, JAXtari exposes structured object states natively from every environment. This removes the need for post-hoc feature extractors and enables clean, rigorous studies of relational, compositional, and logic-augmented RL. (iii) ALE Fidelity and Learning Task Parity: Rather than constraining simulation to bit-level 1:1 emulation, JAXtari targets behavioral and learning-task parity, preserving the human-designed mechanics and challenge of ALE. We validate this consistency by comparing learning dynamics and performance rankings across platforms using standard CleanRL and Gymnasium-compatible baselines. (iv) Accelerator-Native Throughput: We provide pure-JAX, JIT-compiled, vmap-able environments for high-throughput execution. By eliminating the CPU-GPU communication bottleneck, JAXtari allows population-scale evaluation. The remainder of the paper introduces the design of JAXtari, describes its object-centric and modification interfaces, evaluates fidelity against the ALE, and benchmarks accelerator-native throughput across games and training settings. 1 Code provided at https://github.com/k4ntz/JAXtari

2

A. Initialization & Registration

B. The Core Step Loop

C. The Wrapper Stack

(Factory Pipeline)

(Functional & Stateless, JIT Compiled)

(The User Interface) LogWrapper

User Call:

JIT Compiled Block

Game Registry

Observation Wrappers

RNG Split: state.key

AtariWrapper

Statet

Frame Skip (x4)

_player_step(state,action)

Action

Modding Injection

_ball_step(state)

Frame Stack (k)

AutoReset

Raw JAX Env

Statet+1 Pixel Obs Wrapper

Factory Method

Obs Score & Reset Check

Internal Mods (inject logic) Post-Step Mods (modify state)

_enemy_step(state)

Reconstruct

{"pong": "jax_pong"}

Object Centric Wrapper

ball:{x,y, } player:{x,y, } enemy:{x,y, } score1:{x,y, } score2:{x,y, }

Agent Interface (Obs, State, Reward, Done, Info)

jaxtari.make('pong', mods=...)

Reward

Instance JaxPong

Figure 2: JAXtari Architecture. (A) Environment creation adheres to the gym API, but additionally enables injecting modifications. (B) The functional and fully JIT-compiled transition function produces the next state, observation and reward given the previous state and action. (C) The comprehensive wrapper suite implements best practices and simplifies training.

2

JAXtari: The JAX-native Arcade Framework

This section introduces the JAXtari framework. We first provide the technical details of the JAXtari implementation, then describe the unique features it offers compared to the ALE. 2.1

Architecture

Implementing JAXtari requires a fundamental shift from the stateful design of the ALE to a purely functional, stateless paradigm compatible with JAX’s Just-In-Time (JIT) compilation. Unlike the ALE, that relies on the Atari binary ROMs (e.g. via hardware emulation or automated ROM-tosource transpilation), JAXtari offers JAX-native, bottom-up reimplementations of Atari environments. Concretely, game physics, collision handling, object updates, scoring rules, and termination conditions are manually reimplemented using high-level JAX array operations. The environment transitions are placed inside a single XLA (Accelerated Linear Algebra) compiled computation graph, removing the need for a CPU-bound emulator and enabling direct execution on accelerators. We provide an overview of the JAXtari architecture in Figure 2. Let us now detail the different JAXtari components. States. In JAXtari, the transition function follows the signature step(state, action) → (observation, next_state, reward, done, info), ensuring full compatibility with the gymnax API (Lange, 2026). The returned values are either primitive types (reward, done) or defined as immutable PyTreeNode objects (e.g. PongState) that contain all necessary variables (e.g. player positions, ball velocity, scores) and enable direct use within JAX transformations. This allows future algorithms to perform interventions outside the game engine, e.g. a (causal) algorithm may intervene on the symbolic state values between two subsequent calls of the step function, changing the game’s behavior drastically from the outside. Stochasticity is handled via jax.random.PRNGKey, which is defined as part of the game state and thus propagated through all steps. Observations. JAXtari natively supports two distinct observation modes, managed using wrappers: Object-Centric: The core environment state naturally exposes structured data (e.g. positions, velocities). The ObjectCentricWrapper flattens this information into a feature vector, enabling research into relational RL and agents that reason about object interactions directly, without the need for a separate vision backbone. It also only expose the observed objects (depicted on the screen), but hides internal variables (e.g. respawning counters). Pixel renderings: For compatibility with standard deep RL algorithms (e.g. DQN, PPO), we provide a JAXGameRenderer for each environment. This differentiable renderer constructs the frames from the state variables using JAX operations, allowing fast convolutional agents training on GPUs. Notably, both observation types can be retrieved simultaneously for neuro-symbolic approaches. The 3

object-centric observation is JAXtari’s native and thus most efficient instantiation, since it can directly be extracted from the state that is passed between transitions. The pixel-based observations require additional rendering functions that impact performance. Despite this, JAXtari is still significantly more efficient than ALE due to the missing necessity of emulating the hardware architecture, its ability to render game states only when necessary (typically every fourth frame is rendered) and the ability to natively downscale images (cf. Section D.2.1). The implications of the observation mode choice is analyzed in the experimental section.

Vectorized Accelerator Execution. One of the primary advantages of JAXtari is its ability to leverage massive parallelization. Using jax.vmap, we vectorize the step function to simultaneously run thousands of environment instances on a single GPU. This also avoids the data transfer overhead of CPU-based environment and GPU-hosted agent present in ALE workflows.

2.2

The Modification System

A key goal of JAXtari is its modular modification system, designed to test the agent alignment and generalization capabilities. Unlike HackAtari (Delfosse et al., 2024a), which requires complex RAM manipulation, JAXtari’s codebase supports direct, programmatic modifications via a plugin system, making modifications authoring accessible. Furthermore, HackAtari can only modify RAM accessible parts of the state, mostly canceling environments’ dynamics. For example, in Kangaroo, HackAtari e.g. cancels the spawning of the monkeys. While JAXtari also implements such simple modifications, it additionally allows switching the ladders positions (fixed in ALE and not stored in the game’s RAM), changing the sprites, altering the reward system (e.g. punishing the agent for violent actions such as punching), or modifying the environment dynamics (e.g. turning bonuses into penalties). To support this flexibility, JAXtari includes two modification systems that intervene at different environment execution points. Internal modifications are applied before the game environment is built, allowing researchers to efficiently alter core rules and assets—like gravity, enemy speeds, or reward structures—by deeply integrating them into the environment at construction time. Conversely, post-step modifications adjust the environment’s state immediately after each game step. This approach is ideal for making direct interventions that don’t require rewriting the underlying game physics, such as instantly relocating an object, hiding specific variables, or forcing a game to end. Modifications can be freely combined and grouped into reusable collections.

2.3

Using JAXtari for Benchmarking

Having introduced JAXtari as a flexible framework, we focus on its utility as a benchmark substrate to evaluate agent robustness under semantic distribution shifts and to diagnose reliance on spurious shortcuts. Following the intervention-based protocol established by Delfosse et al. (2025), agents are trained exclusively on unmodified JAXtari environments. We subsequently evaluate their zero-shot transfer performance on controlled semantic variations that preserve the core task structure while systematically altering visual representations, object properties, game dynamics, or reward functions. Discrepancies in performance under these distribution shifts serve as a proxy for the agents’ generalization: identifying whether they have internalized robust, task-relevant causal structures or have instead converged on brittle heuristics tied to specific environment configurations. For each game and variation, reports should contain the following metrics (defined formally in Section C). 1. Base Performance (Pbase ). The environmental return achieved by the agent trained on the unmodified environment. 2. Modification Performance (Pmod ). The return achieved by the base-trained agent when evaluated on the modified environment zero-shot. 3. Performance Change (PC). The shift in performance across all evaluated games, quantifying the sensitivity to game modification. 4

3

Experimental Evaluation

We evaluate JAXtari along three axes: fidelity to ALE, accelerator-native scalability, and usefulness for controlled generalization studies. (RQ1) Does JAXtari provide environments that are functionally equivalent to the original ALE? (RQ2) Can JAXtari significantly improve training speeds? (RQ3) Does JAXtari’s modification system allow to probe agent alignment, robustness and shortcut learning behaviors? Experimental Setup As standard deep reinforcement learning baseline, we employ PPO (Schulman et al., 2017), a widely used policy-gradient method. For comparability with the ALE baselines, we employ the training and network configuration from the CleanRL PPO baseline (Huang et al., 2022). Unless stated otherwise, we follow standard Atari training conventions, including frame skipping, frame stacking, and preprocessing to 84 × 84 grayscale observations. We evaluate both raw environmental performance, reported as episodic returns, and performance change, which we analyze using the Interquartile Mean (IQM) with bootstrap confidence intervals. Throughput is measured in total steps per second (SPS). JAXtari experiments were run on a single NVIDIA A100 GPU with 40GB of memory, ALE ones on two AMD EPYC 7742 CPUs (128 cores, 256 threads). Unless specified, results are aggregated over 3 random seeds, using 10M steps per environment. See Appendix D for more details regarding training and environment setup. The JAXtari-15 Benchmark. While JAXtari covers many ALE games (45+, cf. Appendix F), we evaluate on a representative subset of 15 environments selected to span six core real-time control categories: Shooters & Space (Asteroids, Beam Rider, Gravitar, Phoenix, Seaquest), Platformers & Action (Freeway, Frostbite, Kangaroo), Maze & Pathfinding (Ms. Pac-Man), Paddle & Ball (Breakout, Pong), Sports & Racing (Enduro, Skiing, Tennis), and Sparse Reward exploration (Montezuma’s Revenge). This selection isolates distinct RL challenges—from reactive motor control (Pong, Breakout) and delayed credit assignment (Skiing) to sparse multi-room exploration (Montezuma’s Revenge)—while intentionally omitting turn-based or static categories (board and gambling games) to focus on dynamic continuous-time interactions. Each environment is validated against the corresponding ALE version to ensure behavioral consistency. Next to task-fidelity to the original environments and scalability, we position JAXtari as a benchmark for agent robustness. Thus, each game supports a set of pre-implemented modifications, which range from simple visual changes, such as altered colors or background elements, to more substantial gameplay variations. This allows evaluating not only base-task performance, but also the robustness of trained agents under controlled distribution shifts. Details of games and modifications are provided in Section F.3. 3.1

Experimental Results

JAXtari games achieve learning task parity with ALE (RQ1). A central design principle of JAXtari is prioritizing accelerator-native parallel throughput and transparent, source-level modification over bit-level, 1:1 ROM emulation. Exact cycle-level emulation of legacy Atari 2600 hardware (such as TIA register timings and 6502 assembly constraints) imposes rigid execution bottlenecks on GPUs and precludes clean Pythonic interventions on game semantics. Furthermore, deep RL policies are notoriously brittle; minor low-level implementation details or differing pseudo-random number generation prevent direct zero-shot policy transfer between independent simulators. Consequently, JAXtari is designed not as a drop-in emulator replacement, but to achieve behavioral and learning task parity with the ALE. Under this criterion, environments must present equivalent learning dynamics and exploration challenges: tasks solved rapidly in ALE (e.g. Pong, Freeway) are mastered similarly in JAXtari, while sparse-reward exploration bottlenecks (e.g. Montezuma’s Revenge, Skiing) remain equally challenging under standard algorithms. Moreover, agents discover and exploit identical behavioral shortcuts across both suites (e.g. in Kangaroo and Bankheist). As depicted in Figure 3, pixel-based PPO runs (orange vs. green curves) confirm this alignment: learning curves follow comparable trajectories with asymptotic returns in the same order of magnitude. In addition to pixel observations, JAXtari natively 5

Asteroids*

Beamrider*

Breakout

Enduro*

Freeway

1.5K 3K

4K

30

400 1K

20

2K 200

2K

500

10

0

0

1K 0

0

0

Frostbite

Gravitar*

Kangaroo

Montezumarevenge

Episodic return

30K

3K

1.5K

2K

Mspacman

300 200

20K

2K

1K 1K

100

10K

500

1K 0

0

0

0

Phoenix*

Seaquest

20

20K

5M

10 M

10 0

-20K

-10

500

-20 0

Tennis 20

-10K

1K

-10

5K

0

1.5K

0

10K

Skiing

2K

10

15K

0

0

Pong

0

5M

10 M

0

-20

-30K 0

5M

10 M

0

5M

10 M

0

5M

10 M

Environment steps

JAXtari-Pixel

ALE

JAXtari-OC

OCAtari

Not available in OCAtari

Figure 3: JAXtari enables similar learning behavior to ALE. Learning curves of a PPO agent, learning in 15 JAXtari and ALE games. Pixel-based JAXtari (orange) behaves similar to ALE (green), object-centric observations (blue) differ more. Results are aggregated over 3 seeds. provides object-centric (OC) states (Figure 3, blue curves), allowing agents to train directly over structured entities without post-hoc extractors. To rigorously guarantee behavioral consistency across both suites despite independent JAX reimplementation, we employ a systematic Human + Agentic validation pipeline. First, an LLM agent extracts a ground-truth operational specification (active discrete action sets, entity dynamics, collision parameters, and reward formulas) by parsing official ALE documentation and original Atari 2600 manuals. Second, an auditing agent cross-references the Python/JAX source code against this specification to identify implementation gaps (recall) and flag any extraneous logic (precision). Third, human experts verify mapped code blocks and perform side-by-side gameplay validation against the original ALE ROMs to confirm intricate hardware edge cases. Finally, policy training curves are evaluated to detect unintended shortcut exploitation. A comprehensive walkthrough of this validation pipeline, including a case study on Frostbite, is detailed in Appendix B. JAXtari achieves significantly higher throughput, reducing the needed training time (RQ2). The scalability properties of the JAXtari-15 suite are validated against standard Gym-ALE, EnvPool (Weng et al., 2022) and ALE-VectorEnv, a recent EnvPool variation within the ALE. JAXtari-15 environments reach at least a throughput of 1M steps-per-second compared to a maximum of 360k for ALE-VectorEnv and a maximum of 280k for EnvPool (see Figure 4 b). The JAXtari environments significantly outscale CPU-based ALE with a single GPU. Some lower-complexity games, e.g. , pong or freeway, reach numbers exceeding 100M steps per second when using OC observations (cf. Figure 4a). The observed variations in throughput across the JAXtari suite are attributable to heterogeneous visual and gameplay complexities, alongside implementation-specific nuances arising from our ground-up reimplementation of each game. In general, OC observations exhibit significantly higher scalability and throughput, which we attribute to the ability to execute the environment logic exclusively without rendering, significantly reducing memory usage and computational complexity. This is an advantage over ALE, whose emulator does not allow to skip rendering. Additionally, the smaller per-step memory footprint fits roughly 4× as many parallel environments on GPU as in the pixel-based setting. We also validate the utility of scaled-up JAXtari runs by performing PPO experiments with a large number of parallel environments (8192 OC, 2048 RGB, cf. Section E.1). When translated to wallclock time, PPO agents complete training on 100M environment steps in under one hour when 6

10 M 1M

Throughput (SPS)

100 k 10 k

JAXtari-OC JAXtari-Pixel Gym-ALE ALE-VectorEnv EnvPool

10 M 1M 100 k 10 k

2⁰

2⁵ 2⁶ 2⁷ 2⁸ 2 2¹⁹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ 2¹⁶ 2¹⁷ 2¹⁸ 2²⁹ ⁰

2³

2⁰

2⁵ 2⁶ 2⁷ 2⁸ 2 2¹⁹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

1k

1k

2³

Throughput (SPS)

JAXtari-OC JAXtari-Pixel Gym-ALE ALE-VectorEnv EnvPool

100 M

Number of Environments

Number of Environments

(a) Scaling Pong to 1M Environments.

(b) Average throughput comparison.

Figure 4: JAXtari increases the throughput by enabling more parallel environments. (a) The JAX-native JAXtari implementation scales up to 1M parallel environments in Pong before running out of memory. (b) On average, across all games, JAXtari increases throughput by 5× for pixel-based runs compared to ALE-VectorEnv and EnvPool. An in-depth analysis is available in Appendix E. 1.0 PC (IQM)

PC (IQM)

0.5 0.0 -0.5

0.8 0.6 0.4

JAXtari-OC

JAXtari-Pixel

JAXtari-OC

(a) Signed performance change.

JAXtari-Pixel

(b) Absolute performance change.

Figure 5: JAXtari allows the detection of shortcuts by easily modifying the environment. Performance change (PC, defined in Appendix C) of PPO agents trained on the non-modified JAXtari-15 environments and evaluated zero-shot on their modified counterparts, IQM-aggregated across 3 seeds with 95% stratified bootstrap confidence intervals. (a) Signed PC: values below zero indicate that the agent performs worse on the modified environment than on the base version; values above zero indicate (rare) improvement. (b) Absolute PC: magnitude of deviation regardless of direction. These experiments extend the systematic performance drops of standard, model-free RL agents on task simplifications, initially exposed by Delfosse et al. (2025) to a more diverse set of modifications. Per game and mod results are depicted in Appendix E.2.

learning from pixels, and in under seventeen minutes using object-centric observations on a single A100 GPU (cf. Figures 8, 9). This stands in stark contrast to the multi-hour runtimes required by equivalent CleanRL+ALE configurations. This acceleration effectively collapses the experimentiteration loop, e.g. for uncovering novel shortcuts, reducing the experimental cycle from several days to a single working session. Consequently, large-scale sweeps over game variants, which would typically induce long runtimes, have become computationally tractable on modest hardware. JAXtari modification system allows shortcut detection (RQ3). The pre-implemented game modifications (cf. Appendix F) allow us to analyze an agents robustness to various modifications (e.g., visual changes, dynamics simplifications). For this, we follow the proposed benchmark detailed in Section 2.3 and perform zero-shot adaptation evaluation on each modification. The results, depicted in Figure 5, suggest that agents converged on the original games systematically lose performance under modifications, even on simplifications. The signed performance change (PC) is negative across the majority of the game-modifications and the absolute PC stays well above zero, indicating that the drops are consistent collapses. Game simplifications and sprite swaps that preserve gameplay are sufficient to break pixel-based agents on several games, mirroring the brittleness reported in Delfosse et al. (2025), confirming JAXtari’s utility in evaluating robustness and uncovering shortcuts. 7

Beyond exposing these vulnerabilities, JAXtari provides the essential substrate to actively mitigate shortcut learning. Because environment parameters and visual assets are exposed directly in Python and natively vectorizable under jax.vmap, researchers can perform extensive domain randomization across game dynamics and visual features during training with zero simulation overhead. Furthermore, JAXtari’s native object-centric state representations offer the structural scaffolding required to develop and evaluate neuro-symbolic methods—such as NUDGE (Delfosse et al., 2023), BlendRL (Shindo et al., 2025), and Kintsugi (Cao et al., 2026)—which explicitly ground policy learning in relational concepts and allow repairing brittle behavioral shortcuts. Finally, exposing structured semantic states facilitates interfacing multimodal foundation models and code-as-policy architectures (e.g. AtariGPT (Waytowich et al., 2024), RL-GPT (Liu et al., 2024)) to test whether external commonsense priors yield true zero-shot generalization across environmental perturbations.

4

Related Work

JAXtari builds on three lines of RL research: Atari benchmarks, accelerator-native simulators, and generalization benchmarks. Prior work has improved Atari along individual axes such as speed, stochasticity, editability, or state access, but these capabilities have typically remained separate.

ns

Parallelization

Op e

n

rce

bs

er va tio n

fre e

OC o

bservations

sparent states Tran

Tas k

o a ti ific

u so

o el Pix

Accelerator-native RL environments. Recent RL benchmarks increasingly move environment execution into JAX or GPU-native backends to reduce simulation bottlenecks and support largescale vectorized training. Examples include Brax (Freeman et al., 2021) for continuous control, Gymnax (Lange, 2026) for classic control and lightweight environments, Jumanji (Bonnet et al., 2024) for structured decision-making, Pgx (Koyamada et al., 2023) for board games, and JAXMARL (Rutherford et al., 2024) for multi-agent RL. Game-based suites such as Craftax (Matthews et al., 2024) and Octax (Radji et al., 2026) further show the value of scalable JAX-native simulation. These frameworks share JAXtari’s systems motivation, but target domains other than ALE-style Atari gameplay or simplified arcade settings. Atari remains a milestone benchmark because it combines visual input, long-horizon control, diverse game mechanics, and a long history of comparable baselines. Bringing this benchmark class into an accelerator-native, editable, and object-centric setting is therefore a necessary step for testing modern RL agents at Atari-like complexity.

m od

The Arcade Learning Environment and its extensions. The ALE (Bellemare et al., 2013) remains the canonical Atari benchmark, and many extensions address specific limitations of the original setup. Sticky actions and action noise reduce determinism (Machado et al., 2018; Koch et al., 2021), while MinAtar (Young & Tian, 2019) provides simplified Atari-like games for faster experimentation. Other extensions preserve the Atari setting but target different bottlenecks: HackAtari (Delfosse et al., 2024a) enables semantic interventions through RAM manipulation, CuLE (Dalton & Frosio, 2020) and EnvPool (Weng et al., 2022) accelerate environment execution, and AtariARI (Anand et al., 2019) and OCAtari (Delfosse et al., 2024b) expose structured state on top of ALE. These systems expand Atari’s utility, but typically improve one axis at a time. JAXtari instead reimplements Atari-style games directly in JAX, making the environments simultaneously accelerator-native, semantically editable, pixel-compatible, and object-centric by design. We compare JAXtari to ALE and the above-mentioned extensions in Figure 6.

s

as Bi Task variety

Figure 6: Feature coverage across Atari-based reinforcement learning frameworks. JAXtari (ours) is the only library to support the full spectrum of features, including GPU parallelization, diverse observation modes (Pixel and OC), and extensive task modification capabilities.

Generalization and Robustness in RL. Generalization remains a central challenge in RL: agents can achieve high returns in fixed environments while failing under small distribution shifts (Moos 8

et al., 2022). Procedural benchmarks such as CoinRun (Cobbe et al., 2019) and Procgen (Cobbe et al., 2020) expose agents to diverse generated levels, while XLand-MiniGrid (Nikulin et al., 2023) and Craftax (Matthews et al., 2024) provide scalable task diversity. Another line improves robustness through visual augmentation or representation learning, including RAD (Laskin et al., 2020), DrQ (Yarats et al., 2021), and the combination of pixel-based and object-centered representations (Blüml et al., 2026) or environment perturbations (Gu et al., 2025). These approaches broadly improve or evaluate robustness, but do not isolate which semantic assumptions an agent has learned. Targeted semantic interventions address this diagnostic question more directly. HackAtari (Delfosse et al., 2024a) introduced controlled modifications to ALE games through RAM manipulation, and subsequent work showed that deep RL agents can fail even when task variants become easier for humans (Delfosse et al., 2025). JAXtari builds on this intervention-based view, but moves it from opaque RAM manipulation in a CPU-bound emulator to source-level modifications of JAX game logic. This makes counterfactual evaluation easier to implement, compositional across modifications, and scalable to large experimental sweeps. Object-centric representations and neuro-symbolic RL. Learning structured, object-level representations is an established paradigm for compositional generalization and sample efficiency in RL (Diuk et al., 2008; Garnelo et al., 2016). While prior Atari tools such as AtariARI (Anand et al., 2019) and OCAtari (Delfosse et al., 2024b) extracted object states post-hoc from emulator RAM, they remained constrained by CPU-bound stepping and lacked native integration with GPU training loops. Natively exposing object states unlocks methods that incorporate relational inductive biases and neuro-symbolic reasoning. Frameworks such as NUDGE (Delfosse et al., 2023) induce differentiable first-order logic rules over extracted entities, BlendRL (Shindo et al., 2025) unifies neural and symbolic policy learning, and Kintsugi (Cao et al., 2026) repairs misaligned or brittle policies by directly editing executable knowledge bases. Furthermore, structured object representations bridge the semantic gap to foundation models, enabling multimodal large language models and code-as-policy systems (e.g. Atari-GPT (Waytowich et al., 2024), RL-GPT (Liu et al., 2024)) to interpret Atari scenes with human-like commonsense priors. JAXtari provides the first accelerator-native environment suite where such object-centric and neuro-symbolic architectures can be trained, verified, and stress-tested at modern GPU throughput. Positioning of JAXtari. While prior research has independently established the necessity of Atari as a reasoning benchmark, the scalability of accelerator-native simulation, and the diagnostic power of object-centric states and controlled interventions, these capabilities have largely remained fragmented across distinct frameworks. JAXtari resolves this fragmentation by fundamentally reconstructing Atari-style game logic directly in JAX, bypassing the severe bottlenecks of legacy CPU-bound emulation. Through this purely functional paradigm, agents can be trained from standard pixel renderings or evaluated using natively exposed object-centric features, effectively bridging the gap between deep and relational reinforcement learning. Furthermore, the framework leverages massive GPU parallelization for highly efficient scaling, and its modular API allows agents to be tested under direct, source-level semantic modifications—such as shifting dynamics, rules, or visual properties—without relying on complex and opaque RAM manipulation. Ultimately, this unification makes JAXtari a comprehensive benchmark substrate for studying not only how well agents perform, but precisely what causal structures or brittle heuristics they rely on when they succeed.

5

Limitations and Future Work

JAXtari is intended as a scalable and editable benchmark substrate for Atari-style RL, but it has limitations. First, reimplementing Atari games from scratch is substantially more labor-intensive than wrapping an existing emulator. As a result, the current release covers a subset of the full ALE suite rather than the typically used 57 games. We prioritize commonly used games and provide validation for the environments included in the benchmark, but the suite should not yet be interpreted as a complete replacement for ALE. Furthermore, JAXtari does not cover the ALE mode and difficulty switches of the original ROMs which alter core game parameters but instead relies on an independent modification system to introduce similar task-variance. Expanding coverage across the full Atari benchmark remains an important direction for future work. Second, JAXtari aims to reproduce ALE-style gameplay, but it is not an emulator and therefore cannot guarantee bit-level equivalence with the original Atari ROMs. Small divergences in timing, collision 9

handling, object behavior, or rare edge cases may remain, in part due to differences in random number generation. We mitigate this by using visual comparisons, per-game divergence measurements, and downstream learning-curve comparisons, but users should interpret JAXtari as a high-fidelity reimplementation designed for scalable experimentation and controlled intervention, rather than a drop-in substitute for exact emulator-based evaluation. Third, the flexibility of the modification system introduces a benchmark-design challenge. Not every modification defines a meaningful generalization test: some preserve the task structure, while others may change the objective, difficulty, or optimal policy in ways that make comparisons ambiguous. For this reason, benchmark results should distinguish between visual shifts, dynamics shifts, reward shifts, simplifications, and more difficult out-of-distribution variants. Each existing modification is categorized by its intended semantic category and the expected effect. Next to increasing the number of available environments, future releases will include a larger curated modification library. Future work will focus on three directions. First, we plan to expand the number of supported games and broaden coverage of the ALE benchmark. Second, we will extend the library of validated semantic variations and integrate them with benchmark protocols for robustness, shortcut learning, controlled generalization and continual RL. Third, JAXtari’s explicit JAX state and source-level transition functions open opportunities for model-based RL, planning, world modeling, and curriculum generation. Although Atari’s discrete actions and collision-heavy dynamics limit straightforward gradient-based policy optimization, the JAX-native implementation makes it easier to inspect, transform, batch, and potentially differentiate through selected components of the environment.

6

Conclusion

We introduced JAXtari, a JAX-native framework and benchmark substrate for scalable, editable Ataristyle reinforcement learning. By reimplementing game logic directly in JAX, JAXtari removes the CPU-bound emulator bottleneck of standard ALE workflows, enables accelerator-native vectorized execution, and exposes native object-centric state alongside pixel observations. Its source-level modification system further enables controlled semantic variations, allowing agents trained on standard Atari-style tasks to be evaluated under targeted visual, dynamical, reward-level, and objectlevel shifts. Our results show that JAXtari preserves ALE-style learning behavior while providing substantially higher throughput and making counterfactual evaluation practical at scale. Together, these properties make JAXtari a benchmark substrate for studying not only how well RL agents perform, but also whether their behavior reflects robust task understanding or brittle shortcut reliance. By reducing experiments that previously required many hours of emulator-bound execution to substantially shorter accelerator-native runs, JAXtari makes systematic robustness evaluation over many agents, seeds, and environment variants substantially more accessible. Broader Impact. JAXtari is an open-source benchmark substrate for academic reinforcement learning research. Its primary positive impact is democratizing large-scale rl experimentation: by reducing the compute required to run meaningful experiments from many CPU-hours to short GPU-native runs, it lowers the barrier for researchers. The modification system additionally makes systematic robustness and shortcut-learning evaluation accessible, which we expect to benefit the broader community studying agent generalization. We do not foresee significant negative societal impacts; JAXtari releases no pretrained agents, collects no user data, and involves no human subjects. Acknowledgements. This research work has been funded by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Germany’s Excellence Strategy (EXC-3057/1 ”Reasonable Artificial Intelligence“, Project No. 533677015). The authors also thank Lars Teubner and the students at TU Darmstadt who participated in the practical lab implementing JAXtari environments.

References Agarwal, R., Schwarzer, M., Castro, P. S., Courville, A. C., and Bellemare, M. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 2021. Anand, A., Racah, E., Ozair, S., Bengio, Y., Côté, M.-A., and Hjelm, R. D. Unsupervised state representation learning in atari. In Advances in Neural Information Processing Systems, 2019. 10

Badia, A. P., Piot, B., Kapturowski, S., Sprechmann, P., Vitvitskyi, A., Guo, Z. D., and Blundell, C. Agent57: Outperforming the atari human benchmark. In Proceedings of the 37th International Conference on Machine Learning, 2020. Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 2013. Beukman, M., Khetarpal, K., Zheng, Z., Dabney, W., Foerster, J., Dennis, M., and Lyle, C. Preventing Learning Stagnation in PPO by Scaling to 1 Million Parallel Environments. arXiv, 2026. Blüml, J., Derstroff, C., Gregori, B., Dillies, E., Delfosse, Q., and Kersting, K. Do object channels improve robustness in deep reinforcement learning? Trans. Mach. Learn. Res., 2026, 2026. URL https://openreview.net/forum?id=7BFbso4B3R. Bonnet, C., Luo, D., Byrne, D., Surana, S., Abramowitz, S., Duckworth, P., Coyette, V., Midgley, L. I., Tegegn, E., Kalloniatis, T., Mahjoub, O., Macfarlane, M., Smit, A. P., Grinsztajn, N., Boige, R., Waters, C. N., Mimouni, M. A., Sob, U. A. M., de Kock, R., Singh, S., Furelos-Blanco, D., Le, V., Pretorius, A., and Laterre, A. Jumanji: a diverse suite of scalable reinforcement learning environments in jax, 2024. Cao, T., Deng, Y., Shindo, H., Delfosse, Q., Wen, L., Wang, S., Blüml, J., Tauchmann, C., and Kersting, K. Kintsugi: Learning policies by repairing executable knowledge bases. arXiv preprint arXiv:2605.09487, 2026. Cobbe, K., Klimov, O., Hesse, C., Kim, T., and Schulman, J. Quantifying generalization in reinforcement learning. In Proceedings of the 36th International Conference on Machine Learning, 2019. Cobbe, K., Hesse, C., Hilton, J., and Schulman, J. Leveraging procedural generation to benchmark reinforcement learning. In International conference on machine learning, 2020. Dalton, S. and Frosio, I. Accelerating reinforcement learning through GPU atari emulation. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M., and Lin, H. (eds.), Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. Delfosse, Q., Shindo, H., Dhami, D. S., and Kersting, K. Interpretable and explainable logical policies via neurally guided symbolic abstraction. Advances in Neural Information Processing Systems (NeurIPS), 2023. Delfosse, Q., Blüml, J., Gregori, B., and Kersting, K. Hackatari: Atari learning environments for robust and continual reinforcement learning. arXiv preprint, 2024a. Delfosse, Q., Blüml, J., Gregori, B., Sztwiertnia, S., and Kersting, K. OCAtari: Object-centric Atari 2600 reinforcement learning environments. Reinforcement Learning Journal, 2024b. Delfosse, Q., Blüml, J., Tatai, F., Vincent, T., Gregori, B., Dillies, E., Peters, J., Rothkopf, C., and Kersting, K. Deep reinforcement learning agents are not even close to human intelligence. arXiv, 2025. Diuk, C., Cohen, A., and Littman, M. L. An object-oriented representation for efficient reinforcement learning. In Proceedings of the 25th international conference on Machine learning, 2008. Farebrother, J., Machado, M. C., and Bowling, M. Generalization and regularization in dqn. In NeurIPS Deep Reinforcement Learning Workshop, 2018. Freeman, C. D., Frey, E., Raichuk, A., Girgin, S., Mordatch, I., and Bachem, O. Brax–a differentiable physics engine for large scale rigid body simulation. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2021. Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint, 2020. Garnelo, M., Arulkumaran, K., and Shanahan, M. Towards deep symbolic reinforcement learning. arXiv preprint, 2016. 11

Gu, S., Shi, L., Wen, M., Jin, M., Mazumdar, E., Chi, Y., Wierman, A., and Spanos, C. Robust gymnasium: A unified modular benchmark for robust reinforcement learning. ICLR, 2025. Huang, S., Dossa, R. F. J., Ye, C., Braga, J., Chakraborty, D., Mehta, K., and Araújo, J. G. Cleanrl: High-quality single-file implementations of deep reinforcement learning algorithms. Journal of Machine Learning Research, 2022. Koch, J., Langosco, L., Pfau, J., Le, J., and Sharkey, L. Objective robustness in deep reinforcement learning. arXiv, 2021. Koyamada, S., Okano, S., Nishimori, S., Murata, Y., Habara, K., Kita, H., and Ishii, S. Pgx: Hardware-accelerated parallel game simulators for reinforcement learning. In Advances in Neural Information Processing Systems, volume 36, pp. 45716–45743, 2023. Lange, R. T. gymnax: A JAX-based reinforcement learning environment library, 2026. URL http://github.com/RobertTLange/gymnax. Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., and Srinivas, A. Reinforcement learning with augmented data. In Advances in Neural Information Processing Systems, pp. 19884–19895, 2020. Liu, S., Yuan, H., Hu, M., Li, Y., Chen, Y., Liu, S., Lu, Z., and Jia, J. RL-GPT: integrating reinforcement learning and code-as-policy. In Globersons, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J. M., and Zhang, C. (eds.), Advances in Neural Information Processing Systems 37: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10 - 15, 2024, 2024. URL http://papers.nips.cc/paper_files/ paper/2024/hash/31f119089f702e48ecfd138c1bc82c4a-Abstract-Conference.html. Machado, M. C., Bellemare, M. G., Talvitie, E., Veness, J., Hausknecht, M., and Bowling, M. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. Journal of Artificial Intelligence Research, 2018. Matthews, M., Beukman, M., Ellis, B., Samvelyan, M., Jackson, M. T., Coward, S., and Foerster, J. N. Craftax: A lightning-fast benchmark for open-ended reinforcement learning. In Forty-first International Conference on Machine Learning, 2024. Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J., Bellemare, M. G., Graves, A., Riedmiller, M., Fidjeland, A. K., Ostrovski, G., Petersen, S., Beattie, C., Sadik, A., Antonoglou, I., King, H., Kumaran, D., Wierstra, D., Legg, S., and Hassabis, D. Human-level control through deep reinforcement learning. Nature, 518(7540):529–533, 2015. Moos, J., Hansel, K., Abdulsamad, H., Stark, S., Clever, D., and Peters, J. Robust reinforcement learning: A review of foundations and recent advances. Machine Learning and Knowledge Extraction, 2022. Nikulin, A., Kurenkov, V., Zisman, I., Sinii, V., Agarkov, A., and Kolesnikov, S. XLand-minigrid: Scalable meta-reinforcement learning environments in JAX. In Intrinsically-Motivated and OpenEnded Learning Workshop, NeurIPS2023, 2023. Radji, W., Michel, T., and Piteau, H. Octax: Accelerated CHIP-8 arcade environments for reinforcement learning in JAX. In The Fourteenth International Conference on Learning Representations, 2026. Rutherford, A., Ellis, B., Gallici, M., Cook, J., Lupu, A., Ingvarsson, G., Willi, T., Hammond, R., Khan, A., de Witt, C. S., Souly, A., Bandyopadhyay, S., Samvelyan, M., Jiang, M., Lange, R. T., Whiteson, S., Lacerda, B., Hawes, N., Rocktäschel, T., Lu, C., and Foerster, J. N. Jaxmarl: Multi-agent rl environments and algorithms in jax. In The Thirty-eighth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024. Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint, 2017. Shindo, H., Delfosse, Q., Dhami, D. S., and Kersting, K. Blendrl: A framework for merging symbolic and neural policy learning. In International Conference on Learning Representations, 2025. 12

Taiga, A. A., Agarwal, R., Farebrother, J., Courville, A., and Bellemare, M. G. Investigating multi-task pretraining and generalization in reinforcement learning. In The eleventh international conference on learning representations, 2023. Todorov, E., Erez, T., and Tassa, Y. Mujoco: A physics engine for model-based control. In International conference on intelligent robots and systems, 2012. Tunyasuvunakool, S., Muldal, A., Doron, Y., Liu, S., Bohez, S., Merel, J., Erez, T., Lillicrap, T., Heess, N., and Tassa, Y. dm_control: Software and tasks for continuous control. Software Impacts, 2020. Waytowich, N. R., White, D., Sunbeam, M. D., and Goecks, V. G. Atari-GPT: Benchmarking multimodal large language models as low-level policies in Atari games. arXiv preprint arXiv:2408.15950, 2024. Weng, J., Lin, M., Huang, S., Liu, B., Makoviichuk, D., Makoviychuk, V., Liu, Z., Song, Y., Luo, T., Jiang, Y., Xu, Z., and YAN, S. Envpool: A highly parallel reinforcement learning environment execution engine. In Thirty-sixth Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2022. Yarats, D., Kostrikov, I., and Fergus, R. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In International Conference on Learning Representations, 2021. Young, K. and Tian, T. Minatar: An atari-inspired testbed for thorough and reproducible reinforcement learning experiments. arXiv, 2019.

13

Appendix Overview • Appendix A: Code and Data Describes access to JAXtari environments, modifications, data and experimental pipeline. • Appendix B: Human + Agentic Validation Pipeline and Parity Case Study Protocol for guaranteeing behavioral fidelity to ALE, featuring the Frostbite verification checklist and documented hardware divergences. • Appendix C: Metrics Definitions of Interquartile Mean (IQM), and performance change metrics. • Appendix D: Agent Architectures and Training Setup Details hardware setup, GPU usage, training time, and total inference budget required for evaluation. • Appendix E: Extended Results Provides extended results and tables for deep agents, performance comparisons, and modification evaluations. • Appendix F: JAXtari Game Descriptions and Modifications Full documentation of available game environments and all implemented task variances.

14

A

Code and Data Availability

A.1

Artifact Overview

We release JAXtari as an open-source software benchmark for scalable and editable Atari-style reinforcement learning. The artifact includes JAX-native environment implementations, pixelbased and object-centric observation interfaces, Gymnasium-compatible wrappers, a source-level modification API, predefined modification plugins, and scripts for reproducing the main experiments. Detailed evaluation settings are provided in Appendix D, and the full list of supported games and modifications is provided in Appendix F. The full source code, benchmark suite, and documentation are publicly available at: https://github.com/k4ntz/JAXtari A.2

License and Intellectual Property

JAXtari is a bottom-up reimplementation of Atari-style environments in JAX. It does not distribute or execute proprietary Atari ROMs. The released code, wrappers, benchmark configurations, and modification plugins are provided under the MIT license. JAXtari is not affiliated with, endorsed by, or sponsored by Atari or the original game rights holders. We handle visual assets analogously to ALE’s handling of ROMs: JAXtari does not distribute proprietary Atari assets by default, but provides the infrastructure to load user-supplied or licensecompatible assets when available. Users who compare against ALE or original Atari ROMs are responsible for ensuring that their use of third-party ROMs, assets, or emulators complies with the relevant licenses and local regulations. A.3

Reproducibility Statement

To facilitate reproducibility, we make the code publicly available. All scripts used to train and evaluate agents are included in the repository; We also provide a complete description of the directory structure and instructions for running each experiment. The exact configuration files used for each experiment are part of the codebase, allowing exact reproduction of all reported runs. Experiments were averaged over 3 seeds; all scripts enforce fixed seeding to ensure deterministic results. Compute resources are disclosed in Appendix D, which also details the exact hyperparameters, hardware, and training protocols. All experiment scripts log the environment name, observation mode, modification name, agent type, random seed, episode return, episode length, wall-clock time, software version, and hardware configuration. Because JAXtari is a software benchmark, reproducibility also depends on stable releases and regression testing; we therefore provide versioned benchmark releases, fixed configurations for reported baselines, and tests covering environment transitions, rewards, rendering, wrappers, and modification behavior. Substantive changes to game dynamics or evaluation protocols will be documented in the changelog and released under a new version number.

15

A.4

Usage

Installation.

The package can be installed from source as follows:

1 git

clone https://github.com/k4ntz/JAXtari JAXtari 3 pip install -e . 2 cd

A complete list of dependencies is provided in pyproject.toml. Basic usage. 1 import 2 import

A minimal example for creating and stepping an environment is:

jax jaxtari

3 4 env 5 key

= jaxtari.make("pong") = jax.random.PRNGKey(0)

6 7 key,

key_reset, key_action = jax.random.split(key, 3)

8 9 obs,

state = env.reset(key_reset) = env.action_space().sample(key_action) 11 obs, state, reward, done, info = env.step(state, action) 10 action

Modifications can be injected at environment creation: 1 env

= jaxtari.make("pong", mods=["lazy_enemy"])

The same interface can be used with vectorized JAX execution, allowing users to batch environment instances with jax.vmap and compile rollouts with jax.jit.

16

B

Human + Agentic Validation Pipeline and Implementation Verification

Reimplementing Atari 2600 game mechanics from scratch in a pure-JAX functional paradigm requires rigorous quality control to ensure learning task parity with the original Arcade Learning Environment (ALE) (Bellemare et al., 2013). Because Atari ROMs were originally written in 6502 assembly with tight hardware constraints (Television Interface Adapter / TIA clock timings), subtle gameplay mechanics and edge cases can easily be missed in high-level reimplementations. To systematically prevent behavioral discrepancies, we established a four-stage Human + Agentic validation pipeline combining autonomous specification extraction, automated code auditing, humanin-the-loop verification, and empirical policy stress-testing. B.1

The Four-Stage Validation Protocol

1. Stage 1: Autonomous Specification Extraction (Agentic). For every environment, an LLM agent extracts an exhaustive ground-truth operational specification by reading official documentation from the Farama Foundation (ale.farama.org) and canonical AtariAge manuals (atariage.com). This structured specification document formalizes: • Action Space: The active subset of the 18 discrete Atari actions, identifying which inputs induce positional movement, fire projectiles, or trigger special mechanics. • Game Mechanics: Complete entity lifecycle, including spawn conditions, directional velocity, screen boundaries, multi-level progression, and collision consequences. • Reward Structure and Termination: Explicit scoring criteria (points per entity hit or collectable picked up), milestone bonuses, life counters, and game-over triggers. 2. Stage 2: Code Mapping and Gap Auditing (Agentic). A secondary auditing LLM agent parses the JAX implementation against the generated specification: • Recall Audit: Maps each specified game mechanic and scoring rule to exact line numbers and function definitions in the Python/JAX source code, ensuring that no specified rule was omitted. • Precision Audit: Examines all unmapped code paths to identify undocumented logic, hardcoded constants, or custom mechanics not found in official manuals. 3. Stage 3: Dual Ground-Truth Verification (Human-in-the-Loop). Human expert verifiers evaluate the audit outputs along two complementary dimensions: • Recall Verification: Experts review the mapped code blocks to confirm that the mathematical dynamics and coordinate transforms accurately match Atari mechanics. • Precision Verification: For any undocumented code behavior flagged in the precision audit, human verifiers play the original ALE ROM side-by-side with the JAXtari environment. This step confirms whether the logic reflects an authentic assembly-level edge case of the original game or represents an unintended implementation artifact. 4. Stage 4: Empirical Policy Stress-Testing. Finally, standard RL agents (CleanRL PPO) are trained on both JAXtari and ALE under identical hyperparameters. Learning curve trajectories, asymptotic returns, and exploration rates are compared to guarantee that the JAX environment presents an equivalent learning challenge and does not permit unintended shortcut strategies. B.2

Case Study: Validation Checklist for Frostbite

To demonstrate the efficacy of this validation protocol, Tables 1 and 2 present the verified specification mapping and human audit for Frostbite. B.3

Documented Implementation Differences Between JAXtari and ALE

While the validation pipeline guarantees learning task parity, several fundamental differences arise from modernizing Atari for accelerator-native execution: • Explicit PRNG Keying vs. Implicit Hardware Registers: ALE generates pseudo-randomness from uninitialized memory and Stella hardware counter side-effects. JAXtari uses functional JAX PRNG splitting (jax.random.split), ensuring strict mathematical determinism across massive parallel threads. 17

Table 1: Frostbite Action Space and Core Dynamics Validation Checklist. Mapped against the JAXtari implementation and verified against ALE. Mechanic Action Space Coverage Full Discrete Space Directional Controls FIRE Button Core Game Dynamics Igloo Construction Floe Resetting Reversal Penalty

Operational Specification

Source Implementation

Utilizes full Discrete(18) action space; NOOP, FIRE, directional jumps, and combined inputs are supported. Moving joystick triggers Bailey’s jumps across moving ice floes and between shores. Pressing FIRE reverses the movement direction of the ice row currently occupied.

jax_frostbite.py

Landing on a white ice floe turns it blue and adds one block to the igloo (15 blocks required for completion). Once all ice blocks in a single row are turned blue, the entire row reverts back to white. Using FIRE to reverse row direction deducts one uncompleted block from the igloo.

jax_frostbite.py

Precision Audit Finding (Authentic Atari Edge Case) Auditor note: When the igloo is complete (15 blocks), subseFree Reversals quent FIRE reversals incur no block deduction penalty. Human verification: Verified in original ALE ROM; a completed igloo cannot be diminished.

jax_frostbite.py jax_frostbite.py

jax_frostbite.py jax_frostbite.py

jax_frostbite.py

Table 2: Frostbite Reward System Validation Checklist. Scoring Event

Reward Formula

Source Implementation

Coloring an Ice Block Entering Completed Igloo

10 × current level points awarded upon stepping on a white floe. 160×current level bonus points awarded per block upon entering the igloo to clear the level. Level multiplier increases up to level 8; from level 9 onward, scoring caps at 90 points per ice block. Flat bonus of +200 points upon collision with a swimming fish.

jax_frostbite.py jax_frostbite.py

Level Multiplier Cap Catching a Fish

jax_frostbite.py jax_frostbite.py

• Pure Functional State vs. Emulated RAM: In ALE, object coordinates must be inferred by decoding opaque RAM bytes. In JAXtari, state is maintained as typed PyTrees of JAX arrays, guaranteeing exact floating/fixed-point entity bounds. • Floating-Point Coordinate Accumulation: To prevent GPU branching bottlenecks, continuous entity velocities are vectorized across arrays rather than relying on discrete sub-pixel scanline counters.

18

C

Evaluation Metrics and Computational Benchmarking

We assess agent performance and framework efficiency using common metrics designed to provide statistically rigorous insights into learning dynamics and hardware utilization. Reliable Performance Aggregation. To mitigate the impact of outliers and account for the high variance inherent in RL training, we adopt the statistical framework proposed by rliable (Agarwal et al., 2021). We aggregate final performance using the Interquartile Mean (IQM), which computes the mean over the middle 50% of the data distribution (i.e., the 25th to 75th percentiles). This provides a more robust estimate of central tendency than the grand mean or median. Furthermore, we report 95% stratified bootstrap confidence intervals to quantify uncertainty. For longitudinal analysis, learning curves represent the mean in-game rewards accompanied by the standard deviation across 3 independent seeds. Performance Change (PC). To quantify the impact of our environmental perturbations on agent proficiency, we define the Performance Change metric Delfosse et al. (2025). This metric assesses the relative shift in performance when an agent trained on a standard environment is evaluated on a modified variant. It is formulated as: PC =

(Smod − Rmod ) − (Sorg − Rorg ) |Sorg − Rorg |

(1)

where Smod is the score in the modified environment, Sorg is the score in the original environment, and R the random baselines. A value of PC = 0 indicates perfect performance invariance, while PC < 0 signifies a collapse in the perturbation. Throughput and Hardware Efficiency. The computational utility of JAXtari is quantified via environmental throughput, defined as the total number of frames processed per second across all parallel environments. A primary architectural objective of our JAX-native implementation is to maximize hardware efficiency, targeting 100% GPU utilization by eliminating CPU-GPU communication bottlenecks that typically constrain traditional emulators. To provide a transparent account of the computational costs associated with our benchmarks, we track and report total GPU hours. Also, all experiments were done on a single A100 GPU, while scaling with multiple GPUs is possible.

19

D

Agent Architectures and Training Setup

We evaluate the JAXtari-15 environments on PPO (Proximal Policy Optimization) with both pixelbased and object-centric observation types. Whenever possible the implementations are JAX-based for maximum performance (jitted and scanned stepping through the environment). Note that JAXtari can also be employed with non JAX implementations, e.g., pytorch or tensorflow (using DLPack 2 for zero-copy tensor transfer). We also provide a gym-wrapper for maximum compatibility. The network architecture and hyperparameters are adapted from the same sources as the baselines we compare JAXtari’s environments against to ensure that discrepancies in agent performance can, with high certainty, be attributed to environment differences and not the training setup. Therefore PPO is adapted from CleanRL (Huang et al., 2022). D.1

Baseline environments: ALE, EnvPool and ALE-VectorEnv

We evaluate our environments against the ALE (Bellemare et al., 2013) which serves as the standard CPU-based framework for running Atari 2600 games as RL environments. However, the ALE executes environments sequentially, incurring throughput bottlenecks. Since our environments are highly parallelizable, we also evaluate throughput against EnvPool (Weng et al., 2022) and ALEVectorEnv. Both are high-performance, native C++ multi-threaded implementations for running Atari2600 games. Unfortunately, CuLE (Dalton & Frosio, 2020) is currently unmaintained. We made extensive efforts to compile and evaluate the framework locally; however, following the official installation instructions did not result in usable states. Reproducing the environment on modern hardware proved unfeasible due to three well-documented technical blockers: (1) strict dependencies on deeply deprecated software stacks, specifically CUDA 10.0 and PyTorch 1.23 , (2) fatal compilation failures when building the C++ source with modern GCC compilers or the provided Docker containers4 , and (3) unresolved assertion crashes during initialization on multi-GPU systems for users who manage to bypass the build errors5 . Because we could not establish a reliable, non-deprecated baseline for this hardware-accelerated framework, we focus our throughput evaluations on the ALE, EnvPool and ALE-VectorEnv. D.2

Preprocessing

Both agents share the same preprocessing parameters, which directly match the standard ALE preprocessing used in the CleanRL PPO baseline6 , ensuring fair comparison. For pixel-based observations: • Grayscale pixel-wise • Resize via native downscaling to 84 × 84 • Frame skip 4 frames • Max pooling over last two frames • Frame stack of 4 • Normalization of pixels through division by 255 • Reward clipping to {−1, 0, +1} For object-centric observations: • Frame skip 4 frames • Frame stack of 4 • Normalization to [0, 1] using feature specific bounds 2

https://dmlc.github.io/dlpack/latest/ https://github.com/NVlabs/cule/issues/4 and https://github.com/NVlabs/cule/issues/20 4 https://github.com/NVlabs/cule/issues/34 and https://github.com/NVlabs/cule/issues/28 5 https://github.com/NVlabs/cule/issues/5 6 https://docs.cleanrl.dev/rl-algorithms/ppo/ 3

20

• Flatten features into 1D vector • Reward clipping: to {−1, 0, +1} In both modalities: • Noop reset executes up to 30 random no-ops on episode reset • Episodic life terminates current episode on life loss • Sticky actions currently disabled (p = 0.0)

D.2.1

Native Downscaling

The aforementioned native downscaling is a feature of JAXtari’s rendering approach. Since the rendering is not directly tied to the game logic, but a separate computational step that constructs the image from a given state some preprocessing operations can be performed directly inside the renderer instead of afterwards on the finished image. Concretely, rather than rendering the full 210x160 pixel image, the renderer scales down all sprites at initialization and construct the image at the target resolution. This reduces memory consumption and increases throughput significantly, as operating on smaller images lowers the computational cost.

D.3 D.3.1

Network Architecture Pixel-Based Architectures

For pixel-based observations, PPO uses a standard DQN-like convolutional neural network (Mnih et al., 2015) consisting of three convolutional layers followed by a fully connected layer: • Conv1: 32 filters, 8 × 8 kernel, stride 4, ReLU • Conv2: 64 filters, 4 × 4 kernel, stride 2, ReLU • Conv3: 64 filters, 3 × 3 kernel, stride 1, ReLU • FC: 512 units, ReLU

D.3.2

Object-Centric Architecture

For object-centric observations the agents use a fully-connected architecture instead of convolutional layers used for pixel-based observations. For PPO, we use 461 nodes in the first layer and 512 in the second. The first layer size is specifically chosen to match the convolutional layer’s parameter count, enabling a fair comparison between the two input variants. Separate linear actor and critic heads are applied to the output. 21

D.4

Algorithm Details

D.4.1

PPO

Table 3: PPO hyperparameters. Identical for pixel-based and object-centric training runs Hyperparameter Value Total timesteps Parallel environments Rollout length Update epochs Minibatches Minibatch size Learning rate Discount factor γ GAE λ Clip ε Value function coeff. Entropy coeff. Max gradient norm Advantage normalization Value loss clipping Optimizer D.5

107 8 128 steps 4 4 8 × 128/4 = 256 2.5 × 10−4 (linear decay) 0.99 0.95 0.1 0.5 0.01 0.5 Yes No Adam (ε = 10−5 )

Training Protocol

All JAXtari agents reported in this paper were trained from scratch and without using pretrained weights. Furthermore, each agent is trained on environments initialized with three independent random seeds and evaluated on greedy (ε = 0) policies. We report IQM and confidence intervals across seeds. D.6

Computational Resources

We perform all our experiments on NVIDIA DGX hardware with A100 GPUs (40GB memory) and two AMD EPYC 7742 CPUs (128 cores, 256 threads). Each experiments uses the maximum of one GPU per experiment. The scaling experiment took a total of 32.5 GPU hours (including 5 ALE/vec-env runs for an average). The learning curve computation (PPO) took 180h GPU hours for the main results and 20h for appendix results. The evaluation on mods took a total of 2 GPU hours. With additional non-reported development runs, we estimate the total number of GPU hours to be roughly 500.

22

E

Extended Results

This section provides an extension to the evaluation of the JAXtari framework (cf. Section 3, specifically detailing the empirical performance across our benchmark suite. We first present an extension of the throughput analysis in Section E.1 to quantify the computational advantages of our JAX-native implementation, followed by performance comparisons across 15 environments (cf. Section E.2). E.1

Throughput and Scaling Analysis

We extend our work by reporting three complementary views of scaling. First, we present per-game scaling curves, quantifying how throughput, measured in environment steps per second, evolves as the number of parallel environments N increases, and revealing the point at which each title becomes bandwidth- or compute-bound. Second, we report wall-clock time to 100 × 106 environment steps, which serves as a practical, hardware-grounded measure of how quickly researchers can complete long-horizon training runs on a single accelerator. Third, we report the corresponding agent performance, i.e., final and intermediate returns achieved within these extended training budgets, allowing the reader to relate raw throughput to learning outcomes rather than treating speed in isolation. Together, these three axes, scaling, wall-clock cost, and resulting performance, are intended to give a faithful picture of what JAXtari enables in practice and to support reproducibility for follow-up work. Note that the reported ALE-VenctorEnv and EnvPool curves are the same over 5 games averaged throughput values for all runs. Beamrider

10 M

Breakout

Enduro 1M

1M

10 k

10 k

1k

1k

100 k 10 k

10 k

10 k 1k

Number of Environments

Number of Environments

Number of Environments

Frostbite

Gravitar

Kangaroo

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2³

1k

2⁰

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2³

2⁰

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2³

2⁰

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2³

2⁰

1k

Number of Environments

Number of Environments

Montezumarevenge

Mspacman

10 M

10 M

10 M

Number of Environments

10 k

1k

1k

Number of Environments

Number of Environments

Number of Environments

Pong

Seaquest

Skiing

Phoenix

1M

1M

100 k

100 k

10 k

10 k 1k

JAXtari-OC

1k

Number of Environments

JAXtari-Pixel

Gym-ALE

ALE-VectorEnv

Number of Environments

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2³

2⁰

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2⁰

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2³

2⁰

2³

Number of Environments

2³

1k

1k

2⁰

10 k 1k

1M 100 k

10 k

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ ⁰ 2¹ ¹ 2¹ ² 2¹ ³ 2¹ ⁴ 2¹ ⁵

100 k

2³

1M

Number of Environments

Tennis

10 M

10 M

2⁰

Throughput (SPS)

1M

10 k

Number of Environments

100 M

100 M

100 k

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

10 k

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2³

2⁰

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2³

2⁰

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

2³

2⁰

1k

100 k

2³

10 k 1k

1k

100 k

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

100 k 10 k

10 k

1M

1M

1M 100 k

2³

1M

100 k

2⁰

1M Throughput (SPS)

1M

100 k

100 k

2⁵ 2⁶ 2⁷ 2⁸ 2⁹ 2¹ 2¹⁰ 2¹¹ 2¹² 2¹³ 2¹⁴ 2¹⁵ ⁶

100 k

2³

1M

100 k

Freeway 10 M

2⁰

Throughput (SPS)

10 M

2⁰

Asteroids

1M

Number of Environments

EnvPool

Figure 7: JAXtari increased the number of possible environments and consequently the steps per second. While EnvPool and ALE-VectorEnv scale with the number of CPU threads (in our experiments, 256), our approach is not limited by it. Agents reported until reaching out-of-memory errors or 65K environments. ALE-VectorEnv and EnvPool curves are the same for all runs, averaged over 5 games.

23

Wall-Clock Throughput to 100M Steps. To probe the high-throughput regime of JAXtari, we measure the wall-clock time required to reach 100 × 106 environment steps (400 × 106 frames at frame-skip 4), well beyond the 10M–50M step horizon typically reported in the ALE. All runs execute on a single GPU and exploit JAX’s vmap and JIT compilation, keeping the rollout fully on-device. To maximize GPU utilization, we set the number of parallel environments N to 2048 for pixel-based PPO and 8192 for OC-based PPO, which is commonly the maximum for most games. All pixel-based agents reached 100M steps in under 1h, and OC even faster at a maximum of 12min (cf. Figure 8 and Figure 9). For scaling up PPO, we follow the recommendation from Beukman et al. (2026) and simply increase the the number of minibatches to match the original size, as well as the total number of steps. Asteroids

Breakout 400

3K

Reward

Beamrider

20K

0 0m

32m 17m

0 0m

43m 17m

33m

Wall-clock time

Wall-clock time

Frostbite

Gravitar

Reward

17m

0 0m

30m 17m

Wall-clock time

Wall-clock time

Kangaroo

Montezumarevenge

1K

32m 17m

0 0m

44m 33m

26m

0 0m

17m

Wall-clock time

Wall-clock time

Phoenix

Pong

Seaquest

100K

15

50K

17m

0

33m

0m

Wall-clock time

52m 33m

8m

17m

Wall-clock time

0 0m

500

29m

0 0m

17m

Wall-clock time

Skiing

Tennis 0 -5 -10

50K

18m

Mspacman

Wall-clock time 0

-10K

100K

5

45m

0 0m

150K

10

17m

Wall-clock time

1K

500

Wall-clock time 20

19m 8m

1K

10K

17m

0 0m

2K

20K

4K 2K

Reward

8m

2K

6K

0 0m

20m

30K

8K

0 0m

10

50

100 0 0m

20

100

200

10K

Freeway 30

300

2K 1K

Enduro 150

-15

-20K

42m 17m

33m

Wall-clock time

26m 0m

17m

Wall-clock time

-20 0m

21m 8m

17m

Wall-clock time

JAXtari-Pixel

Figure 8: On average, pixel-based approaches need 32 minutes to reach 100M steps on a single GPU. Training curves (reward vs. wall-clock time) for PPO with pixel-based observations across 15 Atari games, using 2048 parallel environments. Orange labels indicate total training time per game, including JAX compilation time.

24

Asteroids

3K

Beamrider

Breakout

Enduro

Reward

200 1K

2K

500

1K

0 0m

9m 3m

7m

0 0m

Wall-clock time

Reward

Frostbite

15

12m 4m

8m

0 0m

10

5

50

5m 2m

3m

Wall-clock time

Wall-clock time

Gravitar

Kangaroo

0 0m

10m 4m

8m

5 0 0m

Wall-clock time 1

4m 2m

3m

Wall-clock time

Montezumarevenge

Mspacman 2K

1K

1K

500

7m 3m

7m

500

250 0 0m

12m 4m

8m

6m

0 0m

3m

0 0m

8m 3m

7m

Wall-clock time

Wall-clock time

Wall-clock time

Wall-clock time

Phoenix

Pong

Seaquest

Skiing

4K

Reward

100

750 1K

20

3K

10

2K

0

1K 0 0m

10

150

1K

2K

0 0m

Freeway 20

11m 4m

8m

Wall-clock time

2K

0m

3m 2m

Wall-clock time

500 0 0m

0 0m

9m 3m

7m

Wall-clock time 1

Tennis

-5K

1K

-10

0

500

11m 4m

8m

Wall-clock time

-10K 0m

6m 3m

Wall-clock time

0 0m

4m 2m

3m

Wall-clock time

JAXtari-OC

Figure 9: On average, OC-based approaches need 8 minutes to reach 100M on a single GPU. Training curves (reward vs. wall-clock time) for PPO with object-centric observations across 15 Atari games, using 8192 parallel environments. Blue labels indicate total training time per game, including JAX compilation time.

25

Reward

Does More Training Time Yield Better PerforSeaquest mance? Whether extended training budgets 160K translate into improved agent performance is not 140K a question with a universal answer. On a substantial fraction of ALE titles, learning curves 120K saturate well before the 100M-step mark, and additional environment interactions yield dimin100K ishing, or even negligible, returns. However, 80K this picture changes markedly for the subset of games that are notoriously difficult to solve or 60K converge upon, most prominently Seaquest and Montezuma’s Revenge. These environments are 40K characterized by sparse, delayed, or hierarchi20K cally structured rewards, which require the agent to execute long sequences of coordinated actions 0 0 20 M 40 M 60 M 80 M before any learning signal is produced. In such Steps regimes, a longer training horizon directly enlarges the agent’s effective exploration budget: JAXtari-OC JAXtari-Pixel it increases the probability of stumbling upon rare reward-bearing trajectories, allows replay Figure 10: Performance after 100M steps. PPO buffers to accumulate a more diverse set of tran- reaches over 150k points in Seaquest before trunsitions, and gives policy- and value-estimates cation by reaching 108000 frames. additional time to escape local optima. Empirically, this manifests as higher final scores and qualitatively richer behaviors that simply do not emerge under shorter budgets. Our results in Figure 10 show that, with a training budget of 100M (or just 43min), we can achieve very high scores. The throughput characterized in the preceding section opens a concrete avenue for future investigation. Because JAXtari renders 100M-step (and beyond) training runs tractable on a single GPU, it becomes feasible to systematically revisit the hard-exploration subset of the ALE under substantially extended compute budgets, 200M, 500M, or even 1B steps, without resorting to large-scale distributed infrastructure. We see several promising directions enabled by this regime, which can be found in our future work section. E.2

Can JAXtari variations be used to evaluate robustness?

Regarding RQ3 we evaluated the performance of our PPO agents in variations of their trained environments. Figure 5 aggregates these over games and mods. Figure 11 shows the performance change per game and per mod.

26

Asteroids

Beamrider

PC (IQM)

4.0

Breakout

5.0

2.0

0.0

0.0

2.0 0.0 matrix theme

slow asteroids

double enemy speed

Enduro

0.1

fog of war

three lanes

Freeway

PC (IQM)

0.0

0.0 0.0

-0.5 filled road

snow

start in max curve

black cars

Gravitar

invert speed

static cars

Kangaroo

PC (IQM)

static ice

dark night

no enemies

Montezumarevenge

0.5 0.0

0.0

-0.5

0.0

-1.0 hyper gravity

inverted colors

zero gravity

no danger

Mspacman PC (IQM)

speed mode

Frostbite

5.0

spawn on second floor

tanks

debug hud

Phoenix

infinite amulet

no fall damage

Pong 0.0

0.0 0.0 -0.5

-0.5

-1.0 caged ghosts

constant fruits

matrix theme

fast enemy bullets

Seaquest PC (IQM)

big paddle

10.0

0.5

-0.1

ball color

fast player

grayscale

change background color

Skiing

lazy enemy

shift enemy

Tennis 0.0

0.0

0.0

-1.0

-0.5 -1.0

-0.5 disable enemies

gravity

JAXtari-OC

mines

dangerous moguls

more trees

JAXtari-Pixel

Improves

classic trees

hard court

Degrades

lazy enemy

random ball speed

Neutral

Figure 11: JAXtari allows users to evaluate performance changes when the environment changes, making it a valuable tool in researching shortcuts. Per game evaluation measuring the drop or improvements of the agents. Evaluations were done over 10 runs per seed. Values show IQM with a 95% CI.

27

F

JAXtari Environment Descriptions and Modifications

F.1

JAXtari with open-source sprite-replacements

Users that do not own a license to the original Atari ROMs can instead utilize our open-source sprite replacement package, available for all games. The replacement sprites are inspired by MinAtar (Young & Tian, 2019) and are single-color rectangles instead of the actual sprites (cf. Figure 12).

(a) Frostbite

(b) Skiing

(c) Tennis

(d) Freeway

(e) Seaquest

Figure 12: The replacement sprites only change the games visually but preserve game mechanics.

F.2

Creating new Modifications

JAXtari features a modular modification system that allows users to edit game semantics directly at the source-code level (as described in Section 2). Modifications are built around the JaxtariInternalModPlugin and JaxtariPostStepModPlugin. When creating a new mod, you can intervene at two distinct points in the environment’s lifecycle: • Internal Modifications: These interventions are applied before Just-In-Time (JIT) compilation. Because they become part of the compiled transition function, they execute highly efficiently. Use this method to override core game constants, collision rules, environment methods, enemy speeds, or object spawning behaviors. You can also use the declarative asset system here to swap visual sprites by overriding values in the environment’s original ASSET_CONFIG. • Post-Step Modifications: These interventions operate on the environment state after the original transition function has evaluated. They are ideal for changes that do not require rewriting the core game dynamics, such as relocating an object, masking a state variable, changing a reward, or triggering an alternative episode termination. Once your modification is defined, you can inject it directly during environment initialization by passing it to the mods argument: env = jaxtari.make("game_name", mods=["custom_mod"]).

Example: Modifying Kangaroo In the original Kangaroo environment, the player controls Mother Kangaroo on a mission to rescue her baby. This involves navigating obstacles, climbing ladders, collecting bonus fruit, and punching interfering monkeys. To evaluate whether an agent has learned robust causal mechanics (like climbing) rather than overfitting to punching the monkeys on the first floor, we can introduce a modification. Let’s create a conceptual implementation of the no_monkey mod of HackAtari, which removes the monkeys from the game. 28

1 import jax 2 import jax.numpy as jnp 3 from jaxtari.modifications import JaxtariInternalModPlugin 4 5 class NoMonkeyMod(JaxtariInternalModPlugin): 6 """ 7 Internal mod to disable monkeys. 8 This patches the environment’s ’_monkey_controller’ method. 9 """ 10 11 @partial(jax.jit, static_argnums=(0,), donate_argnums=(1,)) 12 def _monkey_controller(self, state: KangarooState, punching: chex.Array): 13 """ 14 No-op override for _monkey_controller. 15 """ 16 score_addition = jnp.zeros((), dtype=jnp.int32) 17 18 return ( 19 state.level.monkey_states, 20 state.level.monkey_positions, 21 state.level.monkey_throw_timers, 22 score_addition, 23 state.level.coco_positions, 24 state.level.coco_states, 25 jnp.array(False), 26 ) 27 28 29 # Injecting the mod during environment initialization 30 env = jaxtari.make("kangaroo", mods=[NoMonkeyMod()])

Figure 13: Implementation of an InternalModPlugin in JAXtari. This example demonstrates a structural intervention where the environment’s internal _monkey_controller is patched to a no-op, effectively removing a specific game mechanic while maintaining JIT-compatibility and end-to-end execution. Depending on your design, you could implement this as either an internal (cf. Figure 13) or post-step modification (cf. Figure 14). This mod intercepts the state after the transition function evaluates and simply forces the monkeys’ active status to be disabled. Further, we can use the mod system to swap assets or constants within the game implementation. Lets say the agent learned to dodge coconuts. Now we swap the coconuts with a slightly larger fireball and see what is happening (cf. Figure 15). While both JAXtari and HackAtari allow for environment modifications, they differ fundamentally in their implementation and performance goals. HackAtari is built on top of the standard ALE and Stella emulators. In HackAtari, modifications involve intercepting the emulator’s RAM. In JAXtari, modifications are applied directly to the high-level JAX state objects or by overriding Python methods before JIT compilation. By keeping the environment functional and stateless, these modifications remain compatible with JAX transformations like jax.vmap and jax.jit, allowing you to run thousands of modified environment instances simultaneously on a single GPU.

29

1 import jax 2 import jax.numpy as jnp 3 from jaxtari.modifications import JaxtariPostStepModPlugin 4 5 class NoMonkeyMod(JaxtariPostStepModPlugin): 6 @partial(jax.jit, static_argnums=(0,)) 7 def run(self, prev_state: KangarooState, new_state: KangarooState): 8 """ 9 Called *after* the main step. Overwrites the monkeys. 10 """ 11 level_constants = self._env._get_level_constants(new_state.current_level) 12 13 pinned_level_state = new_state.level.replace( 14 monkey_states=jnp.array(0), 15 monkey_positions=jnp.array(0), 16 monkey_throw_timers=jnp.array(0), 17 ... 18 ) 19 20 return new_state.replace(level=pinned_level_state) 21 22 23 # Injecting the mod during environment initialization 24 env = jaxtari.make("kangaroo", mods=[NoMonkeyMod()])

Figure 14: State-level modification using the PostStepModPlugin. Unlike internal logic patching, this plugin performs post-hoc state manipulation after each transition, allowing researchers to pin specific environmental variables or perform causal interventions on the observation and state space. This however have to be done in every step.

1 class ReplaceCoconutWithFireball(JaxtariInternalModPlugin): 2 asset_overrides = { 3 "coconut": { 4 ’name’: ’coconut’, 5 ’type’: ’single’, 6 ’file’: ’fireball.npy’ 7 } 8 } 9 constants_overrides = { 10 "THROWN_COCONUT_WIDTH": 16, 11 "THROWN_COCONUT_HEIGHT": 12, 12 }

Figure 15: Example of an asset and constant injection for robustness evaluation. By overriding sprite assets and physical constants (e.g., object dimensions or speed), JAXtari enables the rapid generation of Out-of-Distribution (OOD) scenarios to test the generalization capabilities of RL agents.

30

F.3

List of Games and Modifications

Figure 16 provides a visual overview of all 45 environments currently supported by JAXtari. Full environment descriptions and available modifications follow below.

Alien

Amidar

Asterix

Asteroids

Atlantis

Bank Heist

Beam Rider

Berzerk

Blackjack

Breakout

Centipede

Chopper Cmd.

Enduro

Fishing Derby

Flag Capture

Freeway

Frostbite

Galaxian

Gravitar

Haunted House

H. Cannonball

Kangaroo

King Kong

Laser Gates

Montezuma

Ms. Pac-Man

Name This Game

Phoenix

Pong

Q*bert

River Raid

Seaquest

Sir Lancelot

Skiing

Space Invaders

Spacewar

Tennis

Tetris

Time Pilot

Tron

Turmoil

Venture

Video Cube

Video Pinball

Word Zapper

Figure 16: ALE (left) vs. JAXtari (right) renderings for all 45 supported games. 31

F.4 F.4.1

Alien Alien details

The player is trapped aboard a maze-like spaceship alongside three hostile aliens, whose eggs are scattered throughout the vessel and must be destroyed. The aliens actively pursue and attack the player, but can be temporarily repelled using a flamethrower. Occasionally a pulsar power-up can be collected, granting the ability to kill aliens outright for a limited time.

F.4.2

Available Modifications for Alien

Modification

Effect

aggressive_swarm dont_kill end_game extra_lives last_egg matrix_theme pacifist_mode short_circuit

Aliens spend almost all their time actively chasing the player. Punishes for killing and shooting enemies. Most eggs are already eaten. Start with extra lives. Only one egg is remaining. A Matrix-themed mod: black background, green walls/enemies/eggs. Aliens are never frightened, forcing a pure evasion playstyle. Vulnerability window is drastically reduced.

F.4.3

Selection of Available Visual Modifications

Figure 17: A selection of visual modifications available for Alien, with each frame showing a different mod applied to the base environment.

32

F.5 F.5.1

Amidar Amidar details

The player traverses a 2-dimensional grid, aiming to visit every cell while evading pursuing enemies. At certain moments the dynamic reverses: enemies transform into chickens that the player can actively chase and catch.

F.5.2

Available Modifications for Amidar

Modification

Effect

No modifications implemented yet.

33

F.6 F.6.1

Asterix Asterix details

The player controls Asterix, who moves continuously along the horizontal axis and discretely along the vertical. A stream of objects scrolls across the screen, including harmful lyres and various collectibles. The goal is to collect as many useful objects as possible while avoiding lyres, which cost a life on contact. The game starts with three lives, and bonus points are awarded upon reaching certain score thresholds.

F.6.2

Available Modifications for Asterix

Modification

Effect

No modifications implemented yet.

34

F.7 F.7.1

Asteroids Asteroids details

The player pilots a spaceship through an asteroid field, destroying asteroids by shooting them. Large asteroids break into smaller fragments upon being hit, all of which must be cleared to advance to the next level. Flying saucers occasionally appear as an additional threat.

F.7.2

Available Modifications for Asteroids

Modification

Effect

dont_shoot matrix_theme slow_asteroids instant_turn

Internal mod that punishes shooting A Matrix-themed mod: black background, green asteroids The asteroids are slower Player instantly turns in the pressed direction

F.7.3

Selection of Available Visual Modifications

Figure 18: A selection of visual modifications available for Asteroids, with each frame showing a different mod applied to the base environment.

35

F.8 F.8.1

Atlantis Atlantis details

The player defends the submerged city of Atlantis against waves of enemies that slowly descend toward it. Three defense posts are at the player’s disposal to shoot down attackers before they reach striking distance. The game is lost if all seven of Atlantis’ installations are destroyed. After successfully repelling a wave and accumulating enough points, damaged installations can be rebuilt.

F.8.2

Available Modifications for Atlantis

Modification

Effect

jets_only no_last_line random_enemies speed_mode_fast speed_mode_medium speed_mode_slow

Overrides enemy spawn probabilities to only spawn Jet-type enmies. Disable the lowest lane so no enemies can spawn there. Overrides enemy spawn probabilities to make spawning random. Increases overall speed of the game to 6. Increases overall speed of the game to 4. Increases overall speed of the game to 2.

36

F.9

F.9.1

Bankheist

Bankheist details

The player takes the role of a bank robber navigating a getaway car through maze-like cities. Banks can be robbed along the way, though each robbery triggers a police pursuit. Police cars can be neutralized by dropping dynamite, though driving over one’s own explosives is equally fatal. Entering a new city refills the gas tank. The game starts with four lives, which are lost by running out of fuel, getting caught, or falling victim to one’s own dynamite.

F.9.2

Available Modifications for Bankheist

Modification

Effect

2_police_cars blue_player double_speed dynamite_penalty fuel_for_banks golden_banks grey_road moving_banks no_police random_city random_spawns red_police_cars revisit_city unlimited_gas

Replaces 2 banks with police cars. Each bank robbed gives 50 points. Changes the color of the player to light blue. Doubles the player speed. Changes reward for killing a police car with dynamite (-500). Augments the player’s fuel when 3 banks have been robbed. Changes the color of banks to golden. Changes the road color to grey. Makes banks move like police cars. Removes police cars from the game and automatically respawns banks. Randomizes which city is entered next. Restores the procedural, fully random bank spawns. Changes the color of police cars to red. Allows player to go to the previous city by going through the left portal. Prevents gas consumption. Only dropping dynamite consumes gas.

37

F.9.3

Selection of Available Visual Modifications

Figure 19: A selection of visual modifications available for Bankheist, with each frame showing a different mod applied to the base environment.

38

F.10 F.10.1

Beamrider Beamrider details

The player pilots a spaceship moving forward at constant speed, steerable only laterally across a set of discrete positions. The objective is to shoot down enemy vessels while evading their projectiles and navigating around space debris.

F.10.2

Available Modifications for Beamrider

Modification

Effect

dont_kill double_enemy_speed fog_of_war hardcore mothership_laser same_enemies teleport_ufos three_lanes toaster

Internal mod that punishes killing and shooting. Doubles enemy movement speed. Hide the upper two thirds of the screen behind opaque fog. Start with one life and never allow extra lives. Turn the mothership sequence into a stationary lane laser attack. Render all enemy types using the white UFO visuals. Adds a rare attack pattern for the UFOs. Collapses strack to the center three lanes. Disable heavy visual effects while keeping gameplay unchanged.

F.10.3

Selection of Available Visual Modifications

Figure 20: A selection of visual modifications available for Beamrider, with each frame showing a different mod applied to the base environment.

39

F.11 F.11.1

Berzerk Berzerk details

The player navigates a maze populated by evil robots that must be destroyed, while also avoiding contact with the maze walls, which is instantly fatal. Enough points can earn extra lives depending on the selected game mode. In certain modes, the indestructible Evil Otto appears and pursues the player, adding an additional threat that cannot be eliminated and must simply be outrun.

F.11.2

Available Modifications for Berzerk

Modification

Effect

No modifications implemented yet.

40

F.12 F.12.1

Blackjack Blackjack details

The player competes against a dealer in a game of Blackjack, drawing cards with the aim of reaching a total as close to 21 as possible without exceeding it and going bust.

F.12.2

Available Modifications for Blackjack

Modification

Effect

No modifications implemented yet.

41

F.13 F.13.1

Breakout Breakout details

The player controls a paddle at the bottom of the screen, bouncing a ball against a wall of bricks above. Each hit destroys a brick, and breaking through to the other side allows the ball to clear bricks freely on its own. The game provides five lives before it ends.

F.13.2

Available Modifications for Breakout

Modification

Effect

ball_color ball_drift ball_gravity big_paddle block_color player_color small_paddle speed_mode

Changes the balls color to a set color. Consistently drift the ball to the right. Pulls the ball down by 1 every step of the gravity_buffer. Always use a bigger paddle (Size 40). Changes the blocks color to a set color. Changes the player color to a set color. Always use a small paddle (Size 4). Increase speed to maximum at all time steps.

F.13.3

Selection of Available Visual Modifications

Figure 21: A selection of visual modifications available for Breakout, with each frame showing a different mod applied to the base environment.

42

F.14 F.14.1

Centipede Centipede details

The player takes the role of an elf defending an enchanted forest, using magic wands to fend off waves of centipedes, spiders, and fleas that threaten the mushrooms. Being bitten by any creature causes temporary paralysis and costs one wand. Additional wands can be earned by reaching certain score thresholds, but the game ends once all wands are lost.

F.14.2

Available Modifications for Centipede

Modification

Effect

No modifications implemented yet.

43

F.15 F.15.1

ChopperCommand ChopperCommand details

The player pilots a helicopter tasked with protecting ground-based truck convoys by shooting down enemy aircraft. A mini-map at the bottom of the screen provides an overview of the surrounding area.

F.15.2

Available Modifications for ChopperCommand

Modification

Effect

No modifications implemented yet.

44

F.16

F.16.1

Enduro

Enduro details

The player competes in the National Enduro, a long-distance endurance race. Each day a minimum number of cars must be overtaken to remain in the race: 200 on the first day and 300 on each subsequent day. Failing to meet the daily quota ends the game.

F.16.2

Available Modifications for Enduro

Modification

Effect

dawn day dont_max_speed filled_road fog hud max_speed_reward night no_opponents short_days snow start_in_curve start_in_max_curve sunset

Forces the weather to dawn. Forces the weather to day. Does not reward the player if at max. speed. Displays the full filled road instead of only the side. Forces the weather to fog. Renders a HUD that displays sped and x-position. Provides +2 every 20th frame when at maximum speed. Forces the weather to night. Disables all opponent cars. Iterates weather state every 16km driven instead of based on time. Forces the weather to snow. Starts the game directly in the curve. Starts the game in a full curve. Forces the weather to sunset.

45

F.16.3

Selection of Available Visual Modifications

Figure 22: A selection of visual modifications available for Enduro, with each frame showing a different mod applied to the base environment.

46

F.17 F.17.1

Fishingderby Fishingderby details

The player competes against an opponent to catch more sunfish, each fishing from opposite sides of a lake. Lines are cast downward through several layers of water, where fish of varying sizes swim by at different depths. The catch must be reeled in quickly, as sharks patrol the water and will steal the fish right off the line if given the chance. The first player to reach 99 points wins.

F.17.2

Available Modifications for Fishingderby

Modification

Effect

fish_in_middle fish_on_different_sides fish_on_player_side shark_no_movement_easy shark_no_movement_middle

All fish are bound to a small area in the middle (x-position). Half the fish are bound to an area on the right, half on the left. All fish are bound to an area on the players side. Shark is stuck to the right pier. Shark is stuck to the center (x-position).

47

F.18 F.18.1

Flagcapture Flagcapture details

The player explores a grid-based world, moving from square to square while collecting clues scattered throughout. The ultimate objective is to locate and reach a hidden flag.

F.18.2

Available Modifications for Flagcapture

Modification

Effect

No modifications implemented yet.

48

F.19

F.19.1

Freeway

Freeway details

The player guides a chicken across multiple lanes of dense rush hour traffic, aiming to reach the top of the screen. Each successful crossing earns a point, while collisions with vehicles send the chicken back toward the starting side.

F.19.2

Available Modifications for Freeway

Modification

Effect

_hall_of_fame_start bikes black_cars center_cars_on_reset change_sprites frog green_score hall_of_fame invert_speed new_lane_colors slow_cars static_cars stop_all_cars too_close_penalty

Spawns the cars to make a "hall of fame" formation. Replaces all cars with uniquely colored bike sprites. Makes all cars black by overriding CAR_COLORS constant. Positions all cars in the center of the screen when the environment resets. Bundles frog, bikes, new_lane_colors and green_score Replaces the player sprites with frog sprites. Makes the score digits green. Bundles _hall_of_fame_start and static_cars. Inverts the speed of all cars by overriding CAR_UPDATES constant. Makes the lane separation black and the double lane separation red. Halves the speed of all cars. Stops all cars after spawning. Stops all cars randomly with probability 0.4. Provides -1 reward if the chicken is too close to a car.

49

F.19.3

Selection of Available Visual Modifications

Figure 23: A selection of visual modifications available for Freeway, with each frame showing a different mod applied to the base environment.

50

F.20

F.20.1

Frostbite

Frostbite details

The player controls "Frostbite Bailey", hopping across an Arctic river and turning ice blocks from white to blue. Each converted block contributes to the construction of an igloo.

F.20.2

Available Modifications for Frostbite

Modification

Effect

_aligned_ice _misaligned_ice _static_ice bad_fishes_reward change_sprites dark_night early_bear left_igloo lightblue_ice no_enemies recolored_obstacles static_aligned_ice static_misaligned_ice tiger white_igloo

Makes the ice blocks perfectly centered. Makes the ice blocks vertically aligned with each other but off-center. Sets ice block horizontal speeds to zero. Punishes the player for collecting fish. Bundles tiger, white_igloo, recolored_obstacles and lightblue_ice. Spawns constant night and makes the sky darker. Spawns the bear from the start (Level 1). Places the igloo on the left side of the screen. Changes the color of the ice blocks to light blue. Disables all enemies and only keep fishes. Changes the colors of the floating obstacles. Bundles _static_ice and _aligned_ice. Bundles _static_ice and _misaligned_ice. Switches the bear sprites to tiger sprites. Makes the igloo white.

51

F.20.3

Selection of Available Visual Modifications

Figure 24: A selection of visual modifications available for Frostbite, with each frame showing a different mod applied to the base environment.

52

F.21 F.21.1

Galaxian Galaxian details

The player defends a laser base against successive waves of Galaxian invaders. Each wave must be cleared by shooting down all attackers before they reach the base.

F.21.2

Available Modifications for Galaxian

Modification

Effect

No modifications implemented yet.

53

F.22

F.22.1

Gravitar

Gravitar details

The player pilots a small spacecraft through a fictional solar system, free to explore its planets. Flying into a planet transitions the game to a side-view surface landscape.

F.22.2

Available Modifications for Gravitar

Modification

Effect

anti_gravity fuel_crisis grayscale harmless_enemies high_speed hyper_gravity infinite_fuel inverted_colors long_range_tractor neon_mode rapid_fire red_alert sharpshooter slow_enemies valuable_reactor zero_gravity

Reverse gravity from sun, planets, and reactors. Increase fuel consumption rate by 5x. Converts the visual palette to grayscale. Make all enemies harmless by disabling their bullets. Make the ship faster by increasing thrust power and max speed. Increase all gravity from sun, planets, and reactors. Disable fuel consumption. Inverts the primary colors. Increase the range of the tractor beam. Changes colors to bright neon variants. Increase active bullet caps for ship, saucer, and enemies to 4. Makes all terrain red/orange for a high-alert aesthetic. Punishes shooting (-100) and rewards killing enemies (+500). Decrease the movement speed of saucers and bullets. Populate reactor level with 3 enemies and 2 fuel tanks. Disable all gravity from sun, planets, and reactors.

54

F.22.3

Selection of Available Visual Modifications

Figure 25: A selection of visual modifications available for Gravitar, with each frame showing a different mod applied to the base environment.

55

F.23 F.23.1

HauntedHouse HauntedHouse details

The player explores a mansion haunted by the ghost of the vengeful Samuel Graves, searching for three pieces of a magic urn. The mansion must be escaped with all three pieces before exhausting the nine available lives.

F.23.2

Available Modifications for HauntedHouse

Modification

Effect

No modifications implemented yet.

56

F.24 F.24.1

HumanCannonball HumanCannonball details

The player fires a human cannonball and must angle the shot to land the performer precisely into a water tower target.

F.24.2

Available Modifications for HumanCannonball

Modification

Effect

No modifications implemented yet.

57

F.25 F.25.1

Kangaroo Kangaroo details

The player controls Mother Kangaroo on a mission to rescue her baby, starting with three lives. The rescue involves navigating a series of obstacles: climbing ladders, collecting bonus fruit, and throwing punches at interfering monkeys. The goal is to complete the rescue while accumulating as many points as possible.

F.25.2

Available Modifications for Kangaroo

Modification

Effect

_chickens _lethal_bell center_ladders chains chickens dont_punch first_level_only four_ladders high_thrown_coconuts invert_ladders no_bell no_danger no_falling_coconut no_fruit no_monkey no_thrown_coconut pin_child render_debug_info replace_bell_with_danger_sign replace_child_with_monkey replace_coconut_fireball ropes second_level_only spawn_on_second_floor third_level_only

Replaces the monkey sprite with a chicken sprite. Post-step mod that kills the player if they touch the bell. Internal mod to center all ladder positions horizontally. Replaces the ladder sprites with a chain. Bundles _chickens and no_thrown_coconut. Provides negative reward for punching monkeys. Internal mod to force the game to always stay on level 1. Internal mod to add a fourth ladder to each level. Forces coconuts to always spawn at the high position. Internal mod to invert all ladder positions (horizontally). Internal mod to disable the Bell. Bundles no_monkey and no_falling_coconut. Internal mod to disable the single falling coconut. Internal mod to remove Fruits. Internal mod to disable monkeys. Internal mod to disable thrown coconuts. Post-step mod to pin the child kangaroo in place. Renders the player’s X,Y coordinates. Replaces the bell sprite with a danger sign. Replaces the child sprite with a monkey sprite. Replaces the coconut sprite with a fireball sprite. Replaces the ladder sprites with a rope. Internal mod to force the game to always stay on level 2. Spawns the player on the second level position. Internal mod to force the game to always stay on level 3.

58

F.25.3

Selection of Available Visual Modifications

Figure 26: A selection of visual modifications available for Kangaroo, with each frame showing a different mod applied to the base environment.

59

F.26 F.26.1

KingKong KingKong details

The player climbs the Empire State Building to rescue a person kidnapped by King Kong, who hurls bombs downward to impede the ascent. The bombs must be dodged throughout the climb to reach the top safely.

F.26.2

Available Modifications for KingKong

Modification

Effect

No modifications implemented yet.

60

F.27 F.27.1

LaserGates LaserGates details

The player navigates through a malfunctioning computer system using the Dante Dart, with the goal of locating and destroying four Failsafe Detonators scattered within.

F.27.2

Available Modifications for LaserGates

Modification

Effect

No modifications implemented yet.

61

F.28 F.28.1

MontezumaRevenge MontezumaRevenge details

The player ventures through a sprawling maze of chambers inside the emperor’s fortress in search of Montezuma’s treasure. Along the way, deadly creatures must be avoided while collecting valuables and useful tools that aid in the escape with the treasure.

F.28.2

Available Modifications for MontezumaRevenge

Modification

Effect

center_bouncing_skull debug_hud fast_player infinite_amulet jumping_spiders moving_snakes no_enemies no_fall_damage reveal_map rolling_skulls super_jump sword_kill_bonus three_swords

Make the rolling skull jump vertically at the center of the screen. Internal mod to display debug info (Room ID, X, Y) on the HUD. Internal mod to increase player speed. Post-step mod to keep the amulet active forever. Post-step mod to make all spiders jump. Post-step mod to make snakes move horizontally in all rooms. Post-step mod to immediately remove any enemies in the current room. Internal mod to disable fall damage. Internal mod to make dark rooms visible, by granting the torch effect. Make the two skulls that usually bump roll instead, Internal mod to increase jump height. Provide bonus points if the player kills an enemy with the sword. Post-step mod to start the player with 3 swords.

62

F.29

F.29.1

MsPacman

MsPacman details

The player guides Ms. Pac-Man through a maze, collecting all pellets while evading four pursuing ghosts. Power pellets located at the corners of the maze temporarily turn the ghosts vulnerable, allowing Ms. Pac-Man to eat them for bonus points. Fruit items occasionally appear in the maze for additional score.

F.29.2

Available Modifications for MsPacman

Modification

Effect

caged_ghosts constant_fruits fruit_ghost_bonus matrix_theme only_1_ghost only_2_ghost only_3_ghost random_ghost_navigation set_maze_1 set_maze_2 set_maze_3 set_maze_4

Keeps the ghosts caged in the center. Always adds a fruit to the maze. Deactivates points for pellets, But multiplies eating ghost reward. A Matrix-themed mod: black background, green elements. Keeps only 1 ghost in the maze. Keeps only 2 ghost in the maze. Keeps only 3 ghost in the maze. Randomizes the navigation of all ghosts. Forces the level to always be maze 1. Forces the level to always be maze 2. Forces the level to always be maze 3. Forces the level to always be maze 4.

63

F.29.3

Selection of Available Visual Modifications

Figure 27: A selection of visual modifications available for MsPacman, with each frame showing a different mod applied to the base environment.

64

F.30 F.30.1

NameThisGame NameThisGame details

The player dives to protect a discovered underwater treasure, fending off attacks from both a shark and an octopus. Oxygen supplies must be carefully monitored throughout, adding an additional layer of urgency to the defense.

F.30.2

Available Modifications for NameThisGame

Modification

Effect

No modifications implemented yet.

65

F.31

F.31.1

Phoenix

Phoenix details

The player fights through successive waves of enemy war birds, shooting them down while dodging their bombs. The ultimate objective is to reach and defeat the alien pilot commanding the assault.

F.31.2

Available Modifications for Phoenix

Modification

Effect

blood_moon boss_late_missiles fast_enemy_bullets fast_player grayscale infinite_lives inverted_colors invincible_player matrix_theme night_mode no_ability_cooldown penalize_shoot_cooldown

A dark red themed mod. Make boss missiles appear later. Increases speed of enemy projectiles. Increases player movement speed. Turns the entire game into grayscale. Set player lives to 99. Inverts all colors in the game. Player is always invincible. A Matrix-themed mod: black background, green elements. Dims the entire screen by 50% for a night mode experience. Removes cooldown for the special ability (shield). Penalize for each time the player shoots (-50).

66

F.31.3

Selection of Available Visual Modifications

Figure 28: A selection of visual modifications available for Phoenix, with each frame showing a different mod applied to the base environment.

67

F.32 F.32.1

Pong Pong details

The player controls the right paddle, competing against a computer-controlled left paddle. Both sides attempt to deflect the ball past the opponent’s paddle and into their goal while keeping it out of their own.

F.32.2

Available Modifications for Pong

Modification

Effect

change_background_color change_player_color lazy_enemy linear_movement no_fire random_enemy reward_acceleration shift_enemy shift_player zero_score

Changes the background color to navy blue. Changes the player color to red. Enemy does not move after returning a shot until player hits ball. Changes the movement to a linear one (no acceleration). Disables the FIRE action. Randomizes the enemy movement. Rewards player when ball is accelerated by player (+2). Shifts the enemy on the x-axis. Shifts the player on the x-axis. Forces the score to always be zero.

F.32.3

Selection of Available Visual Modifications

Figure 29: A selection of visual modifications available for Pong, with each frame showing a different mod applied to the base environment.

68

F.33

F.33.1

Qbert

Qbert details

The player controls Q*bert, whose goal is to change every cube on a pyramid to a target color by hopping on each one individually. Various hostile creatures roam the pyramid and must be avoided throughout.

F.33.2

Available Modifications for Qbert

Modification

Effect

collecting_bonus_only dark_pyramid diagonal_control grayscale ice_pyramid inverted_colors night_mode no_coily no_enemies no_green_ball no_purple_ball no_red_balls no_sam penalize_all_collectables red_coily swap_collectibles_enemies swap_colors

Overrides the rewards to only get reward from collectibles. Changes the pyramid to a dark aesthetic. Enables diagonal controls. Makes the entire game grayscale. Changes the pyramid to an icy aesthetic. Swaps Qbert’s and Coily’s colors. Dims the background and the characters for a night-time aesthetic. Removes Coily from the screen. Removes all enemies from the screen. Removes green balls from the screen. Removes purple balls from the screen. Removes red balls from the screen. Removes Sam from the screen. Penalizes collecting items. Changes the color of Coily to red. Swaps the colors of the collectibles and Enemies. Swaps the colors of the tiles.

69

F.33.3

Selection of Available Visual Modifications

Figure 30: A selection of visual modifications available for Qbert, with each frame showing a different mod applied to the base environment.

70

F.34 F.34.1

Riverraid Riverraid details

The player pilots a jet flying along a river, steering laterally and firing missiles to destroy enemy targets for points. Fuel depletes continuously and must be replenished by flying over fuel depots before running dry, as an empty tank costs a jet. Collisions with the riverbank or any enemy object, except fuel depots, are equally fatal. The game starts with three jets in reserve, with an additional jet awarded for every 10,000 points scored, up to a maximum of nine.

F.34.2

Available Modifications for Riverraid

Modification

Effect

No modifications implemented yet.

71

F.35 F.35.1

Seaquest Seaquest details

The player commands a submarine capable of moving in all directions and firing torpedoes. The objective is to rescue as many divers as possible while destroying enemy submarines and sharks. The game begins with one active sub and three in reserve, with an additional sub awarded for every 10,000 points scored, up to a maximum of six reserves. The sub is destroyed upon colliding with anything other than friendly divers. Oxygen depletes continuously and must be replenished by surfacing in time. Failing to do so costs a sub. Surfacing with fewer than six divers on board also results in losing one diver.

F.35.2

Available Modifications for Seaquest

Modification

Effect

disable_enemies dont_kill gravity mines no_divers unlimited_oxygen

Disable enemies in the environment. Internal mod that punishes killing and shooting. Adds gravity to that pulls the player downwards. Replaces both sharks and enemy submarines with mine sprites. Internal mod to remove divers from the game. Player always has full oxygen.

F.35.3

Selection of Available Visual Modifications

Figure 31: A selection of visual modifications available for Seaquest, with each frame showing a different mod applied to the base environment.

72

F.36 F.36.1

SirLancelot SirLancelot details

The player controls Sir Lancelot, mounted on Pegasus, on a mission to rescue a prisoner held within a castle and guarded by a fire-breathing dragon.

F.36.2

Available Modifications for SirLancelot

Modification

Effect

No modifications implemented yet.

73

F.37

F.37.1

Skiing

Skiing details

The player controls a skier moving laterally down a slope, aiming to pass through all gates, marked by pairs of poles, in the shortest time possible. Missing a gate incurs a five-second penalty. Collisions with gates or trees briefly knock the skier down before the run continues.

F.37.2

Available Modifications for Skiing

Modification

Effect

_dangerous_moguls _more_moguls _more_trees _trees_everywhere blue_skiier classic_trees flag_flurry green_flags hall_of_fame invert_flag_colors invert_flags jump_to_break moguls_to_trees moving_flags random_flags reward_at_gate speed_burst thin_moguls

Makes colliding with moguls cause the skier to fall. Spawns more moguls during the race. Spawns more trees during the race. Allows trees to spawn anywhere across the entire horizontal axis. Replaces the skier sprites with blue versions. Replaces the default tree sprites with classic versions. Dramatically increases the number of flags on the mountain. Recolors all flags to green. Places the gates dead center and creates a corridor of trees. Inverts the colors of the slalom flags. Flips the orientation or state of the slalom flags on the course. Allows the skier to jump over moguls using the FIRE action. Transforms all moguls into solid trees. Causes the flags to slide horizontally across the screen while you ski. Randomizes the horizontal placement of every flag. Changes the reward function to give +1 reward for passing a gate, Allows the skier to accelerate beyond max. speed using DOWN. Replaces the standard moguls with thinner ones.

74

F.37.3

Selection of Available Visual Modifications

Figure 32: A selection of visual modifications available for Skiing, with each frame showing a different mod applied to the base environment.

75

F.38 F.38.1

SpaceInvaders SpaceInvaders details

The player operates a laser cannon tasked with shooting down waves of descending space invaders before they reach Earth. The game ends either when all lives are lost to enemy fire or when the invaders successfully reach the ground.

F.38.2

Available Modifications for SpaceInvaders

Modification

Effect

controllable_missile disable_shield_left disable_shield_middle disable_shield_right no_danger shift_shields

Forces the fired missile’s horizontal position to match the player’s tank. Erases the left bunker from the screen. Erases the middle bunker from the screen. Erases the right bunker from the screen. Removes all player shields and neutralizes incoming enemy projectiles. Teleports all bunkers to new horizontal positions.

76

F.39 F.39.1

Spacewar Spacewar details

The player pilots a Star Ship through a series of dueling games, aiming to land as many shots on the opponent as possible while avoiding being hit in return.

F.39.2

Available Modifications for Spacewar

Modification

Effect

No modifications implemented yet.

77

F.40

F.40.1

Tennis

Tennis details

The player controls the orange player competing against a computer-controlled blue opponent under standard tennis rules. A set is won by reaching at least six games with a margin of two. If the score reaches 6-6, the first player to pull two games ahead takes the match.

F.40.2

Available Modifications for Tennis

Modification

Effect

clay_court fast_enemy fast_player grass_court grayscale hard_court high_bounce inverted_colors lazy_enemy night_mode random_ball_speed random_walk_speed super_gravity

Changes the court colors to a clay court aesthetic. Enemy moves faster. Increases player walk speed. Changes the court colors to a grass court aesthetic. Makes the court grayscale. Changes the court colors to a hard court aesthetic. Increases ball bounce velocity. Swaps the player and opponent colors. Enemy moves slower. Darkens the court for a night-time aesthetic. Ball has random speed after every hit. Player has random walk speed everytime they start moving. Increases ball gravity.

78

F.40.3

Selection of Available Visual Modifications

Figure 33: A selection of visual modifications available for Tennis, with each frame showing a different mod applied to the base environment.

79

F.41 F.41.1

Tetris Tetris details

The player directs a sequence of falling tile-based shapes, maneuvering and rotating them to fill horizontal rows on the board without leaving any gaps.

F.41.2

Available Modifications for Tetris

Modification

Effect

No modifications implemented yet.

80

F.42 F.42.1

Timepilot Timepilot details

The player pilots an aircraft through a series of increasingly challenging encounters, facing enemies whose technology grows progressively more advanced and futuristic as the game progresses.

F.42.2

Available Modifications for Timepilot

Modification

Effect

No modifications implemented yet.

81

F.43 F.43.1

Tron Tron details

The player wields a flying saucer as a weapon, using it to knock out opponents that close in from all directions before they can reach the player.

F.43.2

Available Modifications for Tron

Modification

Effect

No modifications implemented yet.

82

F.44 F.44.1

Turmoil Turmoil details

The player battles waves of descending aliens, shooting them down while weaving through their projectiles and avoiding direct collisions. Enemy formations grow more aggressive and faster as waves are cleared, demanding increasingly precise movement and quick reflexes to survive.

F.44.2

Available Modifications for Turmoil

Modification

Effect

No modifications implemented yet.

83

F.45

F.45.1

Venture

Venture details

The player navigates a dungeon made up of multiple chambers, each containing treasure to collect and monsters to defeat. Armed with a bow, the player must clear each room of its monsters and claim the treasure within, while also avoiding the Hallway Lurkers that roam the corridors between chambers. Lingering too long in a room causes the monsters to become increasingly aggressive.

F.45.2

Available Modifications for Venture

Modification

Effect

blood_moon fast_enemies grayscale inverted_colors matrix_theme night_mode remove_overworld_mobs reward_for_kill slow_enemies

A dark red themed mod. Increases the speed of all enemies. Turns the entire game into grayscale. Inverts all colors in the game. A Matrix-themed mod: black background, green walls/monsters. Dims the entire screen by 50% for a night mode experience. Removes all monsters from the main map. Adds a reward bonus of 100 per kill in the rooms. Reduces the speed of all enemies.

84

F.45.3

Selection of Available Visual Modifications

Figure 34: A selection of visual modifications available for Venture, with each frame showing a different mod applied to the base environment.

85

F.46 F.46.1

VideoCube VideoCube details

The player guides Hubie around the faces of a Rubik’s cube, swapping adjacent tiles with one another rather than rotating entire rows or columns. The goal is to arrange each face of the cube so that it consists entirely of a single color, offering a distinctly unconventional take on the classic puzzle.

F.46.2

Available Modifications for VideoCube

Modification

Effect

No modifications implemented yet.

86

F.47 F.47.1

VideoPinball VideoPinball details

The player controls a digital pinball table, using flippers to keep the ball in play as long as possible while racking up points. Targets, bumpers, and other classic pinball elements reward precise flipper control and strategic ball placement.

F.47.2

Available Modifications for VideoPinball

Modification

Effect

constant_ball_dynamics low_tilt_effect never_activate_tilt_mode no_scoring_cooldown

Removes most of the ability to change the trajectory. Decreases the tilt effect when nudging. Prevent entering tilt mode by resetting the tilt counter. Remove the cooldown period after hitting a target.

87

F.48 F.48.1

WordZapper WordZapper details

The player shoots letters scrolling across the screen, racing against the clock to hit them in a specific prescribed order. Speed and accuracy are both essential as the pace increases over time.

F.48.2

Available Modifications for WordZapper

Modification

Effect

No modifications implemented yet.

88

Related documents

• LOLLA: Deep Reinforcement Learning for Closed-Loop Link Adaptation Towards a GPU-Accelerated AI-RAN#299860
• Orak: A Foundational Benchmark for Training and Evaluating LLM Agents on Diverse Video Games#144303
• Molt: A Scalable PyTorch-Native Training Framework for Agentic Reinforcement Learning#406536
• Molt: A Scalable PyTorch-Native Training Framework for Agentic Reinforcement Learning#405635
• Official Review · ICLR 2025#787025