Two Rooms lab notebook
World models · 2026 · Lab: Two Rooms
The newcomer’s guide to the whole experiment. Start here if you only open one. The door lab in section 5 is the heart of it: pick a model, move the door, and see where the model believes the dot can cross against where it really can.
Can a machine learn a room just by watching it?
We trained a small world model on pictures of a dot moving between two rooms, then asked it to plan routes, moved the door, and tried a dozen fixes. This page explains each idea from scratch and reports every result, including the failures.
A square room, a wall, one door
A dot lives in a square split by a wall. Each step it gets a small push. A push that would carry it through the wall is stopped, unless the dot is at the door.
The machine never sees coordinates. It sees a picture 32 squares wide and 32 tall, so 1,024 numbers per frame: the bright dot covers about 43 squares, the wall fills columns 15–16, and the door is a gap of rows 13–18. Nobody tells it which numbers are the dot, the wall or the door.
Predict the next moment, but in a smaller language
A world model answers one question: if the world looks like this and I do that, what happens next? Predicting the next picture pixel by pixel gives a blur, because the future is uncertain and the average of many futures is fuzzy. A JEPA (joint-embedding predictive architecture) predicts a short code instead.
How a goal is given, and how a route is found
The world model itself has no goal. We give one as a picture of the dot at the target, encode it, and get a goal code. "Go there" now means "make your imagined code match the goal code".
The planner draws 512 random sequences of 50 pushes and plays each one forward inside the code, never touching pixels. It keeps the 48 whose imagined endings are closest to the goal, samples new plans around them, and repeats 8 times (the cross-entropy method). Paths that hit the wall in imagination score badly and drop out: by the sixth round every kept plan goes through the door.
Then it takes only 4 real steps, looks again and replans. This loop of predict, optimise, act a little, re-observe is model predictive control. It corrects small prediction errors before they pile up.
A world model is only as good as the worlds it has seen
Move the door from the middle to the bottom. The encoder still finds the dot (position error 0.025, against 0.017 normally), but the predictor still expects the door in the middle. Pushed right where the old door was, it predicts the dot reaching x = 0.63; the real dot stops at 0.49. Planning falls from 14 to 7 of 24, and the planner drives straight at the old door.
Everything below is an attempt to answer one question: can the model learn the rule "walls block, gaps let you through", or only the one room it saw?
Ask any model where it thinks it can cross
For each door height we place the dot at 45 heights just left of the wall, imagine eight pushes to the right, and check whether the model believes the dot gets through. Pick a model and move the door.
What helped, door by door
24 cross-room trips per door, 216 in all. Doors 0.1 and 0.9 lie outside every model's training range, so they test extrapolation. Doors in between, but not trained on, test interpolation.
Same start and goal pairs for every model (seed 7). Fix C and the IDM models trained on any door from 0.15 to 0.85 except a held-out band 0.42–0.58, so for them 0.5 is unseen.
What the numbers say
Training longer imagined rollouts (24 steps instead of 8) roughly doubled planning, because planning imagines 50 steps ahead and errors compound. Spreading the code within each clip (fix C) or the inverse-dynamics loss (recover the push from two codes) stopped the any-door shortcut, where the changing door kept the code spread and the dot got dropped: 10 trips became 117 to 135.
What none of them did
Reach past their examples. At doors 0.1 and 0.9 every model gets 1 to 6 of 24. Telling the model the door height as a number did not help either: trained on one door, the number never varies, so the model never learns what it means.
One training door, every door predicted
Our main model squeezes the whole room into 32 numbers, so each door layout needs its own memorised answer. The spatial model keeps a 16 × 16 grid of features and updates it with small 3 × 3 convolutions. Near the door it cannot know where it is, only what is next to it, so the cheapest thing to learn is a local rule: a wall cell blocks, a gap lets you through.
Trained on the middle door only, it predicted crossings correctly at 89–100% of heights at every door from 0.1 to 0.9, in two training runs. Try it in the door lab above.
Can it plan at any door? Yes, once its code measures progress.
With plain distance between grids as the goal score it reached the goal in 7 of 216 trips: in a grid code, two dots that do not overlap look equally far apart, so the planner had no slope to follow. Training with time-similarity (consecutive frames get similar codes) and across-time variance fixed that. Trained on door 0.5 only, the model reached the goal in 193 of 216 trips across all nine doors, including 16 of 24 at door 0.1 and 17 of 24 at door 0.9, where every vector model scored 2 to 6. A second training run gave 18, 24 and 18 of 24 at doors 0.1, 0.5 and 0.9. These runs use a reduced planner (128 routes, 6 rounds, 40 steps ahead); the one-door vector model scored 59 of 216 with the same planner.
Was it the predictions or the measure of progress?
A planner needs good predictions and a good measure of progress, and a failure could come from either. To separate them we froze each model and changed only the measure: every imagined code was matched to the nearest real image of the room, and the route was scored with the true walking distance to the goal, around the wall and through the door. Starts, goals and the planner's random draws stayed the same, so each pair of rows is a like-for-like comparison.
| Model (frozen) | Measure of progress | 0.1 | 0.5 | 0.9 | All 216 | 95% range |
|---|---|---|---|---|---|---|
| No model: real simulator | true walking distance | 24 | 24 | 24 | 216 | 98–100% |
| Spatial, original | its own code distance | 0 | 1 | 1 | 7 | 2–7% |
| Spatial, original | true walking distance | 24 | 24 | 22 | 199 | 88–95% |
| Spatial + time-similarity | its own code distance | 16 | 24 | 17 | 193 | 85–93% |
| Spatial + time-similarity | true walking distance | 17 | 24 | 17 | 184 | 80–89% |
| Vector, one door | its own code distance | 3 | 12 | 1 | 59 | 22–34% |
| Vector, one door | true walking distance | 1 | 24 | 3 | 87 | 34–47% |
What it shows
The search is not the limit: with the real simulator it reaches every goal. The original spatial model's predictions were right at every door, including 0.1 and 0.9; with a correct measure it goes from 7 to 199 trips. The vector model reaches every goal at the door it trained on and almost none at the far doors, because it still imagines the old door, so a perfect measure cannot repair wrong predictions. Time-similarity's own measure is as good as the true walking distance for this planner (193 against 184, inside each other's 95% ranges).
What limits it
The true walking distance is privileged: it needs labelled images of the room and its real layout, so it is a diagnostic, not a method. Each model is one training run with 24 trips per door. Matching imagined codes to a grid of real images adds about 0.01 to position error, against 0.02–0.03 of prediction drift over 40 imagined steps and a 0.07 success radius.
Bump, surprise, update
A child who walks into a new wall is surprised, and soon expects it. We let the one-door model wander the moved-door room and fine-tuned its predictor every 25 steps on what it had just seen. No labels: the prediction error is the lesson.
Two ways to get from here to the goal
Reinforcement learning (RL) learns a policy, a direct mapping from what it sees and the goal to the next action, by trial and reward. A world model with planning learns how the world responds to actions, then searches its own predictions at decision time. Both can solve Two Rooms; they differ in what they pay and what they transfer.
| World model + planning (ours) | Model-free RL policy | |
|---|---|---|
| What it learns | How actions change the world, from reward-free video | Which action to take, from rewards or goal labels |
| Cost at decision time | High: about 205,000 predictor calls per replan | One network pass per step |
| New goal or new task | Change the cost; no retraining | Goal-conditioned policies handle new goals; new tasks need retraining |
| Bad or random data | Robust: it only needs to learn what actions do | Varies; several methods degrade sharply |
| Unseen layouts | Transfers best in published comparisons, but only near its training layouts (our extrapolation fails too) | Generally degrades as layouts move away from training |
| Main risk | Slow; a misleading code distance misguides the planner | Trained inside a world model, it learns to exploit the model's mistakes (walking through walls that are only imagined away) |
What published studies measured
On reward-free Two Rooms data, the PLDM study (Sobal, Zhang, LeCun and others) found its JEPA planner the most robust to poor data and the best at new maze layouts, while goal-conditioned RL methods such as GCIQL and HIQL did very well when data was plentiful and good. Meta's example reaches 97% on Two Rooms with a redrawn wall per episode, and drops to 1% without the inverse-dynamics loss. A learned controller trained on a JEPA's codes (GC-IDM) reached 100% against CEM's 84% on Two-Room at about 100× lower decision cost.
How we would combine them
The usual compromise keeps the world model and adds RL around it: a small policy proposes starting plans and a learned value estimates "steps left", and the planner refines them against the model (the TD-MPC2 recipe). That makes planning fast without handing door transfer to a policy that may not generalise. We have not run an RL baseline ourselves yet; it is the natural next comparison.
1. A JEPA can learn a usable map and the wall rule from 1,024 raw numbers per frame, with no labels, and plan through the door by imagining.
2. It does not adapt on its own. Moved door, wrong predictions, confident planner.
3. With varied doors and the right training pressure (long rollouts, within-clip spread, or inverse dynamics), it fills the gaps between training doors. It never reaches past them.
4. Architecture matters more than hints: a local, spatial predictor learned the wall rule from one door. With time-similarity added it also plans at doors it never saw, including 0.1 and 0.9. Scored with the true walking distance, even the model without time-similarity plans 199 of 216 trips: its predictions were right everywhere, and the missing piece was a measure of progress.
5. Allowed to keep learning, it adapts within a few hundred to a thousand steps, but forgets the old room unless it replays it.
- Sobal, Zhang, Cho, Balestriero, Rudner, LeCun. Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models (PLDM). arXiv 2502.14819
- Meta FAIR. eb_jepa, action-conditioned video JEPA on Two Rooms. GitHub README
- LeWorldModel: Stable End-to-End JEPA from Pixels. arXiv 2603.19312
- SCALE: State-Calibrated Latent Embeddings for JEPA Planning. arXiv 2608.16287
- Latent Geometry Beyond Search: Amortizing Planning in World Models (GC-IDM). arXiv 2605.08732
- Ha and Schmidhuber, World Models (2018). arXiv 1803.10122
- Our code and logs: jepa_two_rooms/video3 (multi_door.py, door_rules.py, adapt.py, spatial_check.py).
Back to the experiment summary: Lab: Two Rooms