DeepMind
摘要
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.
我们提出了算法蒸馏（AD），这是一种将强化学习（RL）算法蒸馏成神经网络的方法，通过用因果序列模型对其训练历史进行建模。算法蒸馏将强化学习视为一个跨剧情的顺序预测问题。一个学习历史的数据集由源RL算法生成，然后通过自回归预测行动来训练因果转化器，并将其之前的学习历史作为背景。与蒸馏后学习或专家序列的顺序策略预测架构不同，AD能够在不更新网络参数的情况下完全在上下文中改进其策略。我们证明了AD能够在各种具有稀疏奖励、组合任务结构和基于像素的观察的环境中进行语境强化学习，并发现AD学习的RL算法比产生源数据的算法更具有数据效率。
研究动机
- Gato和MGDT这类offline（模仿学习）策略蒸馏方法无法trail-and-error的提升策略，此类方法通过策略蒸馏学习策略，但并非强化学习方法（这里的策略蒸馏指从离线数据蒸馏，而非教师模型）
- 其原因主要是训练数据并不能表示学习的过程
主要贡献
- 是第一个通过序列建模离线数据的in-context强化学习方法
方法描述
POMDP：部分可观测马尔可夫决策过程 -> 需要历史信息辅助观测
In-Context Learning：从上下文推断任务的能力，通常指通过prompt推断任务
- 历史经验history：$\mathcal{H} \ni h_t:=\left(o_0, a_0, r_0, \ldots, o_{t-1}, a_{t-1}, r_{t-1}, o_t, a_t, r_t\right)=\left(o \leq t, r \leq t, a_{\leq t}\right)$
algorithm (long history-conditioned policy)：$P: \mathcal{H} \cup \mathcal{O} \rightarrow \Delta(\mathcal{A})$
- Long需要足够跨越学习更新，例如across episodes，把训练过程的episode拼接在一起
- $\Delta(\mathcal{A})$表示动作空间的概率分布空间
- 算法$P$生成任务$\mathcal{M}$的历史表示为：$\left(O_0, A_0, R_0, \ldots, O_T, A_T, R_T\right) \sim P_{\mathcal{M}}$
- 通过行为克隆将long history-conditioned policies蒸馏到一个神经网络中
数据集：$\mathcal{D}:=\left\{\left(o_0^{(n)}, a_0^{(n)}, r_0^{(n)}, \ldots, o_T^{(n)}, a_T^{(n)}, r_T^{(n)}\right) \sim P_{\mathcal{M}_n}^{\text {source }}\right\}_{n=1}^N$
- $P^{\text {source}}$表示生成数据集的源算法
- 损失函数（negative log likelihood）：$\mathcal{L}(\theta):=-\sum_{n=1}^N \sum_{t=1}^{T-1} \log P_\theta\left(A=a_t^{(n)} \mid h_{t-1}^{(n)}, o_t^{(n)}\right)$
- 直觉上，该序列模型可以学到源算法的包括探索、时间信用分配等复杂行为
数据集生成：训练N个独立的单任务RL算法，为防止过拟合N个任务是随机从任务分布中采样的，保存所有历史
- 可能存在有重复任务？
- 训练序列模型：可以使用任意序列模型进行训练（包括RNN），每次采样长度为$c<T$的子序列进行训练
理论分析
无
实验验证
实验环境
Adversarial Bandit：10 arms的多臂老虎机
- out of distribution：训练集：95%的收益在奇数臂；测试集：95%的收益在偶数臂
Dark Room：二维离散的POMDP，要求智能体到达指定目标位置
- 智能体只知道自身位置，必须通过奖励推测目标位置
- 简单：9*9地图，上下左右不动五个动作，20步，出生在中心点，每次到达目标点都获得奖励
- 困难：17*17地图，只有第一次到达目标时获得奖励
Dark Key-to-Door：与Dark Room类似，要求先找钥匙再开门
- agent只知道自己位置，钥匙和门未知，9*9地图，出生在任意位置，获取钥匙和开门获得奖励，步长50
DMLab Watermaze：三维视觉DMLab环境，基于经典的Morris Watermaze
- 在水下迷宫找出口，迷宫的墙上有利于记住位置的彩色方块
- 观测：72*96*3图像；动作：8个，前后左右移动、视角左右移动、前进并视角左右移动
- 奖励：稀疏奖励，仅到达获取奖励
- 步长50，出生在地图中央
Baselines
- Expert Distillation (ED)：数据只包含专家数据
- Source Algorithm：产生数据的源RL算法from scratch训练，作为data-efficiency对比
- RL^2：online元强化学习算法，作为AD的上界
评测标准 adapt in-context
- 主要测试Fast-Adaptation：将pre-train的model在新环境在不断更新自己的context
- 实验效果
In-context reinforcement learning：
- Credit-assignment：对应DarkRoom每次到达目标都获得奖励的情况，尽管模型输入只有单步奖励，但是智能体能学会如何累计奖励
- Exploration：对应DarkRoom只有第一次到达目标获得奖励的情况，可以从历史episode中推测目标位置
- Generalization：Dark Key-to-Door一共包含6.5k个任务，少于2k在训练中见过，其他任务基本可以泛化到near-optimal
- Data-Efficient
- prompting with demonstrations （好像是图片显示有点问题）
- context size
其他思考
- 本文的方法很简单，主要是将RL^2 和Transformer结合在一起，RL^2是将Episode串在一起训练RNN，AD是将Episode串在一起训练Transformer。
- 文章把该方法说的很高级，不变动网络参数即可实现新环境的Adaptation，从context size的对比试验可以直观理解为如果一旦智能体探索到奖励，并在context的历史信息中有保存，则智能体在新的Episode中可以尽可能模仿之前的有效经验，并在此基础上再进行探索。
- 但是本文只在一些很简单的环境中应用AD取得了较好的效果，单局时长也很短，所以历史探索经验相对简单，且context size可以轻松包含1个或多个episode，应用于复杂环境可能效果不太理想。
- 本文提供的思路其实是来源于RL^2，即想办法利用智能体的历史探索经验，从而实现Learn to learn，但是做法过于暴力，还可以进一步研究。
