[论文解读] Transformers are Meta-Reinforcement Learners

Luckeciano C. Melo|arXiv (Cornell University)|Jun 14, 2022

Reinforcement Learning in Robotics被引用 8

一句话总结

该论文提出 TrMRL，一种基于 Transformer 架构的元强化学习智能体，通过多头自注意力机制递归构建情景记忆，实现快速适应与鲁棒泛化。TrMRL 在 HalfCheetahVel 等连续控制环境及机器人操控任务中，实现了优越或相当的样本效率与分布外泛化能力。

ABSTRACT

The transformer architecture and variants presented remarkable success across many machine learning tasks in recent years. This success is intrinsically related to the capability of handling long sequences and the presence of context-dependent weights from the attention mechanism. We argue that these capabilities suit the central role of a Meta-Reinforcement Learning algorithm. Indeed, a meta-RL agent needs to infer the task from a sequence of trajectories. Furthermore, it requires a fast adaptation strategy to adapt its policy for a new task -- which can be achieved using the self-attention mechanism. In this work, we present TrMRL (Transformers for Meta-Reinforcement Learning), a meta-RL agent that mimics the memory reinstatement mechanism using the transformer architecture. It associates the recent past of working memories to build an episodic memory recursively through the transformer layers. We show that the self-attention computes a consensus representation that minimizes the Bayes Risk at each layer and provides meaningful features to compute the best actions. We conducted experiments in high-dimensional continuous control environments for locomotion and dexterous manipulation. Results show that TrMRL presents comparable or superior asymptotic performance, sample efficiency, and out-of-distribution generalization compared to the baselines in these environments.

研究动机与目标

• 为解决元强化学习中快速适应的挑战，利用注意力机制的上下文相关参数化特性。
• 设计一种基于记忆的元强化学习智能体，通过 Transformer 层递归优化情景记忆，模拟记忆再激活过程。
• 利用稳定的 Transformer 架构，提升连续控制环境中样本效率与分布外泛化能力。
• 研究 T-Fixup 权重初始化在稳定基于 Transformer 的元强化学习智能体训练过程中的有效性。
• 验证自注意力机制可计算最小化贝叶斯风险的一致性表征，从而提升动作预测性能。

提出的方法

• TrMRL 使用因果多头自注意力机制，将近期工作记忆通过 Transformer 层整合为任务感知的情景记忆。
• 每一层 Transformer 通过计算最小化贝叶斯风险的一致性表征，优化情景记忆，提升动作选择的特征质量。
• 最后一层输出的情景记忆作为策略头的输入，实现上下文相关的动作预测。
• 应用 T-Fixup 权重初始化以稳定训练过程，并缓解强化学习早期探索阶段的不稳定性。
• 情景记忆通过各层递归构建，每一层均对表征进行优化，逐步逼近能更好预测最优动作的一致性表征。
• 使用线性回归模型与余弦相似度评估各层表征的质量，衡量表征的优化程度与预测能力。

实验结果

研究问题

• RQ1 Transformer 的自注意力机制是否能作为元强化学习中有效的快速适应策略，通过构建上下文相关策略实现？
• RQ2 通过 Transformer 层递归优化情景记忆是否能生成最小化贝叶斯风险的一致性表征，并提升动作预测性能？
• RQ3 与现有元强化学习基线相比，TrMRL 在样本效率与分布外泛化方面表现如何？
• RQ4 T-Fixup 初始化在多大程度上稳定了基于 Transformer 的元强化学习智能体的训练过程？
• RQ5 中间层的情景记忆表征是否可被有意义地用于预测专家动作，表明其表征的逐步优化？

主要发现

• 在 HalfCheetahVel 环境中，TrMRL 从第一集即表现出高性能，仅需约 20 个时间步即可达到最优性能，表明其具备强大的在线适应能力。
• 在目标速度范围为 [3.0, 4.0] 的 HalfCheetahVel 分布外泛化设置中，TrMRL 显著优于所有基线方法，包括 PEARL 和 VariBAD，表明其具备鲁棒的任务表征学习能力。
• 中间层情景记忆表征与最终层表征之间的表征误差随层数增加而减小，证实了向一致表征的逐步优化过程。
• 使用各层情景记忆特征的线性回归模型在测试集上均表现出递减的均方误差，表明随着网络深度增加，表征对专家动作的预测能力持续提升。
• 消融实验表明，T-Fixup 初始化对训练稳定性与性能至关重要，且对网络深度、序列长度及注意力头数量敏感。
• 情景记忆优化过程与定理 4.1 一致，表明自注意力机制在每一层均能计算最小化贝叶斯风险的一致性表征。