该文章介绍了一种基于模型的强化学习算法中的世界模型结构，名为“STORM：Efficient Stochastic Transformer based World Models for Reinforcement Learning”。该模型旨在解决深度强化学习（DRL）算法在现实环境应用中的样本效率低下问题。通过将Transformer强大的序列建模和生成能力与变分自动编码器的随机性相结合，STORM模型提高了智体的策略效率和样本效率。在Atari 100k基准测试中，达到了人类平均水平的126.7%，并创下了不使用前瞻搜索技术的新纪录。
文章介绍了基于模型的强化学习算法的研究背景，指出该算法在视觉输入环境中表现出显著效果，并概述了相关领域的挑战和现状。
文章提出了一种新的世界模型结构——基于随机Transformer的世界模型（STORM）。该模型结合了Transformer的序列建模和生成能力与变分自动编码器的随机性，旨在提高智体的策略效率和样本效率。
文章详细描述了STORM模型的构建过程，包括与其他方法的比较、模型结构的设计、训练过程以及实验设置等。特别强调了模型在序列建模、价值函数估计、样本效率等方面的优势。
文章展示了STORM模型在Atari 100k基准测试中的表现，达到了人类平均水平，并与其他方法进行了比较。同时，文章还提到了模型在其他领域的应用前景和可能的改进方向。
文章总结了基于模型的强化学习算法的重要性和挑战，强调了STORM模型的优点和创新之处，并展望了未来的研究方向。
23年10月来自北理工和清华的论文“STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning”。
最近，基于模型的强化学习算法，在视觉输入环境中表现出了显著的效果。这些方法首先通过自监督学习构建真实环境的参数化模拟
世界模型
。通过利用世界模型的想象力，智体的策略得到了增强，而不受从真实环境中采样的限制。这些算法的性能在很大程度上依赖于世界模型的序列建模和生成能力。然而，构建一个复杂未知环境的完美准确模型几乎是不可能的。模型和现实之间的差异可能会导致智体追求虚拟目标，而在真实环境中的表现不佳。在基于模型的强化学习中引入随机噪声已被证明是有益的。在这项工作中，引入了基于随机 Transformer 的世界模型 (
STORM
)，这是一种高效的世界模型架构，它将 Transformer 强大的序列建模和生成能力与变分自动编码器的随机性相结合。STORM 在 Atari 100k 基准测试中达到人类平均水平的 126.7%，创下不使用前瞻搜索技术的先进方法新纪录。此外，在单个 NVIDIA GeForce RTX 3090 显卡上训练具有 1.85 小时实时交互经验的智体仅需 4.3 小时，与之前的方法相比，效率有所提高。
深度强化学习 (DRL) 已在不同领域取得了显著成功。然而，实现这种成功需要与环境进行大量交互，这阻碍了它在现实环境中的广泛应用。当处理更广泛的现实环境（例如缺乏可调速度模拟工具的无人驾驶和制造系统 [1, 2]）时，这种限制变得尤为具有挑战性。因此，提高样本效率已成为 DRL 算法面临的一个关键挑战。
流行的 DRL 方法（包括 Rainbow [3] 和 PPO [4]）由于两个主要原因而存在样本效率低下的问题。首先，价值函数的估计是一项具有挑战性的任务。这涉及使用深度神经网络 (DNN) 近似价值函数，并使用 n 步引导时间差（TD）对其进行更新，这自然需要多次迭代才能收敛 [5]。其次，在奖励稀疏的场景中，许多样本在价值函数方面表现出相似性，为 DNN 的训练和泛化提供有限的有用信息 [6, 7]。这进一步加剧了提高 DRL 算法样本效率的挑战。
为了应对这些挑战，
基于模型的 DRL
算法应运而生，成为一种有前途的方法，它可以同时解决这两个问题，同时在样本效率高的环境中表现出显著的性能提升。这些算法首先通过自监督学习构建真实环境的参数化模拟世界模型。自监督学习可以通过多种方式实现，例如使用解码器重建原始输入状态 [8–10]、预测帧之间的动作 [7] 或采用对比学习来捕捉输入状态的内部一致性 [6, 7]。这些方法比传统的
无模型 RL
损失提供了更多的监督信息，增强了 DNN 的特征提取能力。随后，通过利用使用世界模型生成的经验来改进智体的策略，消除采样约束，并与无模型算法相比更快地更新价值函数。
然而，使用世界模型进行想象的过程，涉及一个自回归过程，该过程会随着时间的推移累积预测误差。当想象轨迹与真实轨迹出现差异时，智体可能会无意中追求虚拟目标，导致在真实环境中的表现不佳。为了缓解这个问题，在世界模型中引入随机噪声已被证明是有益的 [9–11, 14]。变分自编码器能够自动学习高维数据的低维潜表示，同时将合理的随机噪声纳入潜空间，为图像编码提供了理想的选择。
人们已经进行了大量努力来构建高效的世界模型。例如，SimPLe [11] 利用 LSTM [15]，而 DreamerV3 [10] 采用 GRU [16] 作为序列模型。LSTM 和 GRU 都是循环神经网络 (RNN) 的变型，在序列建模任务方面表现出色。然而，RNN 的循环性质阻碍并行计算，导致训练速度变慢 [17]。相比之下，Transformer 架构 [17] 最近在各种序列建模和生成任务中都表现出优于 RNN 的性能。它克服了忘记长期依赖关系的挑战，专为高效的并行计算而设计。虽然已经进行了多次尝试将 Transformer 纳入世界模型 [12、13、18]，但这些工作并未充分利用该架构的功能。此外，这些方法需要更长的训练时间，并且无法超越基于 GRU 的 DreamerV3 的性能。
基于随机 Transformer 的世界模型 (STORM)，是一种基于模型的 RL 的高效结构。世界模型最近的方法和STORM的比较：
SimPLe [11] 和 Dreamer [10] 依赖于基于 RNN 的模型，而 STORM 采用类似 GPT 的 Transformer [30] 作为序列模型。
与使用多个 token 的 IRIS [13] 相比，STORM 使用单个随机潜变量来表示图像。
STORM 遵循一个 vanilla Transformer[17] 结构，而 TWM[12] 采用 Transformer-XL [21] 结构。
在 STORM 的序列模型中，观察和动作融合为单个 token，而 TWM [12] 将观察、动作和奖励视为三个同等重要的独立 token。
与包含隐藏状态的 Dreamer [10] 和 TransDreamer [18] 不同，STORM 无需利用此信息即可重建原始图像。
如图所示 STORM 和其他方法在 Atari 100k 上的方法比较：SimPLe [11] 和 DreamerV3 [10] 使用 RNN 作为其世界模型，而 TWM [12]、IRIS [13] 和 STORM 使用 Transformer；单个 NVIDIA V100 GPU 上的每秒训练帧数 (FPS) 结果是从 SimPLe、TWM 和 IRIS 的其他显卡推断出来的，而 DreamerV3 和 STORM 则是直接评估的。
下表STORM与其他方法的比较：
该方法遵循基于模型强化学习算法的既定框架，该算法专注于通过想象力增强智体的策略 [5、9-11、13]。
迭代以下步骤，直到达到规定的真实环境交互次数：
在每个时间 t，数据点包含一个观察 ot、一个动作 at、一个奖励 rt 和一个延续标志 ct（一个布尔变量，指示当前情节是否正在进行）。重放缓冲区维护先进先出的队列结构，从而能够从缓冲区中采样连续的轨迹。
该世界模型的完整结构如图所示。实验中，专注于 Atari 游戏 [31]，它会生成环境的图像观测 ot。直接在原始图像上对环境动态进行建模计算成本高昂且容易出错 [7–11, 13, 23]。
为了解决这个问题，利用 VAE [32] 将 ot 转换为潜随机分类分布 Zt。
与之前的研究 [9, 10, 12] 一致，将 Zt 设置为包含 32 个类别的随机分布，每个类别有 32 个级。
编码器 (qφ) 和解码器 (pφ) 结构实现为卷积神经网络 (CNN) [33]。
随后，从 Zt 中采样一个潜变量 zt 来表示原始观测 ot。
由于从分布中采样缺少用于后向传播的梯度，应用直通梯度技巧 [9, 34] 来保留它们。
在进入序列模型之前，用多层感知器 (MLP) 和连接（concatenation）将潜样本 zt 和动作 at 组合成单个 token et。
此操作表示为 mφ，为序列模型准备输入。
序列模型 fφ 将 et 的序列作为输入并产生隐藏状态 ht。
对序列模型采用类似 GPT Transformer 结构 [30]，其中自注意模块被后续掩码所掩盖，从而允许 et 关注序列 e1, e2, ..., et。
通过利用 MLP gφD、gφR 和 gφC，依靠 ht 来预测当前奖励 rtˆ、延续标志 ctˆ 和下一个分布 Zˆt+1。
