Latent Projection for Latent Planning

Yoo Sung Jang; Michael Ryoo

cs.LGOct 19, 2026

Latent Projection for Latent Planning

Authors: Yoo Sung Jang, Michael Ryoo

Organizations: Stony Brook University

Abstract

Latent world models let an agent plan in latent space. The planner searches for an action sequence that minimizes a cost, and that cost is often the $\ell_2$ distance between the predicted terminal latent and the goal latent. The latent, however, is partly noise: the regularizer that prevents collapse during world model training promotes variance in the directions carrying no signal. We show that the same regularizer also supplies the structure to prune that noise in the latent. SIGReg drives the latent marginal to a standard Gaussian at every step, and assuming further that the latent sequence is a Gauss-Markov process, the transition is linear with its covariance pinned by the map itself, $z' \mid z \sim \mathcal{N}(Az,\, I - AA^\top)$ . Fitting the state dynamics $A$ therefore yields the covariance at no additional cost, and in its eigenbasis each latent direction carries its own uncertainty. Furthermore, constraining $A$ to be a normal matrix fixes that basis across timesteps, so one eigenbasis measures the uncertainty over an arbitrary number of steps and serves as a stable basis to compute the planning cost on. We then restrict the $\ell_2$ cost to the $k$ -dimensional subspace the state dynamics carry forward with the least uncertainty. The correction is post hoc: the encoder, the predictor, the rollout, and the planner are all left untouched, and only a single constrained matrix, $A$ , is fit. We evaluate across six environments and show large gains on the 2D navigation tasks, TwoRoom and PointMaze, raising success rate over the plain $\ell_2$ cost by up to $+54.7$ points under a harder planning setting. Code is available at https://github.com/yjang43/lp2.

Figures & tables

Figure 1: How uncertainty grows across transitions in a

2\times 2

block — one independent plane, given by normality. Starting from

z_{0}

(black), the state dynamics rotate by

\theta

per step while contracting by

r

(orange). The blue circle marks the uncertainty the transition leaves, drawn at one standard deviation; it widens toward the unit circle of the stationary law

\mathcal{N}(0,I)

(grey, dashed).

Panel of Figure 2 — Figure 2: (a) The frozen encoder maps the initial observation and each of the $H$ future observations to latents, and $A^{t}$ is fit to the $t$ -step transition by squared error. $A$ is the only trained component. (b) A candidate action sequence is rolled out through the frozen predictor, and its terminal latent is compared to the goal after both are projected onto the retained directions $U_{k}$ .

Figure 3: Planning success rate (%) by retained rank

k

on validation seed, mean

\pm

std over

3

seeds of

50

episodes. Every task is run in both the default and harder settings. A random basis at the same retained rank is included as a control, drawn in gray. At

k=192

every basis recovers LeWM, drawn as the dashed rule; a star marks the selected rank

k^{\star}

Figure 4: Planning success rate (%) under five planning costs, evaluated on

50

episodes. TD-JEPA blends its distance with the plain

\ell_{2}

cost as

(1-\alpha)\,d_{\psi}+\alpha\,\ell_{2}

, and both

\alpha=0

and

\alpha=0.1

are reported. PCA takes the covariance of the latents and keeps its top eigenvectors. PCA and our cost are reported at the retained rank

k^{\star}

selected on the validation seed, swept over

k\in\{2,4,8,\ldots,128,192\}

. Both the state dynamics

A

and

d_{\psi}

are fit with

3

seeds; bars are the mean over those seeds, and error bars one standard deviation.

Regularizer	Strength	Gaussian	Epps–Pulley $\downarrow$	Gap closed $\uparrow$
SIGReg (LeWM)	$w=0.09$	✓	$2.83$	$92\%$ $(60.0\rightarrow 96.7)$
VISReg	$w=20$	✓	$1.52$	$98\%$ $(40.0\rightarrow 98.7)$
RDMReg	$w=20$	$\times$	$23.80$	$87\%$ $(58.0\rightarrow 94.7)$
Stop-grad	$\mathrm{lr}=3\mathrm{e}{-6}$	$\times$	$42.63$	$0\%$ $(52.0\rightarrow 52.0)$

Table 1: Comparison across anti-collapse methods on the harder TwoRoom setting. Each variant is a retrained LeWM recipe with only the regularizer swapped, at a strength chosen to bring their success rate under the plain

\ell_{2}

cost close together. Gaussian marks whether the method targets a Gaussian marginal. Epps–Pulley measures the Gaussianity of the latent. Gap closed is the gap that our cost closes,

(\mathrm{SR}_{\mathrm{ours}}-\mathrm{SR}_{\mathrm{baseline}})/(100-\mathrm{SR}_{\mathrm{baseline}})

, where

\mathrm{SR}

is the planning success rate, measured under the plain

\ell_{2}

cost (

\mathrm{SR}_{\mathrm{baseline}}

) and under our cost at the best retained rank

k^{\star}

(

\mathrm{SR}_{\mathrm{ours}}

); both appear in parentheses. Latent regularizers that encourage Gaussian structure close the largest share of the gap.

	PSD	Normal	Motion	Multi-step	TwoRoom	PointMaze-UMaze
