Upload folder using huggingface_hub

README.md

@@ -1,20 +1,40 @@
 ---
 library_name: pytorch
 license: apache-2.0
-model_name: space-invaders-world-model
 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.
 
+![Real frames (left) and the model's rollout from the same start and actions (right)](results/rollout.gif)
 
+Left: real game. Right: the model, from the same four starting frames and the same actions, never seeing a real frame again.
 
 ## 🚀 Usage
 
@@ -41,13 +61,24 @@ Actions: 0 NOOP, 1 FIRE, 2 RIGHT, 3 LEFT, 4 RIGHTFIRE, 5 LEFTFIRE.
 
 ## 🏋️ Training
 
-Not trained yet.
+| | |
+|---|---|
+| 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 held-out 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.
+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.
 
-Not evaluated yet.
+| 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