Temporal Discriminative World Models

A JEPA-style world model has three parts: an encoder that turns a frame into a vector, a predictor that guesses the next vector from the current one and an action, and something that stops the encoder from cheating — mapping every frame to the same point, where prediction is perfect and useless. The JEPA chapter built that recipe and the world-models chapter planned with it. In 2026 the cleanest version of it is LeWorldModel (LeWM)²: next-embedding MSE plus one regularizer, SIGReg³, that pushes the embeddings toward an isotropic Gaussian.

In June 2026 William Peng and Holger Molin (Stanford) posted a project page arguing that the third part is the wrong fix¹. Their diagnosis: the regularizer stops collapse but does not make the encoder keep what a planner needs, because the prediction loss itself rewards forgetting whatever is hard to predict. Their proposal, the Temporal Discriminative World Model (TDWM), deletes the regularizer and changes the prediction loss: instead of "predict the next vector", it asks for how similar the prediction should be to every frame of the clip, as a function of how far apart in time they are. One loss, no anti-collapse term.

This chapter takes that loss apart line by line, places it among the methods it borrows from, and measures it in a lab — where, in a world small enough to train in minutes, one of its three claims reproduces and two do not.

Slides

Click the slides and use ← → to move through them; F goes full screen and P opens the presenter view, with the notes. The two simulators of this page run inside the slides. Open the slides in a new tab

The recipe this chapter starts from

Write \(o_t\) for the observation (a frame) at time \(t\) and \(a_t\) for the action taken then. A latent world model learns

an encoder \(f_\theta\) with parameters \(\theta\) and a predictor \(g_\phi\) with parameters \(\phi\), trained together on offline trajectories of observations and actions, with no reward. LeWM trains them with

where \(Z\) is the batch of embeddings and \(\lambda = 0.1\) by default². Without the second term, the first one is minimized by a constant encoder: \(z_t = c\) for every frame makes \(\hat z = c\) exact. SIGReg comes from LeJEPA³: project the batch on \(M\) random unit directions (\(M = 1024\) in LeWM) and, along each, measure how far the one-dimensional distribution is from \(\mathcal{N}(0, 1)\) with the Epps–Pulley statistic, a weighted distance between characteristic functions. By the Cramér–Wold theorem, matching every one-dimensional projection means matching the whole distribution, so the embeddings are pushed toward \(\mathcal{N}(0, I)\) — which a constant cannot be.

Planning then happens in latent space. Encode the goal image, \(z_g = f_\theta(o_g)\); for a candidate action sequence, roll the predictor out \(H\) steps and score the last imagined vector by \(\|\hat z_H - z_g\|^2\); search over sequences with the Cross-Entropy Method¹⁵ (CEM): sample many sequences from a Gaussian, keep the best few (the elites), refit the Gaussian to them, repeat. LeWM samples 300 sequences, keeps 30 elites, plans 5 steps ahead (25 environment steps, with a frame-skip of 5) and executes them before replanning². The world-models chapter covers CEM and model-predictive control in full.

Every JEPA world model has to answer "what stops collapse?", and the answers differ more than the architectures do:

What the regularizer does not buy

The project page starts from a probe. Train LeWM on OGBench Scene¹⁴ — a robot arm with a cube, a drawer, a window and buttons — freeze the encoder and fit a linear (ridge) regressor from the embeddings to the true position of each object, scoring held-out \(R^2\):

The slow, mostly still parts of the scene are almost perfectly recoverable; the fast, controllable ones — the ones a planner needs — are not, and a nonlinear probe does not recover them either, so the information is not there. The authors' explanation is about the objective, not the regularizer: next-embedding MSE is minimized by making the future predictable, and the surest way to do that is to stop representing the parts of it that are hard to anticipate. A regularizer that only constrains the shape of the embedding distribution cannot object, because a Gaussian cloud can be built from the drawer and the window alone. Anti-collapse is necessary, not sufficient.

This is a hypothesis with one probe on one environment behind it. It is also a familiar one: it is the same argument the JEPA chapter made for predicting in latent space (ignore what you cannot predict), turned around.

The objective

TDWM keeps the encoder and the predictor and changes three things¹:

1. Norm constraint. Encoder and predictor outputs are divided by their norm: \(\|z\| = \|\hat z\| = 1\). Only the direction is free, so the loss can no longer be lowered by shrinking vectors.
2. History and horizon. From a 3-frame history the predictor predicts \(k\) steps ahead, given the intermediate actions: \(\hat z_{t+k} = g_\phi(z_{t-2:t},\, a_{t:t+k-1},\, k)\).
3. The loss. For a prediction \(\hat z_i\) whose target is frame \(i\) and an encoded frame \(z_j\) of the same clip, let \(c_{ij} = \cos(\hat z_i, z_j) = \hat z_i \cdot z_j\). Then

