PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Ke Fang, Yupu Yao, Lu Cheng

arXiv:2609.36333v2cs.ROcs.AI

TL;DR

Latent world models use representation geometry to compare predicted outcomes with goals, but marginal regularization alone does not determine state-to-state relationships. ATLAS preserves normalized pairwise structure while calibrating the planning latent, and improves mean success across three benchmarks on both novelty subsets, with its largest reported gain on higher-novelty TwoRoom episodes.

  • Problem

    Latent distribution constraints can prevent collapse without specifying which states remain close or far apart, even though those relationships shape goal-conditioned planning.

    Source: §1 · pp. 1–2

  • Method

    ATLAS transfers normalized pairwise distances from encoder features to the planning latent and calibrates its marginal with Wasserstein embedding matching.

    Source: §1 · p. 2, §3.2 · p. 4, §3.1 · p. 4

  • Results

    ATLAS improves mean planning success over LeWM on ID and OOD episodes across three benchmarks; on TwoRoom, success rises from 74.8% to 87.0% on ID and from 41.2% to 61.0% on OOD episodes.

    Source: §1 · p. 2, §4.1 · p. 6

  • Takeaways & Limitations

    Relational preservation and marginal calibration impose non-redundant constraints, and planning stability is linked to relational distortion, latent-scale mismatch, and terminal prediction error.

    Source: §1 · p. 2, §3.3 · p. 5

    The stability analysis concerns finite-candidate selection and models perturbations through terminal prediction error, relational distortion, and latent-scale mismatch.

    Source: §3.3 · p. 5

Abstract

from arXiv · show

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

1 INTRODUCTION

Latent distribution regularization does not specify which states should remain close or far apart, although those relationships shape goal-conditioned planning. ATLAS preserves normalized pairwise geometry from encoder features while calibrating the planning latent, and its analysis links representation and prediction errors to planning stability.

1 INTRODUCTION

  • Well-behaved global latent distributions do not determine which individual states remain close or far apart, despite those relationships guiding comparison of predicted outcomes with goals.Two latent spaces can have similar global distributions yet preserve substantially different state relationships.Source: §1 · pp. 1–2
  • AUROC falls from 0.76 in encoder features to 0.44 in the planning latent, showing that LeWM weakens an OOD-related signal used to predict planning failure.The diagnostic is reported on PushT.Source: §1 · p. 2
  • ATLAS transfers normalized pairwise distances from informative encoder features to the planning latent and calibrates its global distribution with Wasserstein embedding matching.The objective complements its predictive objective by targeting relational preservation, marginal calibration, and future-state prediction.Source: §1 · p. 2
  • Relational preservation and marginal calibration impose non-redundant constraints, and finite-candidate planning stability is linked to representation distortion and rollout prediction error.The paper evaluates ATLAS across three benchmark datasets and reports improved mean planning success over LeWM on both ID and OOD episodes, with the largest improvement on OOD episodes.Source: §1 · p. 2
  • The paper identifies a gap between anti-collapse regularization and planning-relevant geometry: useful OOD information may persist in the encoder but weaken in the planning latent.ATLAS is presented as a way to preserve that relational structure while calibrating global latent geometry.Source: §1 · p. 2

2 RELATED WORK

Prior latent world models differ in their encoders, training data, and anti-collapse regularizers. ATLAS instead modifies the training objective of a from-scratch model, using Wasserstein matching for calibration and encoder patch features as an anchor for relational structure.

2 RELATED WORK

  • ATLAS changes only the training objective of a from-scratch model, unlike related methods that change the encoder or data.LeWM trains a ViT encoder and latent predictor end-to-end; DINO-WM uses a frozen pretrained encoder, while PLDM learns dynamics from reward-free offline data.Source: §2 · p. 2
  • VICReg uses variance and covariance penalties, while LeJEPA and LeWM use Gaussian-oriented regularization; ATLAS instead uses Wasserstein embedding matching.The paper presents WEMReg as retaining the anti-collapse goal while addressing planning-relevant OOD structure.Source: §2 · pp. 2–3
  • WEMReg matches projected latent distributions to a standard Gaussian through exact one-dimensional quantile transport, penalizing every projected quantile residual.The passage contrasts this with finite-frequency normality tests and describes a closed-form batch objective.Source: §2 · p. 3
  • ATLAS transfers pairwise geometry from the model’s own encoder patch tokens, following relational distillation rather than matching absolute activations.The passage also situates k-nearest-neighbor distance as a nonparametric OOD signal used in the paper’s diagnosis and theory.Source: §2 · p. 3

3 ATLAS

