AMAGO: 适应性智能体的可扩展上下文强化学习之旅

在现代人工智能的舞台上，强化学习（RL）正逐渐成为一颗璀璨的明星。无论是在游戏、机器人控制还是复杂决策任务中，RL 都展现出了惊人的潜力。然而，如何让智能体在面对新环境时快速适应，依然是一个亟待解决的难题。今天，我们将探索一项名为 AMAGO 的创新研究，它通过“上下文”强化学习，为这一难题提供了新的解决方案。

什么是 AMAGO？

AMAGO 是“可扩展上下文强化学习”的缩写，旨在利用记忆能力帮助智能体从测试时的经验中快速适应新环境。这一方法将元学习（meta-RL）、零-shot 泛化（zero-shot generalization）和长期记忆（long-term memory）统一为一个单一问题。尽管这一技术是深度元强化学习的早期方法之一，但在面对更复杂的算法时，AMAGO 的表现并不总是占优。幸运的是，通过适当的离线策略实现细节和参数调优，上下文 RL 可以变得稳定且具有竞争力。

AMAGO 重新设计了基于序列的离线强化学习，突破了智能体在模型大小、记忆长度和规划视野上的瓶颈。AMAGO 是开源的，旨在以最小的调优需求，成为新研究中适应性智能体的默认选择。

优化长上下文 Transformers 的能力

AMAGO 通过优化长上下文的 Transformer 模型，来提升记忆和适应性。这一过程面临许多技术挑战，但 AMAGO 通过三种主要思想来应对。

1. 共享单一序列模型：AMAGO 的演员（actor）和评论家（critic）在单一序列模型的基础上同时更新，从而最大化训练目标的学习输出。这一方法在以往的研究中并不被鼓励，但 AMAGO 通过精心设计的细节，使得这一过程稳定可行。
2. 长期离线更新：AMAGO 的学习更新通过始终给序列模型“提供学习内容”，提高了性能并减少了调优的复杂度。我们计算不同优化景观下的 RL 损失，确保模型在多种规划视野（\gamma）下都有所学习。当一切都失败时，AMAGO 还包含了一个类似监督学习的离线 RL 项，这使得它在稀疏奖励和长时间范围内尤其有效。
3. 稳定长上下文 Transformers：无论是 RL 还是 Transformers 本身都有不稳定性，而将二者结合则带来了更多障碍。尤其在内存密集型 RL 中，注意力熵崩溃问题显得尤为突出，因为 RL 环境中最优的记忆模式可能比语言建模时更具特异性。AMAGO 使用一种稳定的 Transformer 模块来防止崩溃，并通过选择安全的过大模型大小来减少调优。

元强化学习与长期记忆

AMAGO 使我们能够有效利用 Transformers 的强大记忆能力，在强化学习任务中展现出色的性能。我们在 39 个来自 POPGym 套件的环境中评估了 AMAGO，结果显示它在记忆密集型泛化问题中取得了显著进展，为基于序列的强化学习提供了强有力的基准。

AMAGO 将元学习视为零-shot 泛化的简单扩展，并在多个常见的元强化学习基准上展现了其稳定性和灵活性。我们可以轻松地调整记忆长度，以适应问题的适应难度，且其效率足以支持数百或数千时间步的上下文长度。

自适应指令跟随

离线学习的一个重要优势是能够“回溯重新标记”奖励。AMAGO 将“后见经验重放”扩展至“指令”或多个目标的序列。重新标记指令增加了数据集的多样性，充分发挥了数据密集型 Transformers 的优势，同时为更复杂的目标生成自动探索课程。AMAGO 的重新标记、基于记忆的适应和长期学习更新的结合，在目标条件下的泛化任务中表现出色。

我们在程序生成的 Crafter 世界中评估 AMAGO，使用来自 Crafter 成就系统的字符串作为指令，并添加导航和方块放置的目标。随着任务对探索和适应新世界布局的需求增加，AMAGO 的记忆和重新标记变得至关重要。

例如，以下是一些 AMAGO 在 Crafter 环境中执行的任务：

• "收集树苗，放置植物 x2，吃牛"
• "吃牛，制作石镐，收集煤，制作石剑，击败僵尸"
• "制作木镐，收集石头，在(30, 30)构建"
• "前往(10, 10)，放置石头，前往(50, 50)，放置石头"

这些任务在测试时由用户提示，而每个视频代表了智能体经过训练的成千上万条指令中的一条。

使用 AMAGO

上下文强化学习适用于任何记忆、泛化或元学习的问题，我们设计 AMAGO 以支持所有这些情况。我们的代码是完全开源的，并可在 GitHub 上获取。我们希望我们的智能体能在需要长期记忆和适应的新基准开发中，作为一个强有力的基准，并包括许多应用 AMAGO 的示例：

• 标准（无记忆）MDP/gym 环境
• POMDP 和长期记忆任务
• K-shot 元强化学习
• 目标条件环境适应
• 从像素中多任务学习

结语

AMAGO 的研究为我们在强化学习领域的探索提供了一条新的思路。通过优化智能体的记忆和适应能力，AMAGO 不仅推动了上下文强化学习的应用，还为未来的研究奠定了基础。随着这一领域的不断发展，我们期待看到更多创新的出现，以及 AMAGO 在实际应用中的广泛应用。

参考文献

1. Grigsby, J., Fan, L., & Zhu, Y. (2023). AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents. arXiv:2310.09971.