Imagination-Augmented Agents for Deep Reinforcement Learning

文献目录

本篇论文是谷歌投的一篇NIPS 2017的论文，提出了一种想象力增强的model-based强化学习方法，思路非常新颖，GitHub上可以搜到不少对I2A的复现代码。由于这篇论文写得比较抽象，个人时间有限，目前更多的只是翻译，之后有理解总结的话再更新本文。

本文旨在解决model-based方法中模型不准确导致预测存在误差的问题。

I2A既利用了基于模型的学习，也利用了无模型的学习。其架构如图1c所示，观测$O_t$同时输入到model-based路径（左）和model-free路径（右）中，其中model-free路径为常规的无模型强化学习算法。model-based路径的输出结果由多个rollout encoder聚合得到，rollout encoder是agent进行imagination任务的地方，其结果如图1b所示。

Rollout encoder有两层：imagine future和encoder。Imagine future是想象力发生的地方，imagine future中包括imagination core（结构如图1a）。向imagination core输入状态$O_t$，我们得到新的状态${\hat{O}}_{t+1}$和奖励${\hat{r}}_{t+1}$，接着将${\hat{O}}_{t+1}$输入到下一个imagination core得到下一个新状态${\hat{O}}_{t+2}$和奖励${\hat{r}}_{t+2}$。重复上述步骤$n$次，我们得到一个基本上是一对状态和奖励的rollout，然后使用encoder（如LSTM）来编码此rollout，得到一个rollout encoding。这些rollout encoding就是想象出的未来路径。我们为不同的未来想象路径提供多个rollout encoder，我们使用一个aggregator 来聚合此rollout encoder。

Imagination core包括策略网络和环境模型。环境模型从agent的历史信息中学习，它获取状态${\hat{O}}_t$的信息，并根据经验想象所有可能的未来，选择高回报的动作${\hat{a}}_t$。

缺点：I2A将所有不确定性估计和模型使用表征至隐式神经网络训练过程中，继承了无模型方法的低效率。

---

1. Introduction

Model-free方法使用了深度神经网络将原始观测直接映射为值或者动作，已经取得了一定的发展。但是这种方法通常需要大量的训练数据，且产生的策略不容易推广到同一环境中的新任务中，因为它缺乏泛化智能的行为灵活性。

Model-based RL 根据历史经验合成世界模型，通过使用内部模型来推理未来（imagining）。虽然需要先学习模型，但model-based的方法具有更好的泛化能力，且能够利用额外的无监督学习信号，提高数据效率。另外基于模型的方法通过增加内部模拟量来扩展性能和更多的计算能力。

基于模型的RL方法目前的应用仅限于具有精确的转换模型或易于学习的领域（符号环境或低维度系统）。对于复杂系统，采用标准规划方法基于模型的agent的性能通常会受到函数逼近导致的模型误差的影响。这些错误在规划过程中累积，导致过度乐观和很差的agent性能。目前没有能够抵抗模型缺陷的规划或基于模型的方法，这些方法在复杂域中是不可避免的。

本文通过提出想象力增强agent（Imagination-Augmented Agents）来解决这一问题，该方法通过“学习解释”不完美的预测来使用近似的环境模型。此种方法是end-to-end的、不对环境模型进行任何假设，且不完全依赖于模拟的returns。

2. The I2A architecture

I2A架构图如上图所示。通过向环境模型输入当前的信息可以对未来进行预测，得到模拟的想象轨迹（imagined trajectories），这些轨迹由神经网络解释并作为策略网络的附加上下文提供。

通常，环境模型是可以从agent轨迹以无监督方式训练的任何循环架构：给定过去状态和当前动作，环境模型预测下一状态和来自环境的任何数量的信号。本文基于 action-conditional next-step predictors 建立环境模型，它们接收当前观测（或观测历史）和当前行动的输入，并预测下一次观测，也可能是下一次奖励。我们在未来的多个时间步中roll out环境模型，通过用当前时间的真实观测初始化想象轨迹，随后将模拟观测值馈送到模型中。

