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）。该模型采用混合记忆机制，结合少量“要点”（gist）tokens与精心设计的注意机制，实现了高效且有效的决策记忆利用。具体而言，滑动观测窗口保留用于即时控制的高保真短期上下文信息；少量要点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]。基于这些认知启示，提出一种混合记忆机制：它既保留高保真的短时上下文信息，又维护少量概括长程历史的“要点”（gist）tokens，同时还保留任务起始时刻的锚定帧。这种类人的混合记忆机制在显著降低推理成本的同时，使策略能够对长时程依赖关系进行推理。

2 架构

MemoryWAM 遵循近期“世界动作模型”（WAM）中基于视频动作扩散的范式：利用预训练的视频扩散 Transformer (DiT) 提供感知动态的视觉表征，并由独立的动作 DiT 基于所学习的视觉动态预测未来动作。该模型架构如图 2 所示。给定观测值 o_t，为了提高计算效率，首先利用因果视频 VAE [63] 将其编码为紧凑的视频潜向量 z_t。该视频潜向量由视频 DiT Phi_v 处理，而动作片段 a_t:t+h-1 则由动作 DiT Phi_a 生成。这两个分支采用 Transformer 混合（MoT）[64] 架构进行组织。此外，MemoryWAM 继承近期高效 WAM [14, 15] 的一项关键优势：它在训练阶段通过视频预测学习物理动态，同时避免推理阶段高昂的视频生成成本。

在推理过程中，当前观测的干净潜在向量 z_t 仅需通过一次视频 DiT 前向传播，以更新视频侧的键值（KV）缓存 C^v_t。

随后，动作 DiT 通过对动作 Token 进行去噪来预测动作块（action chunk），同时关注已缓存的视频表征。由此可见，视频 DiT 负责提取感知动态特性的特征并维护记忆，而动作 DiT 则将这些特征映射为动作。

3 混合记忆

如前所述，基于完整历史记录的注意机制会导致 GPU 显存开销和推理延迟急剧增加。MemoryWAM 采用混合记忆设计来解决这一问题，该设计灵感源自人类记忆的互补形式：用于即时闭环控制的短期记忆、用于检索任务起始状态的事件边界记忆，以及用于紧凑表征长程历史的要旨记忆（gist memory）。具体而言，在时间步 t，Memory-WAM 维护一个紧凑的时间缓存，其中三个组成部分分别对应近期观测、事件边界帧以及压缩后的长期历史信息。

短时记忆。短时记忆负责即时的闭环控制，通过捕捉物体运动、接触状态和手-物构型等快速变化的交互线索，获取近期观测信息。这种记忆被实现为覆盖最近 N 个视频帧的滑动窗口缓存。这既保留了用于动作生成的高保真局部上下文，又将短时注意计算成本限制在固定的窗口大小范围内。

事件边界记忆。并非所有的历史观测都具有同等的信息价值。在连续的体验过程中，事件边界为记忆组织提供了显著信息。在机器人操作任务中，这些边界往往对应于任务的启动时刻和初始场景配置。因此，在任务开始时保留少量包含完整视觉 Token 的“锚定帧”（anchor frames），因为初始场景状态往往构成了指令中关键信息的参照基础，且该状态在后续过程中可能会被遮挡或超出观测窗口范围。

要点记忆（Gist memory）。尽管短时记忆和事件边界记忆完整地保留选定的帧，但它们无法表征完整的长程历史信息。假设每个视频帧包含 L 个视觉 Token，那么在 N 帧之后，用于全历史注意机制的缓存视频 Token 数量为 |C^v_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，则长期缓存的大小变为 |C^v_gist| = O(NM) = O(NL/d)。

对于给定的潜分辨率，由于 L 是固定的，MemoryWAM 将与序列长度（轨迹长度）相关的存储和注意计算开销从 O(N) 降低到 O(N/d)。在实现中，每个视频帧包含 L = 120 个潜视觉 Token，而 MemoryWAM 每帧仅使用 M = 8 个“要旨”（gist）Token，从而实现了 d = L/M = 15 的压缩比。因此，就长期记忆而言，与全历史注意机制相比，MemoryWAM 将 KV 缓存需求降低 15 倍。

在动作生成过程中，动作 DiT 会对混合视频缓存执行注意力操作。这种统一的注意接口使 MemoryWAM 能够整合高保真局部上下文、保留的任务边界信息以及紧凑的长期历史信息，从而实现依赖记忆的动作生成。

实现细节

模型架构。在预训练的 Wan2.2-TI2V-5B [63] 基础上构建 MemoryWAM，利用其视频 DiT（隐藏层维度 3072，FFN 维度 14336，24 个注意头且头维度为 128，30 个 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）设定为 h=16，这是由帧步长 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] 范围内的均匀随机比例进行线性混合（概率 p=1.0，仅应用于视频侧），此举旨在防止教师强制（teacher-forced）训练过程过拟合于完全无噪声的条件帧。用 AdamW 优化器（学习率 2 10^-4，权重衰减 0.01，beta=(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 秒的控制延迟。