One-Step Prediction Fidelity Is a Poor Proxy for Planning Performance in JEPA World Models
Abstract
Latent world models are routinely selected, checkpointed and ablated on one-step prediction loss, on the assumption that a model which predicts the next latent better is a better model. In model-based RL this assumption is known to be unreliable; we ask whether it also fails for latent world models, where the training target is the model's own embedding and the loss is not expressed in interpretable units. We introduce a frozen-decoder protocol reporting one-step error in physical units, decomposed into the part already present in the encoder's representation and the part the predictor adds. Applying it to a recent end-to-end JEPA world model across three continuous-control environments, we move the metric by fine-tuning and measure goal-conditioned planning on the same checkpoints, with three training seeds and paired planning seeds throughout. Within one environment, a per-epoch ladder shows the relationship saturating after a single epoch: the first epoch improves fidelity by 32% and planning by 4.4 pp, but across the fifteen checkpoints that follow -- a further 15% fidelity improvement, monotone in every training seed -- the correlation between one-step error and planning success is r = -0.010. This is not a task-ceiling artifact: stratifying by a model-independent measure of training-data support, planning in the lowest-support stratum starts at 80%, leaving ample headroom, and is nonetheless flat from epoch 1 onward. Across environments the exchange rate is equally unstable: a 93% fidelity reduction on TwoRoom yields -0.7 pp of planning, while a 40% reduction on PushT yields +4.9 pp [+1.1, +8.9]. All results replicate under a decoder-free measure, latent drift. We also show that a matched continued-training control absorbs the entire fidelity gain we had attributed to an auxiliary objective.
Figures & tables
| Planning success (%) | |||||||
|---|---|---|---|---|---|---|---|
| Epoch | (%) | Drift | Low sup. | Med. sup. | High sup. | All | |
| 0 | 0.003336 | — | 1.123 | 80.00 | 90.14 | 100.00 | 90.00 |
| 1 | 0.002268 | 0.839 | 89.17 | 94.84 | 99.15 | 94.44 | |
| 2 | 0.002109 | 0.814 | 85.83 | 95.78 | 97.44 | 93.56 | |
| 3 | 0.002071 | 0.787 | 86.67 | 95.78 | 99.15 | 94.22 | |
| 4 | 0.002026 | 0.774 | 88.33 | 95.78 | 99.15 | 94.67 | |
| Env | Intervention | (%) | (%) | planning (pp) |
|---|---|---|---|---|
| TwoRoom | Continued | |||
| TwoRoom | State-grounded | |||
| PushT | State-grounded | |||
| PushT | Continued | |||
| Cube | State-grounded | |||
| Cube | Continued |
Latent World Models
Latent world models learn compact predictive state representations to simulate how environments change and support planning, control, and reasoning.
See all 212 papers