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

2026-09-28

AI Summary

This study addresses the limitation that regularizing only marginal distributions in latent world models induces relational geometric distortion, thereby undermining action selection reliability during planning. We propose the ATLAS training objective and provide the first proof that relational preservation and marginal calibration constitute non-redundant constraints, establishing a theoretical link between planning stability and geometric distortion. Methodologically, ATLAS employs Wasserstein embedding matching for transport calibration, combined with normalized pairwise structure transfer to optimize latent space geometry. Evaluations on benchmarks including PushT demonstrate that this approach significantly improves goal achievement rates, with the largest gains observed in high-novelty scenarios. These results validate the effectiveness of enhancing latent structural integrity and reducing multi-step prediction errors for robust model-based planning.

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/.

Problem

• latent world models
• planning reliability
• relational geometry
• marginal calibration
• novelty structure

Innovation

• Latent World Model
• Aligned Transport
• Relational Geometry
• Wasserstein Embedding Matching
• World Model Planning