Decision-Metric Alignment in Latent World Models: Diagnostics and Action-Conditioned Objectives for MPC Planning

Probe scores lie. That is the central finding here, stated plainly: a latent world model can decode state, action, reward, and value with high accuracy while still ranking candidate action sequences in the wrong order when those rankings are what actually drive the planner.

The paper calls this gap decision-metric alignment and distinguishes it sharply from information sufficiency. The distinction matters because JEPA-style world models used for model-predictive control score candidate plans by Euclidean distance to a goal latent. Whether that distance correctly orders plans by real-world outcome is a geometric question, not an informational one. A latent can encode everything relevant and still arrange it so that Euclidean distance reverses the preference between two candidates.

To make this testable, the paper introduces Plan-Real Spearman — rank correlation between latent costs and real environment costs over randomly sampled candidate plans — and CEM-stage Spearman, which tracks the same correlation as the cross-entropy method concentrates its proposal toward elite candidates. The stage-wise version matters because CEM does not act on random plans; a cost that orders random candidates well can still fail to discriminate among near-optimal ones.

The analytical account identifies three controlling quantities: encoder distortion, terminal rollout error, and candidate margins. Margins shrink as CEM converges, which is exactly why elite-stage correlations are near zero across every variant tested — not a failure of any particular method, but a structural feature of how the optimizer concentrates its search.

The practical intervention is DA-LeWM, which adds two lightweight auxiliary heads to the base LeWM model: an inverse-dynamics head and a goal-conditioned action head. Both are discarded at test time; the planner still uses plain Euclidean goal distance. The inverse-dynamics head does most of the measurable work. On PushT, it raises Plan-Real Spearman from the LeWM baseline to a substantially higher value, and the combined DA-LeWM objective pushes online success higher still — while probe scores across all four non-collapsed variants differ by less than a few percentage points. That last point is the sharpest result: the thing practitioners typically measure to validate a representation does not track the thing that determines planning quality.

Across all four environments — PushT, Reacher, Cube, and TwoRoom — DA-LeWM accelerates convergence and reaches higher online success than LeWM under both short and extended training budgets. The gains persist over ten epochs, ruling out a simple early-training artifact.

The limitations are stated honestly: four short-horizon simulated tasks, a single training run per configuration, a ViT-Tiny backbone, and CEM as the only planner. Plan-Real Spearman also loses statistical power when real costs are tied, which is why Cube results rely on online success rather than ranking diagnostics. None of this undermines the core argument; it just marks where the evidence stops.