STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning

📄 arXiv: 2310.09615v1 📥 PDF

作者: Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, Gao Huang

分类: cs.LG

发布日期: 2023-10-14

💡 一句话要点

提出STORM以提高强化学习中的世界模型效率

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 强化学习 世界模型 变换器 变分自编码器 序列建模 效率提升 Atari基准测试

📋 核心要点

  1. 现有的基于模型的强化学习方法在构建复杂环境模型时面临准确性不足的问题，导致代理在真实环境中的表现不佳。
  2. STORM通过结合变换器的序列建模能力与变分自编码器的随机特性，提出了一种高效的世界模型架构，以改善代理的策略学习。
  3. 在Atari 100k基准测试中，STORM达到了126.7%的平均人类表现，并且训练效率显著提升，仅需4.3小时的训练时间。

📝 摘要（中文）

近年来，基于模型的强化学习算法在视觉输入环境中展现出显著的效果。这些方法通过自监督学习构建真实环境的参数化模拟世界模型，利用世界模型的想象力增强代理的策略。然而，构建复杂未知环境的完美模型几乎不可能，模型与现实之间的差异可能导致代理追求虚拟目标，从而在真实环境中表现不佳。本文提出了基于随机变换器的世界模型STORM，结合了变换器的强序列建模和生成能力与变分自编码器的随机特性。STORM在Atari 100k基准测试中实现了126.7%的平均人类表现，创造了不使用前瞻搜索技术的最先进方法的新纪录，同时在单个NVIDIA GeForce RTX 3090显卡上仅需4.3小时的训练时间，显示出相较于以往方法的效率提升。

🔬 方法详解

问题定义：本文旨在解决现有基于模型的强化学习方法在复杂环境中构建准确模型的困难，导致代理在真实环境中表现不佳的问题。

核心思路：STORM的核心思路是结合变换器的强大序列建模能力与变分自编码器的随机特性，通过引入随机噪声来增强模型的表现力，从而提高代理的策略学习效果。

技术框架：STORM的整体架构包括三个主要模块：输入处理模块、变换器序列建模模块和输出生成模块。输入处理模块负责将环境状态转化为模型可接受的格式，变换器模块进行序列建模，而输出生成模块则用于生成代理的策略和动作。

关键创新：STORM的主要创新在于将变换器与变分自编码器结合，形成了一种新的世界模型架构，显著提升了模型的生成能力和适应性，区别于传统的确定性模型。

关键设计：在设计中，STORM采用了特定的损失函数来平衡模型的生成能力与真实环境的匹配度，同时在网络结构上进行了优化，以提高训练效率和模型的稳定性。通过这些设计，STORM能够在较短的时间内实现高效的训练。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

STORM在Atari 100k基准测试中实现了126.7%的平均人类表现，创造了不使用前瞻搜索技术的最先进方法的新纪录。同时，训练一个代理仅需4.3小时，显示出相较于以往方法的显著效率提升。

🎯 应用场景

STORM的研究成果在多个领域具有潜在应用价值，特别是在机器人控制、自动驾驶和游戏智能等领域。通过提高强化学习的效率，STORM能够加速智能体在复杂环境中的学习过程，推动相关技术的进步和应用落地。

📄 摘要（原文）

Recently, model-based reinforcement learning algorithms have demonstrated remarkable efficacy in visual input environments. These approaches begin by constructing a parameterized simulation world model of the real environment through self-supervised learning. By leveraging the imagination of the world model, the agent's policy is enhanced without the constraints of sampling from the real environment. The performance of these algorithms heavily relies on the sequence modeling and generation capabilities of the world model. However, constructing a perfectly accurate model of a complex unknown environment is nearly impossible. Discrepancies between the model and reality may cause the agent to pursue virtual goals, resulting in subpar performance in the real environment. Introducing random noise into model-based reinforcement learning has been proven beneficial. In this work, we introduce Stochastic Transformer-based wORld Model (STORM), an efficient world model architecture that combines the strong sequence modeling and generation capabilities of Transformers with the stochastic nature of variational autoencoders. STORM achieves a mean human performance of $126.7\%$ on the Atari $100$k benchmark, setting a new record among state-of-the-art methods that do not employ lookahead search techniques. Moreover, training an agent with $1.85$ hours of real-time interaction experience on a single NVIDIA GeForce RTX 3090 graphics card requires only $4.3$ hours, showcasing improved efficiency compared to previous methodologies.