Control-Geometry Straightening for Sampling-Based Latent Planning

Ziang Fu; Ning Ning

cs.LGSep 28, 2026

Control-Geometry Straightening for Sampling-Based Latent Planning

Authors: Ziang Fu, Ning Ning

Organizations: Texas A&M University

Abstract

Joint-embedding predictive architectures enable planning with latent world models, but accurate transition prediction alone does not ensure that the planning objective is easy to optimize. We introduce Control-Geometry Straightening (CGS), a single auxiliary loss that learns planner-friendly representations by directly straightening control geometry for sampling-efficient planning. CGS matches pairwise cosine similarities among actions to those among corresponding latent differences only using local transitions from pixel-action pairs. The loss can be applied across world-model architectures using end-to-end learned or pretrained representations. Under linear-dynamics, our theoretical analysis connects this objective to temporal straightening and more balanced terminal-cost curvature across the full planning horizon, yielding finite-budget guarantees for MPPI, local contraction results for CEM, and convergence bounds for gradient descent. Across four control environments and multiple planners, CGS improves planning with fewer sampled candidates and refinement steps, achieving success-rate gains up to 20 and 12.6 percentage points over LeWorldModel (LeWM) and its temporal-straightening variant (LeWM+TS), respectively, with sampling-based planners using 128 candidates per update. Probes, comparisons with DINO-WM architecture, and planner-side ablations clarify how latent motion organization, state dependence, and dynamical context shape planning behavior. Straightening control geometry thus makes good action sequences easier to find under limited planning budgets.

Figures & tables

Figure 1: CGS preserves action-related turns and control-response geometry. PushT (top) and Cube (bottom), with LeWM (blue), LeWM+TS (brown), and LeWM+CGS (green). Middle: transition shaping. CGS makes latent transitions follow action-related turns more closely than LeWM and TS. Right: control-response geometry. The heatmaps show pairwise cosines of predicted endpoint changes induced by action-sequence perturbations; Ideal gives the angular relations among the perturbations themselves. CGS more closely preserves these relations in the tested control plane.

Figure 2: CGS training and planning. During training, we minimize the prediction loss between predicted and encoded target embeddings, together with a CGS loss that matches pairwise cosine similarities of standardized actions and consecutive latent differences; SIGReg is omitted from the diagram. During planning, we roll out candidate action sequences for $H$ steps using the frozen predictor and refine them with a sampling-based planner to minimize the distance between the predicted terminal embedding and the goal embedding.

Figure 3: MPPI contraction and finite-sample improvement. (a) CGS lowers population contraction

q_{\infty}

;

\sigma_{\min}

measures the weakest curvature. (b) CGS increases

\Pr(q_{K}<1)

, where

q_{K}

is the

K

-sample contraction factor. Definitions and protocol: Appendix C.4 .

Panel of Table 1 — Table 1: Goal-reaching success rate (SR) at $K=128$ . Values are reported as means $\pm$ sample standard deviation (SD), in %. Bold marks the best per planner and environment, including ties. Parentheses report gains over LeWM in percentage points; the top two per environment are underlined. Prior-guided (PG) variants ( $\dagger$ ) use action-prior initialization. Sampling planners use $I=30$ optimization steps; gradient descent (GD) uses 100 Adam optimization steps, independent of $K$ .

Panel of Table 2 — Table 2: Physical-motion probing on PushT and Reacher. $P_{A}$ and $P_{AS}$ measure how well physical motion is identified from action alone or from action and state; these environment-level scores are shared across representations. $D_{\rm lin}$ and $D_{\rm MLP}$ decode true motion from $\Delta z_{t}$ , while $C_{A}$ and $B_{AS}$ measure agreement with action-only and state-conditioned control effects. Scores are held-out $R^{2}$ , except the normalized agreement score $B_{AS}$ (Appendix C.3 ). Higher is better; bold marks the best representation per metric using unrounded scores.

Figure 6: CGS improves early refinement and BC-prior quality on PushT. Left: CEM/MPPI success over optimization steps

I

(iterations) at

K=128

