Latent World Models Should Plan with Distances Along the Latent Manifold

Zunzhi You; Linwei Tao; Siyu Xu; Yunke Wang; Chang Xu

cs.LGOct 18, 2026

Latent World Models Should Plan with Distances Along the Latent Manifold

Authors: Zunzhi You, Linwei Tao, Siyu Xu, Yunke Wang, Chang Xu

Organizations: The University of Sydney · Machine Learning Research, Apple

Abstract

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.

Figures & tables

Figure 3: Two-dimensional embeddings of the TwoRoom grid latents, with each point colored by its rendered physical position ( left: color key).

Figure 4: Terminal cost against the physical distance to the goal across all four tasks. Top: the L2 cost. Bottom: NGC. Lines show bin medians, and shaded bands indicate the 10th to 90th percentiles. Details in Appendix B.2 .

Figure 5: NGC on a curved manifold $\mathcal{M}$ . Left: a candidate $q$ is close to the goal $g$ in Euclidean distance (blue, the L2 cost) but far from $g$ along $\mathcal{M}$ . NGC scores $q$ by the shortest path through the graph (orange), whose edges join mutual $15$ -nearest neighbors. Right: the neighborhood of $g$ , flattened. When a candidate $q^{\prime}$ and $g$ lie within each other’s neighborhood, NGC scores $q^{\prime}$ by their Euclidean distance $d=\lVert q^{\prime}-g\rVert_{2}$ . The surface is a schematic; the LeWM latents are 192-dimensional.

Figure 6: Top: Success rate against the goal offset (

300

pairs per offset; whiskers indicate 95% bootstrap intervals over pairs). Numbers above the bars show the mean per-pair difference (NGC minus L2, in percentage points), highlighted in orange when the 95% interval excludes zero. Bottom: The same pairs pooled across offsets and binned into quintiles of the physical distance between the start and goal states (

180

pairs per bin). Shaded regions show the gap between NGC and the L2 cost, and filled NGC markers indicate bins where the 95% interval of the difference excludes zero.

Task	offset	L2	NGC ${}_{\text{Euc}}$	NGC	NGC ${}_{\text{Euc}}$ $-$ L2	NGC $-$ NGC ${}_{\text{Euc}}$
TwoRoom	25	91.4	86.6	99.8	$-4.9$ $[-7.4,-2.3]$	$+13.2$ $[+10.0,+16.6]$
	50	65.9	69.0	100.0	$+3.1$ $[-1.4,+7.8]$	$+31.0$ $[+26.2,+36.0]$
	75	27.3	41.3	94.6	$+14.0$ $[+9.6,+18.6]$	$+53.2$ $[+48.0,+58.3]$
Reacher	25	99.6	99.3	99.4	$-0.2$ $[-0.9,+0.3]$	$+0.1$ $[-0.6,+0.8]$
	50	95.6	96.6	100.0	$+1.0$ $[+0.2,+2.0]$	$+3.4$ $[+1.7,+5.4]$
	100	80.6	82.2	100.0	$+1.7$ $[+0.2,+3.2]$	$+17.8$ $[+13.8,+22.0]$

Table 1: Ablation of the graph shortest-path distance on the pairs and graphs of Figure 6 . NGC

{}_{\text{Euc}}

replaces

D_{ij}

with the Euclidean distance between the same graph nodes. The last two columns report mean per-pair differences (in percentage points) with 95% bootstrap intervals.

		L2		NGC		oracle	NGC
Task	offset	learned	oracle	learned	oracle	NGC $-$ L2	oracle $-$ learned
TwoRoom	25	92	84	100	100	$+16$ $[+9,+23]$	$0$
TwoRoom	50	65	58	100	100	$+42$ $[+32,+52]$	$0$
TwoRoom	75	24	19	94	100	$+81$ $[+73,+88]$	$+6$ $[+2,+11]$
Cube	25	44	56	59	76	$+20$ $[+11,+29]$	$+17$ $[+8,+26]$
Cube	50	44	52	53	71	$+19$ $[+11,+27]$	$+18$ $[+11,+26]$

