File size: 3,424 Bytes
2fa05eb faea3aa 2fa05eb faea3aa 2fa05eb faea3aa 2fa05eb faea3aa 2fa05eb faea3aa 2fa05eb faea3aa 2fa05eb | 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: pytorch
license: apache-2.0
tags:
- world-model
- diffusion
- atari
- space-invaders
- video-generation
model-index:
- name: space-invaders-world-model
results:
- task:
type: other
name: Next-frame rollout
dataset:
name: Space Invaders, held-out random play
type: ale/space-invaders
metrics:
- type: psnr_changed_1
value: 18.17
name: PSNR on changed pixels at 1 frames
- type: psnr_changed_5
value: 18.35
name: PSNR on changed pixels at 5 frames
- type: psnr_changed_15
value: 16.64
name: PSNR on changed pixels at 15 frames
---
# space-invaders-world-model
A diffusion world model of Atari Space Invaders, trained from scratch. Given the last 4 frames and actions, a small U-Net denoises the next 64×64 frame; fed its own frames, it generates a video that your actions steer. After [DIAMOND](https://arxiv.org/abs/2405.12399). `atari.py` in this repo collects the data, trains, plays and doubles as the loader.
Left: real game. Right: the model, from the same four starting frames and the same actions, never seeing a real frame again.
## 🚀 Usage
Play it in the browser. Each button press generates one frame:
```
uv run "$(hf download jgalego/space-invaders-world-model atari.py --quiet)" play --game SpaceInvaders
```
From Python, with `atari.py` on the path:
```python
import torch
from atari import DEVICE, collect, load, rollout, to_float
model = load("jgalego/space-invaders-world-model")
frames, _, _ = collect(4, seed=0, game="SpaceInvaders") # four real frames to start from
context = to_float(frames[:4])[None].to(DEVICE)
moves = torch.tensor([[0, 0, 0] + [1] * 16], device=DEVICE) # FIRE for 16 frames
video = rollout(model, context, moves) # (1, 16, 3, 64, 64) in [-1, 1]
```
Actions: 0 NOOP, 1 FIRE, 2 RIGHT, 3 LEFT, 4 RIGHTFIRE, 5 LEFTFIRE.
## 🏋️ Training
| | |
|---|---|
| Model | U-Net, channels 64, 128, 256, 4.2M parameters; noise level and the last 4 actions modulate every block |
| Diffusion | EDM preconditioning, σ_data 0.5, log-normal training noise; 3 Euler steps per frame from σ=5.0 |
| Data | 200,000 frames (382 episodes) of a random policy that holds each uniformly random action for 1 to 6 steps; frameskip 4, no sticky actions |
| Steps | 50,000, batch size 64 |
| Optimizer | AdamW, lr 0.0001, gradient clipping 1.0 |
| Hardware | NVIDIA A10G, 55.9 min |
## 📊 Results
Rollouts on 256 windows from episodes whose seeds never appear in training. The model gets four real frames and the real actions, then only its own frames. **Repeat** keeps showing the last real frame. Much of an Atari frame never changes, so **changed pixels** scores only the pixels where the real or the predicted frame differs from the last real frame.
| Frames ahead | Model, all pixels (dB) | Repeat, all pixels (dB) | Model, changed pixels (dB) | Repeat, changed pixels (dB) |
|---|---|---|---|---|
| 1 | 35.43 | 36.56 | 18.17 | 16.0 |
| 5 | 31.07 | 30.94 | 18.35 | 15.9 |
| 15 | 27.33 | 26.77 | 16.64 | 14.54 |
## ⚠️ Limitations
- A random policy does not get far, so the model knows the early game best.
- 64×64 frames: small sprites are a pixel or two and can blur or vanish over long rollouts.
- Errors compound: every generated frame becomes input for the next. |