Artificial Neural Networks and Deep Learning · Insper

Temporal Discriminative World Models

One loss instead of two

Read this first

Preliminary work. TDWM (Peng & Molin, Stanford, June 2026) is a project page: no preprint, no code, one environment (OGBench Scene). σ, clip length, architecture and number of seeds are not disclosed.

TD ≠ temporal difference. In RL, "TD" almost always means temporal-difference learning (TD-MPC, TD-MPC2: values bootstrapped from the next state). Here TD = temporal discriminative: no value, no reward, no bootstrapping.

Start with a question

= frame at time · = action · = encoder · = predictor (reads the last 3 embeddings)

Both networks are trained together on this loss alone. What is the cheapest way to reach ?

  • Map every frame to the same vector : then is exact. Zero loss, zero information — collapse.

What stops collapse?

model what stops collapse planning cost
DINO-WM a frozen pretrained encoder (DINOv2) distance to goal features
V-JEPA 2-AC a frozen video encoder; only the predictor trains (L1 + rollout loss) L1 to goal, CEM
PLDM a seven-term VICReg-style loss latent distance
LeWM SIGReg: , squared L2, CEM
TDWM the prediction loss itself — no regularizer latent cost-to-go, CEM

LeWM: ~15M parameters, one GPU, planning up to 48× faster than foundation-model world models (Maes et al., 2026).

SIGReg in one slide

  1. Project the batch of embeddings on random unit directions ( in LeWM).
  2. On each 1-D projection, measure the distance to : the Epps–Pulley statistic, a weighted gap between characteristic functions.
  3. Average over directions.
  • Cramér–Wold: if every 1-D projection is , the whole distribution is .
  • A constant cannot look like → no collapse.
  • One hyperparameter, . From LeJEPA (Balestriero & LeCun, 2025).

SIGReg constrains the shape of the cloud of embeddings. Does it say anything about what the encoder keeps?

What the regularizer does not buy

Probe a LeWM encoder trained on OGBench Scene: held-out from embeddings to each object's true position.

drawer window buttons cube arm (7-DOF)
ridge probe 0.95 0.96 0.97 0.27 0.39
MLP probe 1.00 1.00 0.99 0.26 0.40
  • The slow parts survive; the fast, controllable ones — what a planner needs — do not, even for an MLP.
  • TDWM's reading: next-embedding MSE is minimized by making the future predictable — the surest way is to stop representing what is hard to anticipate.
  • A Gaussian cloud can be built from the drawer and the window alone. Anti-collapse is necessary, not sufficient.

The objective

  1. Norm constraint: — only the direction is free.
  2. History and horizon: .
  3. The loss, over every pair of a clip:

= the frame a prediction aims at · = any frame of the clip · = gap in frames · = decay, in frames (not disclosed)

One training step

clip of frameso₀o₁o₂o₃o₄encoder fθevery framezj / ‖zj‖unit vectorsactionsat, …, at+k−1, kpredictor gφ3-frame historyzt−2:tẑicosines cij = ẑi · zjevery pairmean of (c − target)²target e−|i−j|/20.611.000.610.370.220.370.611.000.610.370.220.370.611.000.610.140.220.370.611.00i = j: prediction on the spherei ≠ j: anti-collapse

  • → target 1. On unit vectors : next-embedding prediction on the sphere.
  • → target below 1, falling with the gap: neighbours partly alike, distant frames nearly orthogonal. This half is the anti-collapse term.

Exercise · a collapsed clip

A clip of 5 frames; the predictor guesses frames 1–4, so the matrix is (20 cells). .

  1. The targets for gaps 0, 1, 2, 3, 4?
  2. How many of the 20 cells have each gap?
  3. Every frame and prediction is the same vector. TDWM loss? MSE?
  • target
  • row (1…4), column (0…4), gap
  • collapsed: in every cell

A collapsed clip · answer

gap target cells subtotal
0 1.000 4 0 0
1 0.607 7 0.155 1.084
2 0.368 5 0.400 1.998
3 0.223 3 0.604 1.811
4 0.135 1 0.748 0.748
20 5.640
  • TDWM: 0.2820. MSE: 0 — the collapsed clip is a perfect answer for it.

The loss, with numbers

0 1 2 3 4
.781 .994 .625 .442 0
target .607 1 .607 .368 .223
error² .0304 .0000 .0003 .0055 .0498
  • All 20 cells:
  • MSE:
  • × 0.1: MSE 0.000175, TDWM 0.0285

