# Atari

TorchWM exposes Atari 2600 environments through Gymnasium and the Arcade Learning Environment (ALE). Atari is the environment family; DIAMOND-style Atari support in TorchWM is an optional preprocessing wrapper/factory for Atari, not a separate environment backend.

## Install

```bash
pip install ale-py
pip install "torchwm[atari]"
```

If ROMs are not already available, install and run AutoROM according to your environment's legal and operational requirements.

## Raw Atari APIs

```python
from torchwm import make_atari_env, list_available_atari_envs

env = make_atari_env("ALE/Pong-v5", obs_type="rgb", frameskip=4)
print(list_available_atari_envs()[:10])
```

`make_atari_env()` registers `ale_py` environments with Gymnasium and delegates to `gymnasium.make()`.

## Raw Atari factory options

| Option | Default | Description |
| --- | --- | --- |
| `env_id` | Required | Gymnasium ID such as `ALE/Pong-v5` or `ALE/Breakout-v5` |
| `obs_type` | `"rgb"` | Use `"rgb"` for image observations or `"ram"` for RAM observations |
| `frameskip` | `4` | Number of frames skipped by ALE between actions |
| `repeat_action_probability` | `0.25` | Sticky-action probability |
| `full_action_space` | `False` | Whether to expose the full Atari action set |
| `max_episode_steps` | `None` | Optional episode cap passed to Gymnasium |

Additional keyword arguments are forwarded to `gymnasium.make()`.

## Atari preprocessing for DIAMOND-style training

This is not a separate environment in TorchWM. It is an Atari preprocessing path implemented by `DiamondAtariWrapper` and `make_diamond_atari_env()` for users who want DIAMOND-compatible Atari rollouts.

```python
from torchwm import make_diamond_atari_env

env = make_diamond_atari_env(
    game="ALE/Breakout-v5",
    frameskip=4,
    max_noop=30,
    terminate_on_life_loss=True,
    reward_clip=True,
    resize=(64, 64),
    seed=0,
)
```

The factory creates the underlying Gymnasium Atari environment with `obs_type="rgb"`, `frameskip=1`, `repeat_action_probability=0.0`, and `full_action_space=False`, then applies Atari-specific preprocessing.

| Preprocessing option | Behavior |
| --- | --- |
| `frameskip` | Repeat each selected discrete Atari action and sum rewards |
| `max_noop` | Run a random no-op reset phase before the episode begins |
| `terminate_on_life_loss` | Mark the step done when an ALE life is lost and lives remain |
| `reward_clip` | Clip summed rewards to `[-1, 1]` |
| `resize` | Resize RGB Atari frames to `(height, width)` |

`DiamondAtariWrapper.step()` returns a legacy four-tuple, `obs, reward, done, info`, while `reset()` returns Gymnasium's `(obs, info)` tuple. The preprocessed observation is an HWC `uint8` RGB frame, so transpose it if your model expects channel-first images.

## Native vectorized Atari

For parallel Atari simulation, use `make_atari_vector_env()`:

```python
from torchwm import make_atari_vector_env

vec_env = make_atari_vector_env(
    game="pong",
    num_envs=8,
    obs_type="rgb",
    frameskip=4,
    seed=0,
)
```

This returns ALE's native `AtariVectorEnv`, not TorchWM's multiprocessing `TorchVectorizedEnv`.

## Model usage

- IRIS and Atari-focused workflows can use ALE IDs discovered by `list_available_atari_envs()`.
- Planet catalog entries include a subset of ALE environments.
- Dreamer `make_env()` does not have a dedicated `env_backend="atari"` branch; wrap Atari as a Gym environment with `env_backend="gym"` if you want Dreamer-style image conversion. Use DIAMOND-style preprocessing only when your Atari training recipe expects that specific preprocessing stack.

## Observation and action contract

The raw ALE factory returns Gymnasium's Atari environment directly. Observation shape and dtype depend on `obs_type`; action space is discrete. DIAMOND-style preprocessing still uses Atari as the environment and returns resized HWC RGB frames. If your model expects `{"image": (3, H, W)}`, wrap raw Atari with `GymImageEnv` or transpose/preprocess DIAMOND-style frames before model input.

## Troubleshooting

- **No environments listed**: ensure `ale-py` is installed and importable.
- **ROM errors**: verify ROM installation for the target games.
- **Different game naming conventions**: Gymnasium IDs usually include `ALE/` and `-v5`; native vectorization takes a game name such as `pong`.
- **Preprocessing mismatch**: use the Atari preprocessing section on this page for frame skip, no-op reset, reward clipping, life-loss termination, and resizing when your training recipe expects DIAMOND-style rollouts.