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

26年6月来自上海交大、上海AI实验室、百度AI和港大的论文“AHA-WAM: Asynchronous Horizon-Adaptive World-Action Modeling with Observation-Guided Context Routing”。

世界动作模型已成为机器人操作的一个有前景的范式，它将视觉场景动态和动作联合建模，以将物理先验注入策略学习中。然而，现有的世界动作模型在相同的时间分辨率下耦合世界预测和动作执行，迫使世界分支去建模那些冗余且信息量弱的近期帧变化。将世界预测和动作执行严格绑定在相同的时间节奏上，可能会未充分利用视频分支在实体控制中的潜力。

提出 AHA-WAM，一种基于双扩散Transformer（DiT）架构的异步范围-自适应世界-动作模型，它围绕这种时间非对称性重新组织世界动作建模。AHA-WAM将视频DiT实例化为低频世界规划器，维护对过去观察的滚动KV记忆，并暴露可重用的分层潜上下文以编码长时段场景演变，而高频动作DiT则通过分层联合注意在闭环中执行短动作块，以查询该上下文。为了支持异步执行，引入范围-自适应偏移训练和观察引导的视频上下文路由（OVCR），这两者共同使动作专家能够利用长范围世界上下文，同时对实时执行状态保持响应，无需重新运行视频DiT。

AHA-WAM概述如图1所示：

通用机器人操作策略。通用机器人策略旨在从大规模多样化的演示中学习广泛适用的操作技能[1-4]。RT-1[5]和RT-2[6]通过吸收异构机器人轨迹并将网络规模的视觉语言知识迁移到机器人控制中，建立了视觉-语言-动作（VLA）范式。后续工作[7-15]通过将流匹配和基于扩散的动作头[16,17]与大型预训练骨干网络相结合，进一步提高泛化能力，建立了扩散Transformer（DiT）[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 之间的核心接口是逐层规划器的视频上下文。此上下文是一种潜在的世界规划表示，由一个视频扩散Transformer（video-DiT）前向传播暴露出来，并被多个后续的动作扩散Transformer（action-DiT）前向传播重用。它不同于滚动KV记忆，后者在视频规划器内部存储跨规划器刷新的历史视频状态。

在训练过程中，视频分支在完全因果的视频掩码下预测未来的视频潜元素，从而鼓励视频动态信息树学习前向场景动态，同时利用视觉动态监督来塑造规划器上下文。动作分支被屏蔽，使其无法关注未来的视频元素，因此在推理过程中，可以移除未来视频预测路径。由此，未来视频预测作为世界建模的训练信号，而规划器视频上下文则作为推理时与动作执行器的接口。

这个耦合保留视觉动态和动作生成之间的 WAM 式交互，同时将昂贵的视频 DiT 计算分摊到多个高频动作更新中。

联合世界-动作训练目标。AHA-WAM 的训练目标是结合动作片段和未来视频潜表示，实现动作流匹配。在部署时，AHA-WAM 取消显式的未来帧解码：视频 DiT 仅刷新规划器视频上下文，动作 DiT 使用此上下文生成闭环动作块。

2 基于观察的视频上下文路由

异步执行会为多个动作块复用同一个规划器视频上下文，这虽然分摊了视频DiT的计算成本，但也带来了一个上下文对齐问题：在下一次规划器刷新之前，机器人状态和视觉场景可能已经发生变化。基于观察的视频上下文路由（OVCR）通过使用最新的视觉观察结果，将共享的规划器上下文转换为特定于动作块的上下文，从而解决了这个问题。OVCR 不是将密集的视觉tokens输入到高频动作 DiT 中，而是将它们压缩成一组路由查询，这些查询在动作去噪之前选择并编辑相关的规划器上下文。

对于时间 t 的每个动作块，将对齐的观察上下文拆分为视觉tokens和本体感觉tokens。本体感觉输入紧凑且与瞬时机器人状态直接相关，因此轻量级编码器将其映射到一个状态token，其直接提供给动作 DiT。视觉反馈通过上下文路由间接注入。给定 Q 个可学习的基础查询 B，OVCR 通过对当前视觉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-WAM，RoboTwin 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 则通过其异步规划器-执行器接口进行部署，其中执行器利用最新的观测结果和重用的规划器上下文执行高频闭环更新。