Table 2: Success rate (%) with the learned predictor and with the oracle predictor, which executes each candidate in the simulator and encodes its terminal frame, on one set of

100

pairs with a single planner seed. The last two columns report mean per-pair differences in % with 95% bootstrap intervals: NGC minus L2 under the oracle predictor, and oracle minus learned predictor under NGC.

Task	L2	NGC	( $\times$ L2)	L2	NGC
	CEM (s)		Evaluation	Objective (ms per pool)
TwoRoom	0.49	0.48 ( $0.98\times$ )	0.86	0.071	0.67
Reacher	0.49	0.52 ( $1.06\times$ )	1.05	0.070	0.66
Cube	0.49	0.53 ( $1.07\times$ )	1.08	0.070	0.67
PushT	0.50	0.52 ( $1.05\times$ )	1.13	0.070	0.67

Table 3: Planning time of NGC at offset

25

. CEM: median time of one CEM optimization, with its ratio to L2. Evaluation: total wall-clock time of a full evaluation relative to L2 on the same pairs. Objective: median time to score one pool of

300

candidates (Appendix A.2 ).

Figure 7: Success rate against the number of CEM candidates at the longest offset of each task. The number of elites scales proportionally;

1\times

corresponds to the default setting in Figure 6 . The elite fraction and the

30

CEM iterations remain fixed, and all multipliers are evaluated on the pairs, planner seeds, and graphs of Figure 6 .

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Success criterion	Physical distance
TwoRoom	Endpoint within 16 px of the goal position	Path length via the doorway
Reacher	Joint error $<0.05$ rad, shoulder angle compared on the circle	Norm of the difference of the two joint angles, shoulder angle compared on the circle
Cube	Cube center within 4 cm of the goal position	Gripper-to-cube plus cube-to-goal distance
PushT	Pusher and block within 20 px of the goal pose and block angle within $20^{\circ}$	Pusher-to-block distance plus the root-mean-square displacement over the block’s area to the goal pose

Appendix

Table 4: Success criteria and physical distances (Appendix A.3 ). An episode succeeds at the first step that meets the criterion.

Figure 8: Cost maps and rollout trajectories for additional TwoRoom pairs, formatted as in Figure 2 , with the L2 cost map above the NGC map of each pair. Upper block: all $8$ pairs from the evaluation of Figure 2 on which the L2 cost fails. Lower block: the first $8$ pairs in pair-index order on which the L2 cost succeeds. NGC solves every pair shown.

Task	Offset	L2	NGC	$\Delta$ [95% CI]	gained/lost
TwoRoom	25	91.4 [88.7, 94.0]	99.8 [99.4, 100.0]	$+8.3$ [ $+5.8$ , $+11.0$ ]	45/1
TwoRoom	50	65.9 [60.9, 70.9]	100.0	$+34.1$ [ $+29.2$ , $+39.2$ ]	121/0
TwoRoom	75	27.3 [22.6, 32.1]	94.6 [92.2, 96.7]	$+67.2$ [ $+62.2$ , $+72.1$ ]	223/0
Reacher	25	99.6 [99.1, 99.9]	99.4 [98.9, 99.9]	$-0.1$ [ $-0.8$ , $+0.6$ ]	4/5
Reacher	50	95.6 [93.2, 97.7]	100.0	$+4.4$ [ $+2.3$ , $+6.9$ ]	15/0
Reacher	100	80.6 [76.1, 84.8]	100.0	$+19.4$ [ $+15.2$ , $+23.8$ ]	68/0

Appendix

Table 5: Success rate (%) of the L2 cost and NGC across the offsets of Figure 6 (our checkpoint,

300

pairs per offset, three planner seeds averaged within each pair). Brackets in small type give the 95% bootstrap confidence interval over pairs (corresponding to the whiskers in Figure 6 ).

\Delta

is the mean per-pair difference (NGC minus L2) with its 95% confidence interval, and gained/lost counts the pairs with a positive versus negative difference (Figure 9 ).

Figure 9: Per-pair differences underlying Table 5 . Each bar represents the

300

