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

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

cs.LGOct 16, 2026

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

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

Organizations: Department of Computer Science, Technical University Darmstadt, Germany · Google Intrinsic AI Research · Hessian Center for Artificial Intelligence (hessian.AI) · Centre for Cognitive Science, Darmstadt ·

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 object-centric observations and a flexible modification system, allowing researchers to, e.g. develop neuro-symbolic methods or discover novel shortcuts and subsequently train and adapt their agents rapidly, sometimes in only minutes.

Figures & tables

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 accelerator-native throughput, direct source-level modifications, and native access to object-centric and symbolic state.

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.

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.

Panel of Figure 4 — 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 $\times$ for pixel-based runs compared to ALE-VectorEnv and EnvPool. An in-depth analysis is available in Appendix E .

Panel of Figure 5 — Figure 5: JAXtari allows the detection of shortcuts by easily modifying the environment. Performance change ( $\mathcal{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 $\mathcal{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 $\mathcal{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 .

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.

Appendix figures & tables76 assets

Supplementary material from the paper’s appendix.

Appendix

Mechanic	Operational Specification	Source Implementation
Action Space Coverage
Full Discrete Space	Utilizes full Discrete(18) action space; NOOP , FIRE , directional jumps, and combined inputs are supported.	jax_frostbite.py
Directional Controls	Moving joystick triggers Bailey’s jumps across moving ice floes and between shores.	jax_frostbite.py
FIRE Button	Pressing FIRE reverses the movement direction of the ice row currently occupied.	jax_frostbite.py
Core Game Dynamics
Igloo Construction	Landing on a white ice floe turns it blue and adds one block to the igloo (15 blocks required for completion).	jax_frostbite.py

Appendix

Table 1: Frostbite Action Space and Core Dynamics Validation Checklist. Mapped against the JAXtari implementation and verified against ALE.

Scoring Event	Reward Formula	Source Implementation
Coloring an Ice Block	$10\times\text{current level}$ points awarded upon stepping on a white floe.	jax_frostbite.py
Entering Completed Igloo	$160\times\text{current level}$ bonus points awarded per block upon entering the igloo to clear the level.	jax_frostbite.py
Level Multiplier Cap	Level multiplier increases up to level 8; from level 9 onward, scoring caps at 90 points per ice block.	jax_frostbite.py
Catching a Fish	Flat bonus of $+200$ points upon collision with a swimming fish.	jax_frostbite.py

Appendix

Table 2: Frostbite Reward System Validation Checklist.

Hyperparameter	Value
Total timesteps	$10^{7}$
Parallel environments	8
Rollout length	128 steps
Update epochs	4
Minibatches	4
Minibatch size	$8\times 128/4=256$

Appendix

Table 3: PPO hyperparameters. Identical for pixel-based and object-centric training runs

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.

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.

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.

Figure 10: Performance after 100M steps. PPO reaches over 150k points in Seaquest before truncation by reaching 108000 frames.

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.

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

Panel of Figure 13 — 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.

Panel of Figure 14 — 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.

Panel of Figure 15 — 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.

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

Modification	Effect
aggressive_swarm	Aliens spend almost all their time actively chasing the player.
dont_kill	Punishes for killing and shooting enemies.
end_game	Most eggs are already eaten.
extra_lives	Start with extra lives.
last_egg	Only one egg is remaining.
matrix_theme	A Matrix-themed mod: black background, green walls/enemies/eggs.

Appendix

Table 20

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 22

Modification	Effect
No modifications implemented yet.

Appendix

Table 23

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

Appendix

Table 24

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

Modification	Effect
jets_only	Overrides enemy spawn probabilities to only spawn Jet-type enmies.
no_last_line	Disable the lowest lane so no enemies can spawn there.
random_enemies	Overrides enemy spawn probabilities to make spawning random.
speed_mode_fast	Increases overall speed of the game to 6.
speed_mode_medium	Increases overall speed of the game to 4.
speed_mode_slow	Increases overall speed of the game to 2.

Appendix

Table 26

Modification	Effect
2_police_cars	Replaces 2 banks with police cars. Each bank robbed gives 50 points.
blue_player	Changes the color of the player to light blue.
double_speed	Doubles the player speed.
dynamite_penalty	Changes reward for killing a police car with dynamite (-500).
fuel_for_banks	Augments the player’s fuel when 3 banks have been robbed.
golden_banks	Changes the color of banks to golden.

Appendix

Table 27

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

Modification	Effect
dont_kill	Internal mod that punishes killing and shooting.
double_enemy_speed	Doubles enemy movement speed.
fog_of_war	Hide the upper two thirds of the screen behind opaque fog.
hardcore	Start with one life and never allow extra lives.
mothership_laser	Turn the mothership sequence into a stationary lane laser attack.
same_enemies	Render all enemy types using the white UFO visuals.

Appendix

Table 29

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 31

Modification	Effect
No modifications implemented yet.

Appendix

Table 32

Modification	Effect
ball_color	Changes the balls color to a set color.
ball_drift	Consistently drift the ball to the right.
ball_gravity	Pulls the ball down by 1 every step of the gravity_buffer.
big_paddle	Always use a bigger paddle (Size 40).
block_color	Changes the blocks color to a set color.
player_color	Changes the player color to a set color.

Appendix

Table 33

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 35

Modification	Effect
No modifications implemented yet.

Appendix

Table 36

Modification	Effect
dawn	Forces the weather to dawn.
day	Forces the weather to day.
dont_max_speed	Does not reward the player if at max. speed.
filled_road	Displays the full filled road instead of only the side.
fog	Forces the weather to fog.
hud	Renders a HUD that displays sped and x-position.

Appendix

Table 37

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

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

Appendix

Table 39

Modification	Effect
No modifications implemented yet.

Appendix

Table 40

Modification	Effect
_hall_of_fame_start	Spawns the cars to make a "hall of fame" formation.
bikes	Replaces all cars with uniquely colored bike sprites.
black_cars	Makes all cars black by overriding CAR_COLORS constant.
center_cars_on_reset	Positions all cars in the center of the screen when the environment resets.
change_sprites	Bundles frog , bikes , new_lane_colors and green_score
frog	Replaces the player sprites with frog sprites.

Appendix

Table 41

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

Modification	Effect
_aligned_ice	Makes the ice blocks perfectly centered.
_misaligned_ice	Makes the ice blocks vertically aligned with each other but off-center.
_static_ice	Sets ice block horizontal speeds to zero.
bad_fishes_reward	Punishes the player for collecting fish.
change_sprites	Bundles tiger , white_igloo , recolored_obstacles and lightblue_ice .
dark_night	Spawns constant night and makes the sky darker.

Appendix

Table 43

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 45

Modification	Effect
anti_gravity	Reverse gravity from sun, planets, and reactors.
fuel_crisis	Increase fuel consumption rate by 5x.
grayscale	Converts the visual palette to grayscale.
harmless_enemies	Make all enemies harmless by disabling their bullets.
high_speed	Make the ship faster by increasing thrust power and max speed.
hyper_gravity	Increase all gravity from sun, planets, and reactors.

Appendix

Table 46

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 48

Modification	Effect
No modifications implemented yet.

Appendix

Table 49

Modification	Effect
_chickens	Replaces the monkey sprite with a chicken sprite.
_lethal_bell	Post-step mod that kills the player if they touch the bell.
center_ladders	Internal mod to center all ladder positions horizontally.
chains	Replaces the ladder sprites with a chain.
chickens	Bundles _chickens and no_thrown_coconut .
dont_punch	Provides negative reward for punching monkeys.

Appendix

Table 50

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 52

Modification	Effect
No modifications implemented yet.

Appendix

Table 53

Modification	Effect
center_bouncing_skull	Make the rolling skull jump vertically at the center of the screen.
debug_hud	Internal mod to display debug info (Room ID, X, Y) on the HUD.
fast_player	Internal mod to increase player speed.
infinite_amulet	Post-step mod to keep the amulet active forever.
jumping_spiders	Post-step mod to make all spiders jump.
moving_snakes	Post-step mod to make snakes move horizontally in all rooms.

Appendix

Table 54

Modification	Effect
caged_ghosts	Keeps the ghosts caged in the center.
constant_fruits	Always adds a fruit to the maze.
fruit_ghost_bonus	Deactivates points for pellets, But multiplies eating ghost reward.
matrix_theme	A Matrix-themed mod: black background, green elements.
only_1_ghost	Keeps only 1 ghost in the maze.
only_2_ghost	Keeps only 2 ghost in the maze.

Appendix

Table 55

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 57

Modification	Effect
blood_moon	A dark red themed mod.
boss_late_missiles	Make boss missiles appear later.
fast_enemy_bullets	Increases speed of enemy projectiles.
fast_player	Increases player movement speed.
grayscale	Turns the entire game into grayscale.
infinite_lives	Set player lives to 99.

Appendix

Table 58

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

Modification	Effect
change_background_color	Changes the background color to navy blue.
change_player_color	Changes the player color to red.
lazy_enemy	Enemy does not move after returning a shot until player hits ball.
linear_movement	Changes the movement to a linear one (no acceleration).
no_fire	Disables the FIRE action.
random_enemy	Randomizes the enemy movement.

Appendix

Table 60

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

Modification	Effect
collecting_bonus_only	Overrides the rewards to only get reward from collectibles.
dark_pyramid	Changes the pyramid to a dark aesthetic.
diagonal_control	Enables diagonal controls.
grayscale	Makes the entire game grayscale.
ice_pyramid	Changes the pyramid to an icy aesthetic.
inverted_colors	Swaps Qbert’s and Coily’s colors.

Appendix

Table 62

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 64

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

Appendix

Table 65

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 67

Modification	Effect
_dangerous_moguls	Makes colliding with moguls cause the skier to fall.
_more_moguls	Spawns more moguls during the race.
_more_trees	Spawns more trees during the race.
_trees_everywhere	Allows trees to spawn anywhere across the entire horizontal axis.
blue_skiier	Replaces the skier sprites with blue versions.
classic_trees	Replaces the default tree sprites with classic versions.

Appendix

Table 68

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

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

Appendix

Table 70

Modification	Effect
No modifications implemented yet.

Appendix

Table 71

Modification	Effect
clay_court	Changes the court colors to a clay court aesthetic.
fast_enemy	Enemy moves faster.
fast_player	Increases player walk speed.
grass_court	Changes the court colors to a grass court aesthetic.
grayscale	Makes the court grayscale.
hard_court	Changes the court colors to a hard court aesthetic.

Appendix

Table 72

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 74

Modification	Effect
No modifications implemented yet.

Appendix

Table 75

Modification	Effect
No modifications implemented yet.

Appendix

Table 76

Modification	Effect
No modifications implemented yet.

Appendix

Table 77

Modification	Effect
blood_moon	A dark red themed mod.
fast_enemies	Increases the speed of all enemies.
grayscale	Turns the entire game into grayscale.
inverted_colors	Inverts all colors in the game.
matrix_theme	A Matrix-themed mod: black background, green walls/monsters.
night_mode	Dims the entire screen by 50% for a night mode experience.

Appendix

Table 78

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

Modification	Effect
No modifications implemented yet.

Appendix

Table 80

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

Appendix

Table 81

Modification	Effect
No modifications implemented yet.

Appendix

Table 82

GPU-Accelerated RL

See all 15 papers

Octax: Accelerated CHIP-8 Arcade Environments for Reinforcement Learning in JAX
Oct 2, 2025Waris Radji, Thomas Michel, Hector PiteauRL Benchmarks Reinforcement Learning
Univ. Lille, Inria, CNRS, Centrale Lille, UMR 9189-CRIStAL, France · Independent Researcher
JaxMARL: Multi-Agent RL Environments and Algorithms in JAX
Nov 16, 2023Alexander Rutherford, Benjamin Ellis, Matteo Gallici +18RL Benchmarks Reinforcement Learning
University of Oxford · Universitat Politècnica de Catalunya · University College London +2
Mahjax: A GPU-Accelerated Mahjong Simulator for Reinforcement Learning in JAX
May 20, 2026Soichiro Nishimori, Shinri Okano, Keigo Habara +3RL Benchmarks Game-Playing Agents
The University of Tokyo, Tokyo, Japan. · Nara Institute of Science and Technology, Nara, Japan. · Independent Researcher. +5

Oct 2, 2025J/K move · Enter open · S save

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

Waris Radji, Thomas Michel, Hector Piteau

Univ. Lille, Inria, CNRS, Centrale Lille, UMR 9189-CRIStAL, France · Independent Researcher

RL Benchmarks Reinforcement Learning GPU-Accelerated RL RL for Game AI

Open paper AskPDF arXiv

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

Abstract

GPU-Accelerated RL

Explore similar work

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

JaxMARL: Multi-Agent RL Environments and Algorithms in JAX

Mahjax: A GPU-Accelerated Mahjong Simulator for Reinforcement Learning in JAX