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

cs.LGAug 19, 2026

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

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

Organizations: Simple AI, Beijing, China · University of Chinese Academy of Sciences, Beijing, China

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

Figures & tables

Figure 1: Information sufficiency does not imply decision-metric alignment. In each panel, the left schematic places candidate plans around the goal latent

z_{g}

, while the right plot compares their real costs (horizontal) with their latent Euclidean costs (vertical). In (a), the two costs induce the same ordering. In (b), candidate 1 is truly low-cost but lies far from

z_{g}

, whereas candidate 3 has higher real cost but lies close. Latent distance therefore reverses their preference even though candidate identity and rank remain decodable.

Figure 2: Plan-Real Spearman measurement procedure. This operationalizes the candidate-order comparison in Figure 1 . For each held-out pair $k$ , the same $N=64$ action sequences are evaluated by the world model and the environment, yielding paired latent- and real-cost vectors. Spearman gives a pair-level rank correlation $\rho_{k}$ ; Plan–Real averages the defined correlations among $n=30$ sampled pairs. CEM-stage Spearman reuses the paired scoring procedure with candidates from the selected CEM population.

Figure 3: Probe accuracy vs. online success on PushT (

3

evaluation seeds). Probe scores are nearly identical across the four non-collapsed variants, while online success differs sharply.

	Success (%)		Plan-Real Sp
Task	LeWM	No-SIG	LeWM	No-SIG
PushT	$49.3\pm 12.2$	$\phantom{0}2.0\pm 2.0$	$+0.280$	$+0.031$
TwoRoom	$98.0\pm\phantom{0}2.0$	$41.3\pm 6.1$	$+0.549$	$+0.012$
Reacher	$82.0\pm\phantom{0}2.0$	$10.7\pm 8.1$	$+0.504$	$+0.001$

Table 1: Effect of removing SIGReg across three environments (

3

evaluation seeds, mean

\pm

std).