pairs for a given task and offset. Under each cost, every pair is evaluated with three planner seeds; colors indicate the number of seeds on which NGC succeeds minus the number on which the L2 cost succeeds. Gray segments (equal success under both costs) are centered at zero; pairs where NGC improves over L2 extend to the right in orange, and pairs where it falls behind extend to the left in blue. Numbers count the pairs in each segment (segments with fewer than

13

pairs are unlabeled). Right:

\Delta

, the mean per-pair difference in % (where one seed corresponds to

\frac{1}{3}

, or

33.3

points, of a single pair’s score), with its 95% confidence interval.

Task	L2	NGC	$\Delta$ [95% CI]	L2	NGC	$\Delta$ [95% CI]	reported
	our checkpoint			released LeWM checkpoint
TwoRoom	$91.3$ $\pm 3.1$	$100.0$ $\pm 0.0$	$+8.7$ $[+4.7,+13.3]$	$85.3$ $\pm 3.1$	$100.0$ $\pm 0.0$	$+14.7$ $[+9.3,+20.7]$	$87$
Reacher	$86.0$ $\pm 2.0$	$84.0$ $\pm 3.5$	$-2.0$ $[-9.3,+5.3]$	$82.7$ $\pm 4.2$	$81.3$ $\pm 11.4$	$-1.3$ $[-9.3,+6.7]$	$86$
Cube	$68.0$ $\pm 10.6$	$78.7$ $\pm 7.6$	$+10.7$ $[+4.7,+17.3]$	$69.3$ $\pm 12.1$	$75.3$ $\pm 5.0$	$+6.0$ $[-0.7,+12.7]$	$74$
PushT	$91.3$ $\pm 2.3$	$94.0$ $\pm 5.3$	$+2.7$ $[0.0,+6.0]$	$89.3$ $\pm 5.0$	$88.7$ $\pm 8.1$	$-0.7$ $[-4.7,+3.3]$	$96$

Appendix

Table 6: Evaluation under the original LeWM protocol. Success (%) is reported as the mean

\pm

SD across three runs of

50

random pairs each. The last column lists the integer-rounded figures reported by Maes et al. (2026b) (their Figure 6).

\Delta

is the mean per-pair difference across all

150

pairs (NGC minus L2) with its 95% confidence interval (Appendix A.4 ).

Task	Offset	L2	NGC	$\Delta$ [95% CI]	$\Delta$ , ours
TwoRoom	25	85.3	99.8	$+14.4$ [ $+11.6$ , $+17.6$ ]	$+8.3$
TwoRoom	50	61.9	99.8	$+37.9$ [ $+32.9$ , $+42.9$ ]	$+34.1$
TwoRoom	75	28.0	95.1	$+67.1$ [ $+62.2$ , $+72.0$ ]	$+67.2$
Reacher	25	97.9	97.4	$-0.4$ [ $-1.8$ , $+1.0$ ]	$-0.1$
Reacher	50	96.2	100.0	$+3.8$ [ $+1.9$ , $+5.9$ ]	$+4.4$
Reacher	100	82.6	100.0	$+17.4$ [ $+13.4$ , $+21.6$ ]	$+19.4$

Appendix

Table 7: Performance of the released LeWM checkpoint on the pairs of Figure 6 (

300

pairs per offset, three planner seeds averaged within each pair).

\Delta

is the mean per-pair difference (NGC minus L2) with its 95% confidence interval; the final column repeats

\Delta

from Table 5 for our checkpoint.

Task	Offset	L2	NGC	$\Delta$ [95% CI]	L2	NGC	$\Delta$ [95% CI]	$\Delta$
		Second seed			Third seed			First seed
TwoRoom	25	92.2	100.0	$+7.8$ $[+5.4,+10.3]$	92.6	100.0	$+7.4$ $[+5.1,+10.0]$	$+8.3$
TwoRoom	50	69.7	99.8	$+30.1$ $[+25.4,+34.9]$	66.3	99.2	$+32.9$ $[+27.9,+38.1]$	$+34.1$
TwoRoom	75	31.8	96.1	$+64.3$ $[+59.3,+69.2]$	31.4	86.4	$+55.0$ $[+49.7,+60.3]$	$+67.2$
Reacher	25	99.4	99.4	$0.0$ $[-0.7,+0.7]$	99.7	99.9	$+0.2$ $[-0.2,+0.7]$	$-0.1$
Reacher	50	95.4	100.0	$+4.6$ $[+2.4,+7.0]$	95.6	99.9	$+4.3$ $[+2.2,+6.7]$	$+4.4$