每一个rollout中的actions都来自于rollout策略$\hat{\pi }$。环境模型与策略$\hat{\pi }$共同构成了 imagination core module，用于预测下一时间步。该模块用于生成n个轨迹${\hat{T}}_1\dots {\hat{T}}_n$，每一个轨迹$\hat{T}$都是一个特征序列(${\hat{f}}_{t+1},\dots ,{\hat{f}}_{t+\tau }$)，$t$是当前时间步，$\tau$是rollout的长度，${\hat{f}}_{t+i}$是环境模型的输出（即预测的observation 和/或 reward）。

I2As解决的一个关键问题是学习模型不能被认为是完美的，它有时可能会做出错误或无意义的预测。 因此，我们不希望仅依赖于预测的奖励（或从预测状态预测的值），这在经典规划中经常这么用。 此外，轨迹可能包含奖励序列之外的信息（轨迹可能包含信息性子序列 - 例如解决子问题 - 这不会导致更高的奖励）。 出于这些原因，我们使用一个rollout encoder $\epsilon$将 imagined rollout 作为一个整体来处理并学会解释它，即提取对agent的决定有用的任何信息，或者甚至在必要时忽略它（图1b）。每个轨迹分别编码为rollout embedding $e_i=\epsilon ({\hat{T}}_i)$。最后，aggregator $A$ 将不同的rollout embeddings转换为单个imagination code $C_{ia}=A(e_1,\dots ,e_n)$。

I2A的最后一个组成部分是策略模块，它是一个网络，它从基于模型的预测中获取信息$C_{ia}$，以及无模型的输出$C_{mf}$（一个只接受真实观测作为输入的网络；参见图1c，右），并输出想象增强的策略向量$\pi$和估计值$V$. 因此，I2A学会结合无模型和想象增强路径的信息；如果没有基于模型的路径，I2As将退化为标准的无模型网络。 因此，I2A可以被认为是通过从基于模型的规划提供附加信息来增强无模型agent性能，并且具有比基础无模型agent更严格的表达能力。

3. Architectural choices and experimental setup

3.1. Rollout strategy

实验中针对环境中每一个可能的action执行一个rollout。定义第$i$个rollout中的第一个action为action集合$A$中的第$i$个action，剩下的所有rollouts的后续actions都是由一个shared rollout policy $\hat{\pi }$ 产生。我们研究了几种类型的rollout策略（随机，预训练），发现一种特别有效的策略是将想象力增强的策略提炼为无模型策略。这种蒸馏策略在于创建一个小的无模型网络$\hat{\pi }(o_t)$，并在总损失中增加了一项交叉熵辅助损失（计算于当前观测的想象力增强策略$\pi (o_t)$和计算于相同观测的rollout策略$\hat{\pi }(o_t)$）。通过模仿想象力增强策略，内部的rollouts将类似于agents在真实环境中的轨迹，这也确保了rollouts对应的轨迹具有高回报。同时，不完美的近似导致具有更高熵的rollout策略，可能在勘探和开发之间取得平衡。

3.2. I2A components and environment models

在本文的实验中encoder是一个带有卷积encoder的LSTM，其按时序对轨迹T进行处理。特征${\hat{f}}_t$ 以逆序（${\hat{f}}_{t+\tau }$至${\hat{f}}_{t+1}$）提供给LSTM，以模仿Bellman类型的backup operations（实验证明向前、向后或双向处理的选择似乎对I2A的性能影响相对较小，然而，不应排除研究不同策略）。Aggregator仅将summaries连接起来。对于I2A的无模型路径，我们选择了卷积层的标准网络加上一个全连接网络。 我们还将此架构单独用作baseline agent。

环境模型（图2）定义了一个使用负对数似然损失$l_{model}$优化的分布。我们可以在将环境模型（使用冻结权重）嵌入I2A体系结构之前对其进行预训练，也可以将$l_{model}$作为辅助损失添加到总损失中与agent一起对其进行培训。在实践中，我们发现环境模型的预训练使得I2A体系结构运行的更快，因此我们采用了这种策略。

对于所有环境，我们环境模型的训练数据都是从经过部分训练的标准无模型agent（定义如下）的轨迹生成的。我们使用部分预训练的agents，因为随机agents在我们的一些领域中很少能够看到回报。但是，这意味着我们必须考虑到预训练数据生成agent以及生成数据所需的计算成本（根据实际环境步骤）。在实验中，我们通过两种方式来解决这一问题：明确说明预训练（对于sokoban）中使用的步骤数，或者演示如何将相同的预训练模型重用到许多任务（对于minipacman）。

