One-Step Prediction Fidelity Is a Poor Proxy for Planning Performance in JEPA World Models

AI summary

Overview

Research area: Model-based reinforcement learning and latent world models, specifically Joint-Embedding Predictive Architectures (JEPAs) trained end-to-end from pixels for goal-conditioned control.

Technical level: Advanced. The paper assumes familiarity with latent dynamics models, model-predictive control, checkpoint selection, and paired bootstrap statistics.

Scope: The paper tests whether one-step latent prediction loss — the standard quantity used to select, checkpoint and ablate JEPA world models — actually tracks downstream planning performance, using a frozen-decoder protocol that reports prediction error in physical units across three continuous-control environments.

What This Paper Is About

Practitioners who train latent world models routinely assume that a model which predicts the next latent better is a better model, and they use one-step prediction loss to choose between checkpoints, justify architectural changes and report progress. This paper asks whether that assumption fails for latent world models, where the training target is the model's own embedding and the loss is not expressed in interpretable units. The authors build an intervention-based measurement to turn the question from correlational into causal, then move the fidelity metric by fine-tuning and measure goal-conditioned planning on the same checkpoints.

Key Contributions

1. A frozen-decoder protocol. A single decoder is fitted only on encoder embeddings of real observations, then frozen and applied both to real next-observation embeddings and to predicted next embeddings. This reports one-step error in ground-truth physical units and decomposes it into E_repr (already present in the encoder's representation), E_pred (the quantity a practitioner has in mind) and E_excess (the part the predictor adds). The protocol detects when it has no power.

2. A within-environment saturation result. A per-epoch checkpoint ladder on PushT shows fidelity improving monotonically across all three training seeds while planning goes flat after the first epoch, with a correlation of r = -0.010 over the fifteen post-epoch-1 checkpoints.

3. A cross-environment dissociation. The largest fidelity improvement the authors induce buys the least planning: 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].

4. A noise-floor calibration and methodological caution. Untouched baselines vary by up to 26 percentage points across evaluation seeds — larger than any effect the authors induce. An auxiliary state-grounding objective introduced as an instrument is indistinguishable from five additional epochs of ordinary training on every measure computed.

Main Findings

• First epoch looks like the proxy works, and then it stops. On PushT the first epoch improves fidelity by 32% and planning by 4.4 pp. 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. Including the baseline checkpoint raises this to r = -0.88 over the full 40% range, but the authors do not rest on that figure because it is carried entirely by the single epoch 0 to 1 step, and the baseline was evaluated once and shared across training seeds.

• Fidelity keeps falling after planning stops moving. Over epochs 1 to 5 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, and the per-seed changes from epoch 1 to epoch 5 are +2.00, 0.00 and -1.33, differing in sign. Across the fifteen post-epoch-1 checkpoints, latent drift correlates with planning success at r = -0.175.

• This is not a task ceiling. Support strata are defined from the offline dataset alone: the 512 × 512 workspace is partitioned into an 8 × 8 grid, dataset states are counted per cell, and weighted quartiles define low, medium and high support covering 25.2%, 48.5% and 26.3% of dataset states. Low-support states are edge-like (mean 160.7 px from the workspace centre), high-support states centre-like (38.7 px). Low-support planning begins at 80.0%, roughly 20 pp below saturation, peaks at epoch 1 (+9.17 pp), then declines and never recovers, ending at +7.50 pp while fidelity continues to improve. High support, where a genuine ceiling does operate, is negative throughout (-0.85 to -2.56 pp from a baseline of 100%).

• The models are worse where the training data are sparse. The pretrained baseline is strongly stratified by this model-independent measure. Under a ten-planning-seed evaluation, high-support planning exceeds low-support planning by 18.28 pp [+2.38, +39.35], and the high-minus-low contrast in one-step block-position error is -0.002127 [-0.002353, -0.001927]. The three-seed ladder gives a 20.0 pp planning gap at epoch 0.

• Flat aggregates are not cancelling churn. Relative to epoch 0, trajectories flipping from success to failure number 8, 10, 7, 7 and 10 across epochs 1–5 out of 405 baseline successes — roughly 2%, with no trend as fidelity improves.