Appendix

Table 8: Evaluation across two additional world model training seeds on the pairs of Figure 6 (

300

pairs per offset, three planner seeds averaged within each pair). Columns report success rates (%) of the L2 cost and NGC under each seed, along with the mean per-pair difference

\Delta

(NGC minus L2) and its 95% confidence interval. The final column repeats

\Delta

from Table 5 for our primary checkpoint (first seed).

			Node sample (mutual)			Neighbor rule
Task	Offset	L2	Default	resample 1	resample 2	union
TwoRoom	25	91.3	100.0	100.0 $0.0$	100.0 $0.0$	100.0 $0.0$
TwoRoom	75	23.7	95.7	95.3 $-0.3\,[-3.0,+2.3]$	94.7 $-1.0\,[-3.7,+1.3]$	94.3 $-1.3\,[-4.0,+1.0]$
Reacher	25	99.0	99.0	99.3 $+0.3\,[-1.0,+1.7]$	98.7 $-0.3\,[-2.0,+1.3]$	100.0 $+1.0\,[0.0,+2.3]$
Reacher	100	82.7	100.0	100.0 $0.0$	100.0 $0.0$	100.0 $0.0$
Cube	25	43.7	57.0	54.3 $-2.7\,[-8.3,+2.7]$	53.0 $-4.0\,[-9.7,+1.3]$	56.3 $-0.7\,[-5.3,+4.0]$

Appendix

Table 9: Sensitivity to node sampling and the neighbor connectivity rule on one pair seed (

100

pairs

\times

3

planner seeds per cell). All columns except L2 evaluate NGC. The L2 and Default columns (mutual rule,

k=15

5{,}000

nodes) correspond to the first pair seed of Table 5 . Resamples 1 and 2 draw new sets of

5{,}000

nodes under the mutual rule; the union column uses the default nodes but joins two nodes when either is among the other’s

15

nearest neighbors. Each variant cell reports the success rate (%) and, in small type, the mean per-pair difference

\Delta

against the default graph with its 95% confidence interval.

	TwoRoom		Reacher
Variant	25	75	25	100
L2	91.3	23.7	99.0	82.7
NGC (default)	100.0	95.7	99.0	100.0
$N=625$	100.0 $0.0$	57.0 $-38.7\,[-48.0,-29.7]$	99.7 $+0.7\,[-0.7,+2.0]$	100.0 $0.0$
$N=1{,}250$	100.0 $0.0$	64.3 $-31.3\,[-40.0,-23.0]$	99.7 $+0.7\,[-0.7,+2.0]$	100.0 $0.0$
$N=2{,}500$	100.0 $0.0$	73.3 $-22.3\,[-30.3,-14.7]$	99.7 $+0.7\,[0.0,+1.7]$	100.0 $0.0$

Appendix

Table 10: Sensitivity to graph size

N

and neighborhood size

k

on one pair seed (

100

pairs

\times

3

planner seeds per cell). Column headings indicate goal offsets in steps. Each cell reports the success rate (%) and, in small type, the mean per-pair difference

\Delta

against the default graph (

5{,}000

nodes,

k=15

) with its 95% confidence interval.

Figure 10: Physical-space visualizations of the TwoRoom graphs from Table 10 : nested node subsets and the default graph at $k=15$ ( top ), and the default $5{,}000$ nodes at $k=5$ , $10$ , $30$ , and $60$ ( bottom ), with each edge plotted between the physical positions of its endpoint nodes. Red edges highlight shortcuts that connect nodes separated by more than $80$ px along the shortest feasible path through the doorway. Titles report the number of such shortcut edges and the resulting NGC success rate at offset $75$ .