3.3. Agent training and baseline agents

使用固定的预训练环境模型，我们使用A3C训练剩余的I2A参数。 如上所述，我们在策略$\pi$上添加了一个熵正则化器以鼓励探索和辅助损失，从而将$\pi$提炼为rollout策略$\hat{\pi }$。 我们使用RMSprop优化器在超过32~64 workers中进行异步训练。

为每个agent architecture执行单独的超参数搜索，以确保最佳性能。 除了I2A之外，我们还运行了以下baseline agents（有关所有代理的体系结构详细信息，请参阅附录B）。

• Standard model-free agent
  对于我们的主要基线代理，我们选择了类似于A3C的无模型标准体系结构，包括卷积层（MiniPacman为2，Sokoban为3），后面是全连接层。 最后一层也是全连接层，输出策略logits和value函数。 对于Sokoban，我们还测试了一个“大型”标准架构，其中我们将所有特征图（对于卷积层）和隐藏单元（对于完全连接的层）的数量加倍， 最后架构的参数数量略大于I2A。

• Copy-model agent
  除了具有内部环境模型之外，I2A架构与标准agent的架构非常不同。 为了验证环境模型rollouts中包含的信息是否有助于提高性能，我们实现了一个baseline，我们将I2A中的环境模型替换了成单返回输入观察的“复制”模型。 缺乏模型，这个agent不使用想象力，但使用相同的架构，具有相同数量的可学习参数（环境模型在I2A中保持不变），并且受益于相同的计算量（在两种情况下均随rollouts的长度线性增加）。 此模型有效地对应于一种体系结构，其中策略logits和value是具有skip connections的LSTM网络的最终输出。

4. Sokoban experiments

4.1. I2A performance vs. baselines on Sokoban

从图4中可以看出rollouts展开步为3时，相较于1个展开步，显著提升了学习速度和性能。5优于3，15优于5。但是，总的来说，我们发现使用I2A的rollouts较长会导致收益递减。值得注意的是，5个steps相对于解决一个关卡所采取的步骤数量来说是相对较小的，我们最好的agent平均需要大约50个steps。这意味着即使是这么短的rollouts也可以提供很多信息。

4.2. Learning with imperfect models

I2As的一个关键优势是能够处理不完美的环境模型。然而，对于Sokoban的任务，我们学到的环境模型在推出想象的轨迹时实际上表现得非常好。为了证明I2A可以处理不太可靠的预测，我们进行了另一个实验，其中I2A使用的环境模型表现出更差的性能（由于参数数量较少），在迭代rollout预测中积累了strong artifacts（图5，左）。如图5（右）所示，即使有这样一个明显有缺陷的环境模型，I2A的表现同样出色。这意味着当错误累积时，I2A可以学会忽略rollout的后半部分，但是当错误不那么严重时仍然使用初始预测。最后，请注意，在我们的实验中，令人惊讶的是，模型较差的I2A代理的表现优于具有良好模型的I2A代理。我们认为这是由于随机初始化，虽然我们不能排除提供某种形式的正则化的噪声模型 - 需要更多的工作来研究这种影响。

学习rollout编码器使I2As能够处理不完美的模型预测。 我们可以通过将它们与没有rollout编码器的设置进行比较来进一步证明这一点：如在Tesauro和Galperin的经典蒙特卡罗搜索算法中，我们现在明确地估计每次行动的价值来自rollouts，而不是学习rollouts的任意编码，如I2A。 然后，我们根据这些值选择操作。 具体地，我们从状态学习值函数$V$，并且使用rollout策略$\hat{\pi }$，对每个初始动作采样轨迹rollout，并计算相应的预估蒙特卡罗回归${\sum }_{t\le T}{\gamma }^t{r}_t^a+V(x_T^a)$，其中$((x_t^a,r_t^a){)}_{t=0..T}$来自用动作$a$初始化的轨迹。 选择动作$a$的概率与$\exp (-({\sum }_{t=0..T}{\gamma }^t{r}_t^a+V(x_t^a))/\delta )$成比例，其中$\delta$是学习温度。 这可以被认为是I2A的一种形式，具有固定的摘要summarizer（计算返回值），没有无模型的路径，以及非常简单的策略。 在该架构中，仅学习$V$，$\hat{\pi }$和$\pi$。

