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 —
AI/ML Lab, TU Darmstadt

Citation • Features • Installation • Quick Start • Wrappers • Environments • Contributing • License

JAXtari is a GPU-accelerated, object-centric Atari environment framework powered by JAX. Inspired by OCAtari, it enables training agents with 100M steps in under 1 hour (pixel-based observations) or under 15 minutes (object-centric observation) through JIT compilation, vectorization, and full GPU parallelization — while exposing structured, object-centric observations alongside standard pixel inputs. Similar to HackAtari, it also supports game modifications for testing agent generalization.

Features

• Object-centric observations — structured game state with per-object positions, types, and attributes
• Full GPU pipeline — end-to-end JAX with JIT compilation, vmap, and lax.scan; no CPU/GPU transfer bottlenecks
• Comprehensive wrapper system — pixel, object-centric, combined, normalized, flattened — all composable
• Game modifications — pre-built mods and a clean API for custom distribution shifts

Installation

Basic

python3 -m venv .venv
source .venv/bin/activate
pip install -U pip
pip install -e .

With development tools (tests + manual play)

Includes pytest, pygame, and testing extras:

pip install -e ".[dev]"

With training scripts

Includes wandb, tensorboard, hydra, and other training dependencies:

pip install -e ".[training]"

GPU acceleration (CUDA)

pip install -U "jax[cuda12]"

For other accelerators see the JAX installation guide.

Download sprites

Before running any environment for the first time you will be asked to confirm ROM ownership of the original Atari ROMs. This is necessary to download sprites that look similar to the original ALE sprites.

If you do not have ownership of the original Atari ROMs, you can continue with replacement/custom sprites. In that case, please decline the ownership and the installer will download the alternative sprites package. You can also use your own sprites by placing them in the ~/.local/share/jaxtari/sprites directory (legacy installs under ~/.local/share/jaxatari/sprites are still detected).

Sprite packs ship a .version file. When a newer pack is required, Jaxtari will ask once whether to download it (opt-in). You can always refresh manually with:

python3 -m jaxtari.install_sprites
# or
.venv/bin/install-sprites

Quick Start

Basic environment creation

import jax
import jaxtari

env = jaxtari.make("pong")

# List all available games
print(jaxtari.list_available_games())

Game modifications

JAXtari ships with pre-built modifications for testing generalization:

import jaxtari

# Single mod
env = jaxtari.make("pong", mods=["lazy_enemy"])

# Multiple mods simultaneously
env = jaxtari.make("pong", mods=["lazy_enemy", "shift_enemy"])

Applying wrappers

Wrappers must be applied in order: AtariWrapper first, then an observation wrapper, then optional utility wrappers.

import jaxtari
from jaxtari.wrappers import (
    AtariWrapper,
    ObjectCentricWrapper,
    PixelObsWrapper,
    PixelAndObjectCentricWrapper,
    FlattenObservationWrapper,
    NormalizeObservationWrapper,
    LogWrapper,
)

base_env = jaxtari.make("pong")
atari_env = AtariWrapper(base_env)

# Choose one observation type:
env = ObjectCentricWrapper(atari_env, frame_stack_size=4, frame_skip=4)          # shape: (frame_stack, features)
# env = PixelObsWrapper(atari_env)             # shape: (frame_stack, H, W, C)
# env = PixelAndObjectCentricWrapper(atari_env) # both

# Optional: flatten to 1D
env = FlattenObservationWrapper(env)

# Optional: normalize observations to [0, 1]
env = NormalizeObservationWrapper(env)

# Optional: track episode returns and lengths
env = LogWrapper(env)

Vectorized stepping

import jax
import jaxtari
from jaxtari.wrappers import AtariWrapper, ObjectCentricWrapper, FlattenObservationWrapper

env = FlattenObservationWrapper(ObjectCentricWrapper(AtariWrapper(jaxtari.make("pong"))))

n_envs = 1024
rng = jax.random.PRNGKey(0)
reset_keys = jax.random.split(rng, n_envs)

# Initialise n_envs parallel environments
obs, env_state = jax.vmap(env.reset)(reset_keys)

# Single parallel step
action = jax.random.randint(rng, (n_envs,), 0, env.action_space().n)
obs, env_state, reward, terminated, truncated, info = jax.vmap(env.step)(env_state, action)

# 100 steps with scan
def step_fn(carry, _):
    obs, state = carry
    new_obs, new_state, reward, terminated, truncated, info = jax.vmap(env.step)(state, action)
    return (new_obs, new_state), (reward, terminated, truncated, info)

