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

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

Sep 29, 2026

Anisotropic Representations Improve Planning in JEPA World Models

Latent world models learn action-conditioned dynamics in representation space and often score candidate actions by Euclidean distance to a goal representation. Joint training typically regularizes the representation to prevent collapse, but the resulting representation geometry also determines how terminal errors are weighted during planning. We show that accurate prediction and noncollapsed representations do not guarantee a task-aligned latent planning cost: isotropic Gaussian regularization can induce a geometry that ranks feasible outcomes differently from the task cost. To address this mismatch, we introduce AnisoWM with

Λ

Reg, which replaces the fixed isotropic Gaussian target with a learnable diagonal covariance under fixed-trace and anisotropy constraints. The prediction objective, predictor architecture, and Euclidean planner remain unchanged; the target is used only during training. Our analysis characterizes the prediction-driven allocation of target variance, its dependence on the training distribution, and the conditions under which the induced metric reduces planning regret. Across four visual control environments, AnisoWM improves planning success over LeWorldModel in all four. Its latent planning cost also shows better agreement with task outcomes. Project website: https://rkdrn79.github.io/AnisoWM-page/

Mingu Kang, Yoori Oh, Sookyung Kim +1

Seoul National University · Ewha Womans University

Representation Learning World Model-Based Planning Action-Conditioned World Models Model-Based Planning Predictive Representation Learning Latent World Models+1

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

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