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

Abstract: JEPA-style latent world models can use Euclidean distance to a goal latent as the cost for model-predictive control (MPC). Strong decoding of task variables, however, does not guarantee that this particular cost ranks candidate action sequences by real task progress. We call the latter property \emph{decision-metric alignment}. We introduce Plan-Real Spearman, which measures latent--real rank agreement on random plans, and CEM-stage Spearman, which measures the same agreement as cross-entropy-method (CEM) search concentrates its proposal. We analyze sufficient conditions under which latent distance preserves real-cost rankings, identifying encoder distortion, terminal rollout error, and candidate margins as the controlling quantities. Guided by the observed empirical alignment gap, DA-LeWM augments LeWM with inverse-dynamics and demonstration-conditioned goal-action heads. Across all our experiments, DA-LeWM accelerates convergence and achieves higher online success than LeWM, while probe scores remain similar. These results show that action-conditioned objectives improve the geometry used by Euclidean-cost, CEM-based latent MPC.

• The paper introduces Plan-Real and CEM-stage Spearman diagnostics to measure whether latent Euclidean costs rank candidate action sequences like real task outcomes.
• The paper derives rank-preservation conditions based on encoder distortion, rollout error, and cost margins, while showing that information sufficiency from linear probes does not guarantee decision-metric alignment.
• The paper presents DA-LeWM, which adds inverse-dynamics and goal-action objectives and improves alignment and short-budget success, including gains of 43.4 percentage points on PushT and 10.6 points on Cube, while elite-stage discrimination remains unresolved.

Motivation and core distinction

JEPA-style latent world models for robotic manipulation typically plan by rolling out an action-conditioned predictor and selecting the candidate sequence whose terminal latent is closest, in squared Euclidean distance, to the encoded goal observation. The paper's central observation is that this recipe rests on a property that is neither trained for nor usually measured: whether the latent cost ranks candidate plans consistently with their real environmental outcomes. The authors formalize this as decision-metric alignment, and distinguish it sharply from information sufficiency (decodability of task quantities via linear probes). The two are logically independent: a latent can decode state, action, reward, and value nearly perfectly while its Euclidean geometry reverses the preference between a genuinely low-cost plan and a nearby-but-poor one. This reframing matters because the standard evaluation question — "what does the latent encode?" — is incomplete when the planner consumes distances rather than decoded values.

Diagnostics

To make alignment testable, the paper introduces two rank-based diagnostics:

• Plan-Real Spearman: for each of $n=30$ held-out (start, goal) pairs, $N=64$ random candidate action sequences are scored both by the model's latent cost $\|\hat z_H - z_g\|^2$ and by environment-rollout task cost; the mean pair-level Spearman correlation measures global latent–real rank agreement.
• CEM-stage Spearman: the same correlation computed at three stages of cross-entropy-method search (random iteration 0, mid iteration 15, elite final top-$E$), reflecting that CEM acts on relative order within progressively concentrated proposal distributions.

Ranking is the relevant quantity because CEM refits its proposal from candidate order rather than calibrated costs.

Sufficient-condition analysis

The analytical contribution identifies three controlling quantities for rank preservation. Under pointwise terminal-rollout consistency ($\|\hat z_H - f_\theta(s_H)\|_2 \le \epsilon_H$) and encoder bi-Lipschitzness with constants $\mu_f \le L_f$, every plan's latent cost is bracketed:

$$\mu_f\, c_\text{real}(\mathbf{a}) - \epsilon_H \;\leq\; c_\text{lat}(\mathbf{a}) \;\leq\; L_f\, c_\text{real}(\mathbf{a}) + \epsilon_H$$

A corollary gives margin-implied rank preservation: if real-cost gaps exceed $(L_f/\mu_f)$-scaled margins plus $2\epsilon_H$, ordering is preserved, and Plan-Real Spearman reaches $+1$ when all pairs satisfy it. The analysis is deliberately modest in scope: the authors explicitly note that SIGReg discourages collapse at the distribution level but does not certify bi-Lipschitz behavior or bound Jacobian singular values, that training losses do not establish the constants off-distribution, and that no numerical estimates of $N=64$0, $N=64$1, or $N=64$2 are attached. An empirical soft-margin check on PushT supports consistency: the fraction $N=64$3 of pairs satisfying the margin condition tracks Plan-Real Spearman with pooled Pearson $N=64$4 over 100 points, and the Kendall lower bound $N=64$5 holds for all measured pairs — though the bound remains loose (empirical gap $N=64$6).

