reverie
An AI that plans by imagining the future.
Not yet
It doesn’t beat random guessing yet, and I traced exactly why: its picture of the world isn’t sharp enough.
The technical write-up
A VQ-VAE encoder, a transformer dynamics model over discrete latents, and cross-entropy-method planning inside it — on a gridworld built to require a model rather than a reactive policy. The headline result is negative, and is reported that way: planning closes only 13.4% of the gap between random and an oracle planner running on the true simulator. The architecture is structured so that “the model is inaccurate” and “the planner is weak” are separately measurable, and they have a definite answer.
- 13.4%
- of the random → oracle gap closed
- 2 / 2
- hypotheses tested, refuted
- Tunnels · 12 × 12 grid · 7 × 7 egocentric window · 80,000 steps/s · PyTorch
- 2026 · World models and search
- Tunnels → VQ-VAE encoder → transformer dynamics → imagined rollout (16 codes) → CEM planner → action
- imagined rollout (16 codes) → sampled futures
- imagined rollout (16 codes) → autoregressive within-frame
The environment is the experiment
A world model only earns its cost when the task defeats the alternatives, so Tunnels is built around three properties. A 7×7 egocentric window on a 12×12 grid, so a reactive policy is insufficient rather than merely worse. Reward gated behind a key and a door, so a random policy solves 0.0% of episodes. Hazards that random-walk, which punish a deterministic model and make rollout fidelity worth measuring at all.
The simulator runs at 80,000 steps per second on one core. That is deliberate: the bottleneck should be the model, not the environment.
The model learns everything it is asked to
| quantity | value |
|---|---|
| encoder occupancy F1 | 0.922 |
| codebook usage / perplexity | 98% / 144 of 256 |
| next-code accuracy | 0.893 |
| reward MSE | 0.0007 |
| termination recall / precision | 1.000 / 1.000 |
Every head converges. Nothing in this table suggests a problem, which is exactly why the rollout measurement had to exist.
And imagined rollouts still degrade
| argmax | 1: 0.635, 2: 0.694, 4: 0.642, 6: 0.599, 8: 0.589, 10: 0.534, 12: 0.53 |
|---|---|
| sampled | 1: 0.618, 2: 0.652, 4: 0.6, 6: 0.591, 8: 0.543, 10: 0.5, 12: 0.448 |
| encoder ceiling | 0.922 |
| horizon | 1 | 2 | 4 | 6 | 8 | 10 | 12 |
|---|---|---|---|---|---|---|---|
| F1 (argmax) | 0.635 | 0.694 | 0.642 | 0.599 | 0.589 | 0.534 | 0.530 |
| F1 (sampled) | 0.618 | 0.652 | 0.600 | 0.591 | 0.543 | 0.500 | 0.448 |
Decoded occupancy F1 against the true simulator under identical action sequences.
The encoder’s own reconstruction ceiling is 0.887–0.922, so one-step prediction already gives up roughly 0.27 of the available fidelity and the rest erodes with horizon. Reward MAE stays at 0.02–0.10 and termination agreement at 0.91–1.00 — the scalar heads are fine. It is state prediction that fails.
Attributing the gap with an oracle
| policy | return | solve rate | s/episode |
|---|---|---|---|
| random | −1.444 | 0.000 | — |
| learned-model MPC | −1.306 | 0.000 | 11.4 |
| oracle MPC (true simulator) | −0.411 | 0.100 | 7.9 |
Planning barely beats random, which admits two readings with opposite fixes: the model is too inaccurate, or CEM at this budget cannot solve the task even with a perfect model. The oracle settles it — the same CEM loop run against the real environment through clone/restore. It is not a deployable method; it is an upper bound on the search.
Random → oracle gap: +1.034. The learned model closes 13.4% of it. The oracle reaches a 10% solve rate where random reaches 0%, so the search is capable of something. The shortfall is model fidelity, and the rollout curve says so independently. Structuring the system so this question had a clean answer is the part worth keeping — a single end-to-end objective would have produced the same poor return with no way to attribute it.
Two hypotheses, both refuted
| mode | F1 sampled | F1 argmax | s/episode |
|---|---|---|---|
| parallel | 0.553 | 0.604 | 0.07 |
| autoregressive | 0.538 | 0.562 | 0.44 |
- Sampling imagined futures should beat argmax
- The standard argument for categorical latents: a discrete distribution can represent “left or right, equally likely”, where a regression to the mean gives a state that never occurs. Measured, argmax wins at nearly every horizon — mean F1 0.604 against 0.553. The argument assumes a calibrated model; at 0.62 one-step fidelity the spread is dominated by model error rather than environment stochasticity, so sampling adds noise on top of error.
- The parallel-token shortcut is the cause
- The dynamics head predicts a frame’s 16 codes simultaneously, treating them as conditionally independent — which would explain incoherent samples. So I implemented autoregressive within-frame generation, drawing from the joint. It is worse (0.538 vs 0.553 sampled, 0.562 vs 0.604 argmax) and 6× slower, and argmax still beats sampling after the assumption is removed. Independence was not the cause. Kept behind a default-off flag, because the hypothesis is worth retesting at higher fidelity.
Three bugs worth recording
- The model never saw a termination
- The replay sampler rejected any window containing a terminal step — the safe-looking choice, since windows must not span a reset. Termination recall sat at exactly 0.000 through 3,000 training steps while code accuracy climbed to 0.95: a model that predicted the world accurately and did not know episodes could end. Windows may now end on a termination, the code loss is masked there, and 25% of each batch is forced terminal because the natural rate is 1.6%.
- Unweighted BCE on a 1.6% positive class
- Even with terminal windows sampled, the termination head predicted “alive” unconditionally: 98.4% accuracy, zero recall. pos_weight is now 20 — deliberately short of the balanced 60, because over-correcting makes the planner see hazards everywhere and freeze.
- Cell accuracy as a reconstruction metric
- Observations are ~92% empty, so predicting nothing scores 0.92 and the metric barely moves during training. The first encoder looked like it was at 98% when its occupancy F1 was 0.44. Everything is now reported as F1, where the empty baseline is 0.0.