AHA-WAM：观察引导上下文路由的异步范围-自适应的世界-动作建模

26年6月来自上海交大、上海AI实验室、百度AI和港大的论文“AHA-WAM: Asynchronous Horizon-Adaptive World-Action Modeling with Observation-Guided Context Routing”。世界动作模型已成为机器人操作的一个有前景的范式它将视觉场景动态和动作联合建模以将物理先验注入策略学习中。然而现有的世界动作模型在相同的时间分辨率下耦合世界预测和动作执行迫使世界分支去建模那些冗余且信息量弱的近期帧变化。将世界预测和动作执行严格绑定在相同的时间节奏上可能会未充分利用视频分支在实体控制中的潜力。提出 AHA-WAM一种基于双扩散TransformerDiT架构的异步范围-自适应世界-动作模型它围绕这种时间非对称性重新组织世界动作建模。AHA-WAM将视频DiT实例化为低频世界规划器维护对过去观察的滚动KV记忆并暴露可重用的分层潜上下文以编码长时段场景演变而高频动作DiT则通过分层联合注意在闭环中执行短动作块以查询该上下文。为了支持异步执行引入范围-自适应偏移训练和观察引导的视频上下文路由OVCR这两者共同使动作专家能够利用长范围世界上下文同时对实时执行状态保持响应无需重新运行视频DiT。AHA-WAM概述如图1所示通用机器人操作策略。通用机器人策略旨在从大规模多样化的演示中学习广泛适用的操作技能[1-4]。RT-1[5]和RT-2[6]通过吸收异构机器人轨迹并将网络规模的视觉语言知识迁移到机器人控制中建立了视觉-语言-动作VLA范式。后续工作[7-15]通过将流匹配和基于扩散的动作头[16,17]与大型预训练骨干网络相结合进一步提高泛化能力建立了扩散TransformerDiT[18]作为表达能力强且可扩展的策略架构。然而这些方法仍然主要以动作为中心物理场景动态仅通过演示隐式地捕获这促使了AHA-WAM扩展的世界动作建模范式的出现。用于机器人控制的世界模型。世界模型[19-22]通过预测性视觉动力学来增强机器人策略超越了单纯的动作模仿。现有的基于视频的机器人世界模型大致可以分为两类。第一类遵循“先想象后行动”的范式例如UniPi[23]、Seer[24]和视频预测策略[25]等方法首先预测未来的视觉状态或预测性视觉表征然后通过逆动力学或基于预测未来的策略来恢复动作这在闭环部署中引入了延迟。第二类方法执行联合世界-动作建模其中未来的视觉动力学和动作在共享的生成架构中建模[26-31]。但所有这些方法都将世界预测和动作执行耦合在相同的时间分辨率下将资源消耗在对控制而言冗余的短时域相邻帧上。AHA-WAM通过将视频DiT作为低频长时域规划器与动作DiT作为高频闭环执行器来解决这个问题。双-系统和异步机器人策略。近期的机器人策略探索双-系统或异步设计旨在将深思熟虑的计算与反应式控制相结合。RoboDual [32] 使用 VLA 通用机器人为高效的扩散专家提供任务理解和粗略的动作指导而 Reactive Diffusion Policy [33] 则将缓慢的潜扩散策略与快速的触觉反馈路径相结合以实现丰富的接触操作。AsyncVLA [34] 异步运行一个大型基础模型以提供延迟的语义指导并通过轻量级边缘适配器对其进行细化从而实现快速执行。虽然这些方法验证双-系统执行的有效性但它们的慢速分支主要为动作策略提供指导、反馈或表征。AHA-WAM 则在联合世界-动作建模中实现慢速-快速分离视频 DiT 和动作 DiT 异步运行但通过逐层联合注意机制保持耦合。AHA-WAM世界动作模型将动作预测与学习到的视觉动态相结合但现有模型通常将视频分支和动作分支组织在相同的短时域执行节奏下。其将世界动作建模构建为一个双时间尺度生成问题。图2展示了整体架构。AHA-WAM将WAM推理重组为一个异步世界动作耦合框架同时保留双DiT架构低频视频规划器生成可重用的长时域规划器上下文而高频动作专家则利用该上下文进行闭环动作去噪。这里“时域自适应”指的是在解耦的时间尺度下对任意规划器-执行器相位偏移的鲁棒性而不是在线时域选择。换句话说AHA-WAM在测试时不会动态选择视频或动作时域相反它训练执行器在可变动作起始偏移下利用长时域规划器上下文从而使异步接口能够应对流式部署引起的相位错位。该设计由三种互补机制支持。观察引导的视频上下文路由OVCR无需重新运行视频DiT即可将缓存的规划器上下文调整到最新观察结果。时域自适应偏移训练使动作专家能够感知异步流式传输引起的规划器-执行器相位偏移而滚动KV存储则扩展了视频规划器的时间感受野使其涵盖过去的观察结果。最后实时推理优化进一步加速了高频动作流从而实现闭环部署。1双DiT规划器-执行器架构模型架构。AHA-WAM 基于双 DiT 架构构建该架构包含一个视频 DiT 世界规划器和一个动作 DiT 执行器。视觉观测由预训练的 VAE 进行编码而来自文本编码器的语言嵌入则为两个分支提供条件。动作分支是一个轻量级的动作专家 DiT其层深度与视频 DiT 相同从而实现两个分支之间的逐层交互。视频 DiT 以视觉潜tokens作为输入并经过训练以预测更长规划周期内的未来视频潜tokens。动作 DiT 接收带噪声的动作tokens和本体感受 tokens并在闭环机器人状态反馈下对动作块进行去噪。高频动作分支的视觉反馈并非简单地拼接密集视觉tokens相反视觉信息通过规划器的视频上下文进行中介并由 OVCR 进行调整。逐层规划器-执行器耦合。两个 DiT 之间的核心接口是逐层规划器的视频上下文。此上下文是一种潜在的世界规划表示由一个视频扩散Transformervideo-DiT前向传播暴露出来并被多个后续的动作扩散Transformeraction-DiT前向传播重用。它不同于滚动KV记忆后者在视频规划器内部存储跨规划器刷新的历史视频状态。在训练过程中视频分支在完全因果的视频掩码下预测未来的视频潜元素从而鼓励视频动态信息树学习前向场景动态同时利用视觉动态监督来塑造规划器上下文。动作分支被屏蔽使其无法关注未来的视频元素因此在推理过程中可以移除未来视频预测路径。由此未来视频预测作为世界建模的训练信号而规划器视频上下文则作为推理时与动作执行器的接口。这个耦合保留视觉动态和动作生成之间的 WAM 式交互同时将昂贵的视频 DiT 计算分摊到多个高频动作更新中。联合世界-动作训练目标。AHA-WAM 的训练目标是结合动作片段和未来视频潜表示实现动作流匹配。在部署时AHA-WAM 取消显式的未来帧解码视频 DiT 仅刷新规划器视频上下文动作 DiT 使用此上下文生成闭环动作块。2 基于观察的视频上下文路由异步执行会为多个动作块复用同一个规划器视频上下文这虽然分摊了视频DiT的计算成本但也带来了一个上下文对齐问题在下一次规划器刷新之前机器人状态和视觉场景可能已经发生变化。基于观察的视频上下文路由OVCR通过使用最新的视觉观察结果将共享的规划器上下文转换为特定于动作块的上下文从而解决了这个问题。OVCR 不是将密集的视觉tokens输入到高频动作 DiT 中而是将它们压缩成一组路由查询这些查询在动作去噪之前选择并编辑相关的规划器上下文。对于时间 t 的每个动作块将对齐的观察上下文拆分为视觉tokens和本体感觉tokens。本体感觉输入紧凑且与瞬时机器人状态直接相关因此轻量级编码器将其映射到一个状态token其直接提供给动作 DiT。视觉反馈通过上下文路由间接注入。给定 Q 个可学习的基础查询 BOVCR 通过对当前视觉tokens进行注意池化来构建观察引导的路由查询。由此产生的查询提供紧凑的、受观察条件约束的槽位用于从规划器视频上下文中检索与块相关的信息。OVCR 将规划器上下文的重用从静态缓存转变为基于观察的检索和自适应。因此视频 DiT 保持在每次更新的关键路径之外而每个动作块仍然会收到与最新视觉证据一致的规划器表示。3. 范围-自适应偏移训练异步流会改变慢速规划器和快速执行器之间的相对时间相位。如果训练始终使用视频规划窗口和动作块之间的固定对齐方式则动作 DiT 可能过拟合于单个规划器-执行器相位并在部署期间的中间阶段重用规划器上下文时变得脆弱。因此引入范围-自适应偏移训练使执行器能够应对异步推理引起的相位偏移。令 h_v 表示视频规划视域h_a 表示动作块视域其中 h_a h_v。如图 3 所示对于从时间 τ 开始的每个训练片段视频规划器对 [τ, τ h_v) 上的未来视频潜表示进行建模。由于规划器-执行器对齐方式随动作块大小周期性变化因此采样 δ ∈ [0, h_a) 可以覆盖在多个执行器更新中重用同一规划器上下文时遇到的所有块级相位。因此范围-自适应偏移训练教会动作 DiT 在可变动作开始偏移下利用长范围规划器上下文而 OVCR 在每个阶段处理观察条件上下文选择。4. 滚动式规划器内存由于视频 DiT 作为低频规划器运行它应该在规划器刷新期间保留历史场景信息而不仅仅依赖于当前观测值。这对于长范围操作至关重要因为已完成的子目标、位移的物体或先前观测的状态可以提供重要的上下文信息。因此在视频规划器内部维护一个固定大小的 FIFO 滚动式 K/V 内存。该内存位于慢速视频规划器内部不会被动作 DiT 直接使用。它会在生成下一个规划器上下文之前扩展规划器的时间感受野而高频执行器仍然只与最新的 OVCR 自适应规划器视频上下文交互。5 流式推理和实时优化异步规划器-执行器架构自然而然地导致了流式streaming推理调度。在部署时视频规划器和动作执行器作为两个非阻塞流执行其中 AHA-WAM 将视频 DiT 从每次更新的关键路径中移除但动作分支仍必须以闭环控制频率运行。因此优化动作块推理路径测量一次动作更新的端到端延迟 L_chunk其中包括图像编码、规划器上下文访问、OVCR 路由和动作去噪。视频 DiT 预填充是异步执行的因此 L_chunk 直接决定了动作更新频率。CUDA 加速。通过尽可能使用 TensorRT 和 CUDA 图捕获来执行动作 DiT、内存和上下文模块以及 VAE 编码器从而将重复的动作阶段计算编译成一个静态部署路径同时从去噪热路径中移除冗余的循环不变计算和重复的缓冲区复制。这些更改不会改变模型架构、权重或采样过程但将 10 步动作推理的延迟从 PyTorch eager execution 中的 415.77 毫秒降低到 41.37 毫秒。ODE 蒸馏。在 CUDA 加速的 10 步路径之上构建 AHA-WAM-Flash将动作采样器从 10 个去噪步骤蒸馏到 2 个步骤同时保持相同的观测、规划器上下文和 OVCR 接口。在蒸馏过程中视频 DiT 被冻结学生模型通过采样中间状态进行训练以直接预测教师模型的最终动作输出采样偏向于噪声较大的状态以支持积极的步骤缩减。 ODE 蒸馏进一步将 Lchunk 从 41.37 毫秒减少到 17.56 毫秒。实验设置模型与训练。AHA-WAM 采用双 DiT 策略实现并具有显式异步的规划器-执行器接口。用预训练的 Wan2.2-5B 视频模型初始化世界规划分支其中包括视频 DiT、文本编码器和视频 VAE。与 Fast-WAM [30] 类似高频动作执行器采用紧凑的 DiT 架构隐藏维度 d_a 1024对应约 10.2 亿个参数。除了 49.9 亿个参数的视频规划器和 10.2 亿个参数的动作 DiT 之外AHA-WAM 还引入了 12.2 亿个参数用于滚动 K/V 记忆和上下文路由模块使得最终实例化的模型总参数量约为 72.3 亿个。对于 AHA-WAM 的时间配置视频分支在一个较长的规划时域 h_v 64 上运行而动作分支则预测时域 h_a 16 的短可执行块。在异步推理期间视频规划器维护一个可重用的逐层 K/V 上下文该上下文通过最多 6 个历史观测帧的 FIFO 滚动内存进行增强。动作执行器通过 OVCR 调整此上下文每个动作块使用 32 个基于观测的路由查询。训练过程中世界建模和动作预测均采用流匹配噪声时间采用 logit 正态采样。动作生成方面在推理阶段使用 10 个去噪步骤并将 CFG 设置为 1.0。所有模型均使用 AdamW 进行训练学习率为 1 × 10⁻⁴权重衰减为 0.01采用余弦调度、混合精度和梯度裁剪。延迟是在单个 NVIDIA RTX 5090D GPU 上报告的。评估范围。我们在 RoboTwin 2.0 仿真 [35] 和真实世界实验中评估 AHA-WAM。仿真基准测试在干净场景和随机场景下进行多任务操作而真实世界实验测试在物理双臂任务上的部署。在两种设置下都将任务成功率作为主要指标并进一步分析闭环推理延迟以量化部署效率。RoboTwin 仿真在 RoboTwin 2.0 [35] 上评估 AHA-WAMRoboTwin 2.0 是一个双臂操作基准测试包含 50 个涵盖各种技能的双臂任务。遵循 [27, 29, 30] 中的多任务训练设置对于每个任务训练集包含 50 个干净场景下的演示和 500 个随机场景下的演示。模型采用多任务方式训练全局批大小为 512。每种方法在干净设置和随机设置下每个任务均进行 100 次试验报告任务平均成功率。将模型与具有竞争力的 VLA 和 WAM 基线模型进行比较包括 π0 [2]、π0.5 [9]、ABot-M0 [36]、Motus [27]、LingBot-VA [29] 和 Fast-WAM [30]。真实世界机器人实验设置。用配备自视RGB摄像头的双手动AgileX Piper平台进行评估策略仅使用头部视角RGB观测值、本体感觉状态和语言指令作为输入。考虑四个任务——折叠毛巾、整理桌面、制作豆浆和存放餐盘——涵盖可变形操作、长时域重排、精细工具使用、接触丰富的控制和空间泛化。对于每个任务平均收集约120个回合。实现。由于Fast-WAM和AHA-WAM默认情况下未在任何机器人数据上进行预训练用选定的RoboCOIN子集[37]对Fast-WAM和AHA-WAM进行预训练该子集包含24,600条轨迹和约165小时的机器人数据以确保部署比较的公平性和稳定性。两个模型均在相同的特定任务演示上进行微调。 Motus 和 Fast-WAM 的部署延迟较大导致动作更新稀疏闭环行为不稳定。采用 RTC 式的非阻塞执行方案 [38] 来部署它们并在预测块之间进行动作插值从而在高推理延迟下实现更平滑的控制。AHA-WAM 则通过其异步规划器-执行器接口进行部署其中执行器利用最新的观测结果和重用的规划器上下文执行高频闭环更新。