Anisotropic Representations Improve Planning in JEPA World Models

Mingu Kang; Yoori Oh; Sookyung Kim; Joonseok Lee

cs.ROSep 29, 2026

Anisotropic Representations Improve Planning in JEPA World Models

Authors: Mingu Kang, Yoori Oh, Sookyung Kim, Joonseok Lee

Organizations: Seoul National University · Ewha Womans University

Abstract

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/

Figures & tables

Figure 1: Prediction–planning separation with a nonlinear encoder. (a) The Euclidean task cost and the SIGReg latent cost select different outcomes from the same reachable set; the learned latent geometry closely follows the $\Sigma^{-1}$ reference. (b) As process noise decreases, held-out conditional-mean prediction error decreases while normalized physical planning regret remains nearly unchanged. Points and error bars show the mean and $95\%$ $t$ -confidence interval over ten training seeds.

Figure 2: Planning success across environments. AnisoWM with

\Lambda

Reg at

\kappa=2

in every environment, compared with LeWM and the reference baselines. Reported AnisoWM and LeWM values are means over three training seeds.

Task	LeWM	AnisoWM	Gap	LeWM	AnisoWM	Gap
	$J_{\rm enc}$ (representation only)			$J_{\rm pred}$ (planning cost)
TwoRoom	0.564	0.661	$+0.097$	0.575	0.646	$+0.071$
Reacher	0.790	0.893	$+0.103$	0.676	0.723	$+0.048$
PushT	0.647	0.641	$-0.005$	0.594	0.621	$+0.028$
Cube	0.555	0.576	$+0.021$	0.537	0.553	$+0.016$

Table 1: Ordering recorded action sequences by the outcome they reached. For each initial–goal pair, both models score the same sequences, and the entry is the fraction of sequence pairs whose cost ordering agrees with their outcome ordering.

J_{\rm enc}

encodes the observation a sequence actually reached;

J_{\rm pred}

is the cost used for planning and additionally includes the predictor rollout.

Figure 3: Latent cost neighborhoods around the goal. For Cube and PushT, the left panels show the rendered arena and evaluated region, and the right panels show positions in the lowest 10% of latent planning cost around the goal ( $\star$ ) for LeWM (blue) and AnisoWM ( $\kappa{=}2$ , red). The dashed circle shows the corresponding neighborhood under physical Euclidean distance.

Figure 4: Learned target spectra. Left: learned targets at $\kappa=2$ , sorted by variance, with the isotropic target marked. The trace and maximum variance ratio are fixed, while the allocation is learned; the legend shows the number of coordinates above the mean variance. Right: target spectrum during one TwoRoom run, showing continued evolution after the condition-number bound is reached.

Figure 5: Planning success across anisotropy bounds. Each point with

\kappa>1

reports one training run evaluated on the same set of initial–goal pairs. The

\kappa=1

marker shows the isotropic LeWM baseline from the primary comparison.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Notation	Description
$o_{t},\ o_{g}$	Observation at time $t$ and goal observation.
$f_{\theta},\ g_{\phi}$	Neural encoder and action-conditioned latent predictor.
$z_{t},\ z_{g}$	Latent representation at time $t$ and goal representation.
$a_{t},\ U$	Per-step action and finite-horizon action sequence $U=(a_{0},\ldots,a_{H-1})$ .
$H$	Planning horizon.
$\hat{z}_{H}(U)$	Predicted terminal representation after rolling out $U$ .

Appendix

Table 2: Notation used throughout the paper.

Figure 6: Evaluation environments. From left to right: TwoRoom, PushT, Reacher, and Cube. All four environments have continuous action spaces, and visual goal planning is performed directly from image observations.

Setting	Value
Epochs	15
Training seeds	3072, 1234, 42
Optimizer / batch size	AdamW / 128
Peak learning rate	$5\times 10^{-5}$
Encoder and predictor weight decay	$10^{-3}$
Target-logit learning-rate schedule	Same as the model

Appendix

Table 3: Implementation settings. Target logits use zero weight decay and are clipped after each update as described in Section 4.2 .