Every symbol: \(i\) indexes the frame a prediction aims at, \(j\) any frame of the clip, \(|i - j|\) the gap between them in frames, and \(\sigma > 0\) — in frames — sets how fast the target similarity decays with the gap. The project page does not give the value of \(\sigma\) it used.

Read the loss in two halves:

• \(i = j\) — the target is \(e^0 = 1\): the prediction must point exactly at the frame it aims at. On unit vectors \(\|\hat z - z\|^2 = 2 - 2c\), so this is next-embedding prediction on the sphere.
• \(i \neq j\) — the target is below 1 and falls with the gap: \(e^{-1/2} = 0.607\) one frame away with \(\sigma = 2\), \(e^{-1} = 0.368\) two frames away, \(0.135\) four frames away. The prediction must be partly similar to its neighbours and nearly orthogonal to distant frames. A collapsed encoder gives \(c_{ij} = 1\) everywhere and pays for every one of these cells. This half is the anti-collapse term.

The targets are also consistent. With \(\tau_{ij} = e^{-|i-j|/\sigma}\) and frames ordered \(i \le j \le k\),

so a long-gap target is the product of the short-gap ones. The project page stops there ("more formal analysis to come"). One more line closes the question of whether every target can be met at once: with \(\rho = e^{-1/\sigma}\), the matrix of targets has entries \(\rho^{|i-j|}\) — the correlation matrix of a first-order autoregressive process, whose determinant is \((1 - \rho^2)^{L-1} > 0\) for \(L\) frames. It is positive definite, so \(L\) unit vectors with exactly those cosines exist in \(L\) dimensions. The loss is not asking for something impossible.

The loss, with numbers

A clip of five frames in three dimensions — small enough to check every number by hand. The encoder gave \(Z_{\text{raw}}\) (one row per frame, \(j = 0 \dots 4\)), and the predictor gave \(P_{\text{raw}}\) (its guesses for frames \(i = 1 \dots 4\)):

Take the first prediction, \(\hat z_1 = (1, 0.8, 0)/\sqrt{1.64} = (0.781, 0.625, 0)\). Its cosines with the five normalized frames, the targets for gaps \(1, 0, 1, 2, 3\) and the squared errors are

The \(j = 1\) cell is where an MSE would look, and it is almost perfect. The errors are elsewhere: the prediction is too similar to frame 0 and not similar enough to frame 4. Doing the same for the other three rows and averaging all 20 cells gives \(\mathcal{L}_{\text{TDWM}} = 0.0285\). The plain MSE between raw predictions and raw frames is \((0.04 + 0.01 + 0.01 + 0.01)/4 = 0.0175\). Now two clips that break the comparison:

• Collapse. Every frame and every prediction is \((1, 0, 0)\). The MSE is 0. TDWM sees \(c = 1\) in every cell: the 4 cells at gap 0 cost nothing, but the 7 at gap 1 cost \((1 - 0.607)^2 = 0.155\) each, the 5 at gap 2 \(0.400\), the 3 at gap 3 \(0.604\) and the 1 at gap 4 \(0.748\), a mean of $5.640/20 = $ 0.2820 — ten times the healthy clip.
• Shrinking. Multiply every vector by 0.1. The MSE drops a hundredfold, to 0.000175, without anything being learned. TDWM normalizes first, and stays at 0.0285.

Step through the same computation below, line by line, edit the vectors, change \(\sigma\) or switch to one of the broken clips:

Why the geometry of the latent matters for planning

The planner never reads the loss; it reads distances in the latent space. The project page's second argument is about what those distances mean¹. In Scene, take a task where the arm already sits where the goal image has it and only the drawer must be opened. In an embedding where most of the scene is already correct, start and goal are close. To open the drawer the arm must leave — and leaving makes the image less like the goal, so the cost goes up before it comes down. A planner that sees 5 steps ahead finds no sequence that lowers the cost and stays put. In a latent ordered by time instead, every step toward the drawer is a step closer in time to the goal, and the cost falls the whole way.

The panel below is that example in one dimension, with the two costs written out: an embedding distance that weighs the arm and the drawer, and an idealized TDWM distance \(1 - e^{-(\text{steps left})/\sigma}\). The planner moves to the cheapest point it can see within \(H\) steps, and gets stuck if nothing in view is cheaper by more than the model's error \(\varepsilon\).

Three things to find in it:

