Flappy Bird world model (DFWorld, 172.5M params)
Nothing here is the game engine. Every frame is drawn by the model from its own previous 15 frames and the key pressed (an auto-pilot presses the keys). One dream from the stage-2 checkpoint, 600 frames (every 2nd frame shown, at 15 frames per second).
A game world model: given past 128x128 frames and the key pressed, it draws the next frame.
Diffusion transformer (16 blocks, width 768, 12 heads, patch 8 -> 256 tokens per frame, adaLN-Zero, RoPE over time)
trained with rectified flow and two-copy teacher forcing. Rebuilt from scratch in flappy-world-model.
| folder | what | init | steps |
|---|---|---|---|
stage1/ |
base training, EMA weights | scratch | 29525 |
stage2/ |
self-forcing (2 frames redrawn by the EMA) + bird x5 in the loss, EMA weights | stage 1 EMA | 17181 |
stage3/ |
as stage 2, plus 1 in 4 windows ends on a frame where a new pipe appears, EMA weights | stage 2 EMA | 10405 |
Each folder: ema.pt (plain state_dict), config.yaml, log.jsonl (training log), checks.json (play checks on held-out games).
Play checks (held-out games)
| model | pipe after 4 own frames | wall frames | free-dream survival (median frames) | bird error 1-100 (px) | bird visible 101-200 |
|---|---|---|---|---|---|
| stage 1 | 84.4% | 0.84% | 118 | 0.135 | 6% |
| stage 2 | 80.1% | 0.01% | 196 | 0.117 | 18% |
| stage 3 | 97.3% | 0.12% | 202 | 0.149 | 16% |
Stage 2 keeps the bird and the world alive much longer in free play, but new pipes after the model's own frames dip.
Stage 3 trains on exactly that moment and fixes it: a new pipe appears after 4 of the model's own frames 97% of the time.
Use stage3/ema.pt unless you want a specific stage.
Load
import torch
from fwm.model.world import WorldModel
m = WorldModel(); m.load_state_dict(torch.load("stage3/ema.pt", map_location="cpu")); m.eval()