MemoryWAM：基于持久记忆的高效世界动作建模

26年6月来自港中文大学、清华和浙大的论文“MemoryWAM: Efficient World Action Modeling with Persistent Memory”。在现实世界中实现稳健的机器人操作不仅需要理解当前的观测信息还需要具备记忆能力和对环境动态的建模能力。世界动作模型World Action Models, WAMs通过结合当前及历史观测信息来联合建模视觉预测与动作从而具备上述能力成为机器人操作领域极具前景的范式。然而现有的 WAM 面临着一个根本性的权衡难题推理高效的方法通常仅依赖于近期观测的有限窗口难以应对非马尔可夫环境而能够保留长时历史信息的方法其时间和空间成本会随序列长度的增加而急剧上升。为解决这一挑战提出 MemoryWAM——一种具备高效持久化记忆功能的世界动作模型。MemoryWAM 采用一种混合记忆设计融合近期帧、事件边界锚定帧anchor frame以及用于概括长时历史信息的紧凑“要点gist”tokens。通过定制的注意机制该模型能够同时检索详细的短期上下文和压缩的长期上下文从而在降低推理延迟和 GPU 显存占用的前提下支持依赖记忆的决策过程。在仿真环境和现实世界的长时程、依赖记忆的操作任务中MemoryWAM 均优于强劲的视觉-语言-动作VLA模型及其他 WAM 基线方法同时保持了良好的计算效率。概述如图1 所示MemoryWAM 是一种具备高效持久化记忆的世界动作模型World Action Model。该模型采用混合记忆机制结合少量“要点”gisttokens与精心设计的注意机制实现了高效且有效的决策记忆利用。具体而言滑动观测窗口保留用于即时控制的高保真短期上下文信息少量要点tokens用于压缩长程历史信息而锚定帧anchor frames则在事件边界处如任务初始观测时刻保留完整的视觉tokens从而维持关键时刻的记忆显著性。得益于此随着历史帧数量的增加MemoryWAM 既能维持持久化记忆又仅导致推理延迟和 GPU 显存占用的微小增长。此外与以往具备持久化记忆功能的 WAM 相比MemoryWAM 的推理延迟和 GPU 显存占用均显著降低。上图 1 展示了不同方法在任务成功率和推理延迟方面的对比情况。1 概述现有的视觉-语言-动作模型VLA[4, 62] 和世界动作模型WAM[8, 14] 通常将近期的观测序列 o_t-N:t} 与任务指令 l 映射为动作a_t。尽管它们在短时程任务中表现有效但在决策依赖于长程历史信息的非马尔可夫环境中却难以应对。保留长程历史信息的一种直接方法是在自回归 Transformer [7, 11, 17] 中保存所有过往观测的完整 KV 缓存。虽然这种设计允许策略访问完整的历史信息 o_1:t但随着时间步 t 的增加GPU 显存占用和推理延迟会迅速攀升导致其在长时程任务中不切实际。为了克服这些效率瓶颈本文从人类认知中汲取灵感(1) 短时记忆支持当前的动作规划但容量有限 [19](2) 长时记忆倾向于编码抽象的要点痕迹而非逐字的细节信息 [20](3) 事件边界记忆侧重于任务起始时刻的状态 [21]。基于这些认知启示提出一种混合记忆机制它既保留高保真的短时上下文信息又维护少量概括长程历史的“要点”gisttokens同时还保留任务起始时刻的锚定帧。这种类人的混合记忆机制在显著降低推理成本的同时使策略能够对长时程依赖关系进行推理。2 架构MemoryWAM 遵循近期“世界动作模型”WAM中基于视频动作扩散的范式利用预训练的视频扩散 Transformer (DiT) 提供感知动态的视觉表征并由独立的动作 DiT 基于所学习的视觉动态预测未来动作。该模型架构如图 2 所示。给定观测值 o_t为了提高计算效率首先利用因果视频 VAE [63] 将其编码为紧凑的视频潜向量 z_t。该视频潜向量由视频 DiT Phi_v 处理而动作片段 a_t:th-1 则由动作 DiT Phi_a 生成。这两个分支采用 Transformer 混合MoT[64] 架构进行组织。此外MemoryWAM 继承近期高效 WAM [14, 15] 的一项关键优势它在训练阶段通过视频预测学习物理动态同时避免推理阶段高昂的视频生成成本。在推理过程中当前观测的干净潜在向量 z_t 仅需通过一次视频 DiT 前向传播以更新视频侧的键值KV缓存 Cv_t。随后动作 DiT 通过对动作 Token 进行去噪来预测动作块action chunk同时关注已缓存的视频表征。由此可见视频 DiT 负责提取感知动态特性的特征并维护记忆而动作 DiT 则将这些特征映射为动作。3 混合记忆如前所述基于完整历史记录的注意机制会导致 GPU 显存开销和推理延迟急剧增加。MemoryWAM 采用混合记忆设计来解决这一问题该设计灵感源自人类记忆的互补形式用于即时闭环控制的短期记忆、用于检索任务起始状态的事件边界记忆以及用于紧凑表征长程历史的要旨记忆gist memory。具体而言在时间步 tMemory-WAM 维护一个紧凑的时间缓存其中三个组成部分分别对应近期观测、事件边界帧以及压缩后的长期历史信息。短时记忆。短时记忆负责即时的闭环控制通过捕捉物体运动、接触状态和手-物构型等快速变化的交互线索获取近期观测信息。这种记忆被实现为覆盖最近 N 个视频帧的滑动窗口缓存。这既保留了用于动作生成的高保真局部上下文又将短时注意计算成本限制在固定的窗口大小范围内。事件边界记忆。并非所有的历史观测都具有同等的信息价值。在连续的体验过程中事件边界为记忆组织提供了显著信息。在机器人操作任务中这些边界往往对应于任务的启动时刻和初始场景配置。因此在任务开始时保留少量包含完整视觉 Token 的“锚定帧”anchor frames因为初始场景状态往往构成了指令中关键信息的参照基础且该状态在后续过程中可能会被遮挡或超出观测窗口范围。要点记忆Gist memory。尽管短时记忆和事件边界记忆完整地保留选定的帧但它们无法表征完整的长程历史信息。假设每个视频帧包含 L 个视觉 Token那么在 N 帧之后用于全历史注意机制的缓存视频 Token 数量为 |Cv_full | O(NL)且 KV 缓存的存储量与注意计算成本均随 N 呈线性增长。为了维持高效的长程记忆MemoryWAM 为每一帧附加 M 个可学习的要点 Tokens其中 M 远小于 L。对于时刻 t 的视频帧 f_t其对应的概要 Token g_t 会同时关注 f_t 及其历史上下文信息。对于既非锚定帧也非近期帧的视频帧 f_i后续的视频 Token 和动作 Token 不会直接关注 f_i而是关注 f_i 对应的概要 Token g_i即 f_i 的压缩表征。MemoryWAM 的注意掩码attention mask如图 3 所示。在推理阶段MemoryWAM 会移除 f_i 的 KV 缓存但保留 g_i 的 KV 缓存。因此长程历史信息得以作为一种紧凑的持久化记忆被保留下来而无需依赖高昂的全 Token KV 缓存。这种设计显著降低长程记忆的时间复杂度和空间复杂度。若将压缩比定义为 d L/M则长期缓存的大小变为 |Cv_gist| O(NM) O(NL/d)。对于给定的潜分辨率由于 L 是固定的MemoryWAM 将与序列长度轨迹长度相关的存储和注意计算开销从 O(N) 降低到 O(N/d)。在实现中每个视频帧包含 L 120 个潜视觉 Token而 MemoryWAM 每帧仅使用 M 8 个“要旨”gistToken从而实现了 d L/M 15 的压缩比。因此就长期记忆而言与全历史注意机制相比MemoryWAM 将 KV 缓存需求降低 15 倍。在动作生成过程中动作 DiT 会对混合视频缓存执行注意力操作。这种统一的注意接口使 MemoryWAM 能够整合高保真局部上下文、保留的任务边界信息以及紧凑的长期历史信息从而实现依赖记忆的动作生成。实现细节模型架构。在预训练的 Wan2.2-TI2V-5B [63] 基础上构建 MemoryWAM利用其视频 DiT隐藏层维度 3072FFN 维度 1433624 个注意头且头维度为 12830 个 Transformer 块针对 48 通道潜空间的 Patch 大小为 1×2×2、T5 文本编码器以及 3D 因果视频 VAE。参考 FastWAM动作专家action expert是一个独立的动作 DiT其深度30 个块和注意力配置24 个头头维度 128与视频 DiT 保持一致但采用较小的隐层维度d_a 1024和 FFN 维度4096从而构成一个 1B 参数量的动作专家使模型总参数量约为 6B。参考 LingBot-VA [7]通过在隐藏层维度上对预训练视频 DiT 的权重进行插值来初始化动作 DiT 的权重。动作时域跨度action horizo​​n设定为 h16这是由帧步长 4 和时间维度 VAE 步长 4 共同决定的确保每个潜帧对应一个包含 16 步的动作片段。对于 RMBench [1]来自头部、左腕和右腕相机的图像首先被拼接成一个 384×320 的组合图像底部为 256×320 的头部图像顶部为左右腕图像各 128×160沿宽度方向拼接成 128×320随后由 Wan2.2 VAE 联合编码经 Patch 化处理后每视频帧生成 120 个 Token。机器人状态和动作均为 14 维关节向量双臂本体感知数据通过一个可学习的线性层映射到文本 Token 的维度并附加到动作专家的文本上下文中。对于混合记忆模块为每帧保留 M_v 8 个可学习的视频概要 Token设置包含 N_init 2 个初始帧的汇聚窗口sink window以及包含 N_recent 4 个近期清晰帧的滑动窗口。 要点token 被实现为可学习参数并被放置在与其关联的视频帧相同的 3D RoPE 坐标系中其中 (h, w) 坐标固定于一个恒定标记点marker对于动作专家action expert分支共享视频的 3D RoPE 基底从而使动作查询action queries和缓存的视频键cached video keys处于统一的位置参考系中。训练设置。在视频和动作分支中均采用相同的连续流匹配continuous flow-matching公式并设定 1000 个训练时间步。采用基于 t 的平移 Logit-正态分布shifted logit-normal distribution作为噪声调度策略其中视频分支的平移量设为 5.0动作分支设为 1.0。对于每个训练片段episode构建一个逐帧自回归序列其中交替排列着无噪声clean与含噪noisy的视频 token 以及相应的动作片段action chunks同时应用混合记忆注意掩码hybrid memory attention mask以确保训练时的可见性完全复现推理阶段的 KV 缓存状态。参考 [65] 的做法进一步对每个无噪声条件潜变量clean conditioning latent进行增强将其与高斯噪声按 [0, 1] 范围内的均匀随机比例进行线性混合概率 p1.0仅应用于视频侧此举旨在防止教师强制teacher-forced训练过程过拟合于完全无噪声的条件帧。用 AdamW 优化器学习率 2 10-4权重衰减 0.01beta(0.9, 0.95)在 8 块 GPU 上进行训练单卡批大小batch size设为 1。总损失函数定义为 L lambda_v L_video lambda_a * L_action其中 lambda_v lambda_a 1.0且各项均方误差MSE均根据调度器的 Logit-正态训练权重进行重新加权。仿真和真实实验仿真评估。在 RMBench [1] 上评估 MemoryWAM这是一个针对长时程、依赖记忆的机器人操作任务的具有挑战性的仿真基准。与常见的长时程操作基准不同在那些基准中任务相关信息通常可从当前观测中直接获取RMBench 要求策略能够保留并检索历史观测信息因此非常适合评估机器人操作中的持久记忆能力。RMBench 包含九个双臂操作任务涵盖不同级别的任务记忆复杂度。遵循该基准的协议用每个任务 50 次专家演示来训练所有方法并报告在 100 次试运行rollouts中的成功率。将 MemoryWAM 与具有竞争力的 VLA 和 WAM 基线方法进行比较包括 π0.5 [62]、FastWAM [14] 和 LingBot-VA [7]。这些基线方法涵盖三种代表性范式直接从观测-到-动作的映射、具有有限观测窗口的高效 WAM以及利用完整历史信息的 WAM。硬件设置。实验中使用的机器人平台由一台 ARX 双臂机器人组成每只机械臂均配备一个平行夹爪。一台 RealSense D455 相机用于捕捉工作空间的 RGB 图像。完整的硬件设置如图 6 所示。模仿学习细节。为了验证方法在依赖记忆场景中的应用效果设计两个具有挑战性的任务“猜杯游戏”Shell Game和“观察并按压”Look and Press如图 5 所示。在“猜杯游戏”任务中当人类操作员随机交换杯子位置后机器人需要识别并拾起那个扣住小方块的特定杯子。该任务专门用于评估策略在时间推移过程中追踪被遮挡物体的能力。针对该任务收集 50 次演示数据。在“观察并按压”任务中机器人需观察放置在桌面上的两个数字范围为 1 到 5。随后它必须根据观察到的数字分别按压左侧和右侧按钮相应的次数最后按压一次后方按钮以表示任务完成。该任务旨在评估模型的计数能力和工作记忆能力。针对该任务收集 100 次演示数据。关于硬件与部署细节由摄像头采集的输入图像会被裁剪并调整大小至 256 × 352 的分辨率。训练好的模型部署在单块 NVIDIA RTX 4090 GPU 上。在实际运行过程中机器人在每个动作片段action chunk内的控制频率为 10 Hz。此外考虑到模型的推理耗时各动作片段之间存在约 0.3 秒的控制延迟。