One push apart

World models · 2026 · Lab: Two Rooms

Time-similarity is one extra line in the training loss. In our spatial model it took planning from 1 successful trip in 72 to 64. This page shows, on real frames and real model codes, what it changed inside the model and why that gave the planner a slope to follow.

Two Rooms · world models · time-similarity

Frames one push apart should get codes one step apart

Time-similarity is one extra line in a world model's training loss. It asks the model to give consecutive video frames similar codes, while another rule keeps the codes from freezing. In our spatial world model it took planning from 1 successful trip in 72 to 64, and helped a model trained on a single door plan at doors it never saw.

ModelOriginalTime-similarity 1Time-similarity 10

Eight real frames, one push apart, with the dot pushed toward the door. Under each frame: the patch of the model's 16 × 16 grid that the dot lights up, and how far that frame's code is from the code of the last frame, which plays the goal. Measured on our trained models.

The problem

The model knew how the dot moves, but not which way was closer

A planner needs two things from a world model. It needs to know what happens after a push, and it needs a sense of progress: is this imagined future closer to the goal than that one? Our planner measures progress as the distance between the imagined code and the goal's code.

Our first spatial model was excellent at the first job. Trained on one door, it predicted whether the dot gets through at 89–100% of heights at every door from 0.1 to 0.9. It failed the second. Its code distance said almost nothing about how far the dot really was from the goal, so the planner reached the goal in 1 of 72 trips.

ShowOriginalTime-sim 1Time-sim 10Vector modelWalking distanceStraight line

Follows walking distance (rank correlation, whole room)

–

Follows straight-line distance (whole room)

–

Follows walking distance (left room)

–

Follows straight-line distance (left room)

–

Each map shades every spot by how far its code is from the goal's code (the star), in ten bands from near (dark) to far. A useful map slopes steadily toward the goal and around the wall. The correlations rank all spots by code distance and compare that order with the true order; 1 means the same order, 0 means unrelated. The left room matters most, because there the walk through the door is longer than the straight line.

The rule

What time-similarity is

Think of a flip-book. Each page differs from the last by a small movement. If you describe every page with a list of numbers, time-similarity asks that the lists for neighbouring pages be close. On its own that rule has a lazy solution: give every page the same list. So a second rule says that within each clip, every number in the list must keep changing. Together they force the code to change a little at every step, like marks on a ruler, so the distance between two codes starts to count how many pushes separate them.

Rule 1 · Time-similarity

mean ‖z_t+1 − z_t‖²

Codes of consecutive frames should be close. Weight 1 or 10 in our runs.

Rule 2 · Keep changing

mean max(0, 0.5 − spread of z over the clip)

Within each clip, every one of the code's numbers must vary by at least 0.5. Without this, rule 1 is solved by freezing the code.

# video3/door_rules.py: added to the usual JEPA loss (prediction + variance + covariance)
if tsim > 0:   # rule 1: consecutive codes close
    loss = loss + tsim * F.mse_loss(z[:, 1:], z[:, :-1])
if clipvar > 0:   # rule 2: each number keeps changing within the clip
    zc = z - z.mean(1, keepdim=True)
    loss = loss + clipvar * F.relu(0.5 - torch.sqrt(zc.flatten(0, 1).var(0) + 1e-4)).mean()

z holds the codes of one clip, frame by frame. The spatial model's code is its whole 16 × 16 grid of 8 features, 2,048 numbers.

What it did

It spread the dot out, and the planner got a slope

The strip at the top shows what changed inside the model. The original model writes each dot as a tight spot. Once two spots stop overlapping, their codes are as far apart as they will ever get, so every frame but the last looks equally far from the goal. With time-similarity at weight 10 the model writes a wider patch, consecutive codes overlap, and the distance to the goal falls frame by frame.

Code distance to the last frame, scaled to each model's largest, for the eight frames in the strip.

Spatial model, trained on door 0.5	Original	Time-sim 1	Time-sim 10
90% of a dot's code lies within	0.099	0.133	0.159
Code distance follows walking distance	−0.003	0.08	0.87
Planning at doors 0.2, 0.5, 0.8	1 / 72	26 / 72	64 / 72
Crossing predictions correct at all nine doors	89–100%	91–100%	91–100%

Planning at every door, after training on one

With time-similarity at weight 10, the spatial model trained only on door 0.5 reached the goal in 193 of 216 trips across all nine door heights, including 16 of 24 at door 0.1 and 17 of 24 at door 0.9. The vector model trained on the same door managed 59 of 216 with the same planner, and every vector model we trained, including those trained on several doors, reached only 2–6 of 24 at doors 0.1 and 0.9.