Simulator · the TDWM loss, line by line

Z_raw · encoded frames 0…4
P_raw · predictions of frames 1…4

Can every target be met at once?

  • A long-gap target is the product of the short-gap ones (the project page stops here).
  • With the target matrix has entries : the correlation matrix of an AR(1) process.
  • → positive definite → unit vectors with exactly those cosines exist in dimensions.

The loss asks for something possible. It does not say anything about pairs from different clips — they are never compared.

Why geometry matters: the planner reads distances

CEM (LeWM's settings): sample 300 action sequences of length , roll the predictor out, score , keep the 30 best, refit the Gaussian, repeat 10×; execute, replan.

The drawer task. The arm already sits where the goal wants it; only the drawer must open.

  • Embedding distance: start ≈ goal. Leaving to reach the handle makes the image less like the goal → cost rises first.
  • A 5-step planner finds nothing cheaper and stays put.
  • Temporal distance: every step toward the handle is a step closer in time → cost falls all the way.

Simulator · the cost-to-go along the demonstration

Exercise · how far does σ reach?

The idealized TDWM cost is . The goal is 30 steps away, the planner looks 5 steps ahead and the model error is .

How much does the cost fall over the first 5 steps with ? And with ? Does the planner move in each case?

How far does σ reach · answer

σ cost at 30 steps left cost at 25 steps left drop over 5 steps vs ε = 0.003
3 stuck — a plateau
12 moves

The targets say nothing beyond a few σ: a TDWM cost-to-go has a slope only within that range. σ sets the reach of the planner — and the page does not disclose it.

What the project page reports

Same data for both (expert + action-noise re-rolls), 0.8 exaFLOPs each, best-O25 checkpoint, 1-step autoregressive planning.

TDWM LeWM
probe cube · arm 0.56 · 0.69 0.27 · 0.39
probe drawer · window · buttons 0.93 · 0.91 · 0.96 0.95 · 0.96 · 0.97
SIGReg loss (TDWM does not optimize it) 2.4 (from ~30) 0.6
effective rank (24 DOF in the scene) ~40 ~150
CEM success O25 · O50 · O75 63% · 38% · 15% 50% · 28% · 4%

O25 = goal 25 steps ahead along a dataset trajectory. Error bars: standard error across seeds; number of seeds not given.

The results next to the headline

  • Expert-only data reverses it: LeWM plans better (numbers not given). TDWM's predictor fidelity falls to ≈ 0.70 under an action perturbation of 1.8 standard deviations (LeWM ≈ 0.94).
  • Expert actions predict the state with → TDWM may learn intent, not physics. Adding noisy re-rolls: .
  • Upper bounds at O25: expert-trained TDWM encoder + true simulator 66%; true state + true simulator 77%. Noise-trained TDWM: 63%.
  • Engineering: multi-horizon prediction (97% of FLOPs in the encoder), fused Muon–AdamW.

One environment · seeds undisclosed · checkpoint chosen by the reported metric · σ, k, clip length, architecture undisclosed · no code. O75: 15% vs 4% is a few episodes.

Where TDWM sits

method signal from time role
Time-Contrastive Networks (2018) temporal neighbours pushed apart, simultaneous views pulled together representation for imitation
CPC (2018) predict future latents contrastively (InfoNCE) representation
Contrastive RL, QRL, HILP, TMD (2022–25) a temporal distance, often a quasimetric (asymmetric) value / goal reaching
Temporal-Distance JEPA (Jul 2026) LeWM + SIGReg plus a directed cost head on step counts the planning cost, or a training signal
TDWM (Jun 2026) cosine targets replacing loss and regularizer the representation

TD-JEPA and TDWM: complementary bets

Temporal-Distance JEPA (Bai & Xiong)

  • keeps LeWM's encoder, predictor and SIGReg
  • learns a directed cost :
  • locked protocol, 10 seeds, planning with : Two-Room 100.0 (LeWM 97.4) · Reacher 97.0 (96.0)
  • planning with plain L2 on the latent it shaped: Push-T 86.0 (83.6) · OGB-Cube 82.2 (68.0)

TDWM

  • changes the representation itself
  • keeps the plain latent cost
  • target symmetric and saturating

LeWM's own paper reports 96% on Push-T under its native protocol; TD-JEPA's locked protocol gives 83.6%. Same method, 12 points from the evaluation protocol alone.

Lab · what we measured

A 16×16 world: agent (the arm), puck it pushes (the cube), landmark that never moves (the drawer). Same networks, 2,000 AdamW steps each; only the loss changes.

mse mse+sigreg tdwm σ = 4 diag (i = j only)
spread (0 = collapse) 0.0015 0.1746 0.1700 0.0033
effective rank (of 32) 26.0 19.7 29.3 25.1
MLP probe agent · puck — 0.969 · 0.958 0.937 · 0.925 —
CEM success O10 · O20 · O30 0 · 0 · 0% 63 · 47 · 43% 20 · 13 · 23% —
  • Reproduced: no collapse without a regularizer — and diag shows it comes from the cells.
  • Not reproduced: better probes, better planning. Trap: a collapsed model has a high effective rank.

Lab · why, measured

Predictor rollout (cosine with the truth, steps 1 · 3 · 5)

  • mse+sigreg: 0.99 · 0.95 · 0.89
  • tdwm σ = 4: 0.74 · 0.00 · −0.19

The predictor's share: in a 16-frame clip, is 13 of cells — 6.25% of the loss.

TDWM variant rollout step 5 O10 · O20 · O30
σ = 4 −0.19 20 · 13 · 23%
σ = 12 −0.01 13 · 23 · 30%
σ = 12, balanced (ours) 0.57 50 · 37 · 47%

Balanced = i = j cells weigh as much as all others together. Close to MSE + SIGReg, within the noise of 30 tasks — not ahead.

A toy only tests the fix if it reproduces the failure: here MSE + SIGReg keeps the puck (MLP 0.958).

Lab · the objective, measured

Colab notebook, ~2 h, CPU only (under 10 minutes of compute). Each section ends in a checkpoint.

  1. a world with an agent, a puck, a landmark
  2. one encoder, one predictor
  3. three objectives from their formulas
  4. four models, one budget: collapse
  5. probes: what the embeddings know
  1. the geometry the planner uses
  2. planning with CEM, O10 / O20 / O30
  3. the undisclosed knob: σ
  4. where the gradient goes
  5. open task: make TDWM win, or explain why it cannot

open in Colab →

Takeaways

  1. Every JEPA world model must answer "what stops collapse?" — frozen encoders, VICReg-style terms, SIGReg, or TDWM's loss.
  2. TDWM's diagnosis: anti-collapse is necessary, not sufficient; MSE rewards dropping what is hard to predict.
  3. The loss regresses every onto : is prediction on the sphere, is anti-collapse. Collapse: 0.2820 vs MSE 0.
  4. Consistent targets (positive definite), but symmetric and saturating: σ sets how far the cost-to-go reaches.
  5. TD = temporal discriminative, not temporal difference.
  6. Reported: cube 0.56 vs 0.27, CEM 63/38/15% vs 50/28/4% — one environment, reversed on expert-only data.
  7. Our toy: anti-collapse reproduced; probes and planning did not. A larger share of the loss for the predictor closes most of the planning gap, without reversing it.

Outline: (1) the JEPA world-model recipe and what stops collapse; (2) TDWM's diagnosis; (3) the loss, line by line, with numbers; (4) why the latent geometry matters for planning; (5) what the project page reports and how much it carries; (6) where it sits; (7) what our lab measured. Closes with the Colab lab.

Say both things before any formula. Students who know TD-MPC will otherwise read the whole class through the wrong word.

Let them answer before revealing. This is the JEPA chapter's collapse, and every world model of this class is an answer to it.

One probe, one environment: a hypothesis, not a theorem.

The j = 1 cell is where the MSE looks and it is almost perfect; the errors are elsewhere.

Default case = the worked example. Walk the eight lines. Then: case "collapse" (MSE 0, TDWM 0.2820); case "shrunk" (MSE ÷ 100, TDWM unchanged). Move σ and watch the collapse score change.

Defaults: the embedding cost is stuck at step 0 (the bump); the temporal cost reaches the goal. Raise H to 12: the embedding planner sees past the bump. Then lower σ to 3: the temporal planner is stuck on a plateau.

Oracle: 100% at every offset. 30 tasks per offset: one task = 3.3 points, SE of a 50% rate ≈ 9 points.