1. The bump. With the default settings the embedding cost needs a horizon of 12 steps to see past the bump; at \(H = 5\) it is stuck at step 0. The temporal cost reaches the goal at any horizon.
2. The flat far field. Set \(\sigma = 3\). Thirty steps from the goal the temporal cost is \(1 - e^{-10}\), and five steps of progress change it by \(2 \times 10^{-4}\), below a model error of \(0.003\): the planner is stuck again, on a plateau. The targets \(e^{-|i-j|/\sigma}\) say nothing about pairs more than a few \(\sigma\) apart, so a TDWM cost-to-go has a slope only within that range. \(\sigma\) is not just a smoothing constant; it sets the reach of the planner.
3. The idealization. The orange curve assumes the encoder realizes the targets exactly along the demonstration. A real encoder is only told about pairs inside a training clip; two states from different clips are never compared, and nothing forces their cosine to mean anything.

The same mismatch between what the model is trained on and what the planner scores is what Temporal-Distance JEPA⁴ calls the train–plan gap: on Push-T, LeWM's latent L2 distance correlates with the true temporal separation of held-out pairs at Spearman \(\rho = 0.65\), against \(\rho = 0.95\) for a cost trained on step counts. The world-models chapter treats the gap in general.

What the project page reports

The comparison is TDWM against LeWM, both trained on the same data — expert trajectories plus the same trajectories re-rolled with action noise and re-rendered — for 0.8 exaFLOPs each, keeping the checkpoint with the highest O25 CEM score, and planning with 1-step autoregressive prediction for both¹.

"O25" is a goal 25 steps ahead along a dataset trajectory; the error bars on the page are standard errors across seeds, and the number of seeds is not given. The probe figure does not say whether it is the ridge or the MLP probe; LeWM's numbers in it match its ridge probe above. Effective rank is the exponential of the entropy of the normalized singular values of the embeddings¹⁶ — how many directions they actually use; the environment has 24 degrees of freedom.

Three more results on the page are as informative as the headline:

• Expert-only data reverses the planning result. Trained only on expert trajectories, LeWM outperforms TDWM at planning (the numbers are not given). TDWM's predictor degrades sharply when the planner tries actions outside the expert distribution: its normalized fidelity falls to about 0.70 at an action perturbation of 1.8 standard deviations, against about 0.94 for LeWM (trained on the noisy mix, TDWM stays near 0.89). In the expert data, actions predict the state with \(R^2 = 0.71\); the authors suspect TDWM learned intent rather than physics, and add noisy trajectories (\(R^2 = 0.31\)). The table above is only valid for that mix.
• The upper bounds. At O25, the TDWM encoder trained on expert data, with the true MuJoCo simulator in place of the learned predictor, reaches 66%; CEM on the true state with the true simulator reaches 77%, the ceiling of the evaluation itself. The noise-trained TDWM with its own predictor (63%) comes within 3 points of the first bound. On our reading, the 11 points between 66% and 77% are what planning in a learned latent rather than in the true state costs, even with perfect dynamics.
• Two engineering choices the comparison depends on: multi-horizon prediction (97% of the training FLOPs go to the encoder, so predicting several \(k\) per clip adds little compute) and a fused Muon–AdamW optimizer.

Where TDWM sits

Using time as the supervisory signal for a representation is an old idea; what is new is using it as the only objective of an action-conditioned JEPA world model, with graded rather than binary targets.

The closest neighbour is Temporal-Distance JEPA, and the two make complementary bets. TD-JEPA keeps LeWM's encoder, predictor and losses and adds a separate, directed cost \(d_\psi(z_s, z_g)\) — going from \(A\) to \(B\) need not cost what \(B\) to \(A\) does, which a cosine cannot express. It reports, under a locked protocol over ten seeds, 100.0% on Two-Room (LeWM 97.4%) and 97.0% on Reacher (96.0%) planning with \(d_\psi\), and 86.0% on Push-T (83.6%) and 82.2% on OGBench-Cube (68.0%, a gain of 14.2 points) planning with the plain latent L2 — there the mined cost helps as a training signal for the representation, not as the planning cost⁴. Note in passing that LeWM's own paper reports 96% on Push-T under its native protocol: the same model, a 12-point difference from the evaluation protocol alone. TDWM instead changes the representation and keeps the plain latent cost; its target is symmetric and saturates beyond a few \(\sigma\), which TD-JEPA's regression on step counts does not.

Honest limits

