[论文解读] Towards mental time travel: a hierarchical memory for reinforcement learning agents

Andrew K. Lampinen, Stephanie C. Y. Chan+2 more|arXiv (Cornell University)|May 28, 2021

Reinforcement Learning in Robotics被引用 9

一句话总结

该论文提出了一种新型记忆架构——分层分块注意力记忆（Hierarchical Chunk Attention Memory, HCAM），专为强化学习智能体设计，通过时间分块记忆上的分层注意力机制，实现对过去事件的详细、稀疏回忆。HCAM使智能体能够通过先关注记忆分块的粗粒度摘要，再聚焦于最相关分块的细节，实现心理上的时间旅行，显著提升长期记忆保持能力、样本效率和泛化性能，尤其在需要对长序列进行推理或在不同任务间实现零样本迁移的任务中表现突出。

ABSTRACT

Reinforcement learning agents often forget details of the past, especially after delays or distractor tasks. Agents with common memory architectures struggle to recall and integrate across multiple timesteps of a past event, or even to recall the details of a single timestep that is followed by distractor tasks. To address these limitations, we propose a Hierarchical Chunk Attention Memory (HCAM), which helps agents to remember the past in detail. HCAM stores memories by dividing the past into chunks, and recalls by first performing high-level attention over coarse summaries of the chunks, and then performing detailed attention within only the most relevant chunks. An agent with HCAM can therefore "mentally time-travel" -- remember past events in detail without attending to all intervening events. We show that agents with HCAM substantially outperform agents with other memory architectures at tasks requiring long-term recall, retention, or reasoning over memory. These include recalling where an object is hidden in a 3D environment, rapidly learning to navigate efficiently in a new neighborhood, and rapidly learning and retaining new object names. Agents with HCAM can extrapolate to task sequences much longer than they were trained on, and can even generalize zero-shot from a meta-learning setting to maintaining knowledge across episodes. HCAM improves agent sample efficiency, generalization, and generality (by solving tasks that previously required specialized architectures). Our work is a step towards agents that can learn, interact, and adapt in complex and temporally-extended environments.

研究动机与目标

• 解决现有记忆架构在长延迟或干扰任务下难以保留详细情景记忆的局限性。
• 使强化学习智能体能够实现稀疏、详细的过去事件回忆——类似于人类的心理时间旅行——且不受无关中间事件的干扰。
• 提升记忆增强型强化学习智能体在样本效率、泛化能力以及对超参数变化的鲁棒性方面的能力。
• 使智能体能够外推至远超训练期间所见长度的任务序列。
• 支持从元学习设置到跨多个任务周期保持知识的零样本泛化。

提出的方法

• HCAM将智能体的过去经历划分为固定长度的时间分块，每个分块作为独立的记忆单元存储。
• 在回忆阶段，HCAM首先对分块摘要应用粗粒度注意力机制，以识别相关记忆分块。
• 随后，仅在最相关的分块内执行细粒度注意力机制，以检索详细信息。
• 该架构采用两级注意力机制：先对分块摘要进行高层注意力，再在选定分块内执行低层注意力。
• HCAM被集成到强化学习智能体的策略网络中，实现记忆增强型决策。
• 该方法设计为对超参数变化具有鲁棒性，且无需对分块边界进行端到端训练。

实验结果

研究问题

• RQ1 分层记忆架构能否使强化学习智能体在长延迟或干扰任务后仍能回忆起详细的历史事件？
• RQ2 HCAM在记忆密集型强化学习任务中，能在多大程度上提升样本效率和泛化能力？
• RQ3 具备HCAM的智能体能否外推至远超训练期间所见长度的任务序列？
• RQ4 HCAM能否实现从元学习到多任务周期记忆保持的零样本泛化？
• RQ5 在超参数变化下，HCAM与TransformerXL及基于LSTM的记忆架构相比，在鲁棒性和性能上表现如何？

主要发现

• 在15亿训练步时，HCAM智能体在最困难的任务（如8舞步、长延迟）上达到80–90%的性能，显著优于最佳的TrXL智能体，后者在同一任务上仅达到40–50%。
• 75%的HCAM智能体在最困难任务上表现优于最佳TrXL智能体，且差异具有统计学显著性（配对t检验p = 0.001 和 p = 0.01）。
• HCAM智能体在含20个干扰项或4个任务周期（每周期1个干扰项）的任务中仍保持高性能，而TrXL智能体在复制实验中未表现出显著高于随机水平的性能。
• 在快速词汇学习任务中，50%的HCAM智能体在50亿步内达到高性能，而TrXL智能体中仅17%达到类似水平。
• HCAM在超参数变化下表现出更优的鲁棒性，学习率和熵权重变化下性能保持稳定，而TrXL则表现出高度敏感性。
• HCAM实现了从元学习到多任务周期记忆保持的零样本泛化，这是先前架构未能实现的能力。