打破瓶颈，释放记忆洪荒之力：AMAGO 如何用长序列 Transformer 驯服上下文强化学习？

强化学习（RL）领域近年来发展迅速，涌现出许多能够在特定环境中出色完成单一目标任务的“专家级”决策智能体。然而，构建能够适应多样化环境、具备长期记忆和不确定性推理能力的“通才型”智能体，成为了新的挑战与目标。上下文强化学习（in-context RL）作为通才型 RL 的有力候选者，利用序列模型学习历史经验，在测试阶段动态地适应新情况。试想一下，如果智能体拥有强大的记忆力，能够像福尔摩斯一样，从过往的经历中汲取经验教训，那它在面对新环境时，就能迅速找到应对策略，成为真正的“百事通”！

然而，理想很丰满，现实很骨感。早期的上下文 RL 智能体受限于循环神经网络和在线学习方法，在记忆容量和规划范围上都难以突破瓶颈。Transformer 的出现，将记忆问题巧妙地转化为信息检索，为解决这一难题带来了曙光。然而，将 Transformer 与离线 RL 结合，却像是两位武林高手过招，各自的看家本领反而相互掣肘，难以施展。

AMAGO 横空出世，上下文 RL 迎来新突破！

AMAGO (Adaptive Memory Agent for achieving GOals) 这一全新算法的提出，为上下文 RL 领域带来了两大重要突破。首先，AMAGO 对离线策略 actor-critic 更新机制进行了彻底的重新设计，使其能够支持长序列 Transformer，并实现对整个 rollout 的并行学习（如图 1 所示）。 AMAGO 巧妙地打破了离线策略上下文 RL 的最大瓶颈，将记忆长度、模型规模和规划范围提升到了前所未有的高度。更令人惊喜的是，AMAGO 开源且易于使用，只需少量调整即可应用于新环境，堪称“开箱即用”的利器！

图 1：上下文 RL 技术通过序列模型，从测试时的经验中推断未知环境的特征，从而解决记忆和元学习问题。AMAGO 解决了核心技术挑战，将端到端离线策略 RL 与长序列 Transformer 的性能相结合，将记忆和适应能力推向新的极限。

AMAGO 的第二个贡献在于开辟了上下文 RL 的新方向，并引入了全新的基准测试。AMAGO 利用离线数据，专注于长范围和稀疏奖励，使其能够独特地扩展到探索难度更大的目标条件问题。研究人员为多步目标添加了一种新的 hindsight relabeling 方案，可以在多任务领域生成有效的探索计划。AMAGO 可以学习在程序生成的环境中完成多种可能的指令，同时利用其记忆来适应不熟悉的环境。

AMAGO 的秘密武器：如何化解 Transformer 与离线 RL 的“矛盾”？

AMAGO 的成功并非偶然，其背后隐藏着许多精妙的技术细节。

1. 共享序列模型，化解参数爆炸危机： 传统的离线策略 RL 通常需要为 actor 和 critic 网络分别优化不同的序列模型，并使用额外的目标模型来计算 critic 损失。这种做法虽然有效，但会导致训练过程中需要进行大量的序列模型前向/反向传播，限制了模型规模和上下文长度。AMAGO 反其道而行之，将 actor 和 critic 目标函数整合到一个损失函数中，并使用同一个优化器进行训练。为了防止不同损失项的权重在不同实验中需要反复调整，AMAGO 精心设计了 critic 与 actor 损失项之间的分离机制，并使用 REDQ [74] 的定制变体来计算损失项，从而在最大程度上减少了超参数的调整。

2. 稳定长上下文 Transformer，驯服“脱缰野马”： Transformer 虽然功能强大，但在 RL 中的应用却面临着比监督学习更大的挑战。AMAGO 通过修改 Transformer 架构，借鉴 Normformer [88] 和 σReparam [87] 的方法，并使用 Leaky ReLU 替换饱和激活函数，有效地解决了长序列 RL 中常见的注意力熵崩溃问题。这种架构有效地稳定了训练过程，并减少了超参数调整，使得我们可以放心地选择比所需规模更大的模型，而无需担心训练不稳定。

3. 多伽马学习，破解长视野规划难题： 长范围规划和稀疏奖励是上下文 RL 面临的另一大挑战。AMAGO 采用“多伽马”更新机制，并行优化多个不同的折扣因子 γ，每个 γ 对应一个独立的 Q 值表面，从而提高了共享 Transformer 在训练过程中获得强 actor-critic 学习信号的可能性。此外，AMAGO 还引入了一种“过滤行为克隆” (BC) 项，在估计动作具有正优势时执行监督学习，进一步增强了学习效果。

4. Hindsight 指令重标记，探索未知的宝藏： AMAGO 利用离线数据和高折扣因子的优势，提出了一种针对多步目标的 HER 变体，可以对长轨迹进行 hindsight relabeling。这种方法不仅可以创建更复杂的多阶段任务，还可以在 Crafter [33] 等开放世界领域中有效地创建自动探索计划。

AMAGO 大显身手，挑战高难度任务！

研究人员在一系列长短期记忆、泛化和元学习环境中对 AMAGO 进行了测试，并在多任务领域和程序生成环境中探索了 AMAGO 的自适应记忆和 hindsight 指令重标记的组合效果。实验结果表明，AMAGO 在多个基准测试中均取得了优异的成绩，特别是在 POPGym 套件 [30] 中，AMAGO 的召回率显著提升了性能。

展望未来，上下文 RL 的星辰大海！

AMAGO 的出现，为上下文 RL 领域注入了新的活力。其高效的长序列处理能力，为未来研究开辟了新的方向，也为开发更复杂、更接近真实世界的 RL 基准测试提供了强有力的工具。相信在不久的将来，AMAGO 将帮助我们构建更加智能、更具适应能力的智能体，在更广阔的领域中发挥重要作用！

参考文献：

[8] Goyal, A., Lamb, A., Hoffmann, J., Sodhani, S., Levine, S., Bengio, Y., & Schölkopf, B. (2022). Recurrent independent mechanisms for reinforcement learning. International Conference on Learning Representations.
[10] Duan, Y., Schulman, J., Chen, X., Bartlett, P. L., Sutskever, I., & Abbeel, P. (2016). RL2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:1611.02779.
[11] Wang, J. X., Kurth-Nelson, Z., Tirumala, D., Soyer, H., Leibo, J. Z., Munos, R., … & Botvinick, M. (2016). Learning to reinforcement learn. arXiv preprint arXiv:1611.05763.[22] Ni, A., Agarwal, V., Schwarzer, M., Castro, P. S., Courville, A., & Bellemare, M. G. (2022). Discovering algorithmic leverage for reinforcement learning via modular metalearning. International Conference on Learning Representations.[30] Ni, A., Agarwal, V., Schwarzer, M., Castro, P. S., Courville, A., & Bellemare, M. G. (2023). Popgym: Benchmarking partially observable reinforcement learning in procedurally generated gyms. arXiv preprint arXiv:2302.01561.
[33] Hafner, D., Synnaeve, G., Norouzi, M., & Lillicrap, T. (2020). Mastering atari with discrete world models. International Conference on Machine Learning, 3865-3875. PMLR.