PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Jiawei Wang, Ke Rui, Yushen Zuo, Yichun Feng, Minglei Li

arXiv:2608.18746v1cs.LGcs.CV

TL;DR

Latent world models may encode task information without making Euclidean goal distance a reliable ranking of candidate plans. This paper formalizes and diagnoses decision-metric alignment, then introduces action-conditioned supervision; DA-LeWM accelerates convergence and achieves higher online success than LeWM across four environments despite similar probe scores.

  • Problem

    Linear probes can show task information while leaving unanswered whether latent Euclidean distances rank candidate plans consistently with environmental outcomes.

  • Method

    The paper introduces Plan-Real and CEM-stage Spearman diagnostics, analyzes rank-preservation conditions, and augments LeWM with action-conditioned auxiliary supervision in DA-LeWM.

  • Results

    DA-LeWM accelerates convergence and achieves higher online success than LeWM across all four environments, despite similar linear-probe scores.

  • Takeaways & Limitations

    Decision-metric alignment matters for latent MPC, and action-conditioned supervision improves the cost geometry exposed to the planner.

  • Takeaways & Limitations

    The proposed mechanistic link between inverse-prediction error and norm alignment is empirical rather than a theorem, because action information could be encoded directionally.

Abstract

from arXiv · show

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.

1 Introduction

The introduction argues that latent information sufficiency does not ensure decision-metric alignment for Euclidean-cost latent MPC. It introduces rank-agreement diagnostics and action-conditioned objectives to improve planning geometry and online success.

  • Motivation: Small training-objective changes can substantially alter online planning success while state, action, reward, and value probe scores remain nearly unchanged.This motivates evaluating latent geometry for planning rather than asking only what task information the latent encodes.
  • Decision-Metric Alignment: Decision-metric alignment requires latent Euclidean-cost rankings of candidate plans to track their real-world outcomes, beyond information sufficiency.The distinction is geometric: latent metric structure must be compatible with action consequences in the environment.
  • Diagnostics: Plan-Real Spearman measures latent–real rank correlation across random plans, while CEM-stage Spearman measures it among converged elite candidates.These diagnostics target the rankings used during latent MPC, including the candidates the agent actually executes.
  • Rank-Preservation Analysis: Rank preservation depends on encoder distortion, terminal rollout error, and candidate margins.The paper analyzes sufficient conditions under which latent distance preserves real-cost rankings.
  • DA-LeWM: DA-LeWM adds action-conditioned auxiliary supervision, improving latent–real rank agreement and achieving higher online success than LeWM despite similar linear-probe scores.The reported improvements hold across all four environments.

2 Background and Problem Setup

The paper builds on LeWM, a JEPA-style latent world model whose encoder and action-conditioned predictor support CEM-based latent MPC. It frames auxiliary objectives around a central question: whether they improve the planner’s fixed Euclidean cost, not merely probe or downstream-control performance.

  • LeWM world model: LeWM encodes observations with a Vision Transformer and linear projector, while an autoregressive predictor estimates future latents conditioned on actions.The encoder is fθ, and the predictor is gϕ.
  • Representation regularization: SIGReg matches aggregate embeddings to an isotropic Gaussian and discourages collapse, but does not ensure local metric isotropy, injectivity, or bi-Lipschitz behavior.It also does not bound encoder-Jacobian singular values, so distribution-level regularization alone leaves local geometry unconstrained.
  • Latent MPC: At test time, MPC uses CEM to optimize latent cost by rolling out the predictor and refitting proposals toward the top-E lowest-cost candidate sequences.The rollout produces ˆzH, and CEM iteratively refits its proposal distribution to elite candidates.
  • Goal acquisition: Evaluation goals come from held-out demonstrations at a fixed offset ∆ahead, ensuring in-distribution and reachable goal latents, while the analysis remains independent of goal acquisition.Deployment goals would typically come from a single human demonstration or sub-goal generation.
  • Auxiliary objectives: The study evaluates inverse-dynamics, goal-conditioned action, reward-proxy, and value-proxy heads by asking how they affect the planner’s fixed Euclidean cost.These auxiliary objectives are commonly judged through probes or downstream control, but this work focuses on their impact on planning.

