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

ATLAS. Aligned Transport of Latent Structure for Reliable World Model Planning. By Ke Fang and two others, from Pennsylvania State University. Published on the thirtieth of September. Its code repository is public but slight: seven source files.

A latent world model plans inside its own representation, so the geometry of that representation matters. A common approach regularizes the latent marginal, the overall distribution of latent points. The authors observe that this alone does not pin down the relationships between one state and another, and those relationships are what action selection uses. They show this gap can let planning-relevant novelty structure weaken as representations are transformed into the final latent the planner works with. In other words, a latent space can look well behaved as a whole while the particular differences that tell one situation from a new one have been worn down on the way to the planner.

Aligned Transport of Latent Structure, or ATLAS, is a training objective that does two jobs at once. It explicitly preserves relational geometry, and it calibrates the global latent distribution. Their analysis shows these are not the same constraint twice: relational preservation and marginal calibration impose non-redundant constraints, so neither one can stand in for the other.

For the relational part, ATLAS takes normalized pairwise structure from an informative encoder representation and transfers it to the planning latent. For the marginal part, it uses Wasserstein embedding matching, called WEMReg, which calibrates the marginal through one-dimensional Wasserstein-2 transport. One component looks after how states relate to each other, the other after how the whole cloud of latents is spread.

The analysis also ties finite-candidate planning stability to three things: relational distortion, a mismatch in latent scale, and prediction error. Two of those concern the geometry of the latent the planner searches in, and the third concerns how well the model predicts what happens next. Planning stability, on this account, depends on the shape of the space as well as on the accuracy of the predictions made inside it.

ATLAS is instantiated in LeWM and tested on PushT, TwoRoom and OGBench-Cube. It improves mean goal-reaching success on both the lower-novelty and the higher-novelty evaluation subsets, and the largest gain is on higher-novelty TwoRoom episodes. Diagnostics of the representations and of rollouts point the same way: stronger novelty-related structure in the planning latent, better marginal calibration, and lower multi-step prediction error.

Getting the overall spread of a latent space right is one requirement. Keeping the relationships between particular states intact is another, and the planner leans on the second. The authors conclude that preserving planning-relevant latent geometry is an important ingredient for reliable world-model planning.