ATLAS:面向可靠世界模型规划的对齐潜在结构传输

TL;DR ATLAS 通过显式保留规划潜在空间的关系几何并校准边缘分布，防止新颖性结构在表征转换中丢失，显著提升世界模型规划在低/高新颖任务上的成功率。

问题

问题背景

Latent world models 在模型预测控制中日益受到关注，其规划质量高度依赖 latent representation 的几何结构。当前研究聚焦于如何让 learning latent 分布更利于 planner 的 action selection。

现有方法局限

许多方法仅对 latent marginal 施加正则，例如通过 KL divergence 或 Wasserstein 距离校准全局分布。然而，marginal calibration 不约束 state-to-state 的 pairwise 关系。当 informative encoder representation 被转换为 planner 使用的 final latent 时，归一化或投影等变换可能削弱 planning-relevant novelty structure，导致高新颖度状态下动作选择退化。论文指出，这种关系几何的失真（relational distortion）与 latent-scale mismatch 直接关联有限候选规划的不稳定性。

为何困难且重要

难点在于需要同时满足两个非冗余约束：保持关系几何与校准边际分布。二者可能冲突，因为分布校准倾向于改变尺度，而关系保持需要保留相对距离。在 higher-novelty evaluation 中，这种冲突更明显；若仅优化一项，规划器在 unseen states 上容易失败。业界对此高度关注，因为可靠 world-model planning 是模型预测控制、机器人操作等应用的核心。

行业类比

类似自动驾驶中，感知网络提取的鸟瞰图特征经过多模态融合与归一化后，若丢失了车辆间相对位置关系的几何一致性，下游 planner 在未见过复杂路口时会产生错误轨迹选择。

核心洞察

• ATLAS 的核心洞察是：可靠规划需要同时约束潜空间的**关系几何** 与 **边缘分布**，且二者是**非冗余约束**。以往方法常用 VAE 或分布匹配正则化潜边缘，但这只约束单个状态的分布，无法决定状态间相对距离和排序，导致 planner 选动作时可能被无关的表征变换误导。ATLAS 通过将 encoder 的规范化成对结构显式迁移到 planning latent，并用一维 Wasserstein-2 匹配校准边缘，从理论和实验证明两者缺一不可。这个区分对工程实践很关键：如果你的 world model 只做了 latent prior/regularization，但规划性能在新奇场景下降，很可能是关系几何未保留。
• 从有限候选规划稳定性的角度，ATLAS 揭示了**关系失真**、**潜尺度不匹配** 和 **预测误差** 如何共同影响规划可靠性。现有工作常把规划失败归因于单步预测误差，或者用 rollout accuracy 评估，但 ATLAS 的分析指出，即使预测误差不大，若潜空间的距离排序被破坏或尺度不一致，有限数量候选动作的 argmax 也可能不稳定。这提供了一种新的诊断思路：通过分析 planning latent 的成对距离是否保留 encoder 中的 novelty 结构，以及边缘分布是否校准，来预测 OOD 下的规划失败。对实际系统来说，这种可诊断性比单纯提高平均成功率更有价值，能指导表征层面的修复。

方法

ATLAS 的训练目标由两个核心模块构成，作用于规划网络输出的 planning latent。输入来自一个信息量较高的 encoder 表征（如原始观测的卷积特征），经过编码器映射到 planning latent 空间。

• WEMReg（Wasserstein embedding matching）：通过一维 Wasserstein-2 运输，将 planning latent 的全局边际分布校准到预设的先验分布（如单位高斯）。它显式控制 latent 的尺度与分布形状，避免 latent 幅值漂移。
• 关系几何保持（relational geometry preservation）：从 encoder 表征中提取归一化的 pairwise 结构（如余弦相似度或欧氏距离的相对顺序），要求 planning latent 尽可能保留这些结构关系。这一约束不依赖于边际分布，单独作用于状态之间的关系。

两个约束分别处理 分布形状 与 相对结构，理论上互相独立且不可替代。训练时通过联合优化将它们加入原有 world model 损失。输出是既符合先验分布、又保留状态间拓扑关系的 planning latent，用于后续基于采样的有限候选规划。

