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

Ke Fang; Yupu Yao; Lu Cheng

cs.ROSep 28, 2026

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

Authors: Ke Fang, Yupu Yao, Lu Cheng

Organizations: Pennsylvania State University

Abstract

Latent world models rely on representation geometry for planning, yet regularizing the latent marginal alone does not determine the state-to-state relationships used for action selection. We show that this can cause planning-relevant novelty structure to be weakened as representations are transformed into the final latent used by the planner. We introduce Aligned Transport of Latent Structure (ATLAS), a training objective that explicitly preserves relational geometry while calibrating the global latent distribution. ATLAS transfers normalized pairwise structure from an informative encoder representation to the planning latent and uses Wasserstein embedding matching (WEMReg) to calibrate its marginal through one-dimensional Wasserstein-2 transport. Our analysis shows that relational preservation and marginal calibration impose non-redundant constraints, and connects finite-candidate planning stability to relational distortion, latent-scale mismatch, and prediction error. Instantiated in LeWM, ATLAS improves mean goal-reaching success across PushT, TwoRoom, and OGBench-Cube on both lower- and higher-novelty evaluation subsets, with the largest gain on higher-novelty TwoRoom episodes. Representation and rollout diagnostics further show stronger novelty-related structure in the planning latent, improved marginal calibration, and lower multi-step prediction error. Together, these results highlight preservation of planning-relevant latent geometry as an important ingredient for reliable world-model planning. Code is available at https://anonymous.4open.science/r/atlas-world-model-72C4/.

Figures & tables

Figure 1: An OOD-related diagnostic gap in latent planning on PushT.

Figure 3: Planning success at goal offset

50

on lower- and higher-novelty evaluation subsets. We refer to the two median-split subsets as ID and OOD, respectively. Bars show mean success over five evaluation seeds. The split is determined by a model-independent expert-path novelty score based on distance to training states. Exact values are reported in Table 3 .

Arm	PushT	TwoRoom	OGBench-Cube
LeWM (base)	$45.10\pm 3.65$	$58.00\pm 3.36$	$48.10\pm 2.94$
ATLAS w/o WEMReg	$50.30\pm 3.85$	$71.80\pm 2.99$	$50.50\pm 2.55$
ATLAS w/o OOD-recovery	$47.00\pm 2.00$	$73.90\pm 2.46$	$46.90\pm 2.60$
ATLAS	$\mathbf{50.80\pm 0.68}$	$\mathbf{74.00\pm 1.92}$	$\mathbf{50.70\pm 2.69}$

Table 1:

2\times 2

ablation at goal offset

50

over marginal calibration (SIGReg vs. WEMReg) and relational preservation (OOD-recovery off vs. on). Planning success rate (%, mean

\pm

std over 5 evaluation seeds); highest mean per task in bold.

Figure 4: Planning-latent geometry (PushT). Two-dimensional PCA projections of $z$ . Each point is a state, colored by its model-independent true-state novelty percentile.

Figure 5: Latent-marginal diagnostics under distribution shift. Planning-latent diagnostics on fresh random-policy simulator frames. From left to right, the panels report mean norm, normalized covariance error, median projected CvM discrepancy, and median whitened projected CvM discrepancy.

Figure 6: Multi-step rollout prediction error. Results on PushT (a) and TwoRoom (b) use

100

shared queries and

128

shared candidates per query. Lines show the across-query median of candidate-median latent prediction errors, normalized by the fixed checkpoint-specific reference-bank scale

\ell

. Shading denotes

95\%

percentile intervals from

1{,}000

paired query-bootstrap resamples. One prediction step spans five physical steps.

Figure 7: Selected closed-loop planning trajectory on TwoRoom. One episode in which LeWM fails and ATLAS succeeds (LeWM above ATLAS; frame indices indicate elapsed steps, goal in the rightmost column). The PushT and OGBench-Cube cases are shown in Figure 13 in Appendix.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Group	Hyperparameter	Value
Model	Latent / embedding dimension	$192$
	Predictor	6-layer causal Transformer
	Attention heads / MLP dimension	$16$ / $2048$
	History length	$3$ frames
	Frameskip	$5$
Training	Epoch budget (all tasks)	$10$

Appendix

Table 2: Training, model, and planning hyperparameters.

	PushT		TwoRoom		OGBench-Cube