3 Information Sufficiency vs. Decision-Metric Alignment

Information sufficiency concerns probe-predictable task quantities, whereas decision-metric alignment concerns whether latent planning costs preserve real-cost rankings. The section distinguishes these logically independent properties and evaluates rank agreement on random and CEM-concentrated candidate plans.

  • Definitions: Information sufficiency requires probe regressors to predict task quantities within held-out error ϵ, tested for state, inverse-action, reward, value, and goal-action.The probes are linear, and each task quantity uses a specified input x_k.
  • Definitions: Decision-metric alignment requires latent and real planning costs to induce the same ordering for every pair of candidate action sequences.This is an ordinal requirement: numerical equality between latent cost and reward is unnecessary.
  • Definitions: A latent can encode all task-relevant information yet remain decision-metric misaligned when its geometry poorly ranks candidates by Euclidean distance.Thus, information sufficiency and decision-metric alignment are logically independent.
  • Alignment diagnostics: Plan-Real Spearman computes mean Spearman rank correlation across N = 64 candidate action sequences for each of n = 30 held-out (start, goal) pairs.On PushT, task distance is Euclidean distance over all seven simulator coordinates, including agent and block pose and agent velocity.
  • Alignment diagnostics: CEM-stage Spearman measures rank agreement at random, mid, and elite stages, testing global ordering, proposal-shaping ordering, and local discrimination respectively.This addresses the possibility that agreement on random plans fails as CEM concentrates candidates near low-cost regions.

4 DA-LeWM: A Decision-Aligned Latent World Model

DA-LeWM augments LeWM with inverse-dynamics and goal-conditioned action heads to shape action-relevant latent geometry. An analytical account identifies encoder distortion, terminal rollout error, and candidate margins as determinants of latent-cost rank preservation, while emphasizing that these conditions are assumptions rather than training guarantees.

  • DA-LeWM objective: DA-LeWM adds inverse-dynamics and goal-conditioned action losses to the base LeWM objective, with α = 0.1 and β = 0.1 as cross-environment defaults.Both auxiliary heads are two-layer MLPs with hidden width 256.
  • DA-LeWM objective: Goal-action supervision uses four-frame demonstrations and predicts normalized five-action blocks for goals 5, 10, or 15 raw steps ahead.The targets reflect demonstrated progress rather than optimizer-computed globally optimal actions; squared-error training may average incompatible multimodal actions.
  • Analytical rank preservation: Under pointwise rollout consistency and encoder bi-Lipschitz assumptions, latent and real costs differ within a bracket combining encoder distortion and terminal rollout error.The analysis defines real cost as ||sH(a) − sg||Σ and latent cost as ||ẑH(a) − zg||2.
  • Limitations and mechanism: Training losses and SIGReg diagnose alignment-related behavior but do not establish the analytical assumptions or their off-distribution constants.Inverse-action prediction is presented as a mechanistic claim, not a theorem, and need not force norm alignment.
  • Analytical rank preservation: Rank preservation requires low encoder distortion and sufficiently large candidate margins, which may shrink as CEM concentrates its proposal.This motivates measuring agreement across random, intermediate, and elite candidate populations.

5 Experiments

Across four environments, experiments show that action-conditioned supervision improves the latent geometry used by unchanged Euclidean-cost MPC, increasing decision-metric alignment, online success, and convergence speed. Probe scores remain nearly unchanged, separating control-relevant alignment from information sufficiency.

  • Collapse sanity check: Removing SIGReg flattens latent-cost variation from 3−30× to ≈ 1.005× and coincides with near-zero Plan-Real Spearman and large success drops.This collapse pattern appears on PushT, TwoRoom, and Reacher.
  • Decision-metric alignment: Action-supervised variants raise Plan-Real Spearman from +0.280 for LeWM to +0.410–+0.420 while probe R2 scores vary by at most 0.03.Online success spans 49.3%→92.7%, isolating decision-metric alignment from information sufficiency.
  • Decision-metric alignment: Inverse-only raises PushT Plan-Real Spearman from +0.280 to +0.420 and success from 49.3% to 64.0%, while DA-LeWM reaches 92.7% success with +0.412 Spearman.Held-out transition geometry also improves from −0.03 for LeWM to +0.38 for inverse-only and +0.43 for DA-LeWM.
  • CEM-stage alignment: DA-LeWM achieves the highest random-stage CEM Spearman at +0.536, while action-supervised variants gain +0.112 to +0.133 over LeWM and elite-stage correlations approach zero.Mid-stage correlations remain +0.249–+0.261 for action-supervised variants versus +0.227 for LeWM.
  • Cross-task control: DA-LeWM improves one-epoch success by 43.4 pp on PushT and 10.6 pp on Cube, then maintains higher success throughout ten-epoch training across all four tasks.Against published baselines, it ranks first on PushT at 98.7% and Reacher at 87.3%, while auxiliary heads are discarded at evaluation.

