ATLAS: Aligned Transport of Latent Structure for Reliable World Model Planning
On the map: World models → Latent predictive world models
Abstract (arXiv)
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/.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Placed in Latent predictive world models, one of 11 directions in World models. Builds on 22 of its 28 references that are in the map. First posted on arXiv 28 Sep 2026.
Builds on
22 of 28 references.
Closest ancestors · 6
Same direction, newest first
- 2026Fast LeWorldModel24 Jun
- 2026LeWorldModel: Stable End-to-End Joint-Embedding Predictive Architecture from Pixels13 Mar
- 2024DINO-WM: World Models on Pre-trained Visual Features enable Zero-shot Planning7 Nov
- 2024Revisiting Feature Prediction for Learning Visual Representations from Video15 Feb
- 2023Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture19 Jan
- 2018Learning Latent Dynamics for Planning from Pixels12 Nov
Other references · 16
- 2020Bootstrap your own latent: A new approach to self-supervised Learningcited by 4,750
- 2020Exploring Simple Siamese Representation Learningcited by 2,866
- 2023Diffusion Policy: Visuomotor Policy Learning via Action Diffusioncited by 2,838
- 2021Barlow Twins: Self-Supervised Learning via Redundancy Reductioncited by 1,838
- 2021VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learningcited by 1,057
- 2019Relational Knowledge Distillationcited by 889
- 2022Out-of-Distribution Detection with Deep Nearest Neighborscited by 600
- 2019Similarity-Preserving Knowledge Distillationcited by 516
- 2023TD-MPC2: Scalable, Robust World Models for Continuous Controlcited by 512
- 2021Implicit Behavioral Cloningcited by 452
- 2019Correlation Congruence for Knowledge Distillationcited by 285
- 2025LeJEPA: Provable and Scalable Self-Supervised Learning Without the Heuristicscited by 254
4 more
- 2024OGBench: Benchmarking Offline Goal-Conditioned RLcited by 238
- 2025Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Modelscited by 89
- 2018Sliced-Wasserstein Autoencoder: An Embarrassingly Simple Generative Modelcited by 35
- 2026stable-worldmodel: A Platform for Reproducible World Modeling Research and Evaluationcited by 20
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026SCALE: State-Calibrated Latent Embeddings for JEPA Planning in the Right Geometry
- 2026Latent Geometry Beyond Search: Amortizing Planning in World Models
- 2026Beyond a single latent space: a dual-latent world model for long-horizon planning
- 2026One-Step Prediction Fidelity Is a Poor Proxy for Planning Performance in JEPA World Models
- 2026Predictive but Not Plannable: RC-aux for Latent World Models
- 2026The Objective Is the Bottleneck: Latent World Models Encode What Their Planners Cannot Use
- 2026EpicWorldModel: Exploration-driven Planning with Latent World Models
- 2026Control-Geometry Straightening for Sampling-Based Latent Planning