Method	ID	OOD	ID	OOD	ID	OOD
ATLAS (ours)	$\mathbf{54.00\pm 3.63}$	$\mathbf{47.60\pm 2.58}$	$87.00\pm 2.28$	$61.00\pm 2.19$	$58.40\pm 3.72$	$43.00\pm 4.94$
LeWM	$48.20\pm 4.62$	$42.00\pm 5.76$	$74.80\pm 2.79$	$41.20\pm 4.35$	$55.80\pm 4.31$	$40.40\pm 4.22$
DINO-WM	$51.00\pm 5.93$	$34.40\pm 4.67$	$\mathbf{97.80\pm 1.17}$	$\mathbf{98.20\pm 0.98}$	$59.40\pm 4.50$	$41.60\pm 7.28$
FAST-WM	$40.40\pm 5.08$	$34.60\pm 4.96$	$77.00\pm 4.00$	$47.20\pm 3.87$	$\mathbf{61.60\pm 4.63}$	$\mathbf{43.60\pm 5.71}$
PLDM	$28.80\pm 4.87$	$24.20\pm 5.19$	$85.60\pm 3.72$	$72.40\pm 1.85$	$53.80\pm 4.62$	$36.20\pm 6.97$

Appendix

Table 3: Main results at goal offset

50

, split into in-distribution (ID) and out-of-distribution (OOD) episodes (%, mean

\pm

std over 5 seeds). Best per column in bold; exact data of Figure 3 .

Task	L3	L6	L9	patch	cls	$z$ (SIGReg)	oracle
PushT	0.59	0.81	0.83	0.76	0.72	0.44	0.88
TwoRoom	0.69	0.70	0.60	0.79	0.71	0.55	0.69
OGBench-Cube	0.82	0.86	0.85	0.71	0.79	0.49	0.76

Appendix

Table 4: OOD diagnostic at goal offset

50

(

n{=}200

episodes, baseline LeWM).

\mathrm{AUROC}(\text{OOD score}\rightarrow\text{planning failure})

per representation (shallow

\rightarrow

deep, ending at the planning latent

z

). “oracle” is the true-state reference. Chance

=0.5

Task	L3	L6	L9	patch	cls	$z$	oracle
PushT	0.55	0.63	0.74	0.77	0.77	0.75	0.82
TwoRoom	0.30	0.38	0.61	0.65	0.69	0.64	0.59
OGBench-Cube	0.82	0.89	0.88	0.69	0.84	0.69	0.70

Appendix

Table 5: OOD diagnostic at goal offset

50

for ATLAS (

n{=}200

episodes, same episode set and planner as Table 4 ).

\mathrm{AUROC}(\text{OOD score}\rightarrow\text{planning failure})

per representation.

Figure 8: OOD-related diagnostic signal across representations. AUROC for predicting planning failure from a

k

-NN novelty score at goal offset

50

. Hatched bars denote LeWM and solid bars ATLAS. Scores are computed from mean-pooled patch features, the planning latent

z

, and true environment states (oracle reference). This is a visualization of the per-representation values in Tables 4 and 5 .

Figure 9: Planning-latent geometry on TwoRoom and OGBench-Cube. Two-dimensional PCA projections of $z$ for LeWM (top) and ATLAS (bottom), companion to the PushT panels in Figure 4 . Each point is a state, colored by its model-independent true-state novelty percentile, from lower novelty in blue to higher novelty in red. The separation between low- and high-novelty states is less pronounced than on PushT, but the ATLAS projections still place high-novelty states farther from the in-distribution cluster than LeWM does.

Figure 10: Bound versus regret. Each point is one arm-query evaluation; the dashed line is equality. The conservative bound and absent sufficient-margin events limit quantitative interpretation.

Figure 11: Relational metrics. Across-query medians and paired-bootstrap percentile intervals for within-query pairwise RMSE, 95th-percentile distortion, and maximum distortion. PushT/TwoRoom each use 100 queries, Cube 16 exploratory queries. All four arms use controlled checkpoints.

Figure 12: Theorem 3 decomposition. Bars show mean prediction, relational, and scale contributions to epsilon plan on the original queries. This is a structural decomposition, not a quantitative planning-success predictor.

Panel of Figure 13 — Figure 13: Selected closed-loop planning trajectories on PushT (top) and OGBench-Cube (bottom). One episode per task in which LeWM fails and ATLAS succeeds (LeWM above ATLAS in each panel; frame indices indicate elapsed steps, goal in the rightmost column). On PushT, ATLAS rotates and aligns the T-shaped block with the target pose; on OGBench-Cube, ATLAS carries the cube to the target configuration.

Explore similar work

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
ATM: Why Latent World Models Can Fail to Plan
Jun 8, 2026Jiaheng Chen, Tinghe Zhang, Yucheng Xiao +5Latent World Models Action-Conditioned Planning
School of Software, Northeastern University
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

Sep 29, 2026, cs.ROJ/K move · Enter open · S save

Anisotropic Representations Improve Planning in JEPA World Models

Mingu Kang, Yoori Oh, Sookyung Kim, Joonseok Lee

Seoul National University · Ewha Womans University

Latent World Models Low-Latency Latent Planning World Planning Representation Space Gaussian Mixture Models

Open paper AskPDF arXiv

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/

Abstract

Explore similar work

Anisotropic Representations Improve Planning in JEPA World Models

ATM: Why Latent World Models Can Fail to Plan

Latent Geometry Beyond Search: Amortizing Planning in World Models