我们在Sokoban上运行了这个无编码器的rollout agent，既有精确的环境模型，也有带噪声的环境模型。我们选择rollout的长度为每个环境模型的最佳长度（范围与I2A相同，即从1到5）。如图5（右）、5所示，当使用高精度环境模型时，无编码器agent的性能与基线标准agent的性能相似。然而，与I2A不同，它的性能在使用差模型时会发生灾难性的下降，显示出对模型错误指定的敏感性。

4.3. Further insights into the workings of the I2A architecture

到目前为止，我们已经研究了rollout编码器的作用。为了证明I2A的各种其他成分的重要性，我们进行了额外的控制实验。结果绘制在图4（左）中进行比较。首先，复制模型的I2A（第3.3节）的性能要差得多，这表明环境模型确实是至关重要的。其次，我们训练了一个I2A，其中环境模型预测没有回报，只有观察。这也表现得更糟。然而，经过更长时间的培训（3E9步骤），这些代理确实恢复了接近原始I2A的性能（见附录D.2），即使有这么多步骤，基线代理也从未恢复过。因此，在这项任务中，奖励预测是有帮助的，但不是绝对必要的，仅凭想象的观察就足以获得sokoban的高回报。注意，这与许多经典的规划和基于模型的强化学习方法不同，后者通常依赖于奖励预测。

4.4. Imagination efficiency and comparison with perfect-model planning methods

在前面的章节中，我们说明了I2As可以有效地解决规划问题，并且在模型错误的情况下可以保持健壮。这里，我们问一个不同的问题——如果我们假设一个近乎完美的模型，那么I2A与有竞争力的规划方法相比如何？除了原始性能外，我们还特别关注规划的效率，即解决固定比例水平所需的想象步骤的数量。我们将我们的正则I2A agent与蒙特卡罗树搜索（MCTS）的变体进行了比较，后者是一种现代的引导树搜索算法。对于我们的MCTS实现，我们的目标是通过使用最新的思想建立一个强大的基线：我们包括换位表，并通过使用价值网络（在本例中，是一个使用与I2A相同的总数据量进行训练的深度残差网络，见附录D.3）来评估叶节点的回报。

如果我们假设访问高精度环境模型（包括奖励预测），我们还可以通过使用已训练的I2A执行基本的蒙特卡罗搜索来推动I2A性能：我们让agent播放整集模拟（其中I2A本身使用环境模型进行短期rollouts，因此对应于使用模型内的模型），并且如果找到则执行成功的动作序列，直到达到最多重试次数；这让人想起嵌套的rollouts。固定最多10次重试，得分为95％（I2A本身为87％）。解决一个级别所需的模型仿真步骤的总平均数量，包括在外部循环中运行模型，现在是4k，再次远低于相应的MCTS运行，步长为100k。再次注意，这种方法需要近乎完美的模型;我们不希望I2A与MC搜索在近似模型中表现良好。有关不同方法的想象效率的总结，请参阅表1。

4.5. Generalization experiments

最后，除了处理Sokoban中的随机级别布局之外，我们还探讨了I2As的泛化能力。 我们的agent接受了4个盒子的训练。 表2显示了当在具有不同数量的盒子的水平上测试时的I2A性能，以及用于比较的标准无模型的性能。 我们发现I2As很好地推广; 在7个盒子中，I2A代理仍然能够解决超过一半的水平，几乎与4个盒子上的标准agent一样多。

5. Learning one model for many tasks in MiniPacman

在我们的最后一组实验中，我们演示了如何使用单个模型来为I2A提供对环境动态的一般理解，以解决不同任务的集合。 我们设计了一个名为MiniPacman的简单轻量级域，它允许我们在具有共享状态转换的环境中轻松定义多个任务，这使我们能够进行快速实验。

为了说明基于模型的方法在该多任务设置中的益处，我们训练单个环境模型以预测观察（帧）和事件（如上所定义，例如“吃鬼”）。 请注意，环境模型在所有任务中有效共享，因此学习模型的边际成本为零。 在训练和测试期间，I2A可以访问模型生成的框架和奖励预测; 后者是从模型事件预测和任务奖励向量wrew计算得出的。 因此，奖励向量wrew可以被解释为关于在同一环境中要解决哪个任务的“指令”。 为了公平比较，我们还为所有基线代理提供事件变量作为输入。