6 Related Work

Prior work spans joint-embedding world models, auxiliary representation objectives, learned-model planning, and decision-aware representations. This work complements these directions by auditing whether Euclidean latent goal distance preserves the rankings consumed by MPC and CEM.

  • JEPA and latent world models: JEPA-style methods predict in embedding space, while LeWM applies joint-embedding prediction to action-conditioned visual dynamics for manipulation.Masked autoencoding and self-distillation provide complementary self-supervised approaches; PLDM learns latent dynamics for offline reward-free control.
  • Auxiliary objectives in RL and world models: Inverse dynamics, contrastive learning, self-prediction, goal-conditioned imitation, and hindsight relabeling motivate the auxiliary objectives used in related reinforcement-learning and world-model work.These methods connect representation learning to action relevance, pixel-based RL, demonstrations, and relabeled goals.
  • MPC and planning with learned models: PETS, Dreamer, MuZero, TD-MPC, and DayDreamer demonstrate learned-model planning through rewards, latent imagination, value-equivalent dynamics, and physical-robot control.The paper isolates a complementary requirement: accurate prediction or decoding need not ensure correct relative ordering when Euclidean goal distance is the planning cost.
  • Decision-aware and bisimulation-based representations: Bisimulation metrics and value equivalence ground similarity in rewards, transitions, or predictive consequences, whereas Plan-Real and CEM-stage Spearman audit the fixed metric consumed by the planner.The proposed audits measure ordering globally and as CEM concentrates its proposal, rather than learning a new task-specific metric.

7 Conclusion · Appendix

The paper distinguishes information sufficiency from decision-metric alignment in latent MPC and introduces rank-agreement diagnostics for Euclidean goal costs. Its evidence is limited to short-horizon simulated settings and specific model, cost, and planning configurations.

  • 7 Conclusion: The study examines when Euclidean goal distance is a reliable objective for latent MPC.
  • 7 Conclusion: Plan-Real and CEM-stage Spearman operationalize the distinction between information sufficiency and decision-metric alignment.
  • 7 Conclusion: Encoder distortion, terminal rollout error, and candidate margins control whether latent costs preserve real-cost rankings.
  • 7 Conclusion: Similar linear-probe scores can coexist with substantial differences in Plan-Real agreement across non-collapsed variants.
  • 7 Conclusion: The evidence covers four short-horizon simulated tasks, LeWM-family checkpoints with ViT-Tiny, Euclidean goal costs, and CEM.
  • 7 Conclusion: Each configuration has one training run, so uncertainty reflects evaluation variation rather than training-initialization variation.
  • 7 Conclusion: Plan-Real requires simulator rollouts and loses statistical power when exact real-cost ties occur.
  • 7 Conclusion: Generative AI assisted language editing and code review but was not part of the method or evaluation.

A Experimental Details … Information Probes

