File size: 3,197 Bytes
f942573 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | ---
library_name: simulusion
license: mit
pipeline_tag: reinforcement-learning
tags:
- simulusion
- pytorch
- iris
- model-based-rl
- atari
- transformer
- world-model
---
# IRIS on Atari Boxing
A demo checkpoint from [Simulusion](https://github.com/ParamThakkar123/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)](https://arxiv.org/abs/2209.00588)**.
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](https://paramthakkar123.github.io/simulusion/gallery.html).
## Usage
```bash
pip install simulusion huggingface_hub
```
```python
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:
```python
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:
```bash
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:
```bibtex
@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:
```bibtex
@software{simulusion,
title = {Simulusion: a PyTorch library for world models},
author = {Thakkar, Param},
year = {2026},
url = {https://github.com/ParamThakkar123/simulusion},
}
```
|