Gated Memory Policy

1. 研究背景与问题 (Problem)

机器人操作任务对“记忆”的需求差异巨大：

• 马尔可夫任务 (Markovian)：如简单的抓取放置，仅需当前状态，无需历史记忆。
• 非马尔可夫任务 (Non-Markovian)：
  • 单次试验内记忆 (In-trial)：需要在单次执行过程中回忆过去的动作和视觉观测（例如：物体被遮挡后需回忆其位置）。
  • 跨试验记忆 (Cross-trial)：需要聚合多次交互的信息来推断不可观测的物理属性（例如：通过多次推挤推断物体的摩擦系数）。

现有方法的局限性：

1. ** naive 延长历史 (Naive History Extension)**：简单地增加观测历史长度（如将 Transformer 的输入窗口拉长）往往会导致性能显著下降。
   • 分布偏移 (Distribution Shift)：训练时历史数据是干净的，但推理时若历史帧包含噪声或未知观测，会导致整个输入序列偏离分布。
   • 过拟合 (Overfitting)：输入维度增加导致模型更容易过拟合，泛化能力变差。
   • 计算开销大：Transformer 的自注意力机制随历史长度呈二次方增长 ($O(H^2)$)，导致推理缓慢。
2. 结构化记忆 (Structured Memory)：依赖人工设计的规则（如关键点跟踪、文本规划），缺乏通用性且部署耗时。

核心挑战：如何设计一种策略，既能按需回忆历史（解决马尔可夫任务中的干扰），又能高效地提取关键历史信息（解决非马尔可夫任务），同时保持鲁棒性和低计算成本。

---

2. 方法论 (Methodology)

作者提出了 门控记忆策略 (Gated Memory Policy, GMP)，基于 Transformer 的扩散策略 (Diffusion Policy)，包含三个核心创新：

A. 按需回忆：学习“何时”回忆 (When to Recall)

• 机制：引入一个学习到的记忆门 (Memory Gate) $\mu_t$。
• 原理：
  • 使用 MLP 结合 Sigmoid 激活函数，根据当前观测 $I_t$ 和本体感知 $P_t$ 输出二元门控值 $\mu_t \in \{0, 1\}$。
  • 若 $\mu_t=1$，则激活跨注意力模块，引入历史上下文；若 $\mu_t=0$，则跳过历史，仅基于当前状态决策。
• 门控校准 (Gate Calibration)：
  • 由于端到端训练容易导致门控过度开启（过拟合），作者采用了一种自监督校准方法：
    1. 在验证集上分别训练两个策略：一个始终无记忆 ($\pi$)，一个始终有记忆 ($\pi_{mem}$)。
    2. 比较两者在验证集上的动作预测误差。
    3. 若 $\pi$ 的误差显著大于 $\pi_{mem}$（超过阈值 $\theta$），则标记该时刻为“需要记忆” ($\mu_t=1$)。
    4. 冻结门控权重，在全量数据上重新训练最终策略。
  • 效果：确保策略仅在真正需要历史信息的时刻（如关键决策点）调用记忆，在马尔可夫任务中自动忽略历史，保持反应速度。

B. 高效回忆：学习“回忆什么” (What to Recall)

• 机制：引入轻量级的交叉注意力 (Cross-Attention) 模块。
• 原理：
  • 不将历史直接拼接进主序列（避免 $O(H^2)$ 复杂度），而是将历史动作和视觉特征缓存为 Key/Value (KV Cache)。
  • 当前状态作为 Query，通过交叉注意力机制从缓存的历史中检索最相关的信息。
  • 优势：计算复杂度随历史长度线性增长 ($O(H)$)，且无需额外监督即可自动聚焦关键帧。

C. 增强鲁棒性：扩散噪声注入 (Diffusion Noise Injection)