我们在每项任务中分别训练基线agents和I2As。 图6（右）中的结果表明了I2A架构的优势，在所有任务中都优于标准agent，并且除了一个任务外，其他所有任务都具有复制模型基线。 此外，我们发现I2As和基线之间的性能差距对于任务4和5来说特别高，其中奖励特别稀疏，并且对鬼魂动态的预期尤其重要。 我们假设I2A代理可以利用其环境和奖励模型更有效地探索环境。

6. Discussion

我们提出了I2A，这是一种将无模型和基于模型的想法结合起来实现想象力增强RL的方法：学习解释环境模型以增强无模型决策。 I2A优于MiniPacman以及Sokoban具有挑战性的组合领域的无模型基线。 我们证明，与基于经典模型的RL和规划方法不同，I2A能够成功地使用不完美模型（包括没有奖励预测的模型），因此显着扩展了基于模型的RL概念和想法的适用性。

作为所有基于模型的RL方法，I2As在行动之前通过思考进行权衡环境交互以进行计算。 这在不可逆转的领域至关重要，在这些领域，行动可能会产生灾难性后果，例如在Sokoban。 在我们的实验中，I2A每次交互总是比无模型基线慢一个数量级。 计算量可以变化（它随着推出的数量和深度线性增长）; 因此，我们期望I2As能够从动态计算资源分配的进步中获益。 未来研究的另一个途径是抽象环境模型：在“正确”的复杂性水平上学习预测模型，并且可以在测试时有效地评估，这将有助于将I2As扩展到更丰富的领域。

值得注意的是，在Sokoban上，I2As与具有完美环境模型的强大规划基线（MCTS）相比毫不逊色：在性能相当的情况下，I2A对模型的函数调用要求远远少于MCTS，因为它们的模型rollout被已学习的rollout策略引导到状态空间的相关部分。 这通过训练rollout政策来指出进一步的潜在改进，这些政策以任务相关的方式“学会查询”不完美的模型。

7. Supplementary material for: Imagination-Augmented Agents for Deep Reinforcement Learning

A. Training and rollout policy distillation details

本文使用的每个agent都定义了一个随机策略，即离散actions $a$上的类别分布$\pi (a_t\mid o_t;\theta )$。以时间步$t$的观测$o_t$为输入，$\pi (a_t\mid o_t;\theta )$的logits由参数$\theta$的神经网络计算。在训练期间，为了增加采取有回报的action的概率，A3C使用策略梯度$g(\theta )$对参数$\theta$进行更新：
$$g(\theta )={▽}_{\theta }\log \pi (a_t\mid o_t;\theta )A(o_t,a_t)$$
其中$A(o_t,a_t)$是估计的优势函数。实际上，我们学习了一个值函数$V(o_t;{\theta }_v)$并用它来计算作为$k$步bootstrap返回值和当前值估计的差异的优势：
$$A(o_t,a_t)=(\sum _{t<t^{\prime }\le t+k}{\gamma }^{t^{\prime }-t}{r}_{t^{\prime }})+{\gamma }^{k+1}V(o_{t+k+1};{\theta }_v)-V(o_t;{\theta }_v)$$
值函数$V(o_t;{\theta }_v)$也计算为参数${\theta }_v$的神经网络的输出。选择值函数网络的输入是计算$\pi$的策略网络的倒数第二层。参数${\theta }_v$通过$k$步bootstrap返回值进行更新：
$$g({\theta }_v)=-A(o_t,a_t){\partial }_{{\theta }_v}V(o_t;{\theta }_v)$$
在实现过程中，我们将上述更新表示为相应的替代损失的梯度。在替代损失中我们加入了熵正则化${\lambda }_{ent}{\sum }_{a_t}\pi (a_t\mid o_t;\theta )\log \pi (a_t\mid o_t;\theta )$以鼓励探索，其中在所有实验中${\lambda }_{ent}=10^{-2}$。在适用的情况下，我们为策略distillation添加一个损失（$\pi$和$\hat{\pi }$之间的交叉熵）：
$$l_{dist}(\pi ,\hat{\pi })(o_t)={\lambda }_{dist}\sum _a\pi (a\mid o_t)\log \hat{\pi }(a\mid o_t)$$
其中${\lambda }_{dist}$是尺度参数。这里$\hat{\pi }$表示我们不会通过行为策略$\pi$反向传播$l_{dist}$的梯度关于rollout策略的参数（Here $\hat{\pi }$ denotes that we do not backpropagate gradients of $l_{dist}$ wrt. to the parameters of the rollout policy through the behavioral policy $\pi$.）。最后，即使我们预先训练了我们的环境模型，原则上我们也可以通过在模型下添加适当的observations的对数似然项来与I2A agent 共同学习它。

