Metadata-Version: 2.5
Name: JAXtari
Version: 0.2.0
Summary: GPU-accelerated, object-centric Atari environments for reinforcement learning with JAX.
Project-URL: Homepage, https://github.com/k4ntz/Jaxtari
Project-URL: Repository, https://github.com/k4ntz/Jaxtari
License: MIT License
        
        Copyright (c) 2025 Quentin Delfosse, Jannis Blüml
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
License-File: LICENSE
Requires-Python: >=3.10
Requires-Dist: absl-py>=2.3
Requires-Dist: ale-py>=0.11.1
Requires-Dist: chex>=0.1.87
Requires-Dist: flax
Requires-Dist: gymnasium>=1.2.0
Requires-Dist: gymnax>=0.0.8
Requires-Dist: jax
Requires-Dist: ml-dtypes
Requires-Dist: numpy
Requires-Dist: opt-einsum>=3.4.0
Requires-Dist: platformdirs>=4.5.1
Requires-Dist: requests>=2.32.5
Requires-Dist: scipy>=1.15.3
Requires-Dist: toolz>=1.0.0
Requires-Dist: typing-extensions>=4.14.0
Provides-Extra: dev
Requires-Dist: gymnasium[other]==1.2.0; extra == 'dev'
Requires-Dist: pygame==2.5.0; extra == 'dev'
Requires-Dist: pytest; extra == 'dev'
Requires-Dist: syrupy==4.9.1; extra == 'dev'
Provides-Extra: gh-ci
Requires-Dist: opencv-python-headless; extra == 'gh-ci'
Requires-Dist: pygame==2.5.0; extra == 'gh-ci'
Requires-Dist: pytest; extra == 'gh-ci'
Requires-Dist: pytest-github-actions-annotate-failures; extra == 'gh-ci'
Requires-Dist: pytest-sugar; extra == 'gh-ci'
Requires-Dist: pytest-xdist; extra == 'gh-ci'
Requires-Dist: syrupy==4.9.1; extra == 'gh-ci'
Provides-Extra: training
Requires-Dist: hydra-core>=1.3.2; extra == 'training'
Requires-Dist: omegaconf>=2.3.0; extra == 'training'
Requires-Dist: rtpt>=0.0.4; extra == 'training'
Requires-Dist: safetensors>=0.7.0; extra == 'training'
Requires-Dist: tensorboard; extra == 'training'
Requires-Dist: torch; extra == 'training'
Requires-Dist: tyro; extra == 'training'
Requires-Dist: wandb[media]>=0.24.0; extra == 'training'
Description-Content-Type: text/markdown

# 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](https://www.aiml.informatik.tu-darmstadt.de/)

[Citation](#citation) • [Features](#features) • [Installation](#installation) • [Quick Start](#quick-start) • [Wrappers](#wrapper-reference) • [Environments](#available-environments) • [Contributing](#contributing) • [License](LICENSE)


**JAXtari** is a GPU-accelerated, object-centric Atari environment framework powered by [JAX](https://github.com/google/jax). Inspired by [OCAtari](https://github.com/k4ntz/OC_Atari), 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](https://github.com/k4ntz/HackAtari), it also supports game modifications for testing agent generalization.

---


<div class="collage">
  <div class="row" align="center">
    <img src="./docs/source/_static/gifs/pong.gif" alt="Pong" width="24%">
    <img src="./docs/source/_static/gifs/beamrider.gif" alt="Beamrider" width="24%">
    <img src="./docs/source/_static/gifs/phoenix.gif" alt="Phoenix" width="24%">
    <img src="./docs/source/_static/gifs/tennis.gif" alt="Tennis" width="24%">
  </div>
  <div class="row" align="center">
    <img src="./docs/source/_static/gifs/skiing.gif" alt="Skiing" width="24%">
    <img src="./docs/source/_static/gifs/montezumarevenge.gif" alt="Montezuma" width="24%">
    <img src="./docs/source/_static/gifs/seaquest.gif" alt="Seaquest" width="24%">
    <img src="./docs/source/_static/gifs/kangaroo.gif" alt="Kangaroo" width="24%">
  </div>
  <div class="row" align="center">
    <img src="./docs/source/_static/gifs/freeway.gif" alt="Freeway" width="24%">
    <img src="./docs/source/_static/gifs/venture.gif" alt="Venture" width="24%">
    <img src="./docs/source/_static/gifs/qbert.gif" alt="Qbert" width="24%">
    <img src="./docs/source/_static/gifs/frostbite.gif" alt="Frostbite" width="24%">
  </div>
  <div class="row" align="center">
    <img src="./docs/source/_static/gifs/bankheist.gif" alt="Bankheist" width="24%">
    <img src="./docs/source/_static/gifs/mspacman.gif" alt="Ms. PacMan" width="24%">
    <img src="./docs/source/_static/gifs/gravitar.gif" alt="Gravitar" width="24%">
    <img src="./docs/source/_static/gifs/enduro.gif" alt="Enduro" width="24%">
  </div>
</div>

---

## 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

```bash
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:

```bash
pip install -e ".[dev]"
```

### With training scripts

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

```bash
pip install -e ".[training]"
```

### GPU acceleration (CUDA)

```bash
pip install -U "jax[cuda12]"
```

For other accelerators see the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html).

### 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:

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

---

## Quick Start

### Basic environment creation

```python
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:

```python
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.

```python
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

```python
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.

```python
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):

```python
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

```bash
# 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]
```


| Wrapper                       | Description                                                                                                                            |
|-------------------------------|----------------------------------------------------------------------------------------------------------------------------------------|
| `AtariWrapper`                | Atari-specific pre-processing: sticky actions, episodic life, noop reset, frame-skip config. Must come before any observation wrapper. |
| `ObjectCentricWrapper`        | Stacked object-centric features. Output shape: `(frame_stack, features)`.                                                              |
| `PixelObsWrapper`             | Stacked pixel frames with max-pooling. Output shape: `(frame_stack, H, W, C)`.                                                         |
| `PixelAndObjectCentricWrapper` | Both pixel and object-centric observations as a tuple.                                                                                 |
| `PixelAndObjectObsWrapper`    | Same as above but returns structured (non-flattened) OC observations.                                                                  |
| `FlattenObservationWrapper`   | Flattens any observation pytree to a single 1D array.                                                                                  |
| `NormalizeObservationWrapper` | Normalizes observations to `[0, 1]` (or `[-1, 1]` with `to_neg_one=True`). Compatible with any pytree structure.                       |
| `LogWrapper`                  | Tracks episode returns and lengths.                                                                                                    |
| `MultiRewardWrapper`          | Computes multiple reward functions at every step. Apply before `AtariWrapper`.                                                         |
| `MultiRewardLogWrapper`       | Tracks multiple reward components separately. Use with `MultiRewardWrapper`.                                                           |
| `ContinuousActionWrapper`     | Converts a 3D continuous action (r, theta, fire) into discrete Atari actions. Requires full_action_space=True in AtariWrapper. Place this wrapper after all other wrappers (i.e., at the outermost layer)        |

---

## 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](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](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

```bibtex
@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](LICENSE) for details.