without an action prior. Middle: BC-prior/no-op latent goal-distance ratios for the first plan (lower is better). Right: prior-only success using the prior mean as the sole candidate without additional candidate sampling (

K=1

I=1

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

World model	Planner	$K=1$	$K=2$	$K=4$	$K=8$	$K=16$	$K=32$	$K=64$	$K=128$	$K=256$
LeWM	MPPI	$\mathbf{0.0}{\scriptstyle\,\pm 0.0}$	$30.7{\scriptstyle\,\pm 6.1}$	$42.7{\scriptstyle\,\pm 9.0}$	$56.7{\scriptstyle\,\pm 1.1}$	$58.7{\scriptstyle\,\pm 10.1}$	$61.3{\scriptstyle\,\pm 1.1}$	$59.3{\scriptstyle\,\pm 10.1}$	$64.0{\scriptstyle\,\pm 5.3}$	$60.7{\scriptstyle\,\pm 4.2}$
	PG-MPPI	$84.7{\scriptstyle\,\pm 6.4}$	$82.0{\scriptstyle\,\pm 4.0}$	$82.7{\scriptstyle\,\pm 6.1}$	$77.3{\scriptstyle\,\pm 5.0}$	$76.0{\scriptstyle\,\pm 6.0}$	$78.7{\scriptstyle\,\pm 6.4}$	$72.7{\scriptstyle\,\pm 8.3}$	$72.0{\scriptstyle\,\pm 6.0}$	$62.0{\scriptstyle\,\pm 3.5}$
	CEM	$\mathbf{0.0}{\scriptstyle\,\pm 0.0}$	$\mathbf{35.3}{\scriptstyle\,\pm 3.1}$	$37.3{\scriptstyle\,\pm 6.4}$	$56.0{\scriptstyle\,\pm 5.3}$	$74.0{\scriptstyle\,\pm 7.2}$	$83.3{\scriptstyle\,\pm 4.2}$	$83.3{\scriptstyle\,\pm 5.0}$	$89.3{\scriptstyle\,\pm 6.1}$	$87.3{\scriptstyle\,\pm 4.2}$
	PG-CEM	$85.3{\scriptstyle\,\pm 6.1}$	$84.0{\scriptstyle\,\pm 3.5}$	$80.7{\scriptstyle\,\pm 5.8}$	$91.3{\scriptstyle\,\pm 1.1}$	$89.3{\scriptstyle\,\pm 1.1}$	$93.3{\scriptstyle\,\pm 1.1}$	$94.0{\scriptstyle\,\pm 4.0}$	$92.0{\scriptstyle\,\pm 5.3}$	$92.0{\scriptstyle\,\pm 2.0}$
	GD	$85.3{\scriptstyle\,\pm 5.0}$
LeWM + TS	MPPI	$\mathbf{0.0}{\scriptstyle\,\pm 0.0}$	$38.0{\scriptstyle\,\pm 5.3}$	$40.7{\scriptstyle\,\pm 6.4}$	$50.0{\scriptstyle\,\pm 5.3}$	$51.3{\scriptstyle\,\pm 1.1}$	$56.7{\scriptstyle\,\pm 5.0}$	$65.3{\scriptstyle\,\pm 3.1}$	$64.7{\scriptstyle\,\pm 8.3}$	$66.7{\scriptstyle\,\pm 7.0}$

Appendix

Table 3: PushT sample-budget sweep. Success rate (%, mean

\pm

sample SD over three evaluation seeds, 50 episodes each). Bold marks the best representation for each planner and budget, including ties. GD is independent of

K

World model	Planner	$K=1$	$K=2$	$K=4$	$K=8$	$K=16$	$K=32$	$K=64$	$K=128$	$K=256$
LeWM	MPPI	$\mathbf{44.7}{\scriptstyle\,\pm 7.0}$	$46.0{\scriptstyle\,\pm 5.3}$	$44.7{\scriptstyle\,\pm 11.0}$	$46.7{\scriptstyle\,\pm 6.4}$	$51.3{\scriptstyle\,\pm 4.2}$	$51.3{\scriptstyle\,\pm 5.0}$	$56.7{\scriptstyle\,\pm 5.0}$	$59.3{\scriptstyle\,\pm 4.6}$	$64.0{\scriptstyle\,\pm 9.2}$
	PG-MPPI	$\mathbf{100.0}{\scriptstyle\,\pm 0.0}$	$\mathbf{75.3}{\scriptstyle\,\pm 8.3}$	$\mathbf{70.0}{\scriptstyle\,\pm 5.3}$	$60.7{\scriptstyle\,\pm 4.2}$	$61.3{\scriptstyle\,\pm 6.1}$	$76.7{\scriptstyle\,\pm 2.3}$	$82.7{\scriptstyle\,\pm 6.4}$	$89.3{\scriptstyle\,\pm 4.2}$	$90.0{\scriptstyle\,\pm 2.0}$
	CEM	$\mathbf{44.7}{\scriptstyle\,\pm 7.0}$	$\mathbf{50.0}{\scriptstyle\,\pm 3.5}$	$45.3{\scriptstyle\,\pm 10.1}$	$\mathbf{55.3}{\scriptstyle\,\pm 9.0}$	$57.3{\scriptstyle\,\pm 7.0}$	$62.7{\scriptstyle\,\pm 10.1}$	$64.0{\scriptstyle\,\pm 6.0}$	$63.3{\scriptstyle\,\pm 2.3}$	$64.0{\scriptstyle\,\pm 11.1}$
	PG-CEM	$\mathbf{100.0}{\scriptstyle\,\pm 0.0}$	$94.0{\scriptstyle\,\pm 0.0}$	$88.0{\scriptstyle\,\pm 5.3}$	$91.3{\scriptstyle\,\pm 3.1}$	$\mathbf{94.7}{\scriptstyle\,\pm 1.1}$	$89.3{\scriptstyle\,\pm 4.2}$	$88.0{\scriptstyle\,\pm 4.0}$	$90.0{\scriptstyle\,\pm 2.0}$	$89.3{\scriptstyle\,\pm 1.1}$
	GD	$59.3{\scriptstyle\,\pm 8.1}$
LeWM + TS	MPPI	$\mathbf{44.7}{\scriptstyle\,\pm 7.0}$	$\mathbf{48.0}{\scriptstyle\,\pm 5.3}$	$\mathbf{46.7}{\scriptstyle\,\pm 3.1}$	$48.0{\scriptstyle\,\pm 2.0}$	$46.0{\scriptstyle\,\pm 8.7}$	$56.7{\scriptstyle\,\pm 9.0}$	$68.0{\scriptstyle\,\pm 5.3}$	$75.3{\scriptstyle\,\pm 5.0}$	$75.3{\scriptstyle\,\pm 6.1}$

Appendix

Table 4: Cube sample-budget sweep. Success rate (%, mean

\pm

sample SD over three evaluation seeds, 50 episodes each). Bold marks the best representation for each planner and budget, including ties. GD is independent of

K

World model	Planner	$K=1$	$K=2$	$K=4$	$K=8$	$K=16$	$K=32$	$K=64$	$K=128$	$K=256$
LeWM	MPPI	$\mathbf{13.3}{\scriptstyle\,\pm 1.1}$	$40.7{\scriptstyle\,\pm 7.0}$	$53.3{\scriptstyle\,\pm 7.0}$	$60.0{\scriptstyle\,\pm 10.6}$	$70.7{\scriptstyle\,\pm 11.7}$	$82.0{\scriptstyle\,\pm 10.0}$	$88.7{\scriptstyle\,\pm 4.2}$	$92.0{\scriptstyle\,\pm 3.5}$	$94.7{\scriptstyle\,\pm 1.1}$
	PG-MPPI	$\mathbf{100.0}{\scriptstyle\,\pm 0.0}$	$77.3{\scriptstyle\,\pm 9.9}$	$68.7{\scriptstyle\,\pm 4.2}$	$79.3{\scriptstyle\,\pm 1.1}$	$87.3{\scriptstyle\,\pm 4.6}$	$91.3{\scriptstyle\,\pm 6.4}$	$\mathbf{98.7}{\scriptstyle\,\pm 1.1}$	$98.0{\scriptstyle\,\pm 2.0}$	$\mathbf{100.0}{\scriptstyle\,\pm 0.0}$
	CEM	$\mathbf{13.3}{\scriptstyle\,\pm 1.1}$	$55.3{\scriptstyle\,\pm 9.0}$	$57.3{\scriptstyle\,\pm 6.4}$	$56.7{\scriptstyle\,\pm 3.1}$	$68.0{\scriptstyle\,\pm 5.3}$	$79.3{\scriptstyle\,\pm 4.2}$	$79.3{\scriptstyle\,\pm 9.0}$	$76.7{\scriptstyle\,\pm 10.1}$	$78.0{\scriptstyle\,\pm 5.3}$
	PG-CEM	$\mathbf{100.0}{\scriptstyle\,\pm 0.0}$	$98.0{\scriptstyle\,\pm 3.5}$	$97.3{\scriptstyle\,\pm 2.3}$	$99.3{\scriptstyle\,\pm 1.1}$	$96.0{\scriptstyle\,\pm 4.0}$	$96.7{\scriptstyle\,\pm 2.3}$	$96.0{\scriptstyle\,\pm 3.5}$	$93.3{\scriptstyle\,\pm 3.1}$	$94.7{\scriptstyle\,\pm 2.3}$
	GD	$41.3{\scriptstyle\,\pm 7.0}$
LeWM + TS	MPPI	$\mathbf{13.3}{\scriptstyle\,\pm 1.1}$	$\mathbf{56.0}{\scriptstyle\,\pm 3.5}$	$\mathbf{64.0}{\scriptstyle\,\pm 9.2}$	$74.7{\scriptstyle\,\pm 1.1}$	$81.3{\scriptstyle\,\pm 1.1}$	$88.0{\scriptstyle\,\pm 5.3}$	$95.3{\scriptstyle\,\pm 4.2}$	$\mathbf{97.3}{\scriptstyle\,\pm 3.1}$	$96.0{\scriptstyle\,\pm 2.0}$

Appendix

Table 5: TwoRooms sample-budget sweep. Success rate (%, mean

\pm

sample SD over three evaluation seeds, 50 episodes each). Bold marks the best representation for each planner and budget, including ties. GD is independent of

K

World model	Planner	$K=1$	$K=2$	$K=4$	$K=8$	$K=16$	$K=32$	$K=64$	$K=128$	$K=256$
LeWM	MPPI	$\mathbf{2.0}{\scriptstyle\,\pm 2.0}$	$46.0{\scriptstyle\,\pm 7.2}$	$\mathbf{40.7}{\scriptstyle\,\pm 8.3}$	$\mathbf{44.7}{\scriptstyle\,\pm 3.1}$	$36.7{\scriptstyle\,\pm 9.0}$	$36.7{\scriptstyle\,\pm 6.1}$	$49.3{\scriptstyle\,\pm 11.4}$	$65.3{\scriptstyle\,\pm 4.2}$	$\mathbf{73.3}{\scriptstyle\,\pm 4.6}$
	PG-MPPI	$\mathbf{78.7}{\scriptstyle\,\pm 2.3}$	$\mathbf{67.3}{\scriptstyle\,\pm 5.0}$	$\mathbf{49.3}{\scriptstyle\,\pm 1.1}$	$46.0{\scriptstyle\,\pm 7.2}$	$42.7{\scriptstyle\,\pm 5.0}$	$46.0{\scriptstyle\,\pm 5.3}$	$46.7{\scriptstyle\,\pm 10.1}$	$62.7{\scriptstyle\,\pm 11.0}$	$\mathbf{72.0}{\scriptstyle\,\pm 7.2}$
	CEM	$\mathbf{2.0}{\scriptstyle\,\pm 2.0}$	$\mathbf{44.0}{\scriptstyle\,\pm 10.0}$	$\mathbf{41.3}{\scriptstyle\,\pm 4.2}$	$65.3{\scriptstyle\,\pm 5.8}$	$\mathbf{72.0}{\scriptstyle\,\pm 5.3}$	$74.7{\scriptstyle\,\pm 9.2}$	$\mathbf{83.3}{\scriptstyle\,\pm 7.0}$	$\mathbf{83.3}{\scriptstyle\,\pm 8.3}$	$85.3{\scriptstyle\,\pm 3.1}$
	PG-CEM	$\mathbf{78.7}{\scriptstyle\,\pm 2.3}$	$\mathbf{82.0}{\scriptstyle\,\pm 2.0}$	$\mathbf{73.3}{\scriptstyle\,\pm 6.1}$	$\mathbf{79.3}{\scriptstyle\,\pm 4.2}$	$\mathbf{85.3}{\scriptstyle\,\pm 3.1}$	$80.0{\scriptstyle\,\pm 3.5}$	$81.3{\scriptstyle\,\pm 9.4}$	$\mathbf{86.7}{\scriptstyle\,\pm 6.1}$	$86.7{\scriptstyle\,\pm 10.1}$
	GD	$76.7{\scriptstyle\,\pm 4.2}$
LeWM + TS	MPPI	$\mathbf{2.0}{\scriptstyle\,\pm 2.0}$	$\mathbf{47.3}{\scriptstyle\,\pm 8.1}$	$40.0{\scriptstyle\,\pm 10.6}$	$36.7{\scriptstyle\,\pm 9.9}$	$\mathbf{44.7}{\scriptstyle\,\pm 4.2}$	$40.0{\scriptstyle\,\pm 9.2}$	$53.3{\scriptstyle\,\pm 6.1}$	$64.0{\scriptstyle\,\pm 2.0}$	$68.7{\scriptstyle\,\pm 6.4}$

Appendix

Table 6: Reacher sample-budget sweep. Success rate (%, mean

\pm

sample SD over three evaluation seeds, 50 episodes each). Bold marks the best representation for each planner and budget, including ties. GD is independent of

K

Planner	Model	PushT	Cube	TwoRooms	Reacher
CEM	LeWM	$62.79$	$56.54$	$64.79$	$\mathbf{63.46}$
	LeWM + TS	$62.17\,{\scriptstyle(-0.62)}$	$59.96\,{\scriptstyle(+3.42)}$	$\mathbf{80.96}\,{\scriptstyle(+16.17)}$	$61.71\,{\scriptstyle(-1.75)}$
	LeWM + CGS	$\mathbf{67.75}\,{\scriptstyle(+4.96)}$	$\mathbf{62.58}\,{\scriptstyle(+6.04)}$	$79.58\,{\scriptstyle(+14.79)}$	$63.04\,{\scriptstyle(-0.42)}$
PG-CEM	LeWM	$89.17$	$91.25$	$96.75$	$\mathbf{81.33}$
	LeWM + TS	$90.38\,{\scriptstyle(+1.21)}$	$89.88\,{\scriptstyle(-1.38)}$	$\mathbf{99.83}\,{\scriptstyle(+3.08)}$	$76.88\,{\scriptstyle(-4.46)}$
	LeWM + CGS	$\mathbf{93.17}\,{\scriptstyle(+4.00)}$	$\mathbf{94.46}\,{\scriptstyle(+3.21)}$	$99.58\,{\scriptstyle(+2.83)}$	$74.50\,{\scriptstyle(-6.83)}$

Appendix

Table 7: Complete sample-budget sweep nAUC.

K\in\{1,2,4,8,16,32,64,128,256\}

; sampling planners use

I=30

. nAUC uses

\log_{2}K

. Parentheses give percentage-point differences from LeWM for the same environment and planner. Bold marks the best representation for each planner and environment. Results average three evaluation seeds (50 episodes each); GD repeats its

K

-independent SR.

Figure 7: Cube reference geometry: sample-budget and refinement curves with nAUC summaries. Lines show mean success and bands show sample SD over three evaluation seeds (50 episodes each). The table reports mean normalized AUC (nAUC) over the complete sample-budget and refinement sweeps:

K=1

–

256

on a

\log_{2}K

axis and

I=1

–

30

on a linear

I

axis. Higher is better; bold marks the best reference for each planner and sweep, including ties. Parentheses give percentage-point differences from grasp-active XYZ15D for the same sweep and planner; negative values indicate decreases. Both planners use no action prior.

Figure 8: Reacher sample-budget sweeps under three observation-history protocols. Top: MPPI; bottom: CEM. Columns show single-frame input, dataset-initialized history, and online history, respectively. All curves use

I=30

without an action prior; bands show

\pm

one sample SD across three evaluation seeds, 50 episodes each. MPPI uses

\tau=64

. State labels are used only to train transition-CGS.

History	Model	MPPI		CEM		GD
		SR	nAUC	SR	nAUC	SR
Single frame	LeWM	$65.3\,{\scriptstyle\pm 4.2}$	$44.63$	$\mathbf{83.3}\,{\scriptstyle\pm 8.3}$	$\mathbf{63.46}$	$76.7\,{\scriptstyle\pm 4.2}$
	action-CGS	$\mathbf{66.7}\,{\scriptstyle\pm 2.3}\,{\scriptstyle(+1.3)}$	$\mathbf{45.79}\,{\scriptstyle(+1.17)}$	$80.7\,{\scriptstyle\pm 3.1}\,{\scriptstyle(-2.7)}$	$63.04\,{\scriptstyle(-0.42)}$	$82.0\,{\scriptstyle\pm 5.3}\,{\scriptstyle(+5.3)}$
	transition-CGS	$66.0\,{\scriptstyle\pm 7.2}\,{\scriptstyle(+0.7)}$	$\mathbf{45.79}\,{\scriptstyle(+1.17)}$	$79.3\,{\scriptstyle\pm 7.0}\,{\scriptstyle(-4.0)}$	$62.88\,{\scriptstyle(-0.58)}$	$\mathbf{85.3}\,{\scriptstyle\pm 7.0}\,{\scriptstyle(+8.7)}$
Dataset-initialized history	LeWM	$68.7\,{\scriptstyle\pm 2.3}$	$49.38$	$96.0\,{\scriptstyle\pm 5.3}$	$71.75$	$84.0\,{\scriptstyle\pm 3.5}$
	action-CGS	$72.7\,{\scriptstyle\pm 3.1}\,{\scriptstyle(+4.0)}$	$49.75\,{\scriptstyle(+0.37)}$	$\mathbf{98.7}\,{\scriptstyle\pm 1.1}\,{\scriptstyle(+2.7)}$	$\mathbf{73.83}\,{\scriptstyle(+2.08)}$	$90.7\,{\scriptstyle\pm 4.2}\,{\scriptstyle(+6.7)}$

Appendix

Table 8: Reacher reference and history ablation: success rate and sample-budget nAUC. SR is mean

\pm

sample SD (%) over three evaluation seeds, 50 episodes each. MPPI/CEM SR uses

K=128

I=30

; GD uses 100 Adam steps and is independent of

K

. nAUC summarizes the complete sample-budget sweep. Parentheses give percentage-point differences from LeWM for both CGS models, within the same history, planner, and metric. Bold marks the highest value within each history/metric group. MPPI uses

\tau=64

Figure 9: Action–effect geometry across four environments. From top to bottom: two examples each of PushT, Cube, TwoRooms, and Reacher. Left: environment illustrations. Middle: complete 30-frame action references (purple dashed) and encoded transition paths for LeWM (blue), TS (brown), and CGS (green), with shared action PCA in 2D or 3D; squares and stars mark starts and ends. Shape agreement indicates preservation of action-related directions. Right: cosine Grams of predicted endpoint responses to circular action perturbations; Ideal is the unit-circle reference. Protocol and interpretation are given in Appendix C.2 .

Figure 10: Selected qualitative PushT planning rollout. Simulator and decoded trajectories for a LeWM failure (top) and a LeWM+CGS success (bottom). Columns show the rollout over environment steps with the target configuration. The CGS example reaches the target while its decoded trajectory remains qualitatively aligned with the simulator rollout, illustrating that the learned representation retains sufficient task-relevant information for successful planning in this example.

Environment	Physical change	$P_{A}^{\rm lin}$	$P_{A}^{\rm MLP}$	$P_{S}^{\rm lin}$	$P_{S}^{\rm MLP}$	$P_{AS}^{\rm lin}$	$P_{AS}^{\rm MLP}$	$R^{2}_{A\mid S}$
PushT	Agent position	0.999	0.999	0.122	0.266	0.999	0.999	0.999
	Block position	0.121	0.317	0.092	0.372	0.175	0.613	0.384
	Block angle	-0.012	0.074	0.045	0.271	0.044	0.543	0.372
Cube	Joint position	0.946	0.966	0.718	0.834	0.960	0.997	0.981
	End-effector position	0.992	0.994	0.753	0.870	0.993	0.999	0.990
	End-effector yaw	0.991	0.992	0.528	0.628	0.992	0.998	0.995

Appendix

Table 9: Physical-motion identification.

A

S

, and

AS

use the action, current state, and both.

R^{2}_{A\mid S}

uses the MLP probes and the state-only residual as its denominator. Negative values denote performance below the held-out mean predictor.

Model	Target	$D_{\rm lin}$	$D_{\rm MLP}$	$C_{A}$	$C_{AS}$	$B_{A}$	$B_{AS}$
LeWM	Agent position	0.809	0.888	0.807	0.808	0.807	0.808
TS	Agent position	0.831	0.891	0.829	0.830	0.829	0.830
CGS	Agent position	0.903	0.959	0.901	0.902	0.901	0.902
LeWM	Block position	0.813	0.864	0.791	0.602	0.256	0.603
TS	Block position	0.816	0.830	0.817	0.670	0.314	0.678
CGS	Block position	0.847	0.873	0.871	0.655	0.276	0.664

Appendix

Table 10: Complete PushT representation probes.

D

denotes dynamic decoding,

C

control-effect accessibility, and

B

bridge agreement. Subscripts

A

and

AS

denote action-only and state-conditioned physical targets.

Model	Target	$D_{\rm lin}$	$D_{\rm MLP}$	$C_{A}$	$C_{AS}$	$B_{A}$	$B_{AS}$
LeWM	Joint position	0.742	0.747	0.741	0.743	0.723	0.743
TS	Joint position	0.622	0.643	0.629	0.622	0.631	0.621
CGS	Joint position	0.743	0.895	0.748	0.741	0.729	0.741
LeWM	End-effector position	0.895	0.969	0.893	0.893	0.891	0.893
TS	End-effector position	0.710	0.776	0.706	0.708	0.700	0.709
CGS	End-effector position	0.931	0.982	0.930	0.930	0.927	0.930

Appendix

Table 11: Complete Cube representation probes. CGS uses grasp-active transitions and XYZ controls for its geometry objective. All three models share the dataset, episode split, horizon, latent dimension, and probe protocol.

Environment	Model	Target	$D_{\rm lin}$	$D_{\rm MLP}$	$C_{A}$	$C_{AS}$	$B_{A}$	$B_{AS}$
Reacher	LeWM	Joint position	0.355	0.988	0.368	0.357	0.447	0.364
	TS	Joint position	0.347	0.988	0.371	0.348	0.447	0.353
	CGS	Joint position	0.123	0.988	0.135	0.127	0.248	0.134
	LeWM	Fingertip position	0.951	0.994	0.155	0.932	0.028	0.932
	TS	Fingertip position	0.974	0.994	0.157	0.960	0.011	0.959
	CGS	Fingertip position	0.899	0.993	0.041	0.883	0.074	0.884

Appendix

Table 12: Complete Reacher and TwoRoom representation probes, using the same notation and protocol as Tables 10 and 11 .

World model / objective	Planning method	Wall			PointMaze–UMaze			PointMaze–Medium			PushT
		Sampling budget			Sampling budget			Sampling budget			Sampling budget
		$K{=}32$	$K{=}64$	$K{=}128$	$K{=}32$	$K{=}64$	$K{=}128$	$K{=}32$	$K{=}64$	$K{=}128$	$K{=}32$	$K{=}64$	$K{=}128$
DINO-WM (CLS)	CEM	$74.7{\pm}3.1$	$76.0{\pm}5.3$	$74.7{\pm}6.1$	$88.0{\pm}3.5$	$92.7{\pm}3.1$	$90.7{\pm}5.0$	$85.3{\pm}4.2$	$86.0{\pm}2.0$	$84.7{\pm}7.0$	$47.3{\pm}6.4$	$56.0{\pm}8.7$	$63.3{\pm}5.0$
	MPPI	$84.0{\pm}5.3$	$\mathbf{86.7{\pm}6.4}$	$\mathbf{89.3{\pm}5.0}$	$92.0{\pm}2.0$	$93.3{\pm}2.3$	$96.7{\pm}4.2$	$90.7{\pm}5.0$	$90.7{\pm}2.3$	$92.7{\pm}3.1$	$31.3{\pm}6.1$	$35.3{\pm}4.6$	$46.7{\pm}3.1$
	GD	$55.3{\pm}2.3$			$62.0{\pm}5.3$			$56.0{\pm}4.0$			$46.0{\pm}2.0$
+TS (CLS)	CEM	$79.3{\pm}6.1$	$78.7{\pm}4.2$	$77.3{\pm}8.3$	$86.0{\pm}2.0$	$89.3{\pm}5.0$	$90.0{\pm}2.0$	$87.3{\pm}1.2$	$84.7{\pm}3.1$	$89.3{\pm}4.6$	$41.3{\pm}5.0$	$38.0{\pm}2.0$	$47.3{\pm}4.2$

Appendix

Table 13: Additional DINO-WM benchmark results. Goal-reaching success rate (%) for Base, TS, and CGS using frozen DINOv2 CLS or Patch representations on the DINO-WM benchmark settings used by Temporal Straightening. CEM and MPPI use

I=30

optimization steps, and GD uses 100 Adam steps. A single environment-level MPPI temperature is shared across representation types and training objectives. Values are means

\pm

sample SD over three evaluation seeds with 50 evaluation episodes per seed; PushT additionally uses proprioception. Bold marks the best objective for a fixed representation type, planner, environment, and sampling budget.

Model	CEM	MPPI	GD
DINO-WM (Patch)+CGS	86.852	86.959	6.516
LeWM+CGS	5.163	1.850	0.161
DINO/LeWM	$16.8\times$	$47.0\times$	$40.5\times$

Appendix

Table 14: PushT planner-solve time at

K=128

. Seconds per episode measured in a timing experiment. Times include both planner solves and exclude loading, environment execution, and rendering. Ratios characterize the stated implementations and batching settings.

Figure 11: PushT training ablations. Ramp fraction, history length, SIGReg weight, and CGS weight, evaluated with MPPI, CEM, and GD. Means use three evaluation seeds with 50 episodes each; GD is independent of

K

. Reference points reuse the main PushT checkpoint with

\lambda_{\mathrm{CGS}}=0.01

, a 5% ramp,

N=3

, and

\lambda_{\mathrm{SIGReg}}=0.09

Figure 12: Temporal straightening in LeWM. Mean cosine similarity of consecutive latent displacements on fixed validation windows. Lines follow one training run per method and environment; bands show

\pm 1.96

standard errors across probe episodes. TS produces higher final mean cosine similarity on all four environments.

Figure 13: MPPI temperature sensitivity and temperature-selection nAUC at

K=128

. Top: Mean success over three evaluation seeds (50 episodes each,

I=30

), using one shared temperature per environment and a 0–100% axis. Bottom: nAUC over the displayed

\tau

values on a

\log_{2}

scale; parentheses give percentage-point differences from LeWM and bold marks the best representation.

Figure 14: Refinement efficiency and matched BC prior quality on Cube, TwoRooms, and Reacher. Left: CEM and MPPI success over

I

(iterations) at

K=128

; bands show sample SD over three evaluation seeds. Middle: first-solve prior endpoint distance divided by zero-action endpoint distance, on the same 150 starts; diamonds and labels show means. Right: prior-only success at

K=1

I=1

, with the prior mean as the sole candidate; error bars show sample SD.

Planner	Model	PushT	Cube	TwoRooms	Reacher
CEM	LeWM	$73.93$	$64.00$	$76.87$	$75.16$
	LeWM + TS	$72.23\,{\scriptstyle(-1.70)}$	$71.98\,{\scriptstyle(+7.98)}$	$\mathbf{94.85}\,{\scriptstyle(+17.98)}$	$73.92\,{\scriptstyle(-1.24)}$
	LeWM + CGS	$\mathbf{80.25}\,{\scriptstyle(+6.32)}$	$\mathbf{77.40}\,{\scriptstyle(+13.40)}$	$93.34\,{\scriptstyle(+16.47)}$	$\mathbf{77.60}\,{\scriptstyle(+2.44)}$
MPPI	LeWM	$64.85$	$64.36$	$86.06$	$67.23$
	LeWM + TS	$67.06\,{\scriptstyle(+2.21)}$	$70.70\,{\scriptstyle(+6.34)}$	$92.89\,{\scriptstyle(+6.83)}$	$68.39\,{\scriptstyle(+1.16)}$
	LeWM + CGS	$\mathbf{75.71}\,{\scriptstyle(+10.86)}$	$\mathbf{73.72}\,{\scriptstyle(+9.37)}$	$\mathbf{94.49}\,{\scriptstyle(+8.44)}$	$\mathbf{71.44}\,{\scriptstyle(+4.21)}$

Appendix

Table 15: Refinement-iteration nAUC at

K=128

I\in\{1,2,5,10,15,20,25,30\}

. nAUC integrates over linear

I

. Parentheses give percentage-point differences from the matching LeWM baseline. Bold marks the best representation for each planner and environment. All entries summarize three evaluation seeds.

Explore similar work

Latent Geometry Beyond Search: Amortizing Planning in World Models
May 9, 2026Hoang Nguyen, Xiaohao Xu, Xiaonan HuangLatent World Models World Models
Department of Robotics, University of Michigan, Ann Arbor
Anisotropic Representations Improve Planning in JEPA World Models
Sep 29, 2026Mingu Kang, Yoori Oh, Sookyung Kim +1Latent World Models Low-Latency Latent Planning
Seoul National University · Ewha Womans University
Representation World Model: Learning States, Transition and Executable Plans in Representation
Sep 24, 2026Yijun Yuan, Weicheng Zheng, Weibang Wang +7World Model Planning World Models
IIIS, Tsinghua University

May 9, 2026, cs.ROJ/K move · Enter open · S save

Latent Geometry Beyond Search: Amortizing Planning in World Models

Hoang Nguyen, Xiaohao Xu, Xiaonan Huang

Department of Robotics, University of Michigan, Ann Arbor

Latent World Models World Models Inverse Dynamics Planning Intermediate Latent States

Open paper AskPDF arXiv

Modern vision-based world models can represent observations as compact yet expressive latent manifolds, but fast goal-oriented planning in these spaces remains challenging. This raises a central question: when does a learned representation simplify control, rather than merely enabling prediction? We study this question in a pretrained LeWorldModel, whose latent geometry is regularized for smoothness and uniformity. Our key insight is that, under such geometry, planning can be amortized into a latent inverse-dynamics mapping instead of requiring online search. We therefore replace iterative planning with a lightweight Goal-Conditioned Inverse Dynamics Model (GC-IDM) that maps the current latent state, goal latent state, and remaining horizon directly to the next action. Empirically, across four benchmark environments spanning navigation, contact-rich manipulation, and continuous control, our controller matches or exceeds CEM in seven of eight environment-protocol settings while reducing per-decision cost by 100-130x. A broader sweep over test-time planners (CEM, MPPI, iCEM, and gradient-based methods) shows that this result is not specific to a particular optimizer. These findings suggest that much of the structure recovered by test-time planning is already locally encoded in the latent representation. More broadly, our results indicate that sufficiently structured latent spaces can shift part of the planning burden from online optimization to learned inference. Our code is publicly available at https://github.com/hdnndh/Latent-Geometry-Beyond-Search-Amortizing-Planning-in-World-Models .

Abstract

Explore similar work

Latent Geometry Beyond Search: Amortizing Planning in World Models

Anisotropic Representations Improve Planning in JEPA World Models

Representation World Model: Learning States, Transition and Executable Plans in Representation