Lab: Two Rooms

World models · 2026 · our experiment

The question

Every page before this one reports somebody else’s results. This one is ours. Can a small JEPA world model learn a room just by watching it - and when the room changes, did it learn the rule (“walls block, gaps let you through”) or only the one room it saw?

The world is deliberately tiny. A dot moves in a square split by a wall with one door. The model never sees coordinates, only 32 × 32 pictures. It learns to predict the next picture’s code, not its pixels, the JEPA bet, then a planner searches inside those codes for a route to a goal picture, as in LeWorldModel and PLDM. Then we move the door.

The process, in order

1. Baseline JEPA. A vector model squeezes each frame into 32 numbers. Its code is a clean map of the room (x and y read back at R² 0.997), and a cross-entropy-method planner reaches the goal in 14 of 24 cross-room trips. Move the door and it falls to 7: the predictor still expects the old door.
2. Many doors and generalisation fixes. Train on several doors, roll out 24 steps instead of 8, spread the code within each clip, add an inverse-dynamics loss. The gaps between training doors fill in; doors beyond them never do.
3. Bump, surprise, update. Let the model wander the new room and fine-tune on its own prediction error. It learns the new door within a few hundred to a thousand steps, and forgets the old room unless it replays it.
4. Door as an input, a spatial model, time-similarity. Telling the model the door height does not help. Changing the architecture does: a spatial model (a 16 × 16 feature grid updated by 3 × 3 convolutions) trained on the middle door predicts crossings at 89-100% of heights at every door.
5. Spatial planning and the distance diagnosis. That spatial model still could not plan (7 of 216): in a grid code, two dots that do not overlap look equally far apart, so the planner has no slope. Adding time-similarity fixed it.
6. The oracle test. Freeze each model and swap only its measure of progress for the true walking distance, to separate bad predictions from a bad sense of progress.
7. Robustness pack. Three training runs of each model, re-planned with three measures of progress, including one that needs no labels at all. See the table below.
8. Deep-research review of the extrapolation literature, which pointed at walking distance (the eikonal cost) as the ideal planning cost.

Headline numbers

Trips reaching the goal, 24 per door, nine doors from 0.1 to 0.9, 216 in total. Doors 0.1 and 0.9 are outside every model’s training range, so they test extrapolation.

• Vector (MLP) models fail beyond their training range: 2-6 of 24 at doors 0.1 and 0.9, whatever they were trained on.
• Spatial JEPA with time-similarity, trained on door 0.5 only: across 3 training runs it reached 47–62 of 72 trips at doors 0.1, 0.5 and 0.9 (24 of 24 at the training door every time, 11–19 of 24 at each far door). One run, tested at all nine doors, reached 193 of 216.
• Oracle test (frozen model; only the measure of progress changes):

Model	Own code distance	True walking distance
Ceiling (real simulator)	-	216
Original spatial model	7	199
Vector model	59	87
Time-similarity model	193	184 (no real difference)

Read the table across. The original spatial model’s predictions were right at every door; with a correct measure it goes from 7 to 199. The vector model’s predictions are wrong at the far doors, so even a perfect measure only lifts it to 87. Time-similarity’s own measure is as good as the privileged one.

Robustness pack

Trips reaching the goal out of 72 (24 at each of doors 0.1, 0.5 and 0.9), for three training runs of each model. Each entry is one run, with its 95% range in brackets.

Model	Its own code distance	True walking distance (privileged)	Experience map (no labels)
Original spatial	2 (1–10%), 2 (1–10%), 4 (2–13%)	68 (87–98%), 71 (93–100%), 71 (93–100%)	46 (52–74%), 46 (52–74%), 51 (59–80%)
Time-similarity	62 (76–92%), 55 (65–85%), 47 (54–75%)	55 (65–85%), 57 (68–87%), 49 (57–78%)	43 (48–70%), 39 (43–65%), 40 (44–66%)
Vector	16 (14–33%), 17 (15–35%), 14 (12–30%)	29 (30–52%), 29 (30–52%), 28 (28–50%)	26 (26–48%), 27 (27–49%), 24 (24–45%)

1. The original spatial model’s predictions are right; its measure of progress is what fails. In all three runs the frozen model plans well with walking distance and fails with its own code distance (paired McNemar p < 0.0001).
2. Time-similarity makes code distance as good as walking distance. The difference is not significant in any run (p = 0.12, 0.79, 0.69). Its predictions are a little less accurate than the original spatial model’s: a walking-distance mean of 54 against 70.
3. The label-free experience map recovers most of the gap. It reaches about 48 of 72 for the spatial model, but stays below walking distance.
4. The vector model’s predictions are wrong at the far doors. It fails there whatever measure of progress it uses.
5. Checks. A finer decoding grid changes nothing (67 vs 68, p = 1.0). Straight-line distance drops to 41, so knowing where the wall is matters.

How sure is this? There are only 3 runs per model and 72 trips per entry, so each entry’s 95% range is about ±11 percentage points mid-range, and tighter near 0 or 72. The same trained model also moved by about 5 trips between two evaluations with different start and goal pairs.

Verdict

What it fixedA spatial JEPA trained on one door, with time-similarity in its loss, reached the goal in 47-62 of 72 trips across 3 training runs, including doors outside its training range where every vector model reached at most 6 of 24.

What it still costsOne small toy world, three training runs per model, 24 trips per door. The strongest diagnostic (true walking distance) is privileged information a real robot would not have.

The three explainers

Each one is a complete interactive page from the experiment, placed inline on its page.

• Two Rooms lab notebook - every experiment from scratch, with a door lab you can drive.
• The planner’s path integral - Feynman’s sum over paths beside our sampling planner.
• One push apart - what time-similarity is and why it made planning work.

PaperLearning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models (PLDM)Sobal, Zhang, Cho, Balestriero, Rudner, LeCun · 2025arXiv:2502.14819 →

Where it sits

Worth having first: LeWorldModel · JEPA

The code, logs and raw results live in the physicalAILabs project (jepa_two_rooms/video3/), which is not public yet.