Pub-AI: AI in Science

Latent Dynamics · 2026-09-30

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

Ke Fang, Yupu Yao, Lu Cheng

arXiv:2609.36333PDF

Auto-summarized: this summary was generated by a language model from the paper’s text and has not been reviewed by an editor. Check the paper before relying on it. Benchmark numbers appear only after a human has verified them.

TL;DR

ATLAS addresses a gap in latent world-model planning: regularizing only the latent marginal (anti-collapse) does not guarantee preservation of the relational geometry needed for goal-conditioned action selection. ATLAS transfers normalized pairwise structure from an informative encoder representation to the planning latent and calibrates the planning latent’s marginal using Wasserstein embedding matching (WEMReg). Instantiated in LeWM, it improves mean goal-reaching success, especially on higher-novelty TwoRoom episodes.

Why it matters

For latent planning, what matters is not only that representations are non-degenerate, but that state-to-state relationships in the representation used for planning preserve the geometry present earlier (e.g., in encoder patch features). The authors show that marginal regularization alone can wash out OOD-related novelty structure, weakening planning reliability, and propose an objective that enforces complementary constraints: relational preservation plus Wasserstein-based marginal calibration.

Method

  • ATLAS adds a relational preservation loss that transfers normalized pairwise distances from the encoder’s mean-pooled patch tokens (anchor) to the planning latent (with stop-gradient on the anchor).
  • ATLAS uses WEMReg to calibrate the global planning-latent marginal by matching one-dimensional projected distributions to a standard Gaussian via Wasserstein-2 distance (averaged over random directions).
  • ATLAS is trained with a combined latent prediction objective plus λ·WEMReg and μ·OOD-recovery, and evaluated with goal-conditioned CEM planning in the LeWM-style latent planner.

Limitation

The guarantees concern the evaluated pool and finite candidate set; small training losses alone do not control errors on unseen states. (Scope of Theorem 3 / Appendix C.1.)

Abstract (from 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/.

Related papers