Figure 11: Behavior of the L2 cost with the oracle predictor on TwoRoom at offset $25$ . Left: the first four pairs (in pair-index order) where the L2 cost fails under the oracle predictor but succeeds under the learned predictor, cropped around the stopping point and the goal. Top row: L2 cost maps overlaid with rollouts from both predictors. Bottom row: corresponding NGC cost maps. Each panel is normalized to its own 2nd–98th percentile within the crop; the dashed circle marks the goal success radius, and white dots indicate strict local minima on the grid (points cheaper than all $8$ neighbors). Right: L2 cost profiles along the straight line from the oracle rollout’s stopping point to the goal, normalized by the cost at the stopping point, across all $16$ pairs where the L2 cost fails with the oracle predictor (the four pairs shown on the left are highlighted in black).

Latent World Models

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

See all 204 papers

Oct 4, 2026

How Long, Not How Close: A Learned Temporal Metric for Planning in Latent World Models

Latent world models plan by rolling a frozen predictor forward under candidate action sequences and ranking the candidates by the latent distance between their imagined end state and the goal. However, this ranking breaks down when the goal lies several plans away, because the latent distance measures how closely an end state resembles the goal rather than how far it remains from reaching it. To address this, we propose TEMPO, a temporal-distance planning objective that leaves the world model untouched, learns only from the recorded trajectories already used to train it, and adds negligible cost to the planner's search. TEMPO learns a small map of the frozen latent in which the distance between two states of an episode reflects the number of environment steps between them, and blends this distance into the planner's cost. It requires no rewards, policies or success labels and, being a cost rather than a model, applies to frozen world models with one latent vector per state that plan by a latent distance. We evaluate TEMPO on eleven simulated environments (e.g., maze navigation, tabletop pushing, robotic arm control and three-dimensional manipulation) with the LeWM and PLDM planners. With a small MLP that adds at most 0.3% to a plan's arithmetic, TEMPO improves both planners at every goal distance, including the one-plan setting of their evaluations, raises LeWM from 36% to 99% on TwoRoom three plans from the goal, and remains competitive on a broad range of 2D and 3D navigation, reaching and manipulation tasks.

Lama Moukheiber, Haotian Xue, Yongxin Chen

Georgia Institute of Technology

World Model-Based Planning Long-Horizon Planning Latent World Models Latent-Space Planning

Ask PDF

Aug 13, 2026

The Objective Is the Bottleneck: Latent World Models Encode What Their Planners Cannot Use

Latent world models are judged by how well they predict, so when planning fails at long horizons the natural reading is that the predictor degrades. On a reproduction of LeWorldModel on TwoRoom we show the binding constraint is the planner's objective instead. The predictor is not the limit: its imagined state seventy-five environment steps ahead is still only 0.189 as wrong as assuming the world froze, while the planner never imagines beyond twenty-five. The objective is. Cross-entropy-method planning minimises squared latent distance, which tracks true distance at r = 0.426, saturates by about eighty arena units and decreases beyond a hundred and twenty, so moving away from the goal can lower the cost. The information is present throughout: a ridge probe recovers position from the frozen embedding at R^2 0.9922. The pathology is the method's, not one reimplementation's. It is present in the authors' released weights, and across four checkpoints long-horizon success rank-orders exactly with metric quality and inversely with prediction accuracy. Replacing only the objective, with nothing retrained and no GPU, lifts goals reached at offset 100 from 26.0% to 98.0%, equals the 98.0% at offset 25, and reaches 92.0% under a third of the budget: planning stops depending on the horizon. The best cost is not the most accurate. A head learned from frame separation alone predicts spatial distance worse than a position probe (r = 0.819 against 0.9897) yet plans better, charging 24% more to cross the environment's dividing wall where squared latent distance charges 4% less. It has learned reachability, not proximity.

Joyjeet Singh

Independent researcher

World Model-Based Planning Model-Based Planning Latent World Models World Model Evaluation Latent-Space Planning

Ask PDF

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

Abstract

Latent World Models

Explore similar work

How Long, Not How Close: A Learned Temporal Metric for Planning in Latent World Models

The Objective Is the Bottleneck: Latent World Models Encode What Their Planners Cannot Use

World Model Control by Trajectory Reachability Metrics