与同类方法的差异点：以往工作只正则化 latent 的边际分布（如 VAE 的 KL 项），忽略了配对关系；ATLAS 首次显式地将关系几何从 encoder 迁移到 planning latent，并通过 Wasserstein 运输同时校准边际。

实验

实验设计

论文在 LeWM 世界模型框架上实例化 ATLAS，并在三个环境 PushT、TwoRoom、OGBench-Cube 上评估目标达成成功率。每个环境都划分了 lower-novelty 和 higher-novelty 评估子集，用于检验模型在分布内与分布外泛化下的规划可靠性。对比基线包括基础 LeWM 以及仅做潜在边际正则化的变体，重点观察关系几何保持和边际校准的独立贡献。

关键发现

• ATLAS 在三个环境的 lower-novelty 和 higher-novelty 子集上均提升平均目标达成成功率，其中 higher-novelty TwoRoom 提升最大。
• 表征诊断显示：规划潜在空间中保留了更强的 novelty 相关结构，边际校准 得到改善，同时 多步预测误差 下降。
• 消融实验表明，关系保持与边际校准是 非冗余约束：仅校准边际不足以维持状态间关系，而仅保持关系可能导致尺度不匹配。

与基线的深度解读

传统世界模型训练常通过正则化潜在边际分布来稳定 Latent，但这不约束状态间相对几何，导致规划时用于动作选择的 pairwise 关系被弱化。ATLAS 显式地将编码器表征中的规范化 pairwise 结构迁移到规划潜在空间，同时用一维 Wasserstein-2 传输校准边际，解决了“全局分布对齐但局部关系失真”的隐患。与其他仅做分布对齐的方法相比，ATLAS 同时覆盖关系几何与尺度校准，因此在更高 novel 的评测子集上优势更明显，说明保留 planning-relevant latent geometry 对可靠规划至关重要。

行业影响

落地场景

ATLAS 针对 latent world model 训练中表示几何与边缘分布失配问题，适用于 model-based RL 的机器人操作与导航任务。典型场景包括：

• 仓储机器人 在货架变化、新环境下的 pick-and-place；
• 自动驾驶 中基于 world model 的 motion planning，需处理未见过路口或天气；
• 家庭服务机器人 在用户家中执行 manipulation 任务时应对 novelty。

商业价值

• 降本：提升 novel 场景下成功率，减少人工重启/干预成本；
• 增收：提高自动化设备在长尾场景的可用性，扩大可部署范围；
• 体验提升：减少因表示退化导致的规划失败，用户端更稳定可靠。

与现有工作流集成

ATLAS 是训练目标层面的正则化，可插入 LeWM、Dreamer、TD-MPC 等现有 latent world model 训练循环，无需改动推理架构。只需在 encoder 和 planning latent 之间增加 relational preservation loss 与 WEMReg，即可提升潜在几何质量。工程团队可在现有 RL 训练脚本中加入该 loss，作为辅助项微调。

局限

• **验证范围有限**：方法仅在单一世界模型架构 LeWM 上实例化，未在 Dreamer、TD-MPC 等主流架构上验证泛化性。实验环境限于 PushT、TwoRoom、OGBench-Cube，缺少高维视觉输入或更复杂连续控制任务。此外，未报告计算开销与超参数敏感性；关系保持需计算成对结构，潜在 O(N²) 成本可能限制在大规模或高维潜在空间中的应用。
• **理论假设与对齐方式的局限**：WEMReg 使用一维 Wasserstein-2 传输校准边缘分布，无法显式对齐高维联合分布，可能遗漏跨维度依赖。理论分析基于有限候选规划和特定假设，若实际候选生成策略与设定不符，稳定性结论可能失效。与对比学习方法（如 CURL、InfoNCE）或直接潜在归一化的关系未充分讨论，存在与已有正则化技术重叠的风险。
• **缺少消融与对比**：实验仅与 LeWM 基线比较，未纳入其他潜在空间正则化方法（如 distribution matching、contrastive regularization）的直接对比，无法证明 ATLAS 优于更简单的替代方案。对新颖性提升的归因可能部分来自额外正则化带来的隐式泛化，而非明确的关系保持机制。作者未提供失败案例分析或对极端新颖性场景的鲁棒性测试，限制了结论的适用范围。