基于 Transformer 的世界模型可以满足 100K 次交互的要求

23年3月来自多特蒙德大学的论文“Transformer-based world models are happy with 100k interactions”。

深度神经网络已在许多强化学习环境中取得了成功。然而，与人类学习者相比，它们过于渴求数据。为了构建一个样本高效的世界模型，以自回归的方式将Transformer应用于现实世界的事件：不仅将紧凑的潜状态和采取的行动，而且将经验或预测的奖励输入到Transformer中，以便它可以灵活地在不同时间步骤中关注所有三种模态。Transformer允许世界模型直接访问先前的状态，而不是通过压缩的循环状态查看它们。通过利用 Transformer-XL 架构，它能够学习长期依赖关系，同时保持计算效率。基于Transformer的世界模型 (TWM) 生成有意义的新经验，用于训练策略，该策略在 Atari 100k 基准上的表现，优于以前的无模型和基于模型的强化学习算法。

世界模型由一个观察模型和一个动态模型组成，它们不共享参数。如图所示TWM世界模型架构。使用 CNN 对观测 ot-l:t 进行编码。随机离散潜状态 zt-l:t、动作 at-l:t 和奖励 rt-l:t 的线性嵌入被输入到Transformer中，Transformer在每个时间步骤计算确定性隐状态 ht。使用 MLP 基于 ht 计算奖励 rt、折扣因子 γt 和下一个潜状态 zt+1 的预测。

观测模型是一个变分自动编码器（Kingma & Welling，2014），它将观测值 ot 编码为紧凑的随机潜状态 zt，并用解码器重建观测值，这里只需要获得 zt 的一个学习信号。

采用 DreamerV2（Hafner，2021）的神经网络架构，并对观察模型进行轻微修改。因此，潜状态 zt 是离散的，由 32 个分类变量的向量组成，每个类别有 32 个级。观察解码器，重建观察并预测所有像素的独立标准正态分布均值。观察模型的作用，是仅捕获有关当前时间步骤的非时间信息，这与 （Hafner 2021）不同。但是，包括短时时域信息，因为单个观察 ot 由四帧组成（又名帧堆叠）。

动态模型根据其过去预测的历史来预测下一个时间步。主干是一个确定性聚合模型 fψ，它根据 l 个先前生成的潜状态、动作和奖励的历史计算确定性隐状态 ht。奖励、折扣和下一个潜状态的预测器，取决于隐状态。

聚合模型实现为因果掩码 Transformer-XL（Dai，2019），它通过递归机制和相对位置编码增强 vanilla Transformer（Vaswani，2017）。通过这些编码，世界模型可以学习独立于绝对时间步的动态。按照 （Chen 2021）的做法，潜状态、动作和奖励在传递给 Transformer 之前被发送到特定模态的线性嵌入中。输入tokens的数量为 3l-1，因为三种模态（潜状态、动作、奖励）和最后一个奖励不是输入的一部分。将动作模态的输出视为隐状态，并忽略其他两种模态的输出（上图中的橙色框与灰色框）。

潜状态、奖励和折扣预测器被实现为多层感知器 (MLP)，并分别计算独立分类分布向量、正态分布和伯努利分布的参数，这些参数取决于确定性隐状态。下一个状态从 p (zˆt+1 | ht) 中采样确定。奖励和折扣分别由 pψ (rˆt |ht) 和 pψ (γˆt |ht) 的平均值确定。

由于这些设计选择，世界模型具有以下有益特性：

1. 动态模型是自回归的，可以直接访问其先前的输出。
2. 由于序列是并行处理的（与 RNN 相比），因此训练效率高。
3. 由于输出缓存（与 vanilla Transformers 相比），因此推理效率高。
4. 可以通过递归机制捕获长期依赖关系。

全自回归动态模型的好处：
首先，与 RNN 相比，直接访问先前的潜状态可以对它们之间更复杂的依赖关系进行建模，而 RNN 只能通过压缩的循环状态间接看到它们。这也有可能使推理更加稳健，因为退化的预测更容易被忽略。
其次，由于模型可以看到它之前产生的奖励，因此它可以对自己的预测做出反应。当奖励是从概率分布中采样时，这一点更为重要，因为没有自回归就无法观察到引入的噪声。

优化观察模型的自监督损失函数，该函数是解码器损失、熵正则化器和一致性损失的预期总和

对于平衡的交叉熵损失，最小化动态模型损失中的交叉熵，这就是如何训练潜状态预测器的方式。奖励和折扣预测器通过负对数似然进行优化。这个导致一个动态模型的自监督损失如下：

该策略 πθ (at | zˆt ) 是在想象的轨迹上训练的，主要使用标准优势 A-C (Mnih，2016) 方法。训练两个独立的网络：一个演员 at ∼ π (at | zˆt )，具有一个参数 θ 和一个评论家 vξ (zˆt )。通过广义 ξt 优势估计 (Schulman，2016) 计算优势，同时使用由世界模型 γˆt 预测的折扣因子，而不是所有时间步骤的固定折扣因子。与 DreamerV2 (Hafneretal.，2021) 一样，通过折扣因子的累积乘积来加权 A-C 的损失，以便软解释情节结束。

该策略在给定状态某个视图 xt 的情况下计算动作分布 πθ(at | xt)。例如，在推理时（即应用于真实环境时），xt 可以是 ot、zt 或 [zt, ht]，或者在训练时是世界模型 oˆt 、zˆt 或 [zˆt, ht] 的相应预测。必须谨慎选择此视图，因为它会对策略的性能产生重大影响，并影响世界模型的设计选择。即使 oˆt 的重建不完善，使用 xt = ot（或 oˆt ）也相对稳定，因为观察 p(o) 的底层分布在训练期间不会发生变化。但是，它的计算效率也较低，因为它需要在想象期间重建观察结果，并为策略添加额外的卷积层。使用 xt = zt（或 zˆt ）稍微不太稳定，因为策略必须在训练期间适应分布 pφ（zt | ot ）和 pψ（zˆt | ht ）的变化。

尽管如此，等式 (3) 中的熵正则化器和一致性损失，可以稳定这些分布。使用 xt = [zt , ht ]（或 [zˆt, ht ]）为智体提供了经验历史的摘要，但它也增加了在推理时运行Transformer的负担。当使用一个最新帧的堆叠时，无模型智体已经在大多数 Atari 游戏中表现良好（例如，Mnih 2015；Schulman 2017）。因此，选择 xt = zt 并在推理时应用帧堆叠，以便将短时信息直接合并到潜状态中。在训练时，用 xt = zˆt，即预测的潜状态，这意味着不应用帧堆叠。因此，策略在训练时（想象期间无需重建）和推理时（在真实环境中运行时无需Transformer）都具有计算效率。

与使用世界模型进行学习一样，反复以下步骤： (i) 使用当前策略在真实环境中收集经验，(ii) 使用过去的经验改进世界模型，(iii) 使用世界模型生成的新经验改进策略。

在训练期间，构建一个收集的经验数据集 D = [(o1, a1, r1, d1),…,(oT, aT, rT, dT)]。在使用当前策略收集新经验后，通过从 D 中采样长度为 l 的 N 个序列并使用随机梯度下降优化方程 (3) 和 (4) 中的损失函数来改进世界模型。在执行世界模型更新后，从 N × l 个观测中选择 M 个，并将其编码为潜状态，作为新轨迹的初始状态。根据策略提供的动作，动态模型迭代生成长度为 H 的 M 条轨迹。随后，使用标准无模型目标改进策略。在如下算法 1 中，提供用于训练世界模型和策略的伪代码。