• No stable exchange rate across environments. The two largest fidelity improvements, both above 91%, produce the only two negative planning estimates. The single resolvable planning improvement, +4.9 pp [+1.1, +8.9], arises from a fidelity change less than half that size and belongs to the control arm, not the intervention. OGBench-Cube's -0.3% produces a larger planning point estimate than TwoRoom's -93.1%. Only the bolded planning interval in the cross-environment table excludes zero.

• The protocol detects when it has no power. Replacing the true action with one sampled from a different transition raises E_pred by 29.1× on PushT, 15.7× on TwoRoom and 2.1× on OGBench-Cube, so the predictors are genuinely action-conditioned. PushT and TwoRoom exceed a persistence and constant-velocity skill level of 0.87. OGBench-Cube is the exception: its constant-velocity skill is -0.74, meaning the learned predictor is worse than linear extrapolation, and its E_excess is only 0.3% of E_pred, so almost all measured error originates in the decoder.

• The result replicates under a decoder-free measure. Latent drift never passes through the decoder, and it exhibits the same dissociation independently: it falls 77% on TwoRoom for -0.7 pp of planning and 32% on PushT for +4.9 pp. Within the PushT ladder, mean E_excess and mean latent drift correlate at r = 0.85. Per transition the two are essentially uncorrelated (PushT r = 0.012, OGBench-Cube r = -0.007; TwoRoom is the exception at r = 0.34–0.59).

• A third intervention shows the same signature. Holding checkpoint, architecture, objective, sample count and epochs fixed, the authors resampled the 200,000 selected training windows to favour rare block positions, raising the low-support share from 27.1% to 43.6%. Pre-specified outcomes were null: low-support block-position error changed by -0.000036 [-0.000101, +0.000032] and low-support planning by +1.245 pp [-3.611, +5.989]. A secondary, uncorrected outcome moved — low-support mean state-prediction nMSE fell by 0.000207 [-0.000386, -0.000019], a 6.7% relative reduction.

• The planning noise floor is larger than the effects. Evaluating untouched baseline checkpoints across planning seeds 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. Five of six intervention points lie within the ambient variation of doing nothing.

• A matched control absorbs the whole intervention effect. Paired planning differences between state grounding and continued training, resampled over training seeds, are -1.33 [-4.22, +1.78] on PushT, +0.22 [-4.44, +4.00] on TwoRoom and +0.89 [-4.00, +5.11] on OGBench-Cube: no consistent sign, all containing zero. Reported without the matched continued-training arm, the natural conclusion would have been that an auxiliary state-grounding loss reduces one-step physical prediction error by more than 90%. That number is real; its attribution would have been wrong.

Methodology in Plain English

The authors need three things that are individually unremarkable and rarely combined: a fidelity measure that is interpretable and comparable across models, an intervention that moves that measure so the question becomes causal rather than a comparison across an uncontrolled model zoo, and a downstream measurement the intervention cannot directly optimize, evaluated with enough statistical care to separate a real effect from seed noise.

The measure works like this. A model's latent space is its own, so raw latent error norms are not comparable across models. The authors therefore fit one small decoder from latents to physical simulator state using only encoder embeddings of real observations, on held-out training episodes, with the world model frozen. That decoder is then frozen and applied to two different inputs on disjoint test episodes: the encoder's embedding of the real next observation, and the predictor's output given the true action. Squared error per state component is divided by that component's training-set variance and averaged, and normalization statistics are shared by all models within an environment. The difference between the two errors, E_excess, is what they argue should be reported: if the predictor were perfect, predicted latents would equal real latents and E_excess would be exactly zero, whatever the decoder does. Conversely, when E_repr is approximately equal to E_pred, almost all measured error originates before the predictor is involved — a condition the authors observe on one environment and which E_pred alone would have concealed.