_, (rewards, terminations, truncations, infos) = jax.lax.scan(
    step_fn, (obs, env_state), None, length=100
)

Gymnasium compatibility (WIP)

Note: This wrapper is currently work in progress and supports interoperability with CPU-based Gymnasium pipelines (e.g. stable-baselines3). It currently only exposes pixel observations and does not accept JAXtari wrappers. For JAX-native training use the wrapper stack above instead.

from jaxtari.gym_wrapper import GymnasiumJaxtariWrapper
import jaxtari

base_env = jaxtari.make("pong")
gym_env = GymnasiumJaxtariWrapper(base_env)

obs, info = gym_env.reset()
obs, reward, terminated, truncated, info = gym_env.step(gym_env.action_space.sample())

Multiple reward functions

Use MultiRewardWrapper to compute several reward signals in parallel (apply it directly after the base environment, before any other wrapper):

import jaxtari
from jaxtari.wrappers import MultiRewardWrapper, AtariWrapper, ObjectCentricWrapper, MultiRewardLogWrapper

def survival_reward(prev_state, state):
    return 1.0  # reward every surviving step

def score_delta(prev_state, state):
    return state.score - prev_state.score

base_env = jaxtari.make("pong")
env = MultiRewardWrapper(base_env, reward_funcs=[survival_reward, score_delta])
env = ObjectCentricWrapper(AtariWrapper(env))
env = MultiRewardLogWrapper(env)

Manual play

# requires the [dev] extra (pygame)
python3 scripts/play.py -g Pong --mods lazy_enemy

Wrapper Reference

All wrappers live in src/jaxtari/wrappers.py. The standard stack is:

base env  →  [MultiRewardWrapper]  →  AtariWrapper  →  <obs wrapper>  →  [utility wrappers]

Project Structure

Jaxtari/
├── src/jaxtari/
│   ├── core.py              # make() factory, game and mod registries
│   ├── environment.py       # JaxEnvironment base class
│   ├── wrappers.py          # all wrappers
│   ├── modification.py      # mod system (plugins, conflict detection)
│   ├── gym_wrapper.py       # Gymnasium compatibility adapter
│   ├── renderers.py         # JAX rendering utilities
│   ├── spaces.py            # action/observation space definitions
│   ├── install_sprites.py   # sprite download script
│   └── games/
│       ├── jax_<game>.py    # one file per environment
│       └── mods/
│           ├── <game>_mods.py          # mod controller per game
│           └── <game>/
│               └── <game>_mod_plugins.py  # individual mod plugin classes
├── scripts/                 # see scripts/README.md for a full description
│   ├── play.py              # interactive human play
│   └── benchmarks/          # PPO/PQN training and evaluation scripts
├── tests/                   # pytest test suite
├── docs/                    # Sphinx documentation source
├── games_covered.md         # full environment status table
└── pyproject.toml

Contributing

Contributions are welcome! See CONTRIBUTING.md for detailed guides on adding mods, environments, and wrappers. Quick overview below.

Adding a new environment

1. Create src/jaxtari/games/jax_<game>.py implementing JaxEnvironment
2. Register it in GAME_MODULES in src/jaxtari/core.py
3. Add a test in tests/games/
4. Update your game's status in games_covered.md

Adding a mod

1. Create src/jaxtari/games/mods/<game>/<game>_mod_plugins.py with your plugin class(es) extending JaxtariInternalModPlugin or JaxtariPostStepModPlugin
2. Create or update src/jaxtari/games/mods/<game>_mods.py — add your mod key to the REGISTRY dict
3. Register the controller in MOD_MODULES in src/jaxtari/core.py (if not already present)

Adding a wrapper

1. Subclass JaxtariWrapper in src/jaxtari/wrappers.py
2. Implement reset(), step(), and observation_space() / action_space()
3. Export it from src/jaxtari/__init__.py

General

1. Fork this repository
2. Create a feature branch: git checkout -b feature/my-feature
3. Commit your changes and open a pull request

Feel free to share new mods or environments by opening a PR!

Citation

@misc{jaxatari2026,
  author = {Delfosse, Quentin and Emunds, Raban and Seitz, Paul and Wette, Sebastian and Kirn, Daniel and Mandok, Dominik and Bl{\"u}ml, Jannis and Kersting, Kristian},
  title = {JAXtari: High-Throughput and Easy-to-Modify Arcade Learning Environment},
  year = {2026},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {https://github.com/k4ntz/Jaxtari/},
}

License

This project is licensed under the MIT License — see LICENSE for details.