simillusion/iris-atari-boxing

IRIS on Atari Boxing

A demo checkpoint from Simulusion, an open-source PyTorch library that puts world models behind one API. It implements the paper Transformers are Sample-Efficient World Models (Micheli et al., ICLR 2023).

IRIS compresses every frame into 16 discrete tokens with a VQ autoencoder, models the game as a sequence of tokens and actions with a Transformer, and trains its policy only on the Transformer's imagination.

Results

Trained for 300 epochs (60,000 real environment steps) on a 4 GB RTX 3050 laptop GPU. The world model keeps the ring, both boxers and the scoreboard, and the policy has learned to close in and trade punches. In a full round in the real emulator it led 29-21 with 24 seconds left and lost 31-32: level with the game's built-in opponent, not yet better. The paper reaches about +70 with several times more policy training per epoch.

This is a short demo run sized for a laptop GPU, meant to show what the model does. It is not a reproduction of the paper's results. More clips and learning curves are in the gallery.

Usage

pip install simulusion huggingface_hub

from huggingface_hub import snapshot_download
from simulusion.world_model import load_world_model

path = snapshot_download("simillusion/iris-atari-boxing")
wm = load_world_model(path, "iris")  # encode / step / rollout / decode

The same interface works for every world model in Simulusion. To get the full agent:

from simulusion.models.iris_agent import IRISAgent

# The first load also downloads torchvision's VGG16 (528 MB) for the perceptual loss
agent = IRISAgent.from_pretrained(path)

Checkpoints are plain state dicts and load with torch.load(..., weights_only=True). They include optimizer state, so you can continue training from them.

Training

Reproduce this checkpoint from a clone of the repository:

python demos/iris_demo.py train --game ALE/Boxing-v5 --run demos/runs/iris_boxing --until-epoch 300 \
    --autoencoder-steps 200 --transformer-steps 200 --actor-critic-steps 50 --actor-critic-batch 64

Citation

If you use this model, please cite the original paper:

@inproceedings{micheli2023iris,
  title     = {Transformers are Sample-Efficient World Models},
  author    = {Micheli, Vincent and Alonso, Eloi and Fleuret, Fran{\c{c}}ois},
  booktitle = {International Conference on Learning Representations},
  year      = {2023},
  url       = {https://arxiv.org/abs/2209.00588},
}

and Simulusion:

@software{simulusion,
  title  = {Simulusion: a PyTorch library for world models},
  author = {Thakkar, Param},
  year   = {2026},
  url    = {https://github.com/ParamThakkar123/simulusion},
}