异步视界自适应的世界-动作模型:带观测引导上下文路由的AHA-WAM AHA-WAM:Asynchronous Horizon-Adaptive World-Action Modeling with Observation-Guided Context Routing
解耦世界规划与动作执行,异步双系统机器人策略提升闭环控制效率
前置知识
世界-动作模型(World-Action Model)
一种将视觉场景动力学预测与动作生成联合建模的机器人学习范式。不同于传统仅从演示中学习动作标签的方法,WAM通过学习动作与视觉场景演化的共同演化,将物理先验注入到控制策略中,提供更泛化和可迁移的策略表示。核心思想是让机器人理解当前场景以及场景在机器人动作下的演化过程。
本文直接建立在这一范式之上,但重新组织了世界预测和动作执行的时序关系,是理解AHA-WAM异步设计的基础。
扩散模型与Flow Matching
一种生成式建模方法,通过逐步去噪从高斯噪声中恢复目标数据。Flow Matching是扩散模型的一种高效变体,通过学习从噪声分布到数据分布的速度场来生成样本。给定目标变量y、噪声epsilon服从正态分布N(0, I)和流时间rho属于(0, 1),构造插值样本y_rho = (1-rho)y + rho*epsilon,模型预测速度场f_theta(y_rho, rho),最小化损失L_FM(y) = E[||f_theta(y_rho, rho) - (epsilon - y)||_2^2]。
AHA-WAM使用flow matching联合训练动作生成和视频动力学学习,这是模型训练的核心损失函数。
键值记忆与注意力机制
Transformer架构中的核心组件,通过查询、键、值三元组实现信息检索和聚合。给定查询Q、键K和值V,注意力计算为Attention(Q, K, V) = softmax(QK^T/sqrt(d))V。键值记忆存储历史状态的(K, V)对,新查询可以通过注意力机制检索相关信息。在AHA-WAM中,滚动KV记忆用于存储历史视频状态,跨规划器刷新周期维持长时间视野。
OVCR机制通过查询规划器的KV记忆来获取上下文,滚动KV记忆连接历史观察与未来规划,这是异步执行的关键技术基础。
研究动机
现有的世界-动作模型存在一个根本性的设计缺陷:它们将世界预测和动作执行耦合在相同的时间分辨率上。这种强制同步意味着世界分支必须对相邻帧的密集变化进行建模,而这些帧变化通常是高度相关的,对控制来说只有微弱的信息价值。例如,在机器人操作任务中,每30毫秒捕获的连续帧之间往往只有微小的光照变化或相机噪声,世界模型花费大量容量去建模这些冗余信息,却难以捕捉任务相关的长期场景演化。实验数据显示,Fast-WAM在RTX 5090D GPU上的推理延迟高达190毫秒,闭环控制频率仅为5.26 Hz,这对于需要快速物理响应的实时机器人操作来说是严重瓶颈。这种时序耦合反映了抽象层次的结构性错配:视频世界模型应该形成对未来视觉状态的时序扩展潜在规划,而动作模型应该与实时控制循环紧密耦合,及时发出闭环修正。
本文的目标是本文的目标是重新组织世界-动作建模以解决时序耦合问题,提出一个异步视界自适应的世界-动作模型。具体来说,要将视频扩散Transformer实例化为低频世界规划器,它对过去的观察维护滚动键值记忆,并暴露可复用的层级潜在上下文编码长期场景演化;同时让高频动作DiT通过层级联合注意力查询这个上下文来执行短动作块,实现闭环控制。目标是在不牺牲操作性能的前提下,显著提升闭环控制频率,从Fast-WAM的5.26 Hz提升到24 Hz以上,甚至通过蒸馏达到56.95 Hz。
与已有工作不同的是,本文的独特切入角度是从时间不对称性出发,重新审视世界预测和动作执行的功能角色。现有方法假设世界预测和动作执行共享相同的短视界,而AHA-WAM认为两者应该服务于不同的时间尺度:世界模型应该成为慢思考的长期规划器,而动作模型应该是快反应的实时执行器。这种异步设计让视频分支可以专注于捕捉长期场景演化,而不受短周期动作更新的干扰。更关键的是,通过Observation-Guided Video-Context Routing(OVCR),每个动作预测都能获得观察条件的规划器状态视图,而无需将密集视觉token送入高频动作分支。这与之前的双系统方法有本质区别——RoboDual使用VLA提供粗略动作指导,AsyncVLA运行大模型提供延迟语义指导,而AHA-WAM在联合世界-动作建模框架内实现慢快分割,通过层级联合注意力保持耦合。
核心方法
AHA-WAM的整体思路是将世界-动作建模重新组织为两个时间尺度的生成问题。直观来说,就像人类在复杂操作中既有长期规划又有快速反应一样,模型分为两个协同工作的子系统:慢速视频规划器和快速动作执行器。视频规划器像一个战略家,它以低频率运行,维护场景的长期记忆和演化计划;动作执行器像一个战术家,它以高频率运行,根据最新观察快速生成动作。技术路线上,AHA-WAM基于双扩散Transformer架构实现这种异步设计:视频DiT作为世界规划器,对未来视频潜在变量进行建模,产生可复用的层级规划器上下文;动作DiT作为执行器,接收噪声动作token和本体感觉token,在闭环机器人状态反馈下去噪动作块。两个DiT通过层级联合注意力耦合,动作DiT的查询同时关注动作自身的键值和规划器上下文的键值。这种设计保留了联合世界-动作建模的好处,同时让每个分支在它最有信息量的时间尺度上运行。
核心创新点是异步视界自适应的规划-执行解耦,通过三个互补机制实现。第一是Observation-Guided Video-Context Routing(OVCR),它使用最新视觉观察将共享规划器上下文转换为块特定的上下文,而不是简单地将规划器上下文静态缓存。OVCR构造观察引导的路由查询,其中B是可学习基础查询,通过注意力池化从当前视觉token获得查询向量,然后预测残差键值更新,通过门控残差更新得到适应后的上下文。第二是视界自适应偏移训练,通过随机采样偏移delta让动作执行器学习在不同规划-执行相位关系下消费规划器上下文,暴露模型到异步流式部署引起的相位偏移。第三是滚动KV记忆,在视频规划器内部维护固定大小FIFO记忆,存储来自最近规划器刷新的历史视频状态,扩展规划器的时序感受野。这些机制共同解决了异步执行引入的上下文对齐问题。
方法步骤详情
AHA-WAM的完整方法步骤包括训练和推理两个阶段。在训练阶段,首先使用flow matching损失联合训练动作生成和视频动力学学习。给定目标变量y(动作块At或未来视频潜在Z_v:t:t+hv)、高斯噪声epsilon服从标准正态分布和流时间rho属于(0, 1),构造插值样本y_rho = (1-rho)y + rho*epsilon,模型预测速度场,损失为L_FM(y) = E[||f_theta(y_rho, rho, O_v^t, s_t, l) - (epsilon - y)||_2^2]。最终损失为L = L_a + lambda * L_v,其中L_a = L_FM(At),L_v = L_FM(Z_v:t:t+hv)。训练时视频分支使用完全因果掩码学习时序动力学,动作分支被掩码不能关注未来视频token。关键是引入视界自适应偏移,在每个训练段开始时采样偏移delta,将动作块网格在视频规划视界内移动delta。在推理阶段,采用流式调度:视频规划器和动作执行器作为两个非阻塞流运行。视频DiT异步执行prefill生成规划器上下文,不在每次动作更新的关键路径上。每个动作更新时,首先使用OVCR将最新观察转换为路由查询,查询并更新规划器的层级视频上下文,得到块特定的上下文,然后动作DiT通过层级联合注意力消费这个上下文去噪当前动作块。本体感觉反馈直接进入动作DiT,视觉反馈间接通过路由的视频上下文注入。CUDA优化将动作阶段计算编译为静态部署路径,通过TensorRT和CUDA图捕获加速。ODE蒸馏进一步将采样从10步减少到2步,构建AHA-WAM-Flash变体。
技术新颖性
AHA-WAM的技术新颖性体现在三个层面。在架构层面,它是第一个在联合世界-动作建模框架内实现异步执行的工作。与RoboDual使用VLA提供粗略指导、AsyncVLA运行大模型提供延迟语义指导不同,AHA-WAM在双DiT架构内实现慢快分割,通过层级联合注意力保持耦合,而不是让慢分支仅提供指导或反馈。在训练层面,视界自适应偏移训练是一种新颖的训练策略,它通过随机移动动作块网格暴露模型到异步部署中的规划-执行相位偏移,这与传统的固定对齐训练有本质区别。在推理层面,OVCR是一种创新的上下文路由机制,它将规划器上下文重用从静态缓存转变为观察条件的检索和适应,让视频DiT保持在每次更新关键路径之外,同时每个动作块仍然获得与最新视觉证据对齐的规划器表示。这种设计让世界-动作模型首次能够在不牺牲控制频率的情况下享受视频动力学建模的好处。
实验结果
在RoboTwin 2.0模拟环境中,AHA-WAM在50个双臂操作任务上取得了令人印象深刻的性能。在clean设置下成功率为93.40%,在随机设置下为92.20%,平均成功率达到92.80%,超过了所有强基线。特别是与Fast-WAM相比,AHA-WAM在无需机器人数据预训练的情况下提高了0.97个百分点,表明异步设计不仅提高了闭环控制频率,还保持并改进了性能。与需要大规模机器人数据预训练的最强基线LingBot-VA相比,AHA-WAM仍然超过了0.60个百分点。AHA-WAM-Flash变体保持了90.20%的平均成功率,仅带来轻微的性能下降。消融实验揭示了每个组件的重要性:Naive-Async(简单解耦但无滚动KV记忆或OVCR)从91.83%下降到88.60%,表明仅靠更快的动作更新无法补偿过时或相位错位的规划器上下文;添加滚动KV记忆恢复到91.01%,显示持久规划器状态有助于稳定重用的视频上下文;OVCR进一步提升到91.47%,表明观察条件的上下文适应是异步规划-执行错位的更直接 remedy;完整的AHA-WAM达到92.80%,证明记忆和路由是互补的。真实世界机器人实验在AgileX Piper双臂平台上评估了四个任务:折叠毛巾、整理桌面、准备豆浆和存放盘子。在原始设置下,AHA-WAM达到78.33%的成功率,明显超过WAM基线Motus(21.67%)和Fast-WAM(68.33%),与强泛化VLA基线pi0.5(76.67%)相当。在泛化偏移下,pi0.5获得最高成功率,而AHA-WAM保持第二并达到最高进度分数35.00,表明尽管不依赖pi0.5规模的泛化预训练,AHA-WAM获得了相当的真实世界鲁棒性。延迟分析显示,AHA-WAM将延迟从Fast-WAM的190.00毫秒降低到41.37毫秒,闭环频率从5.26 Hz提升到24.17 Hz。通过蒸馏的AHA-WAM-Flash进一步达到17.56毫秒和56.95 Hz,对应10.82倍加速。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| RoboTwin 2.0 (50 tasks, clean setting) | Success Rate (%) | 93.40 | LingBot-VA 92.90 | +0.50 |
| RoboTwin 2.0 (50 tasks, randomized setting) | Success Rate (%) | 92.20 | LingBot-VA 91.50 | +0.70 |
| RoboTwin 2.0 (50 tasks, average) | Success Rate (%) | 92.80 | Fast-WAM 91.83 | +0.97 |
| Real-world tasks (4 tasks, original setting) | Success Rate (%) | 78.33 | pi0.5 76.67 | +1.66 |
| Real-world tasks (generalization shifts) | Progress Score | 35.00 | pi0.5 33.25 | +1.75 |
| Inference latency (single RTX 5090D) | Latency (ms) | 41.37 | Fast-WAM 190.00 | -78.22% |
| Closed-loop control frequency | Frequency (Hz) | 24.17 | Fast-WAM 5.26 | +4.59x |
局限与改进
作者在论文中明确指出了几个局限性。规划器更新频率、视频视界和动作块大小引入了时间超参数,它们的最优分配可能依赖于任务动力学和具身实现。这意味着在不同任务或机器人平台上可能需要调整这些超参数才能获得最佳性能。此外,当前的滚动KV记忆窗口大小固定为6个历史观察帧,这可能不足以支撑非常长期的任务。模型总参数量约为7.23B(4.99B视频规划器 + 1.02B动作DiT + 1.22B记忆和路由模块),这比一些轻量级策略模型更大,部署资源要求更高。真实世界实验显示AHA-WAM在泛化偏移下的性能仍然落后于pi0.5,表明大规模预训练带来的泛化能力仍有价值。此外,论文没有评估在极端动态环境或高干扰环境下的鲁棒性,这些场景可能对异步设计的上下文对齐提出更大挑战。
独立分析的弱点
从独立分析的角度,AHA-WAM存在几个可以改进的弱点。首先,当前的时间超参数(视频视界hv=64、动作块视界ha=16、滚动记忆窗口大小6)是固定的,但不同任务的最佳配置可能差异很大。例如,需要精细操作的任务可能需要更小的动作块和更高的控制频率,而长序列规划任务可能需要更长的视频视界和更大的记忆窗口。改进方向可以是根据任务复杂度在线调整这些超参数,或者学习自适应的时间配置。其次,OVCR机制使用固定数量的路由查询(32个),这可能不足以捕捉所有任务相关的视觉信息。可以引入动态查询数量或层次化查询结构来提高表达能力。第三,真实世界泛化实验显示AHA-WAM在分布偏移下的鲁棒性仍有提升空间,特别是在纹理泛化和新环境测试中性能下降明显。可以通过更多样化的预训练数据、域适应技术或在线学习来改进。第四,当前模型没有显式建模不确定性,在观察与规划器上下文严重错位时可能产生不可靠的动作。引入贝叶斯框架或不确定性估计可以提高安全性和可靠性。
未来方向
作者提出的未来工作方向包括加强慢规划器,使其具有更长期的预测和更丰富的场景表示,并在专门的长视界基准上进行系统评估以更好地量化异步设计如何随任务复杂性扩展。基于AHA-WAM的成果,可以延伸出多个研究方向。一是探索更丰富的世界表示,例如引入物理引擎、3D场景图或对象中心表示,让规划器不仅能预测视觉演化,还能预测物理属性变化。二是研究层级异步架构,例如添加中间时间尺度的战术规划器,形成战略-战术-反应三级系统。三是将异步设计扩展到其他模态,例如触觉反馈或力控制,实现多模态异步融合。四是探索在线学习和适应,让规划器能够根据执行反馈动态调整其表示和预测策略。五是研究更高效的蒸馏和量化方法,进一步降低部署资源要求,使AHA-WAM能够在边缘设备上运行。六是探索在多智能体协作场景中的应用,研究异步世界-动作建模如何支持团队协调和通信。
复现评估
论文提供了项目页面(https://serene-sivy.github.io/aha-wam/),但未明确说明是否开源代码和模型权重。从实现细节来看,复现需要一定的资源:单个NVIDIA RTX 5090D GPU用于延迟测量,训练可能需要更多算力。模型基于预训练的Wan2.2-5B视频模型初始化视频规划器分支,动作DiT采用紧凑架构(隐藏维度da=1024,约1.02B参数)。训练使用AdamW优化器,学习率1e-4,权重衰减0.01,余弦调度,混合精度和梯度裁剪。RoboTwin 2.0评估的设置是每个任务50个clean演示和500个randomized演示,全局批次大小512,每个任务测试100次试验。真实世界实验每个任务收集约120个episodes。CUDA优化包括通过TensorRT和CUDA图捕获编译重复的动作阶段计算,移除冗余循环不变计算和重复缓冲区复制,将10步动作推理延迟从415.77毫秒降低到41.37毫秒。ODE蒸馏将采样步骤从10步减少到2步,进一步将延迟降低到17.56毫秒。这些优化是工程密集的,复现需要深厚的CUDA编程经验。总体来说,虽然论文提供了详细的实现细节和超参数设置,但缺乏开源代码和模型权重会增加复现难度。
论文图表