Task	Seed 3072	Seed 1234	Seed 42	Mean
TwoRoom	90	94	94	92.7
Reacher	90	88	88	88.7
PushT	98	96	96	96.7
Cube	78	80	80	79.3

Appendix

Table 4: Per-seed AnisoWM planning success (%). Each entry reports AnisoWM planning success under the same evaluation protocol used for Figure 2 . The final column reports the mean across the three training seeds, which the main text rounds to an integer.

Task	Outcome scalar	Threshold	$H$	Cases
TwoRoom	agent–goal distance	$16$ px	20	45
Reacher	largest joint-position error	$0.05$ rad	25	50
PushT	$\max\!\big(\text{position}/20\,\mathrm{px},\ \text{angle}/(\pi/9)\big)$	$1$	25	48
Cube	block–target distance	$0.04$ m	15	33

Appendix

Table 5: Outcome scalar and diagnostic horizon. Each scalar is normalized by its threshold, so one is the success boundary and lower is better. Cases are counted out of the fifty initial–goal pairs of the evaluation protocol.

Figure 7: Ranking accuracy, pair by pair. Each point is one initial–goal pair: how well LeWM’s cost orders its sequences on the horizontal axis against how well AnisoWM’s does on the vertical one, both scoring the same sequences. The shaded half-plane above the diagonal is where AnisoWM orders that pair better. Top row:

J_{\mathrm{enc}}

, which leaves the representation on its own. Bottom row:

J_{\mathrm{pred}}

, the cost the planner minimizes. Rows of a column share their axis limits; each panel shows the seed Table 1 reports.

Figure 8: Additional latent-cost neighborhoods around the goal. Four additional Cube/PushT example pairs, shown using the same visualization as Figure 3 . The dashed circle denotes the task-cost neighborhood in physical space; blue and red contours denote the lowest-cost regions under LeWM and AnisoWM ( $\kappa=2$ ), respectively.

Figure 9: Evolution of the learned target condition number. Each curve shows one TwoRoom training run with the indicated anisotropy bound, and the dotted line at each bound is the ratio that bound permits. Filled markers indicate that the learned target reaches its bound.

$\kappa$	TwoRoom	Reacher	PushT	Cube
1.25	90	90	94	82
2	90	90	98	78
4	92	84	92	78
8	92	84	92	78
16	94	86	92	78

Appendix

Table 6: Planning success across anisotropic target bounds (%). Each entry corresponds to the training run used in the sweep of Figure 5 .

Figure 10: Measured feature-covariance spectra in TwoRoom. Eigenvalues of the centered feature covariance are normalized by their mean and sorted independently for each displayed anisotropy bound; the dashed line marks index

100

, the value quoted in the text. These empirical feature covariances are distinct from the learned target covariance

\Lambda

Figure 11: Latent prediction loss and planning success. One marker per trained run, coloured by its anisotropy bound. A run further left predicts better and a run higher plans better. The red ring marks the best-planning run and the grey ring the lowest-loss one; where a single run is both, the two rings are concentric. The MSE is measured in each model’s own learned representation space.

Figure 12: Rollouts in TwoRoom and Reacher. Two rollout pairs per environment compare AnisoWM at $\kappa=2$ with LeWM under the same evaluation planner and seed. Columns show the indicated environment steps and the goal observation; red outlines indicate recorded goal attainment.

Figure 13: Rollouts in PushT and Cube. Two rollout pairs per environment compare AnisoWM at $\kappa=2$ with LeWM under the same evaluation planner and seed. Columns show the indicated environment steps and the goal observation; red outlines indicate recorded goal attainment.

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
ATLAS: Aligned Transport of Latent Structure for Reliable World Model Planning
Sep 28, 2026Ke Fang, Yupu Yao, Lu ChengLatent World Models World Models
Pennsylvania State University
Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control
Jul 28, 2026Jiaxin Bai, Jiaxuan XiongLow-Latency Latent Planning Video Joint Embedding Predictive Architecture
Hong Kong Baptist University Hong Kong, China

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

ATLAS: Aligned Transport of Latent Structure for Reliable World Model Planning

Temporal-Distance JEPA: Plan-Aware Representation Learning for Latent World Model Predictive Control