To move the metric, they warm-start from released checkpoints and continue training for 5 epochs at learning rate 10^-5 and batch size 128, with three training seeds per environment, saving a checkpoint each epoch. One arm uses the original two-term objective only; another adds an auxiliary head predicting standardized physical state from encoder embeddings at weight 0.1, discarded after training. Planning uses the released evaluation unchanged: the cross-entropy method under receding-horizon model-predictive control over 50 goal-conditioned test trajectories, with solver, hyperparameters and test trajectories fixed across variants and the same three planning seeds used at every training seed, so all differences are paired.

The evaluated model is LeWorldModel, an end-to-end JEPA world model trained from pixels with next-embedding prediction plus a regularizer enforcing approximately Gaussian latents. The three environments are PushT (2D manipulation, 5 state dimensions, 41,983 test transitions), TwoRoom (2D navigation through a doorway, 2 dimensions, 16,076) and OGBench-Cube (3D manipulation, 9 dimensions, 38,000); all are continuous-action and trained offline without reward. The decoder is a two-layer MLP with hidden width 512, trained up to 60 epochs with early stopping (patience 10) and batch size 4096; no test episode contributes to fitting it. Statistics use three training seeds per variant per environment, five paired decoder/data-split seeds for fidelity and three paired planning seeds, with 10,000 hierarchical paired bootstrap repetitions resampling at all three levels while preserving variant pairing.

The authors state two limits of the procedure up front: bootstrapping over only three training seeds gives poorly calibrated intervals at the top level, so planning intervals should be read as indicative rather than exact; and planning success is a binomial over 50 trajectories, placing a floor of roughly 4–7 pp on what any single evaluation can resolve.

Why This Matters

Impact on research. The paper does not claim that fidelity interventions never affect planning — one of six effects is resolved and positive, and the first epoch moves both quantities together. What it excludes is a stable, monotone exchange rate of the kind the proxy's use presupposes. Neither the within-task saturation nor the cross-environment pattern licenses ranking checkpoints or architectures by one-step loss. The authors note that the model-based RL community has already named the objective mismatch between model likelihood and control performance, but that this had not been measured for latent world models, where the training target is the model's own embedding. The methodological cautions — decompose a fidelity metric, always run a matched continued-training control, report paired multiple seeds with the seed spread — are portable to any lab evaluating world models.

Real-world applications. These follow from the environments and pipeline the paper studies:

• Robotic manipulation, where a policy or planner is trained against a learned latent dynamics model and deployed on real hardware; choosing a checkpoint by one-step loss may select a model no better at the task.
• Autonomous navigation through constrained spaces such as doorways, the TwoRoom analogue, where planning failure modes are concentrated in regions the offline data under-cover.
• Offline model-based control pipelines generally, where the model-selection step is cheap and the planning evaluation is expensive — the paper shows the cheap signal is not a substitute.
• Benchmarking and evaluation practice in machine learning, where single-run comparisons at these variances can produce an apparent effect of ten or more percentage points from nothing.

Industry relevance. Teams that train world models and report latent prediction loss as evidence of progress are directly affected. The noise-floor result — an unmodified baseline varying by 26 percentage points across evaluation seeds, against a largest observed intervention effect of 4.9 pp — means single-seed comparisons are not informative at this scale, and paired multi-seed evaluation should be the default for this class of comparison. The matched-control result is a warning about attribution: an auxiliary objective with a well-motivated design, a large and tightly-resolved improvement, and reproduction across five evaluation and three training seeds still turned out to be indistinguishable from five extra epochs of ordinary training.

Future Directions

• Rollout-level measurement. The current fidelity measure is one-step and teacher-forced while planning is long-horizon and off-distribution. The authors regard a rollout-level version of the measurement as the experiment that would settle how far the result generalizes, and as the most promising direction; the protocol extends directly.

• A wider within-task sweep. The within-environment ladder covers one environment, and the correlation is estimated over the 15% fidelity range lying beyond the first epoch — about a fifth of the total baseline-to-epoch-5 reduction. An auxiliary-weight or capacity ablation would strengthen the estimate.

• Mechanism. The authors advance no mechanism, having measured none. They name three possibilities to distinguish: one-step teacher-forced error may be dominated

Authors’ 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.