ATLAS combines latent prediction with marginal calibration and relational-geometry preservation because a well-behaved latent distribution alone does not specify which states remain close or far apart. Its analysis connects these distinct representation constraints, prediction error, and latent-scale mismatch to finite-candidate planning stability.

3 ATLAS

  • ATLAS combines latent prediction, WEMReg marginal calibration, and OOD-recovery relational preservation to address planning’s prediction, distribution, and state-relationship requirements.The three objective terms respectively target accurate dynamics prediction, marginal calibration, and preservation of relational geometry.Source: §3 · p. 3, §3.2 · p. 5

3.1 CALIBRATING THE LATENT DISTRIBUTION

  • WEMReg averages one-dimensional projected-distribution Wasserstein costs against standard-Gaussian quantiles; finite-frequency matching can miss non-Gaussian distributions with zero mean and identity covariance.The theorem establishes that such a distribution can be indistinguishable from a standard Gaussian under any specified finite-frequency objective while retaining positive WEMReg.Source: §3.1 · p. 4

3.2 PRESERVING RELATIONAL GEOMETRY

  • OOD-recovery transfers normalized pairwise geometry from mean-pooled patch features, rather than copying features, and Proposition 2 shows relational and marginal losses constrain different properties.Pairwise distances are computed after coordinate standardization and divided by their batch mean to remove overall scale differences.Source: §3.2 · pp. 4–5

3.3 FINITE-CANDIDATE PLANNING STABILITY

  • The stability analysis bounds novelty-score differences and finite-candidate planning regret in relation to geometry distortion, terminal prediction error, and latent-scale mismatch.The analysis links prediction loss to terminal rollout error, relational loss to geometry distortion, and WEMReg to scale mismatch.Source: §3.3 · p. 5

4 EXPERIMENTS

ATLAS improves mean planning success over LeWM across all three tasks and both novelty subsets. Ablations show that WEMReg and OOD-recovery have distinct, partly overlapping benefits whose planning gains vary by task.

4 EXPERIMENTS

  • The experiments compare offline-trained models on navigation and manipulation tasks using a shared goal-conditioned planning protocol and receding-horizon CEM.Evaluation uses a goal 50 steps ahead, 200 episodes per seed, and five evaluation seeds.Source: §4 · pp. 5–6

4.1 MAIN RESULTS

  • Evaluation episodes are split into equal ID and OOD halves using a fixed, model-independent expert-path distance-to-training-states score.For each task and seed, episodes above the median score are labeled OOD.Source: §4.1 · p. 6
  • Across all three tasks, ATLAS improves mean success over LeWM on ID and OOD; on TwoRoom, ID rises from 74.8% to 87.0% and OOD from 41.2% to 61.0%.The largest improvement is on higher-novelty TwoRoom episodes; gains on PushT and OGBench-Cube are modest.Source: §4.1 · p. 6, Figure 3 · p. 7

4.2 ABLATION STUDY

  • ATLAS attains the highest mean success on all three tasks, but its gain over the stronger single-component variant is small and varies across tasks.On TwoRoom either component alone gives a large improvement over LeWM; on PushT, OOD-recovery gives the larger individual gain, while on OGBench-Cube WEMReg alone does not improve over LeWM.Source: §4.2 · p. 7, Table 1 · p. 7

5 ANALYSIS OF LATENT GEOMETRY AND PLANNING RELIABILITY

Diagnostics examine novelty-related geometry, marginal calibration, multi-step prediction, and selected closed-loop behavior. They align with the paper’s planning-stability factors, but are presented as mechanism-oriented diagnostics rather than causal explanations.

Novelty-related structure

  • ATLAS retains more novelty-predictive signal in the planning latent: its k-NN diagnostic reaches 0.75, compared with 0.77 in patch features, while LeWM’s signal weakens substantially across representation depth.The relational objective transfers normalized pairwise geometry from patch features without directly supervising novelty scores.Source: §5 · pp. 7–8
  • On PushT, ATLAS’s PCA projection separates high-novelty states from the low-novelty cluster more clearly than LeWM’s; the two-dimensional plot is qualitative support, not a direct measure of relational preservation.The separation is less pronounced on TwoRoom and OGBench-Cube.Source: §5 · p. 8, Figure 4 · p. 8

Marginal calibration

  • Variants using WEMReg improve several latent-marginal diagnostics, especially covariance and projected-distribution errors on TwoRoom, though improvements are not uniform across tasks and statistics.Diagnostics use fresh random-policy simulator frames, separate from the planning ID/OOD evaluation split.Source: §5 · p. 8, Figure 5 · p. 8

