AHA-WAM: Asynchronous Horizon-Adaptive World-Action Modeling with Observation-Guided Context Routing

AHA-WAM 把世界—动作模型拆成两个异步运行的时间尺度：较重的视频 DiT 利用视觉历史维护长期上下文，较轻的动作 DiT 根据最新图像、机器人状态和规划上下文生成短动作块。每次动作更新先通过 OVCR 读取并残差修正规划器的逐层 K/V，再完成动作去噪；训练时联合学习视频动态与动作，并随机移动动作块起点，部署时移除未来视频生成路径，通过 CUDA 优化及可选蒸馏降低闭环延迟。

关键是让缓存保持可校正，而不要求昂贵的视频网络参与每一次动作更新。最新视觉先被压缩成少量查询，这些查询从旧规划上下文检索相关信息，预测带门控的 K/V 残差，动作网络逐层读取修正后的上下文。滚动记忆负责保存过去，错位训练负责适应规划与动作时钟不同步，OVCR 负责对齐当前场景；三者处理不同问题。未来视频只提供训练监督，动作分支训练时看不到未来视频 token，因此推理时可以删掉未来生成部分。

• 慢规划器：用 Wan2.2-5B 初始化视频分支，以因果视频监督学习动态；输出每层 K/V，上下文跨多个动作更新复用，历史则写入固定大小 FIFO 记忆。
• 快反馈与 OVCR：把最新视觉压成 32 个观测查询，逐层读取旧 K/V，通过门控残差更新生成当前动作块专用上下文；本体状态另经状态编码器直接输入动作网络。
• 动作生成：动作 DiT 的查询联合读取自身动作 token 和已适配的规划 K/V，在语言、本体状态与视觉上下文条件下去噪出 16 步动作块。
• 时序训练：视频规划跨度与短动作块解耦，在训练段中采样 δ∈{0,…,15}，平移动作块网格，让同一规划上下文可服务不同起点的执行。
• 部署优化：两条非阻塞计算流并行；快分支用 TensorRT、CUDA Graph 和循环不变量外提压缩延迟，Flash 进一步把动作采样减为 2 步。
• 沿图 1 看全局：历史、规划上下文与现场反馈： 图 1 阅读提示：过去的观测给慢规划器提供背景；规划器输出面向后续执行的隐式上下文；动作分支持续根据新观测调整执行。这张概念图中的“未来规划”不意味着部署时要把未来像素逐帧解码出来。本文的世界模型主要通过训练中的视频预测塑造表示，再让动作分支使用该表示。
• 沿图 2 从下往上：每条箭头到底传什么： 图 2 阅读提示：左侧蓝色是慢视频分支，中间绿色是快动作分支，横跨两者的色带是上下文路由接口；右侧矩阵解释训练与推理时哪些信息可以互相读取。圈叉表示交叉注意力，圈加表示残差更新，斜纹块表示更新后的 K/V。
• ① 左下：图像进入 VAE，历史进入滚动记忆。 视觉观测先编码为 latent token。视频 DiT 的每一层在因果约束下处理视频 token，结合保留的历史信息形成新的规划上下文。历史不是把整段视频反复拼回快分支，而是由慢分支维护 FIFO K/V 记忆。实验设置最多保留 6 个历史观测帧，避免记忆容量无限增长。
• ② 左中：规划器导出逐层 K/V。 设最新一次慢规划刷新发生在 τ，接口为 Cᵖ_τ={(Kᵖ,ℓ_τ,Vᵖ,ℓ_τ)}，ℓ 是 Transformer 层号。动作网络与视频网络具有相同层数，便于逐层连接；它们不需要相同宽度或参数量。论文实例化的视频 DiT 约 4.99B，动作 DiT 约 1.02B，记忆与路由模块另有约 1.22B，总计 7.23B。所谓轻量执行器是相对于视频骨干而言。
• ③ 中下：最新图像与本体状态分流。 当前观测经视觉编码/投影形成 Xᵛ_t；机器人状态 s_t 经轻量状态编码器成为状态 token，直接进入动作分支。视觉不以大量 dense token 直接拼进动作网络，而是先生成观测查询。语言编码同时为两条分支提供条件，图中未把所有条件连线都展开。
• ④ 中间第一个圈叉：用 32 个查询压缩当前视觉。 令 B 为可学习的基础查询，得到 Zᑫ_t=Attn(B,fᵥ(Xᵛ_t),fᵥ(Xᵛ_t))。可以理解为让少量槽位从当前画面提取与控制有关的信息；虽然 B 是学习到的参数，输出 Zᑫ_t 每次都随最新图像改变，并非固定查询内容。每个 chunk 使用 32 个槽位，控制接口的 token 数量。
• ⑤ 横向色带：先从旧计划读取，再做残差编辑。 在每一层，以 Zᑫ_t 查询规划 K/V，得到 R^ℓ_t=Attn(Zᑫ,ℓ_t,Kᵖ,ℓ_τ,Vᵖ,ℓ_τ)。轻量路由器 g^ℓ_ψ 结合 R 与 Zᑫ 预测 ΔK、ΔV；通过学习到的门 α^ℓ_t 得到 K̃=K+αΔK、Ṽ=V+αΔV。图里的旧缓存到圈加的旁路表示保留基底，圈叉到圈加的支路表示当前观测驱动的修改。OVCR 因此包含检索和适配两个动作，而不仅是给缓存做一个只读 attention。
• ⑥ 右侧绿色 Joint Attention：动作查询读取校正后的上下文。 对动作分支第 ℓ 层，核心计算是 Attn(Qᵃ,[Kᵃ;K̃ᵖ],[Vᵃ;Ṽᵖ])：同一个动作查询同时看动作 token 与路由后的世界上下文。随后继续经过层内调制、交叉注意力与前馈计算。重复 L 层、完成所需去噪步数后，输出当前短动作块。下一次动作更新拿到新的图像和状态，重新做路由；此时慢规划器未必完成新的刷新。
• ⑦ 不要混淆三个缓存概念。 Rolling K/V Memory 是慢规划器跨刷新保存的历史；Cᵖ_τ 是最近一次刷新导出的、可跨动作更新复用的逐层上下文；C̃ᵖ_t 是为当前 chunk 适配后的上下文。动作执行器直接读取第三项，而不是直接查询完整历史 FIFO。把三者分开，才能理解“历史不丢”和“当前观测有效”如何同时成立。
• 图 2 右侧掩码：为什么训练能看视频，部署却不用生成未来： 训练矩阵的左上角是下三角。 视频 token 按时间因果读取，不能任意看更晚的帧；未来视频预测成为学习场景动态的辅助任务。
• 动作块对应各自的局部色块。 每块动作读取本块更新后的规划上下文，以及块内允许的动作信息；不同动作块的观测 o¹、o² 分别驱动自己的上下文适配。不要把图中的两个动作 token 误读为真实 chunk 只有两步，那只是示意。
• 动作不依赖未来视频 token。 正文明确屏蔽动作对未来视频的访问，避免训练时借助部署时不存在的信息。因而推理矩阵保留可用观测形成的上下文及其适配接口，删除未来视频预测路径仍能生成动作。
• 训练目标是联合 flow matching。 对视频 latent 或动作 y，采样噪声 ε 和噪声时间 ρ，构造 yρ=(1−ρ)y+ρε，网络拟合速度 ε−y；总损失 L=Lₐ+λLᵥ，附录给出联合优化时两项等权。部署不需要视觉解码，但训练中的 Lᵥ 会塑造规划表示。
• 沿图 3 看错位训练：让执行起点不必等待规划时钟： 图 3 阅读提示：上方蓝线表示较粗的世界规划节奏，下方绿线表示更密的动作节奏；左侧 δ 标出两者起点的偏移，括号提示一个较长窗口内承载多个动作更新。图中 t₀、t₄ 等位置是时间关系示意，不能据此代替附录的实际调度超参数。
• 如果训练时动作块总从视频窗口边界开始，执行器容易学会一个固定的“上下文—动作位置”对应关系。真实异步部署中，规划刷新完成的时间、动作请求的到达时间不可能永远整齐对齐；相同规划上下文可能在窗口中段被再次使用。训练时作者采样 δ∼Uniform{0,…,hₐ−1}，把动作块网格整体平移后再计算动作损失，使执行器遇到不同的块级相位。
• 正文设置 hᵥ=64、hₐ=16。举例来说，δ=5 意味着动作块从规划窗口内偏后的位置开始；这是帮助理解的示例，不是固定推理超参数，也不是强迫系统等 5 帧才动。覆盖 0–15 的偏移是为了覆盖以 chunk 长度为周期的相位。OVCR 再用各起点对应的最新观测，使当前动作知道自己处于什么场景。
• 这里存在一个重要区分：horizon 表示模型涉及的时间范围，刷新频率表示规划上下文多久重算一次，执行频率表示动作请求多久更新一次。hᵥ/hₐ=4 不能直接推出部署视频线程每 4 次动作刷新一次；附录另列视频/动作 frequency ratio 为 8，两种量需要分别按实现核实。论文也没有给出能无界容忍陈旧缓存的理论保证。
• 部署时一次完整闭环，以及 Flash 改了什么： 慢分支收到可用视觉观测和历史记忆，后台计算新的 Cᵖ_τ；任务语言作为条件输入。起始阶段需要先建立可用上下文。
• 快分支到达更新时刻，读取最新图像 o_t 和状态 s_t，取最近一次完整可用的规划上下文。本文重点是快分支不必等待这一次慢分支重新运行。
• 图像编码、OVCR 适配、动作去噪构成当前动作块的关键路径；生成短动作后交给执行层，下一次更新再纳入反馈。
• 慢分支完成新结果后，后续动作更新改用新上下文；其内部 FIFO 记忆也随刷新维护。以上是论文机制的执行解释，具体跨流同步、缓冲区切换和动作覆盖策略需查看实现，不能仅凭结构图确定。
• 基础 AHA-WAM 使用 10 步动作去噪。工程优化将 Action DiT、上下文模块和 VAE 编码器编译为 TensorRT 路径，在可行处用 CUDA Graph，移出噪声步循环内不变的条件/位置计算，并避免重复拷贝。论文的同系统累计优化把 L_chunk 从 415.77 ms 降到 41.37 ms，约为原来的十分之一；这不应误写为“仅异步就快了十倍”。
• Flash 冻结视频 DiT，蒸馏动作采样路径，部署用 2 步完成采样。附录描述从教师轨迹选择中间锚点，让学生从这些带噪状态回归最终动作，并更偏重高噪声起点。表 5 给教师 16 步与锚点 {0,1,2,4,8,12,16}，正文则以 10→2 步叙述，因此可以确认 2 步部署和 17.56 ms 结果，但严格复现教师设置前需要澄清正文/附录口径。