[论文解读] In-context Reinforcement Learning with Algorithm Distillation

Michael Laskin, Luyu Wang+12 more|arXiv (Cornell University)|Oct 25, 2022

Reinforcement Learning in Robotics被引用 10

一句话总结

本文提出算法蒸馏（AD），一种通过因果Transformer模型从强化学习（RL）智能体的跨episode扩展学习历史中预测动作的方法，实现无需参数更新的上下文内策略改进。AD通过将策略改进建模为完整训练轨迹上的序列预测问题，学习到比源RL算法更高效的数据利用型RL算法。

ABSTRACT

We propose Algorithm Distillation (AD), a method for distilling reinforcement learning (RL) algorithms into neural networks by modeling their training histories with a causal sequence model. Algorithm Distillation treats learning to reinforcement learn as an across-episode sequential prediction problem. A dataset of learning histories is generated by a source RL algorithm, and then a causal transformer is trained by autoregressively predicting actions given their preceding learning histories as context. Unlike sequential policy prediction architectures that distill post-learning or expert sequences, AD is able to improve its policy entirely in-context without updating its network parameters. We demonstrate that AD can reinforcement learn in-context in a variety of environments with sparse rewards, combinatorial task structure, and pixel-based observations, and find that AD learns a more data-efficient RL algorithm than the one that generated the source data.

研究动机与目标

• 实现上下文内强化学习，使模型通过试错改进而无需微调。
• 通过长上下文序列建模，将强化学习算法训练过程建模为序列预测问题。
• 不仅蒸馏固定策略，还从强化学习训练历史中蒸馏出完整的策略改进算子。
• 证明可通过模仿学习在离线强化学习学习历史上实现上下文内强化学习。
• 表明蒸馏后的算法在数据效率方面可优于源RL算法并具备泛化能力。

提出的方法

• 从多个独立任务中收集源RL算法的完整学习历史大规模多任务数据集。
• 将每个学习历史表示为状态-动作-奖励标记的序列，保留时间与回合结构。
• 训练因果Transformer，利用完整的先前学习历史作为上下文，自回归地预测动作。
• 使用仅依赖过去状态、动作和奖励的序列建模损失，不依赖回报。
• 确保上下文长度足够长，以捕捉跨多个回合的策略改进。
• 在推理时以上下文方式应用训练好的Transformer：以先前回合的历史为提示，预测动作而不更新权重。

实验结果

研究问题

• RQ1 序列模型能否通过关注完整学习历史，在上下文中学习到策略改进？
• RQ2 将强化学习训练建模为序列预测问题，是否能实现上下文内探索与时间信用分配？
• RQ3 蒸馏模型能否在新任务上泛化，并在数据效率方面优于源RL算法？
• RQ4 是否可通过模仿学习将强化学习算法的策略改进算子蒸馏为单一神经网络？
• RQ5 在样本效率方面，蒸馏算法与源算法相比表现如何？

主要发现

• AD在无需任何参数更新的情况下实现上下文内强化学习，完全依赖先前回合的上下文信息。
• AD在部分可观察环境（如DMLab的Watermaze）中展现出上下文内探索与时间信用分配能力。
• AD学习到的强化学习算法比生成训练数据的源算法更具数据效率。
• 该方法在稀疏奖励、组合结构及基于像素的观测任务中均表现出良好的泛化能力。
• AD优于以往需要微调或专家演示以适应新任务的离线策略蒸馏方法。
• AD的成功关键在于足够长的上下文窗口，能够跨越多个回合以捕捉策略改进过程。