The experiments use a modular LeWM implementation with shared latent-model backbones, controlled training and evaluation settings, and diagnostics that distinguish latent information from decision-metric alignment. Information probes are standardized across variants, while exploratory reward/value supervision is explicitly limited by target-specification issues.

  • A Experimental Details: The implementation adds auxiliary decision heads and losses modularly on top of the open-source LeWM codebase without changing the base world-model architecture.All variants share the same ViT-Tiny encoder and Transformer predictor configuration.
  • Auxiliary Head Architecture: Auxiliary heads include inverse-action, goal-conditioned action, reward-proxy, and value-proxy predictors, with the canonical PushT DA-LeWM checkpoint using only the first two.Each head is a two-layer MLP; the all-heads configuration instantiates all four.
  • Auxiliary Head Architecture: The all-heads reward/value ablation is exploratory and implementation-specific because its targets do not match the intended PushT state fields.No corrected all-heads checkpoint was trained, so the result cannot establish whether correctly specified reward/value supervision helps or hurts.
  • Training Configuration: Training comparisons control optimizer, data, model settings, and randomness, while the base loss remains identical across variants except No-SIGReg.Matched-budget ablations use one epoch; learning-dynamics comparisons use 100 epochs and analyze the first ten checkpoints.
  • Computing Infrastructure: Experiments run on one 80 GB NVIDIA A100-SXM4-80GB or A800-SXM4-80GB GPU per process using the specified Linux, Python, PyTorch, and Lightning environment.Training uses bfloat16 mixed precision, while diagnostics use float32 inference and double-precision host aggregation where required.
  • Evaluation Configuration: Online evaluation averages 3 × 50 = 150 episodes per reported mean using identical five-block CEM planning and receding-horizon parameters across environments.The controller uses 300 online CEM samples, 30 iterations, and an elite size of 30; each five-block horizon executes 25 raw actions.
  • Plan-Real and CEM-stage Spearman: Plan-Real and CEM-stage Spearman diagnostics compare latent and environment costs on random candidate plans and at random, middle, and elite CEM stages.Plan-Real samples 30 validation start-goal pairs and 64 candidates per pair; CEM-stage uses 300 samples × 30 iterations with a top-30 elite set.

B Geometric Signature of Inverse-Action Training

Inverse-action training produces a measurable geometric signature: latent displacement magnitude becomes moderately positively correlated with action magnitude over 38,400 held-out transitions. This observation supports Claim 1 empirically but is not asserted as a logical consequence of inverse loss alone.

  • Geometric Signature of Inverse-Action Training: Over 38,400 held-out transitions, the study measures global rank correlation between latent displacement magnitude ∥∆z∥ and action magnitude ∥a∥.The correlation is evaluated on held-out transitions rather than training data.
  • Geometric Signature of Inverse-Action Training: Inverse-action training changes a near-zero baseline correlation into a moderate positive correlation between ∥∆z∥ and ∥a∥.This change is presented as a falsifiable empirical signature of Claim 1.
  • Geometric Signature of Inverse-Action Training: The observed correlation is not claimed to follow logically from inverse loss alone, because the inverse head could recover actions through direction rather than norm.The passage distinguishes the empirical signature from a guaranteed consequence of the objective.

C Additional Ablations … D Scope of the Sufficient-Condition Diagnostics

The ablations show that action-conditioned weighting improves global latent–real alignment but can regress across environments, motivating β = 0.1 as the conservative default. The sufficient-condition diagnostics are treated as observable tests rather than numerically identified guarantees.

  • Per-Method Probes and Online Success on PushT: Probes detect No-SIGReg’s information loss but do not distinguish the four non-collapsed methods.No-SIGReg produces negative or near-zero R2 probe scores, whereas the other four methods remain unresolved by probes.
  • Goal-Action Weight Sensitivity: +0.448 ± 0.207 Plan-Real Spearman and (+0.549, +0.232, −0.038) CEM-stage Spearman occur at β = 0.3 on PushT.The Plan-Real and random-stage gains versus LeWM are significant, but mid- and elite-stage differences are not.
  • Goal-Action Weight Sensitivity: The β = 0.3 and separately initialized β = 0.1 checks reproduce global alignment improvement and shared local saturation.Both checks show significant global and random-stage gains over LeWM, with nonsignificant mid- and elite-stage differences.
  • Goal-Action Weight Sensitivity: +0.436 ± 0.224 Plan-Real Spearman and (+0.542, +0.237, +0.028) CEM-stage Spearman occur for a separately initialized β = 0.1 checkpoint.Its Plan-Real and random-stage gains over LeWM are significant, while mid- and elite-stage gains are not.
  • Goal-Action Weight Sensitivity: 78.0% Reacher success at β = 0.3 trails 84.0% at β = 0.1 and 82.0% for the LeWM baseline, motivating β = 0.1 across environments.The β = 0.3 regression supports the conservative cross-environment default.
  • D Scope of the Sufficient-Condition Diagnostics: Proposition 1 provides sufficient conditions, not numerically identified constants, because average prediction loss and finite held-out pairs cannot certify ϵH, µf, or Lf.The paper therefore evaluates observable implications through diagnostics rather than attaching numerical estimates to these constants.