Multi-step prediction

  • ATLAS has lower normalized multi-step latent prediction error than LeWM at every measured horizon on PushT and TwoRoom, with the gap generally increasing over longer rollouts.One prediction step spans five physical steps.Source: §5 · p. 9, Figure 6 · p. 9

Closed-loop trajectories

  • In one selected TwoRoom episode, ATLAS passes through the connecting opening and reaches the goal while LeWM remains across the wall; the example is illustrative, not representative evidence.The figure shows LeWM above ATLAS, with frame indices marking elapsed steps and the goal in the rightmost column.Source: §5 · p. 9, Figure 7 · p. 9

6 CONCLUSION

The conclusion presents ATLAS as jointly preserving relational geometry and calibrating the latent marginal, with improved planning success and supporting representation and rollout diagnostics.

Conclusion

  • ATLAS jointly targets relational preservation and marginal calibration, connecting representation properties to finite-candidate planning through distortion, latent-scale mismatch, and prediction error.OOD-recovery preserves normalized geometry relative to an internal patch anchor, while WEMReg calibrates the global latent distribution.Source: §6 · p. 9
  • Across PushT, TwoRoom, and OGBench-Cube, ATLAS improves mean goal-reaching success over LeWM on both lower- and higher-novelty subsets.Ablations show distinct but partly overlapping benefits; diagnostics show stronger novelty-related structure, improved marginal calibration, and lower multi-step prediction error.Source: §6 · p. 9
  • The results support preserving planning-relevant latent geometry as a useful principle for improving the reliability of latent world-model planning.Source: §6 · p. 9

A EXPERIMENTAL SETUP DETAILS

The experiments use offline demonstration datasets for three control tasks and share model architecture and task-level hyperparameters. ID/OOD labels are assigned by a fixed median split of model-independent expert-path novelty scores.

A EXPERIMENTAL SETUP DETAILS

  • Models use the same architecture, with variants differing only in anti-collapse regularizer and OOD-recovery weight; training uses offline demonstration datasets for TwoRoom, PushT, and OGBench-Cube.The validation split monitors training, while evaluation goals are sampled from valid episodes and methods are evaluated at goal offset 50.Source: A EXPERIMENTAL SETUP DETAILS · p. 11, Table 2 · p. 12

A.1 IN-DISTRIBUTION / OUT-OF-DISTRIBUTION SPLIT

  • Within each task and seed, the 200 evaluation episodes are split evenly at the median expert-path distance to a fixed bank of training states, with above-median episodes labeled OOD.The 50-neighbor novelty score uses standardized state coordinates and a fixed, model-independent bank.Source: §A.1 · p. 12, Table 3 · p. 12

B OOD DIAGNOSTIC: PROTOCOL AND FULL RESULTS

The diagnostic tests whether encoder-level OOD information remains predictive of planning failure in the planner's latent, and ATLAS restores much of the signal lost by LeWM. The accompanying setup defines the diagnostic representations and bounds the stated guarantees to the evaluated pool and finite candidate set.

B OOD DIAGNOSTIC: PROTOCOL AND FULL RESULTS

  • The OOD score is mean distance to the 50 nearest neighbors in a standardized bank of 1.5k training frames; larger scores indicate greater distance from training support.The true-state oracle uses the same k-NN construction on raw state.Source: B OOD DIAGNOSTIC: PROTOCOL AND FULL… · p. 13
  • AUROC measures whether logged OOD scores predict simulator-labeled failure, with chance at 0.5; the oracle estimates how strongly failures are OOD-driven.The unchanged latent CEM planner runs on 200 evaluation episodes, with each episode scored over its last three steps.Source: B OOD DIAGNOSTIC: PROTOCOL AND FULL… · p. 13
  • ATLAS raises planning-latent failure-prediction AUROC to 0.75/0.64/0.69 on PushT/TwoRoom/Cube, matching its patch anchor and oracle after LeWM's values of 0.44/0.55/0.49.The diagnostic uses the same 200-episode set and planner for both models.Source: Table 5 · p. 13, Table 4 · p. 13, B OOD DIAGNOSTIC: PROTOCOL AND FULL… · p. 13

C.1 DEFINITIONS AND SETUP

  • WEMReg is described as an exact directional expectation using quadratic Wasserstein distance, while evaluation fixes a finite pool and shared normalizers.The stated guarantees concern that evaluated pool and finite candidate set; small training losses alone do not control errors on unseen states.Source: §C.1 · pp. 13–15

C.2 EXACT FINITE-SAMPLE SLICED TRANSPORT