Ours	✓	✓	✓	✓	96.7 $\pm$ 2.3	96.7 $\pm$ 1.2
	$\times$	✓	✓	✓	96.0 $\pm$ 2.0	94.7 $\pm$ 1.2
	✓	$\times$	—	✓	88.0 $\pm$ 2.0	90.7 $\pm$ 3.1
	✓	✓	$\times$	✓	95.3 $\pm$ 1.2	96.0 $\pm$ 2.0
	✓	✓	✓	$\times$	94.7 $\pm$ 5.0	92.0 $\pm$ 5.3
	$\times$	$\times$	—	✓	82.0 $\pm$ 2.0	42.0 $\pm$ 0.0

Table 2: Ablating the structure of

A

on the harder setting of TwoRoom and PointMaze-UMaze. Each row keeps (✓) or removes (

\times

) components of the full parameterization (shaded). Success rate is measured at the best retained rank

k^{\star}

, mean

\pm

std over

3

seeds. The last row is LeWM, which plans with the plain

\ell_{2}

cost. PSD :

\Sigma

remains a valid covariance. Normal :

AA^{\top}=A^{\top}A

. Motion : the

2\times 2

rotating planes; without them

A

is a symmetric matrix. Multi-step : fitting multi-step transitions instead of the single step. Every component contributes.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Episodes	Episode length
TwoRoom	$10{,}000$	$92$ (avg.)
PushT	$18{,}685$	$125$ (avg.)
OGBench-Cube	$10{,}000$	$200$
Reacher	$10{,}000$	$200$
PointMaze-UMaze	$10{,}000$	$91$ (avg.)
PointMaze-Medium	$10{,}000$	$98$ (avg.)

Appendix

Table 3: Training data for each world model. Lengths marked “avg.” are averages over variable-length episodes.

Symbol	Quantity	Default	Harder
$N$	candidate sequences per iteration	$300$	$300$
$M$	CEM iterations	$10$ ( $30$ on PushT)	$10$ ( $30$ on PushT)
$E$	retained elites	$30$	$30$
$T$	planning horizon (predictor steps)	$5$	$10$
$c$	frame skip (env. timesteps per predictor step)	$5$	$5$
$\Delta$	goal offset (timesteps)	$25$	$50$

Appendix

Table 4: Planning configuration for the default and harder settings.

Figure 6: Eigenvalues of

\Sigma=I-AA^{\top}

per task, sorted in increasing order. Mean

\pm

std over the

3

seeds

A

is fit with.

		ours				random
Task	State	$2$	$8$	$32$	$128$	$2$	$8$	$32$	$128$	$192$
TwoRoom	$x$	$0.71$	$0.94$	$0.97$	$0.99$	$0.03$	$0.09$	$0.35$	$0.99$	$1.00$
	$y$	$0.23$	$0.72$	$0.90$	$0.98$	$0.01$	$0.06$	$0.29$	$0.94$	$0.99$
	mean	$0.47$	$0.83$	$0.94$	$0.98$	$0.02$	$0.07$	$0.32$	$0.96$	$0.99$
PushT	agent $x$	$0.02$	$0.20$	$0.47$	$0.93$	$0.03$	$0.09$	$0.32$	$0.93$	$0.96$
	agent $y$	$0.03$	$0.19$	$0.41$	$0.92$	$0.02$	$0.07$	$0.33$	$0.93$	$0.96$

Appendix

Table 5: Held-out

R^{2}

of a linear probe reading the ground-truth state from a

k

-dimensional subspace of the frozen

192

-dimensional latent, fit by ordinary least squares on

10

k frames and scored on

10

k held-out frames from disjoint episodes. Ours keeps the

k

least-noisy eigenvectors of

\Sigma=I-AA^{\top}

, the directions the cost ( 10 ) projects onto; random is a random orthonormal

k

-subspace. Rank increases left to right within each group. Each mean averages the rows of its task. Only positions are shown; velocities are omitted.

Latent-Space Planning

See all 36 papers

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
Anisotropic Representations Improve Planning in JEPA World Models
Sep 29, 2026Mingu Kang, Yoori Oh, Sookyung Kim +1Representation Learning World Model-Based Planning
Seoul National University · Ewha Womans University
The Objective Is the Bottleneck: Latent World Models Encode What Their Planners Cannot Use
Aug 13, 2026Joyjeet SinghWorld Model-Based Planning Model-Based Planning
Independent researcher

Oct 18, 2026J/K move · Enter open · S save

Latent World Models Should Plan with Distances Along the Latent Manifold

Zunzhi You, Linwei Tao, Siyu Xu, Yunke Wang+1

The University of Sydney · Machine Learning Research, Apple

Latent World Models Latent-Space Planning

Open paper AskPDF arXiv

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.

Abstract

Latent-Space Planning

Explore similar work

Latent World Models Should Plan with Distances Along the Latent Manifold

Anisotropic Representations Improve Planning in JEPA World Models

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