Dreams the game. Learns to walk in the dream. Plays for real.
A complete walkthrough: the 3D map reconstruction, the diffusion world model's imagination, the System 1 / System 2 game harness with its LLM planner, and the Boulder Badge won in the real game.
PokéDreamer is a research project on world models for game agents - teaching an agent to imagine Pokémon Red, train skills inside that imagination, and measure where the dream matches the emulator's memory.
From symbolic state probes to pixel-level imagination - three generations of world models.
The first world model - a Variational Autoencoder compressing frames to 32-dim latents, with a GRU dynamics model predicting the next latent. An MPC planner imagines action sequences and picks the best path using linear probes for map position.
The major leap - native 160×144 resolution, a 4-layer Residual CNN encoder, and a full Recurrent State-Space Model with 32×32 discrete categorical latents trained via Gumbel-Softmax. The RSSM's prior enables pure imagination rollouts without any emulator steps.
A DIAMOND-style diffusion world model of Pokémon Red, and an agent whose navigation was learned entirely inside it. Ten navigation skills trained in imagination carry the agent to the Boulder Badge in the real game in 30 of 30 runs; battles and goal choice are rules, and six local LLMs were tested as the planner (System 2 can call tools: Pokédex and guide search over Bulbapedia via lab/agent/web_search.py). Whether battles can also be trained inside the world model was measured: not with this model. Next: Path B, battles learned inside the world model (estimated ~50 GPU-hours, planned).
def raw(self, x_noisy, sigma, ctx, actions, aug_bucket):
"""Network output F and EDM scalings; x_noisy (B,1,H,W), sigma (B,), ctx (B,CTX,H,W), actions (B,CTX)."""
s = sigma[:, None, None, None]
c_in = 1 / (s ** 2 + SIGMA_DATA ** 2).sqrt()
c_skip = SIGMA_DATA ** 2 / (s ** 2 + SIGMA_DATA ** 2)
c_out = s * SIGMA_DATA / (s ** 2 + SIGMA_DATA ** 2).sqrt()
cond = self.noise_emb(sigma.log() / 4) + self.act_emb(actions).flatten(1) + self.aug_emb(aug_bucket)
out = self.unet(torch.cat([x_noisy * c_in, ctx], 1), self.cond_mlp(cond))
return out, c_skip, c_outQuantitative results from training and evaluation across all three versions.
Empirical findings grounded in Game Boy RAM state across the diffusion world model, imagination-trained navigation policies, start-tile sweeps, battle fidelity, and local LLM planners.
| Scope / Task | Result | Page |
|---|---|---|
| Boulder Badge, real game, 10-leg agent | 30/30 runs (3 seeds × 10) | v3_gym |
| Runs where every navigation leg is a world-model policy (no fallback anywhere) | 14/30 | v3_gym |
| Navigation legs at 100/100 on all 3 seeds | 7 of 10; to_pewter_pc 100/0/100, forest 100/69/97 | v3_gym |
| Start-tile robustness after fixing a reward hack (move gate) | 2,339/2,340 episodes over 26 start tiles | v3_tile_sweep |
| Can battles be trained inside this world model? | No: imagined battles end early in 74.6% of held-out windows (real-frame probe: 1.1%) | v3_battle_fidelity |
| Local LLMs as System 2 | 5 of 6 reach 10/10 badges | v3_system2 |
The game harness cleanly separates high-level planning from low-level execution. System 2 selects the next sub-goal from verified RAM state and can call tools: Pokédex and guide search over Bulbapedia (lab/agent/web_search.py). Every decision is validated against game memory before execution, and System 1 executes using navigation policies trained entirely in the world model's imagination.
Measured over 15,536 held-out trajectories of length 29 steps. Scheduled sampling (SS) is critical: the SS model's compounding drift stays roughly flat (3.71–4.38 tiles) out to 29 imagined steps, while the pure teacher-forcing (TF) ablation reaches 11.63 tiles - a 2.7× degradation.
| Rollout Step | SS Latent MSE | TF Latent MSE | SS Tile Error | TF Tile Error | Improvement |
|---|---|---|---|---|---|
| Step 1 | 0.10187 | 0.13640 | 4.38 tiles | 4.34 tiles | 0.99× |
| Step 5 | 0.11551 | 0.25249 | 3.81 tiles | 5.11 tiles | 1.34× |
| Step 10 | 0.13915 | 0.46385 | 3.83 tiles | 6.98 tiles | 1.82× |
| Step 15 | 0.15572 | 0.69037 | 3.92 tiles | 8.94 tiles | 2.28× |
| Step 20 | 0.17057 | 0.87813 | 4.02 tiles | 10.36 tiles | 2.58× |
| Step 25 | 0.18885 | 0.99588 | 4.17 tiles | 10.81 tiles | 2.59× |
| Step 29 | 0.20115 | 1.06557 | 4.25 tiles | 11.63 tiles | 2.73× |
4 epochs on 20 NPZ files (~16,000 native-resolution transitions). Batch size 64, sequence length 15. Reconstruction loss steadily decreases - the best checkpoint (epoch 4, val recon = 0.1003) demonstrates pixel-level world modeling at native Game Boy resolution.
| Epoch | Train Loss | Train Recon | Train KL | Val Loss | Val Recon | Val KL |
|---|---|---|---|---|---|---|
| 1 | 0.1476 | 0.1379 | 0.0078 | 0.1266 | 0.1256 | 0.0010 |
| 2 | 0.1207 | 0.1144 | 0.0063 | 0.1172 | 0.1110 | 0.0062 |
| 3 | 0.1490 | 0.1068 | 0.0422 | 0.1228 | 0.1142 | 0.0086 |
| 4 | 0.1021 | 0.1015 | 0.0005 | 0.1651 | 0.1003 ★ | 0.0648 |
How the world model evolved across versions - from continuous latents and teacher forcing to discrete RSSM with KL-balancing.
| Component | v1 (VAE + GRU) | v2 (Discrete RSSM) |
|---|---|---|
| Resolution | 40×36 pixels (PWhiddy downsampled) | 160×144 pixels (native Game Boy) |
| Encoder | Variational Autoencoder (VAE) | 4-layer Residual CNN → 512-dim embed |
| Latent Space | Continuous R³² (reparameterization) | Discrete 32×32 categorical (1024-dim) |
| Dynamics | Autoregressive GRU (scheduled sampling) | RSSM: GRU h_t (512) + Stochastic s_t |
| Gradient Estimator | Reparameterization trick | Gumbel-Softmax straight-through |
| KL Balancing | Standard ELBO | 80% prior / 20% posterior balancing |
| Imagination | Prior-only rollout | Full posterior + prior with KL balancing |
| Decoders | Pixel decoder only | Pixel + Reward predictor + Continue predictor |
| Controller | Lookahead MPC planner (coordinate probe) | Actor-Critic trained in imagination |
Get PokéDreamer running in minutes. You'll need a legally-obtained Pokémon Red ROM.
git clone https://github.com/xoTEMPESTox/PokeDreamer.git
cd PokeDreamer
conda env create -f environment.yml
conda activate pokemon-rlCopy your legally-obtained Pokemon - Red Version (USA, Europe).gb to the project root. Alternatively, download the pre-collected dataset from Hugging Face.
# Collect data (optional - dataset available on HF)
python legacy/v2/scripts/collect_data.py --episodes 20 --out-dir datapython legacy/v2/scripts/train_rssm.py \
--data-dir data \
--epochs 12 \
--batch-size 64 \
--out-dir legacy/v2/checkpoints/rssm_v2Renders a side-by-side video comparing the real emulator vs. the RSSM's imagined frames.
python legacy/v2/scripts/generate_demo_video_v2.py \
--checkpoint legacy/v2/checkpoints/rssm_v2/best_world_model.pt \
--save-state saves/intro_done.state \
--out-video legacy/v2/checkpoints/rssm_v2/side_by_side_demo_v2.mp4RSSM v2 best checkpoint (best_world_model.pt), v1 VAE and dynamics checkpoints - all on Hugging Face.
20 NPZ files of native-resolution (160×144) gameplay transitions with full RAM state annotations. ~340MB total.
HuggingFace Dataset →Full Python source - models, dataset loader, game state extractor, training scripts, and demo video generator.
GitHub Repository →Detailed documentation of every module: models.py, dataset.py, game_state.py, ram_addresses.py, and all scripts.