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

May 21, 2026

World Model Control by Trajectory Reachability Metrics

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.

Liangyu Li, Shengzhi Wang, Libin Qiu +2

Tongji University

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

Ask PDF

Sep 3, 2026

Latent Energy Action Planning with World Models

Latent world models support efficient model predictive control from high-dimensional observations, yet optimizing a single learned latent objective can favor action sequences whose decoder-predicted terminal descriptor does not match the goal descriptor. We introduce Latent Energy Action Planning (LEAP), which treats the complete action horizon as a differentiable variable and optimizes it through a frozen LeWorldModel (LeWM). LEAP couples terminal latent goal matching with a terminal-window state energy. Low energy requires the predicted terminal latent to agree with the goal latent and the decoder-predicted terminal descriptor to agree with the goal descriptor. A frozen goal-conditioned proposal initializes the search, a quasi-Newton solver refines actions through the autoregressive rollout, and post-optimization projection enforces the admissible action range. Across four control domains using the officially released LeWM checkpoints, the complete LEAP planning system raises mean success from 77.5% for LeWM planned with the cross-entropy method (LeWM+CEM) to 94.8% under a matched protocol, a 17.3-percentage-point improvement, while retaining the frozen LeWM representation and predictor.

Phu Pham, Aniket Bera

Department of Computer Science, Purdue University, USA

Trajectory Optimization World Model-Based Planning Model Predictive Control Model-Based Planning Latent World Models

Ask PDF

Oct 18, 2026

Latent World Models Should Plan with Distances Along the Latent Manifold

Latent world models such as LeWM plan by searching for an action sequence whose predicted terminal latent is closest to the goal latent under an ambient norm, such as the Euclidean distance. Planning with such a cost assumes that proximity in the ambient latent space reflects proximity to the goal in the environment. Studying LeWM, we find that this assumption holds only in a small neighborhood around the goal, whereas farther away the L2 cost forms spurious low-cost basins that trap the planner. Yet neighborhood-based embeddings show that the latents retain the underlying physical state. We therefore hypothesize that the latents lie near a curved, low-dimensional manifold, and that the terminal cost should measure the distance along this latent manifold rather than through the ambient space. To estimate this manifold distance, we propose the Neighborhood Graph Cost (NGC), which computes shortest-path distances on a mutual nearest-neighbor graph constructed over training latents. NGC changes only the terminal cost and requires no retraining. This single change improves the planning success rate of LeWM across all four benchmark tasks and nearly solves TwoRoom and Reacher: even at the farthest goal offsets tested, NGC reaches 95% and 100% success, whereas the L2 cost falls to 27% and 81%. These results support the manifold hypothesis and our argument that latent world models should plan with distances along the latent manifold.

Zunzhi You, Linwei Tao, Siyu Xu +2

The University of Sydney · Machine Learning Research, Apple

Latent World Models Latent-Space Planning

Ask PDF

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