One-Step Prediction Fidelity Is a Poor Proxy for Planning Performance in JEPA World Models
Latent world models for control are routinely selected and checkpointed by one-step next-embedding prediction loss, on the assumption that lower loss implies better planning. This paper tests that proxy in JEPA architectures by introducing a frozen-decoder protocol that reports one-step error in interpretable physical units, then applies fine-tuning interventions and measures goal-conditioned planning with paired seeds. It finds the relationship saturates after a single epoch within a task and has no stable exchange rate across tasks, implying that latent prediction loss is an unreliable stand-in for planning competence.
The paper makes a clean, controlled case that one-step latent prediction fidelity is a poor proxy for planning in JEPA world models. Its central claim—that within a task the relationship saturates after coarse initial gains and lacks a stable cross-task exchange rate—is well supported by the PushT ladder and cross-environment dissociation. The authors are admirably candid about limitations, though the scope is narrow: one architecture, three low-dimensional tasks, and warm-started five-epoch fine-tunes rather than full training runs.
The frozen-decoder protocol is methodologically sharp: it converts uninterpretable latent-space norms into physical units and decomposes total error into $E_{\mathrm{repr}}$ (encoder readability) and $E_{\mathrm{excess}} = E_{\mathrm{pred}} - E_{\mathrm{repr}}$, flagging degenerate predictors such as OGBench-Cube where $E_{\mathrm{excess}}$ is negligible. The within-environment PushT result is stark and convincing: after epoch 1, 'fidelity falls a further $11.2\%$ and latent drift a further $8.5\%$, both monotonically and on every individual training seed. Aggregate planning moves $+0.22$ pp,' yielding a post-epoch-1 correlation of $r = -0.010$ that is hard to dismiss as noise.
First, the two interventions—continued training and an auxiliary state-grounding loss—turn out to be indistinguishable, so the paper effectively has only one causal instrument (five extra epochs), not two. As the authors note, 'An auxiliary state-grounding objective we introduced is indistinguishable from five additional epochs of ordinary training on every measure we computed. Reported without a matched control it would have appeared as a large, tightly-resolved improvement.' Second, OGBench-Cube is essentially a failed run: 'Its constant-velocity skill is $-0.74$: the learned predictor is worse than linear extrapolation,' and the protocol has almost no power there, yet it is still included in the cross-environment narrative. Third, the detailed fidelity ladder covers only one environment.
The evidence for within-task saturation is strong, but the cross-environment claim is weaker because the comparison is almost entirely between rather than within environments. The authors are forthright about this: 'Because the two interventions are indistinguishable (Section 8), the fidelity variation in Table 2 is almost entirely between environments.' The link to prior work is fair: Lambert et al. showed that 'the likelihood of one-step ahead predictions is not always correlated with control performance,' and this paper usefully extends that observation to latent JEPA models with a protocol that yields interpretable units. The matched control for the auxiliary objective is a valuable methodological model for the field.
Reproducibility is middling. Hyperparameters and evaluation seeds are reported, but code and pretrained checkpoints are not mentioned as released. More troubling, the planning noise floor is enormous: 'Evaluating the untouched baseline checkpoints across planning seeds — changing nothing about the model — gives 84–96 on PushT (sd 6.0), 82–88 on TwoRoom (sd 3.1) and 58–84 on OGBench-Cube (sd 13.1). The unmodified OGBench-Cube baseline varies by 26 percentage points across evaluation seeds alone, against a largest observed intervention effect anywhere of 4.9 pp.' The authors correctly advocate paired multi-seed evaluation, yet three training seeds barely calibrate bootstrap intervals, and 50 planning trajectories set a resolution floor of roughly 4–7 pp, as they acknowledge.
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.
Pick a starting point or write your own. Challenges run in the background, so you can keep reading while the AI investigates.
No challenges yet. Disagree with the review? Ask the AI to revisit a specific claim.