B. Agent and model architecture details

我们在所有代理的所有隐藏层之间使用了ReLUs非线性单元。对于环境模型，我们使用Leak ReLUs，斜率为0.01。

B.1. Agents

• Standard model-free baseline agent
  标准的无模型基线代理是一个多层卷积神经网络（CNN），以当前观察为输入，然后是完全连接（FC）隐藏层。该FC层分为两个头：进入FC层，每个动作一个输出，计算策略logits $\log \pi (a_t\mid o_t,\theta )$；并进入另一个FC层得到单一输出，计算值函数$V(o_t;{\theta }_v)$。 层的大小选择如下：
  • for MiniPacman： CNN有两层，都有3x3内核，16个输出通道和步幅1和2；以下FC层有256个单元。
  • for Sokoban： CNN有三层，内核大小为8x8,4x4,3x3，步长为4,2,1，输出通道数为32,64,64; 以下FC有512个单位。
• I2A
  I2A的无模型路径由一个与标准的无模型基线相同的CNN组成（无FC层）。Rollout编码器使用另一个相同大小的CNN处理由环境模型生成的每个帧。然后将该CNN的输出与奖励预测（单个标量广播成帧形状）连接起来。此特征是LSTM的输入，具有512（对于Sokoban）或256（对于Minipacman）单元。相同的LSTM用于处理所有5个rollouts（每个action一个）；所有rollouts的LSTM的最后一个输出连接到单个向量$c_{ia}$中，对于sokoban长度为2560，对于minipacman长度为1280。该向量与无模型CNN路径的输出$c_{mf}$连接，并被馈送到完全连接的层，计算策略logits和值函数，如上述基线代理中所示。
• Copy-model
  复制模型代理具有与I2A完全相同的体系结构，但环境模型被恒等函数替换（不断返回输入观察）。

B.2. Environment models

对于I2A，我们为MiniPacman和Sokoban环境的原始像素观测预先训练1阶自动回归模型（见图7和图8）。在这两种情况下，模型的输入都包括最后一次观测$o_t$，一个broadcasted，和最后一个动作$a_t$的one-hot表示。在先前的研究之后，训练模型的输出以通过网络输出和数据$o_{t+1}$之间的伯努利交叉熵的随机梯度下降来预测下一帧$o_{t+1}$。

Sokoban模型是MiniPacman模型的简化案例；Sokoban模型几乎完全是本地的（除了奖励模型），而MiniPacman模型需要处理非本地交互（鬼魂的移动受到Pacman位置的影响，可以任意远离鬼魂）。

• MiniPacman model
  输入和输出帧的大小为15 x 19 x 3（宽x高x RGB）。该模型如图7所示。它由尺寸保持，多尺度CNN架构和额外的完全连接层组成，用于奖励预测。为了捕获像素之间的长期依赖关系，我们还使用了一个称为pool-and-inject的层，该层将全局最大池应用于每个功能图，并将结果值作为相同大小的功能图广播，并将结果连接到输入。因此，Pool-and-inject layers是保持大小的层，其将每个层的最大值全局地传递到下一个卷积层。
• Sokoban model
  Sokoban模型被选择为残差CNN，具有额外的CNN /完全连接的MLP路径以预测奖励。尺寸为80x80x3的输入首先使用大型8x8的核和8的步幅进行卷积处理。在通过8×8卷积层输出预测帧之前，利用两个保持大小的CNN层进一步处理该减少的表示。

References

[1] http://aigradients.com/2019/06/24/imagination-augmented-agents-for-deep-reinforcement-learning/
[2] https://deepmind.com/research/publications/imagination-augmented-agents-deep-reinforcement-learning/