ATLAS: Aligned Transport of Latent Structure for Reliable World Model Planning
Ke Fang, Yupu Yao, Lu Cheng
cs.RO, cs.AI
2026-09-29
Gaussian latent regularization erases relational structure a planner needs; ATLAS restores it, lifting TwoRoom high-novelty planning success from 41% to 61%.
Latent world models pick actions by comparing a predicted latent state against a goal latent. Anti-collapse regularizers like SIGReg keep the latent's overall distribution close to Gaussian, but that says nothing about which states should sit close together and which should stay apart — the relationships a planner actually uses. The paper runs a diagnostic on LeWM, a JEPA-style world model: a k-nearest-neighbor novelty score is used to predict planning failure. In the encoder's patch features this score hits an AUROC of 0.76 on PushT, but in the final planning latent it drops to 0.44, worse than a coin flip. The cue that a state is unfamiliar exists inside the encoder and gets washed out on the way to the representation the planner actually uses.
ATLAS adds two regularizers on top of LeWM's prediction loss:
Together with the prediction loss, the three terms target dynamics accuracy, relational preservation, and marginal calibration separately; Proposition 2 shows neither regularizer determines the other. A theoretical result (Theorem 3) ties the stability of finite-candidate CEM planning to the sum of prediction error, relational distortion, and latent-scale mismatch.
Each task (PushT, TwoRoom, OGBench-Cube) is split into lower- and higher-novelty evaluation halves by each episode's distance to the training data (200 episodes, 5 seeds):
| Task | LeWM low-novelty | ATLAS low-novelty | LeWM high-novelty | ATLAS high-novelty |
| PushT | 48.2% | 54.0% | 42.0% | 47.6% |
| TwoRoom | 74.8% | 87.0% | 41.2% | 61.0% |
| OGBench-Cube | 55.8% | 58.4% | 40.4% | 43.0% |
The biggest gain is on TwoRoom's high-novelty subset. DINO-WM, built on a frozen DINOv2 encoder pretrained on 142M external images, reaches 97.8%/98.2% on TwoRoom, well ahead of ATLAS, but the paper flags this as an unmatched comparison, not an apples-to-apples one. Ablations show WEMReg and OOD-recovery each help on their own; combining them gives the best mean across tasks but only a small edge over the stronger single component. Diagnostics show the planning-latent novelty AUROC recovering from 0.44 to 0.75, and lower multi-step prediction error at every rollout horizon.
The finding worth taking away: keeping a latent distribution Gaussian does not preserve the state relationships a planner relies on. These are separate constraints that need separate objectives. The k-NN AUROC probe used here is a reusable diagnostic for anyone training JEPA-style representations, a cheap way to check whether a regularizer is quietly erasing useful structure at some layer. ATLAS itself only changes the training objective, not the backbone or the data, which makes it easy to try. The gains are modest on PushT and OGBench-Cube and substantial only on the TwoRoom navigation task, so this reads as a regularizer worth testing rather than a guaranteed win.
The paper states the two regularizers' benefits are not additive: they address overlapping representation bottlenecks, and the full ATLAS objective beats the stronger single-component ablation by only a small margin. The theoretical planning-stability bound is labeled "conservative" in the paper's own figure caption, and no evaluated case satisfied the margin condition needed for the bound to translate cleanly into guaranteed action selection. The representation and rollout diagnostics are explicitly called "mechanism-oriented diagnostics rather than causal explanations of planning success." DINO-WM's large lead on TwoRoom, driven by external pretraining, shows how much headroom a better encoder leaves on the table regardless of the training objective. All evaluation is on three small-scale simulated tasks; nothing here speaks to real robots or higher-dimensional observations.