These appendices characterize finite-sample sliced transport and establish that WEMReg calibration, relational structure, and finite-candidate stability address distinct mathematical properties. They also give a quantitative connection between WEMReg and normalization mismatch.

C.2 EXACT FINITE-SAMPLE SLICED TRANSPORT

  • Applying the Gaussian quantile-cell decomposition to sorted projections yields the exact batch WEMReg cost with a non-optimizable residual.A lower bound is attained in one projection, but need not be attainable in every direction of a finite high-dimensional point cloud.Source: §C.2 · p. 15

C.3 PROOF OF FINITE-FREQUENCY AMBIGUITY

  • A finite-frequency characteristic-function loss can be zero for a non-Gaussian distribution even when its WEMReg cost is positive.The proof constructs a distribution distinct from the Gaussian that satisfies the finite-frequency constraints.Source: §C.3 · pp. 15–16

C.4 WEMREG CONTROLS THE NORMALIZATION MISMATCH

  • WEMReg cost s bounds mean and covariance mismatch, including the coordinate-scale mismatch used in relational standardization.The result applies to population or empirical reference distributions; a finite Monte Carlo estimate requires an additional integration-error bound.Source: §C.4 · p. 16

C.5 PROOF OF NON-REDUNDANCY

  • A one-dimensional construction shows that positive scaling can preserve normalized relational distances while changing the sorted values, so relational preservation and marginal transport are non-redundant.The construction works for every ε0 ≥ 0.Source: §C.5 · pp. 16–17

C.6 STABILITY OF FINITE MINIMIZATION

  • For finite candidate sets, a unique minimizer remains the minimizer after a uniform cost perturbation of at most ε when its margin exceeds 2ε.The lemma also bounds the difference between the two minimum costs by ε.Source: §C.6 · p. 17

C.7 PROOF OF RELIABILITY AND PLANNING STABILITY

The appendix connects relational and prediction errors to finite-candidate planning stability. It gives sufficient margin conditions for preserving choices and bounds terminal prediction error from one-step errors under a Lipschitz predictor.

C.7 PROOF OF RELIABILITY AND PLANNING STABILITY

  • Uniformly bounded discrepancies between anchor and latent candidate costs yield a finite-candidate regret bound.The proof applies finite-minimization stability to the anchor goal distance and the predicted terminal-embedding cost.Source: §C.7 · p. 17
  • Under Theorem 3’s assumptions, including η < 1, a unique anchor minimizer i⋆ is guaranteed to be selected whenever each alternative i has anchor-cost margin dP(i,g) − dP(i⋆,g) > 2ϵplan.When ϵrel = 0, all kNN scores agree on the pool.Source: §C.7 · p. 18

C.8 FROM ONE-STEP PREDICTION TO TERMINAL ERROR

  • With bounded one-step prediction error and an L-Lipschitz predictor on compared states, terminal error follows by iterating e_(t+1) ≤ L e_t + ε_dyn from e_0 = 0.The stated iteration also covers L = 1.Source: §C.8 · p. 18

C.9 CONNECTION TO A TASK-ALIGNED TERMINAL COST

Task-cost regret is related to planning regret when task cost is aligned with anchor distance. The accompanying diagnostics distinguish illustrative trajectories and structural decompositions from quantitative planning-success predictions.

C.9 CONNECTION TO A TASK-ALIGNED TERMINAL COST

  • Task-cost regret is bounded under Theorem 3 when task cost differs from anchor goal distance by at most ε_anchor for every candidate.The result requires this anchor-alignment condition in addition to Theorem 3.Source: §C.9 · p. 18
  • Figure 10 compares bound and regret across arm-query evaluations against an equality line, while its conservative bound and absent sufficient-margin events limit quantitative interpretation.Source: Figure 10 · p. 19
  • Figure 11 reports within-query pairwise RMSE, 95th-percentile distortion, and maximum distortion using across-query medians and paired-bootstrap intervals.PushT and TwoRoom use 100 queries each; Cube uses 16 exploratory queries, with controlled checkpoints for all four arms.Source: Figure 11 · p. 20
  • Theorem 3's bars decompose ε_plan into mean prediction, relational, and scale contributions, but the figure is not a quantitative predictor of planning success.Source: Figure 12 · p. 20
  • The selected PushT and OGBench-Cube episodes illustrate cases where LeWM fails and ATLAS succeeds, with ATLAS aligning the block to the target pose and carrying the cube to its target configuration.Figure 13 shows one episode per task, with frame indices marking elapsed steps.Source: Figure 13 · p. 21
Loading 2609.36333v2…