Figure 4: Online success throughout ten training epochs (

3

evaluation seeds per task, mean

\pm

std). Shading shows one standard deviation. Annotations report DA-LeWM minus LeWM averaged over epochs (Avg.

\Delta

) and at epoch 10 (

\Delta_{10}

Model	Plan-Real Sp	Positive pairs (out of $30$ )
LeWM	$+0.280$	$24$
No-SIGReg	$+0.031$	$14$
Inverse-only	$\mathbf{+0.420}$	$30$
All-heads	$+0.410$	$30$
DA-LeWM	$+0.412$	$29$

Table 2: Full-state Plan-Real Spearman on PushT (

n=30

Model	Random	Mid	Elite
LeWM	$+0.403$	$+0.227$	$+0.036$
No-SIGReg	$+0.029$	$-0.017$	$-0.102$
Inverse-only	$+0.523$	$\mathbf{+0.261}$	$-0.089$
All-heads	$+0.515$	$+0.249$	$-0.010$
DA-LeWM	$\mathbf{+0.536}$	$+0.253$	$-0.011$

Table 3: Full-state CEM-stage Spearman on PushT (

n\!=\!15

held-out pairs, CEM budget

300\!\times\!30

, top-

30

elites).

Model	PushT	Reacher	Cube
LeWM	$49.3\pm 12.2$	$82.0\pm 2.0$	$62.7\pm 4.2$
No-SIGReg	$\phantom{0}2.0\pm\phantom{0}2.0$	$10.7\pm 8.1$	$52.7\pm 1.9$
Inverse-only	$64.0\pm\phantom{0}7.2$	$82.7\pm 3.1$	$68.0\pm 4.0$
DA-LeWM	$\mathbf{92.7\pm\phantom{0}1.2}$	$\mathbf{84.0\pm\phantom{0}3.5}$	$\mathbf{73.3\pm 1.2}$

Table 4: Cross-environment online success (matched training budget,

3

evaluation seeds for all variants.)

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 5: Global Spearman correlation between

\|\Delta z\|

and

\|a\|

38{,}400

held-out PushT transitions. Values are read directly from the saved diagnostic output. No synthetic points are shown.

Model	Success	State $R^{2}$	Action $R^{2}$	Goal $R^{2}$
LeWM	$49.3\pm 12.2$	$0.90$	$0.86$	$0.78$
No-SIGReg	$\phantom{0}2.0\pm\phantom{0}2.0$	$-6.14$	$0.14$	$-0.67$
Inverse-only	$64.0\pm\phantom{0}7.2$	$0.89$	$0.88$	$0.77$
All-heads	$71.3\pm\phantom{0}4.2$	$0.89$	$0.88$	$0.78$
DA-LeWM	$\mathbf{92.7\pm\phantom{0}1.2}$	$\mathbf{0.90}$	$\mathbf{0.89}$	$\mathbf{0.80}$

Appendix

Table 6: Information probes and online success on PushT (

3

evaluation seeds). Across the four non-collapsed methods, probe scores change by less than

0.03

R^{2}

while online success spans

43

percentage points. No-SIGReg is shown as a degenerate sanity check.

Variant	$\rho$ mean	$\rho$ CV	Log spread
LeWM (no DS)	$2.13\pm 0.86$	$0.710\pm 0.200$	$1.87\pm 0.26$
Inverse-only	$2.90\pm 1.08$	$0.714\pm 0.172$	$2.03\pm 0.36$
All-heads	$2.92\pm 1.20$	$0.685\pm 0.116$	$1.98\pm 0.41$
DA-LeWM ( $\beta\!=\!0.1$ )	$2.90\pm 1.04$	$0.767\pm 0.118$	$2.12\pm 0.20$

Appendix

Table 7: Elite-neighborhood local latent geometry on PushT (

K\!=\!30

elites/pair,

n\!=\!15

pairs). Mean

\pm

std across pairs. CV is

\sigma(\rho)/\mu(\rho)

. Log spread is

\log(p_{95}/p_{5})

Contrast	$\Delta\mathrm{CV}\ (t)$	$\Delta$ log spread $(t)$
All-heads $-$ Inverse-only	$-0.029\ (-0.73)$	$-0.049\ (-0.87)$
All-heads $-$ DA-LeWM	$-0.081\ (-1.85)$	$-0.135\ (-1.09)$
All-heads $-$ LeWM	$-0.024\ (-0.47)$	$+0.119\ (+0.85)$
DA-LeWM $-$ Inverse-only	$+0.053\ (+0.99)$	$+0.086\ (+0.87)$
DA-LeWM $-$ LeWM	$+0.057\ (+0.93)$	$+0.255\ (+2.50)$
Inverse-only $-$ LeWM	$+0.004\ (+0.06)$	$+0.169\ (+1.34)$

Appendix

Table 8: Paired comparisons of elite-neighborhood local geometry on PushT (

n\!=\!15

pairs, same pair indices across variants). Positive numbers indicate the first variant has higher anisotropy. Parentheses report paired

t

. All all-heads contrasts are non-significant. Among the remaining contrasts, only DA-LeWM versus LeWM log spread is significant (

p=0.026

Variant	$p$	$\rho_{s}$	$\tau_{a}$
LeWM (no DS)	$0.288\pm 0.073$	$+0.338\pm 0.206$	$+0.239\pm 0.148$
No-SIGReg	$0.221\pm 0.052$	$+0.053\pm 0.209$	$+0.037\pm 0.142$
Inverse-only	$0.314\pm 0.088$	$+0.448\pm 0.207$	$+0.324\pm 0.154$
All-heads	$0.317\pm 0.093$	$+0.457\pm 0.236$	$+0.331\pm 0.178$
DA ( $\beta\!=\!0.1$ )	$0.325\pm 0.091$	$+0.479\pm 0.228$	$+0.347\pm 0.174$

Appendix

Table 9: Corollary 1 soft-margin consistency check on PushT (

N\!=\!64

random plans/pair,

n\!=\!20

pairs, mean

\pm

std across pairs).

\tau_{a}\geq 2p-1

holds for every pair.

Latent World Models

Latent world models learn compact predictive state representations to simulate how environments change and support planning, control, and reasoning.

See all 212 papers

World Model Control by Trajectory Reachability Metrics
May 21, 2026Liangyu Li, Shengzhi Wang, Libin Qiu +2Reachability Analysis World Models for Robotics
Tongji University
Latent Energy Action Planning with World Models
Sep 3, 2026Phu Pham, Aniket BeraTrajectory Optimization World Model-Based Planning
Department of Computer Science, Purdue University, USA
Latent World Models Should Plan with Distances Along the Latent Manifold
Oct 18, 2026Zunzhi You, Linwei Tao, Siyu Xu +2Latent World Models Latent-Space Planning
The University of Sydney · Machine Learning Research, Apple

May 21, 2026J/K move · Enter open · S save

World Model Control by Trajectory Reachability Metrics

Liangyu Li, Shengzhi Wang, Libin Qiu, Mingliang Xiong+1

Tongji University

Reachability Analysis World Models for Robotics World Model-Based Planning Model-Based Planning Latent World Models

Open paper AskPDF arXiv

Latent world models can learn representations that contain information needed for control, while the downstream controller may still rank candidate actions poorly when it relies on terminal latent distance alone. We study this failure in a fixed encoder and introduce trajectory reachability metrics (TRM), a small temporal pairwise cost trained from logged trajectories and used to rank predicted endpoints of a candidate action sequence against a goal. In the TwoRoom evaluation on 100 episodes from a high distance range, the original controller with LeWorldModel reaches 7.0% mean success. Temporal TRM trained after excluding all evaluation episodes reaches 96.7%, shuffled label controls stay at 0.0%, and TRM also improves PLDM from 32.7% to 84.0%. TopoNav, a separate pixel navigation task, shows the same link between repaired ranking and closed loop control. The selection audit with shared candidates (SASC) and rowspace interventions show that the gain comes from reweighting latent directions carrying the action decision. The XY rowspace contributes less than 1% of terminal latent MSE but carries most of the information needed for control. Coverage and boundary tests define the scope. Balanced doorway coverage restores a 0.0% coverage failure to 100.0% success, while an unseen wall orientation and contact rich PushT expose layout, dynamics, and recovery limits.

Abstract

Latent World Models

Explore similar work

World Model Control by Trajectory Reachability Metrics

Latent Energy Action Planning with World Models

Latent World Models Should Plan with Distances Along the Latent Manifold