• Evidence. One environment, one comparison, a project page. Nothing yet on navigation, on real robots or on scaling.
• Data dependence. The method needs time to mean progress. On expert-only data it planned worse than LeWM; on aimless data (random walks that revisit the same place) a target that says "frames 20 steps apart are orthogonal" contradicts an image that is identical. The page's fix — expert plus noise — is a requirement on the data, not a detail.
• Clip-local supervision. Only pairs inside a training clip are ever compared. Nothing in the loss relates two states from different clips, so the cost-to-go between them is whatever the network's smoothness makes it.
• Symmetric, saturating cost. \(e^{-|i-j|/\sigma}\) treats "5 steps before" and "5 steps after" alike and goes flat past a few \(\sigma\); goal reaching is directional and long-range. Quasimetric methods¹⁰ and TD-JEPA⁴ were built around exactly that.
• The predictor's share. In a clip of \(L\) frames, the \(i = j\) cells — the only ones that train the predictor to hit its target — are a fraction \(1/L\) of the plain mean. Unless pairs are sampled or weighted differently (the page does not say), the rollout quality the planner needs gets little of the gradient. The lab measures this.

Lab: the objective, measured

The Colab notebook rebuilds TDWM from its formula in a world small enough to train on a CPU in under ten minutes: a \(16 \times 16\) image of an agent moved by the action (the arm), a puck it can push (the cube) and a landmark that never moves (the drawer). Four models with the same encoder, predictor and budget — 2,000 AdamW steps — differ only in the loss: plain MSE, MSE + SIGReg (LeWM's recipe), TDWM and TDWM's \(i = j\) half alone. What it found:

(The oracle — CEM with the true dynamics — succeeds on 100% of the 30 tasks at every offset; one task is 3.3 points and the standard error of a 50% rate is about 9 points.)

• Reproduced: anti-collapse without a regularizer. mse and diag collapse; tdwm keeps the same spread as SIGReg with no regularizer, and the difference between diag and tdwm is exactly the \(i \neq j\) cells.
• Not reproduced: better probes or better planning. TDWM's linear probe is slightly better on the agent and worse on the puck; with an MLP probe both models know everything. TDWM plans worse at every offset.
• Why, measured. TDWM's predictor reaches a cosine of only 0.74 with its target one step ahead, and its 5-step rollout is unrelated to the truth. Two changes improve the rollout: \(\sigma = 12\) (0.88 one step ahead, but success 13% / 23% / 30%, no better than \(\sigma = 4\) within the noise) and, our change, not the paper's, weighing the \(i = j\) cells as much as all the others together — at \(\sigma = 12\) the rollout holds at 0.57 after five steps and success reaches 50% / 37% / 47%: behind the baseline at O10 and O20, ahead at O30, each gap within about one standard error of a difference between two rates (some 13 points).
• Trap. The collapsed models have a high effective rank (26 of 32): rank is scale-invariant, and what remains after collapse is isotropic noise. Read it only next to a spread.

What this does not show is that the project page is wrong: its premise is that MSE + SIGReg drops the hard-to-predict variable, and in this world it does not (its MLP probe recovers the puck at 0.958). A toy can only test the fix if it reproduces the failure. The final open task asks you to build a world where it does.

Open the lab in Colab — save a copy to your Drive (File → Save a copy in Drive) to run and edit it.

Key takeaways

1. A JEPA world model must answer "what stops collapse?" Frozen pretrained encoders, multi-term VICReg-style losses, SIGReg, and now TDWM's loss are four different answers.
2. TDWM's diagnosis: anti-collapse is necessary but not sufficient — next-embedding MSE rewards dropping what is hard to predict, and a distributional regularizer does not object.
3. The TDWM loss regresses every cosine \(c_{ij}\) of a clip onto \(e^{-|i-j|/\sigma}\). The \(i = j\) cells are prediction on the sphere; the \(i \neq j\) cells are the anti-collapse term. Collapse costs it 0.2820 in the worked example, against 0 for the MSE; shrinking costs it nothing.
4. The targets are consistent (\(\tau_{ik} = \tau_{ij}\tau_{jk}\), a positive-definite matrix), but symmetric and saturating: \(\sigma\) sets how far from the goal the cost-to-go still has a slope.
5. TD means temporal discriminative, not temporal difference: no value, no reward, no bootstrapping.
6. The project page reports better probes of the cube (0.56 vs 0.27) and the arm (0.69 vs 0.39) and better CEM success (63/38/15% vs 50/28/4%) on one environment, with noise-augmented data — and the reverse on expert-only data. Preliminary, undisclosed hyperparameters, no code.
7. In this course's toy, the anti-collapse claim reproduced and the probe and planning claims did not; giving the predictor a larger share of the loss (at \(\sigma = 12\)) closed most of the planning gap, but did not reverse it.