E Elite-Neighborhood Local Latent Geometry

The elite-neighborhood diagnostic finds no evidence that R/V auxiliary losses distort local latent geometry. Within-elite latent–real distance association remains weak across variants, including DA-LeWM and LeWM.

  • Diagnostic procedure: The diagnostic compares latent and real-state pairwise geometry within final CEM elite sets across 15 held-out PushT start-goal pairs.It uses 30 elites per pair and computes a diagonal-Mahalanobis real-state distance against predicted final-step latent distance.
  • Results: Across variants, coefficient of variation spans 0.685–0.767 and log spread spans 1.87–2.12 in the elite neighborhood.All-heads does not differ significantly from inverse-only, DA-LeWM, or LeWM on either statistic.
  • Results: All all-heads contrasts are non-significant, providing no evidence that R/V losses produce detectable increases in local anisotropy.The paired-comparison table reports only one significant remaining contrast: DA-LeWM versus LeWM log spread, p = 0.026.
  • Results: Pearson correlation between latent and real-state distances remains uniformly low at +0.050–+0.092 across all four variants.The isolated DA-LeWM–LeWM log-spread contrast does not support an R/V-specific effect because all three contrasts involving all-heads are non-significant.

F Soft-Margin Consistency Check for Corollary 1 … H Training Dynamics and Checkpoint Selection

The soft-margin condition consistently tracks latent–real ranking agreement on PushT, though its Kendall lower bound is conservative. Cross-environment diagnostics show limited headroom on Reacher, are uninformative for contact-sparse Cube, and motivate a common epoch-10 checkpoint for training comparisons.

  • F Soft-Margin Consistency Check for Corollary 1: The tie-aware bound τa ≥ q(2p −1) reduces to τa ≥ 2p −1 for the reported full-state PushT pairs, with ties contributing zero.Here q = 1 because all reported pairs have distinct real costs.
  • F Soft-Margin Consistency Check for Corollary 1: +0.9996 Pearson correlation across five variant means shows that soft-margin satisfaction p tracks Plan-Real Spearman ρs consistently.Pooling all 100 variant–pair points gives Pearson(p, ρs) = +0.895 (p = 3.54×10−36), while the bound holds for every pair.
  • F Soft-Margin Consistency Check for Corollary 1: The empirical gap τa − (2p −1) is 0.670 ± 0.083, indicating that the Kendall lower bound is conservative and loose.The bound τa ≥ 2p −1 holds for all 100 variant–pair observations.
  • G Cross-Environment Spearman Scope: Cross-environment Spearman diagnostics are informative on Reacher but not meaningfully summarized on contact-sparse Cube, where PushT remains the detailed main-body evaluation.The appendix collects Reacher Plan-Real and CEM-stage measurements and explains the narrower Cube interpretation.
  • Reacher: Reacher has high baseline alignment, with LeWM achieving Plan-Real Spearman Sp = +0.504 and auxiliary objectives leaving behavioral Spearman essentially unchanged.These measurements do not estimate the constants or tightness of Proposition 1.
  • Reacher: Alignment-improving variants on Reacher vary within sampling noise, while online success differs by less than 3 percentage points across variants.The comparison uses n = 30 pairs with std ∼0.22.
  • Cube: Tie-Dominated Real Costs: Cube random plans often produce tied real costs, including tie blocks of size 63 among 64 candidates and fully constant cost vectors.Average-rank Spearman is undefined for constant vectors and low-support when only one or two distinct costs remain.
  • H Training Dynamics and Checkpoint Selection: Epoch 10 is used as the common pre-specified endpoint because intermediate checkpoint performance is non-monotonic, especially on Reacher and Cube.Matched 100-epoch schedules evaluate epochs 1–10, with 50 episodes for each of 3 evaluation seeds per task; separate task-specific sweeps are not pooled.
Loading 2608.18746v1…