Three jobs

Moving, measuring and searching are separate skills

A better search cannot rescue a flat measure of progress, and a perfect measure cannot rescue wrong predictions. Our results separate the three jobs cleanly.

Job	The question it answers	What our models showed
Dynamics	What happens if I push this way?	The spatial model, trained on one door, predicts crossings at 89–100% of heights at all nine doors. Scored with the true walking distance, the same frozen model plans 199 of 216 trips.
Measure of progress	Is this future closer to the goal?	The original grid's distance ignores the real distance (−0.003). Time-similarity at weight 10 brings it to 0.87. Swapping the original's code distance for the true walking distance takes it from 7 to 199 of 216 trips.
Search	Can I find a good push sequence?	The film's planner finds the door from 3% of first-round routes, 100% of kept routes by round six. With the real simulator and the true distance, this planner budget reaches 216 of 216.

The decisive test

Was it the predictions or the measure?

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	Door 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%

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).

Trips that reached the goal, out of 24 per door; the totals cover all nine doors. Matching imagined codes to real images on a grid spaced 0.025 apart adds about 0.01 to position error. A grid twice as fine halves that on real frames but barely changes imagined positions after 40 steps (0.020 against 0.022 for the original model), because prediction drift is larger.

Open questions

What this does not prove yet

• The decisive test used a privileged measure. The true walking distance needs labelled images of the room and its real layout. It shows what each model can do with a perfect measure of progress. A robot could not use it as it stands.
• Wider patches may not be the whole story. Across three models, a wider dot patch went with better planning. Time-similarity may also reshape the code in other ways. A model made to write equally wide patches without time-similarity would separate the two causes.
• Smooth is not the same as walking distance. Time-similarity only compares frames one push apart. In the training videos the dot reaches the other room only through the door, which is likely why the time-sim-10 code follows walking distance in the left room (0.91) better than straight-line distance (0.71). It is still not the exact walking distance around the wall.
• Small samples. Each model is one training run, with 24 trips per door and a reduced planner (128 imagined routes, 6 rounds, 40 steps ahead). A second run of the time-similarity model repeated its result at doors 0.1, 0.5 and 0.9. Differences inside the 95% ranges, such as 193 against 184, are not yet real differences.

How we will firm this up

• Take away the privilege one piece at a time. First keep the matching to real images but score straight-line distance, which needs no knowledge of the wall. Then drop the labels too: let the model explore the new room for a few minutes, link consecutive frames into a map, and count steps on that map. Neither version uses the true layout.
• More training runs, compared trip by trip. Two more runs of each model at doors 0.1, 0.5 and 0.9, saving every trip's outcome so paired tests can say which differences are real.
• A finer grid as a check. Repeat the decisive test with a matching grid twice as fine. The drift measurements predict the same results.
• Other layout changes. Different door widths, a shifted wall and two doors, to test whether the spatial model learned a rule about walls in general or only about door height.
• A width-matched control. Train without time-similarity but with wider dot patches built in, to separate the width effect from anything else time-similarity does.
• A learned arrival-time map. Estimate steps-to-goal with a local rule trained only from the order of frames in the videos, so the measure of progress needs no labels at all.

Where it comes from

An old idea about what matters in video

The idea that useful features change slowly is decades old. Slow feature analysis (Wiskott and Sejnowski, 2002) learns features that vary slowly over time, on the theory that the causes behind a scene change more slowly than its pixels. Mobahi, Collobert and Weston (2009) used the temporal coherence of video to train image features. Time-contrastive networks (Sermanet and others, 2018) pull together views of the same moment and push apart different moments.

World models that plan use it for the reason shown here. PLDM, a JEPA planner on the same Two Rooms task, reaches about 98% success; without its time-similarity term, 71%. Meta's eb_jepa example weights time-similarity at 12 and drops from 97% to 61% without it.

Sources

1. Our experiments: jepa_two_rooms/video3 (door_rules.py, spatial_ts_full.py, time_sim_export.py).
2. Sobal, Zhang et al., Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models (PLDM). arXiv 2502.14819
3. Meta FAIR, eb_jepa action-conditioned Two Rooms example. README
4. L. Wiskott and T. Sejnowski, Slow feature analysis: unsupervised learning of invariances, Neural Computation (2002).
5. H. Mobahi, R. Collobert and J. Weston, Deep learning from temporal coherence in video, ICML (2009).
6. P. Sermanet et al., Time-contrastive networks: self-supervised learning from video, ICRA (2018).

Back to the experiment summary: Lab: Two Rooms