ATLAS preserves relational geometry while calibrating the latent distribution in LeWM, improving goal-reaching success on PushT, TwoRoom, and OGBench-Cube, with the largest gain on higher-novelty TwoRoom episodes

The authors show that a latent world model can retain useful out-of-distribution information in its encoder while losing much of that signal in the latent representation used for planning, exposing a gap between anti-collapse regularization and preservation of planning-relevant geometry. Prior work such as SIGReg prevents representation collapse by constraining the latent marginal, but does not specify which states should remain close or far apart; this work uses a k-nearest-neighbor novelty score against a training feature bank as a diagnostic and shows the signal weakens substantially from the patch representation to the planning latent in LeWM. On PushT, failure-prediction AUROC decreases from 0.76 in the encoder representation to 0.44 in the planning latent, below chance; on TwoRoom and OGBench-Cube the planning latent reaches 0.55 and 0.49 respectively, both below the patch anchor and the true-state oracle.

The authors introduce ATLAS, a training objective that explicitly preserves planning-relevant relational geometry while calibrating the global latent distribution. ATLAS transfers normalized pairwise distances from the encoder's mean-pooled patch features into the planning latent and uses WEMReg to match random one-dimensional projections of the latent distribution to a standard Gaussian via Wasserstein-2 transport; unlike LeWM's SIGReg finite-frequency matching, WEMReg penalizes every projected quantile residual. Theoretical analysis provides Theorem 1 showing finite-frequency ambiguity for non-Gaussian distributions and Proposition 2 showing relational preservation and marginal calibration impose non-redundant constraints; instantiated in LeWM with the same architecture, data, and planner, changing only the marginal regularizer and adding relational preservation.

The authors connect representation distortion and prediction error to finite-candidate planning stability and provide a regret bound for candidate selection under anchor geometry. Theorem 3 relates the prediction loss, relational objective, and WEMReg to terminal prediction error, geometry distortion, and latent-scale mismatch respectively, and gives a bound on the difference between k-nearest-neighbor novelty scores computed in the planning and anchor representations. Proofs are provided in Appendix C.1 under assumptions including a finite candidate set and bounded latent-scale mismatch; the guarantees concern the evaluated pool and finite candidate set, and small training losses alone do not control errors on unseen states.

Across PushT, TwoRoom, and OGBench-Cube, ATLAS improves mean goal-reaching success over LeWM on both lower- and higher-novelty evaluation subsets, with the largest gain on higher-novelty TwoRoom episodes. Compared with the closest controlled comparison LeWM, ATLAS uses the same JEPA-style encoder-predictor setup and changes only the training objective; ablations show WEMReg and OOD-recovery constrain different representation properties, but their downstream benefits need not be additive. On TwoRoom, ATLAS improves from LeWM's ID success rate to and OOD from to; gains on PushT and OGBench-Cube are modest; episodes are evaluated per task per evaluation seed, with mean and standard deviation reported across five seeds.

ATLAS as a training objective can be directly applied to other JEPA-style latent world models, suited to offline-data training and goal-conditioned tasks with finite-candidate CEM planning at test time; its theoretical guarantees concern the evaluated pool and finite candidate set, the method is instantiated in LeWM, and evaluation covers PushT, TwoRoom, and OGBench-Cube using the same architecture, data, and planner, so the most directly applicable setting is latent world model planning under similar conditions.

Theoretical guarantees concern the evaluated pool and finite candidate set, and small training losses alone do not control errors on unseen states; benefits of relational preservation and marginal calibration are not uniform across tasks, and ablations show their downstream gains need not be additive; DINO-WM performs better on TwoRoom but uses an externally pretrained encoder and is not a controlled objective-level comparison; PCA visualizations are only two-dimensional projections and serve as qualitative support rather than direct measures; readers of the abstract alone may miss ablation and diagnostic details and should consult the appendix for full settings and values.