DA-LeWM

Guided by the observed alignment gap, DA-LeWM augments LeWM with two lightweight auxiliary heads: an inverse-dynamics head $N=64$7 and a demonstration-conditioned goal-action head $N=64$8, each weighted at $N=64$9. Both heads are discarded at test time, so inference-time MPC is unchanged. A mechanistic claim links inverse-action training to norm structure: the correlation between latent-displacement magnitude and action magnitude rises from $\|\hat z_H - z_g\|^2$0 (LeWM) to $\|\hat z_H - z_g\|^2$1 (inverse-only) and $\|\hat z_H - z_g\|^2$2 (DA-LeWM) over 38,400 held-out transitions. The authors correctly flag this as empirical evidence, not a theorem — action information could in principle be encoded directionally rather than in norms.

Empirical results

The experimental design cleanly separates collapse from alignment. Removing SIGReg contracts the latent-cost dynamic range from $\|\hat z_H - z_g\|^2$3–$\|\hat z_H - z_g\|^2$4 to roughly $\|\hat z_H - z_g\|^2$5, collapsing Plan-Real Spearman to near zero and online success dramatically (PushT $\|\hat z_H - z_g\|^2$6; Reacher $\|\hat z_H - z_g\|^2$7). Among non-collapsed variants, the key dissociation emerges: linear-probe $\|\hat z_H - z_g\|^2$8 varies by less than $\|\hat z_H - z_g\|^2$9 across four variants while Plan-Real Spearman moves from $E$0 to $E$1–$E$2 and PushT online success spans 43 percentage points ($E$3). Inverse-only accounts for most of the ranking lift; the combined objective delivers most of the additional control gain.

CEM-stage results add nuance. Action-supervised variants improve random-stage agreement significantly (paired $E$4), but elite-stage correlations are near zero for every variant, indicating shared local saturation among near-optimal candidates. Notably, the paper reports that mid- and elite-stage gains over LeWM are not statistically significant even for DA-LeWM — so the diagnostic benefit is concentrated at the global-ordering stage.

Cross-environment short-budget results show DA-LeWM gains of $E$5 pp on PushT and $E$6 pp on Cube, with Reacher essentially flat (its baseline alignment is already high at $E$7, leaving little headroom). Under ten-epoch learning curves, DA-LeWM's advantage persists across all four tasks, supporting accelerated convergence rather than a favorable checkpoint. Against published baselines, DA-LeWM ranks first on PushT ($E$8) and Reacher ($E$9), surpasses PLDM on Cube, and reaches $\|\hat z_H - f_\theta(s_H)\|_2 \le \epsilon_H$0 on TwoRoom — above LeWM by 9.0 pp though below DINO-WM by 4.0 pp.

Limitations and open questions

The paper is candid about scope. Evidence covers four short-horizon simulated tasks with ViT-Tiny encoders, Euclidean goal costs, and CEM only; each configuration uses a single training run, so uncertainty reflects evaluation rather than initialization variation. Cube's contact-sparse dynamics produce exact real-cost ties that make Spearman undefined or low-support, so no Cube alignment metric is reported. The all-heads ablation used implementation-specific reward/value proxy targets (with a stated coordinate-layout mismatch on PushT), so it cannot determine whether correctly specified reward/value supervision helps or hurts. Sensitivity to the inverse weight $\|\hat z_H - f_\theta(s_H)\|_2 \le \epsilon_H$1 is untested, and gradient-based planning, partial observability, larger backbones, learned rewards, and generalization to real robots remain open. Whether elite-stage saturation can be addressed by representation changes at all — given that it appears uniform across every variant — is left unresolved.

Conclusion

This paper makes a precise, measurable argument that latent world models intended for Euclidean-cost MPC should be evaluated for the ordering their cost induces, not only for what they encode. The Plan-Real and CEM-stage Spearman diagnostics, the sufficient-condition analysis naming encoder distortion, rollout error, and candidate margins, and the demonstration that probe-invariant representations can differ by tens of success-percentage points together establish decision-metric alignment as a distinct and consequential property. DA-LeWM shows that lightweight action-conditioned supervision improves the planner-facing geometry without altering the planner, yielding consistent convergence and success gains across four environments — with the honest caveat that local, near-optimal discrimination remains unsolved across all configurations studied.