• 问题：直接使用干净的历史数据进行训练会导致模型对历史过度依赖，无法应对推理时的噪声或分布偏移。
• 机制：在训练和推理过程中，对历史动作注入扩散调度噪声 (Diffusion-scheduled noise)。
  • 在扩散步 $k$ 时，模型预测未来动作，条件是基于噪声水平为 $k-1$ 的历史动作（即比未来动作“干净”一步，但仍有噪声）。
  • 效果：强制模型学习在历史不完美（含噪声）的情况下依然能做出正确决策，显著提升了泛化能力和鲁棒性。

---

3. 关键贡献 (Key Contributions)

1. Gated Memory Policy (GMP) 框架：
   • 提出了一种高效、实用的记忆聚焦策略框架。
   • 通过门控机制实现了“按需回忆”，在保持马尔可夫任务高性能的同时，成功处理了复杂的非马尔可夫任务。
   • 通过交叉注意力实现了线性计算复杂度，解决了长历史带来的计算瓶颈。
2. MemMimic 基准测试：
   • 提出了一个新的基准套件，包含模拟和真实世界的任务，明确区分并评估了单次试验内记忆和跨试验记忆能力。
   • 填补了现有基准（如 RoboMimic）主要关注马尔可夫任务，而忽视跨试验记忆任务的空白。
3. 全面的实证研究：
   • 在仿真和真实世界（包括 ARX X5 机器人和 iPhone 视觉）中进行了广泛验证。
   • 深入分析了记忆门控校准、噪声注入等设计选择对性能、鲁棒性和计算效率的影响。

---

4. 实验结果 (Results)

基准测试表现

• MemMimic (非马尔可夫任务)：
  • GMP 在长历史基线方法上取得了 30.1% 的平均成功率提升。
  • 在“迭代推挤 (Iterative Pushing)"、“迭代投掷 (Iterative Flinging)"和“真实世界迭代铸造 (Iterative Casting)"等需要跨试验记忆的任务中，GMP 能够成功推断物理属性（如摩擦系数、质量）并调整策略，而长历史基线往往失败。
• RoboMimic (马尔可夫任务)：
  • GMP 在不需要记忆的简单任务（如 Square, Transport, Tool Hang）上保持了与无历史策略相当甚至更优的性能，证明了门控机制有效避免了历史信息的干扰。
• MIKASA-Robo 基准：
  • 相比现有的 VLA (Vision-Language-Action) 基线（如 MemoryVLA），GMP 平均提升了 26.6% 的成功率。

关键发现 (Findings)

1. 简单延长历史会退化性能：在马尔可夫任务上，历史越长性能越差；GMP 通过门控机制解决了这一问题。
2. 注意力聚焦正确：可视化显示，GMP 的交叉注意力能自动聚焦到关键帧（如第一次看到物体颜色的时刻，或上一次试验的失败结果），无需显式监督。
3. 计算效率高：
   • 自注意力基线的推理时间随历史长度呈二次方增长。
   • GMP 在门控开启时呈线性增长，在门控关闭时（大多数马尔可夫时刻）推理时间恒定且极短。
4. 噪声注入提升鲁棒性：在历史动作中注入扩散噪声的策略，比随机噪声或无噪声策略在真实世界任务中表现更稳定。

---

5. 意义与影响 (Significance)

• 理论突破：解决了机器人策略中“记忆长度”与“任务类型”不匹配的矛盾，证明了选择性记忆优于全量记忆。
• 工程价值：
  • 低成本部署：线性计算复杂度使得在边缘设备（如机器人自带 GPU）上运行长历史策略成为可能。
  • 真实世界适应性：通过噪声注入和门控机制，GMP 在充满噪声的真实世界环境中表现出极强的鲁棒性，能够处理物体物理属性未知、环境动态变化等复杂场景。
• 未来方向：为构建具备长期记忆、能够进行上下文自适应（In-context Adaptation）的通用机器人策略提供了新的范式。

总结：GMP 通过“门控机制决定何时用记忆”、“交叉注意力决定用什么记忆”以及“噪声注入增强鲁棒性”的三重设计，成功实现了在复杂非马尔可夫任务中的高性能，同时避免了在简单任务中的性能退化，是机器人模仿学习领域的重要进展。