Installation

Install

Create a virtual environment and install JAXAtari:

python3 -m venv .venv
source .venv/bin/activate

python3 -m pip install -U pip
pip3 install -e .

Note

This will install JAX without GPU acceleration.

GPU Support (CUDA)

CUDA users should run the following to add GPU support:

pip install -U "jax[cuda12]"

For other accelerator types, please follow the instructions in the JAX installation guide.

Sprite Assets

Next, you need to download the original Atari 2600 sprites. Before downloading, you will be asked to confirm ownership of the original ROMs:

.venv/bin/install_sprites