Dreaming to Dodge — VizDoom take_cover world model + agents
Pretrained weights from Dreaming to Dodge — a from-scratch reproduction of the World Models (Ha & Schmidhuber, 2018) "train the agent inside its own dream" result on VizDoom take_cover, using the IRIS-style recipe (VQ-VAE tokenizer + GPT world model + agents trained purely in imagination), plus a world-model calibration study that goes beyond the paper.
All checkpoints are plain PyTorch state_dicts (the CMA controller is a small dict — see below). The model classes live in the GitHub repo: tokenizer.py, world_model.py, controller.py, modal_apps/train_policy_v6.py.
Files
| File | What it is | Held-out survival (steps) |
|---|---|---|
tokenizer.pt | VQ-VAE tokenizer (64×64×3 → 64 tokens, vocab 512); warmth-weighted recon keeps the fireball (recall 0.953) | — |
world_model_r7.pt | Best world model (calibration round 7): GPT over tokens, predicts next tokens + reward + done; token_acc 0.609, policy-ranking probe fixed | — |
world_model_calibrated.pt | Round-6 world model (first to pass the probe gate, worst |bias| 11.4) | — |
world_model_prepatch.pt | Pre-calibration world model with the no-op blind spot (dream−real +29.8 for standing still) — useful for studying WM exploitation | — |
actor_critic_v6_real.pt | Champion agent: v6 recipe trained in the real env (831k frames) | 131.6 ± 64.7 (beats oracle 98.3) |
actor_critic_v6_dream_r7.pt | v6 agent trained 100% inside world_model_r7.pt — zero real frames | 75.8 ± 39.9 |
actor_critic_v6_dream_cal.pt | v6 agent trained inside the round-6 world model | 61.2 ± 22.9 |
actor_critic_v6_dream.pt | v6 agent trained inside the pre-patch world model — exploits the no-op blind spot (94% no-op) | 52.1 |
controllerc2.pt | Best CMA-ES controller evolved in the dream (dict: theta, thetas, use_ensemble) | 92.7 (79.0 ± 14 adversarial) |
The dream-agent trio is the point of this release: the same agent recipe trained in three world models of increasing calibration quality — 52.1 → 61.2 → 75.8 — against a real-environment twin at 131.6. The gap is the "world-model tax," and the trend is what probe-driven calibration buys.
Loading
import torch, yaml
from tokenizer import Tokenizer, TokenizerConfig
from world_model import WorldModel, WorldModelConfig
cfg = yaml.safe_load(open("configs/doom.yaml"))
tc, wc = cfg["tokenizer"], cfg["world_model"]
tok = Tokenizer(TokenizerConfig(in_channels=3, resolution=64, ch=tc["ch"],
num_tokens=tc["num_tokens"], vocab_size=tc["vocab_size"],
embed_dim=tc["embed_dim"]))
tok.load_state_dict(torch.load("tokenizer.pt", map_location="cpu"))
wm = WorldModel(WorldModelConfig(num_tokens=tc["num_tokens"], vocab_size=tc["vocab_size"],
num_actions=3, max_timesteps=wc["max_timesteps"], embed_dim=wc["embed_dim"],
num_layers=wc["num_layers"], num_heads=wc["num_heads"], dropout=wc["dropout"]))
wm.load_state_dict(torch.load("world_model_r7.pt", map_location="cpu"))
The v6 actor-critics are 3,620-param feature MLPs over Reducer(rows=3, cols=6) features (input 108 = 2×54, hidden 32, 3 actions) — the exact class (FeatAC) is defined in modal_apps/train_policy_v6.py. The CMA controller loads via TinyController + Reducer (see modal_apps/verify_controller.py).
Environment: VizDoom take_cover, 64×64×3 frames, frame_skip 4, 3 actions (no-op/left/right), living reward +1, t_max 525. Note vizdoom requires numpy<2.
Headline results
| Agent | Survival (steps, held-out) |
|---|---|
| v6-real twin (champion) | 131.6 ± 64.7 |
| Heuristic oracle | 98.3 |
| CMA-ES c2 (dream-evolved) | 92.7 |
| MPC in the world model (zero training) | 84.1 ± 16.8 |
| v6-dream on round-7 WM | 75.8 ± 39.9 |
| v6-dream on round-6 WM | 61.2 ± 22.9 |
| v6-dream on pre-patch WM | 52.1 |
Full method, calibration campaign (probe suite, data patching rounds, failure analysis), and reproduction details: project README.