STORM：基于高效随机变换器的强化学习世界模型
STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning
Weipu Zhang, Gang Wang, Jian Sun, Yetian Yuan, Gao Huang
neural information processing systems · 2023-01-01
摘要
近年来，基于模型强化学习算法在视觉输入环境中展现出卓越的效能。这类方法首先通过自监督学习构建真实环境的参数化模拟世界模型，借助世界模型的“想象力”来优化智能体的策略，从而摆脱对真实环境采样的限制。此类算法的性能高度依赖于世界模型的序列建模与生成能力。然而，要构建一个完全精准的复杂未知环境模型几乎是不可能的；模型与现实之间的偏差可能导致智能体追求虚拟目标，进而使其在真实环境中的表现不佳。研究已证实，在基于模型的强化学习中引入随机噪声能够带来积极效果。在本研究中，我们提出了基于随机变换器的世界模型（STORM），这是一种高效的世界模型架构，它将变换器强大的序列建模与生成能力与变分自编码器的随机特性相结合。在Atari 100k基准测试中，STORM的平均人类表现水平达到了126.7%，在无需前瞻搜索技术的最新方法中创下新纪录。此外，仅需4.3小时即可在单块NVIDIA GeForce RTX 3090显卡上完成基于1.85小时实时交互数据的智能体训练，其效率较以往方法有明显提升。
本摘要由 AI（Hy-MT2-Pro）翻译生成，仅供参考，请以原文为准。
AI 综述
基于全文生成- 主要观点
- 论文将问题域定位为视觉输入环境下的基于模型的强化学习，核心论点是：通过自监督学习构建参数化世界模型，并仅用其想象轨迹改进策略，可缓解无模型RL对真实环境交互的依赖；而世界模型的序列建模与生成能力决定算法表现。作者指出复杂未知环境无法被完全准确建模，模型与现实差异会使智能体追逐虚拟目标，因此在世界模型中引入随机噪声有益。STORM的具体主张是以分类变分自编码器提供随机潜变量、以GPT类Transformer作为序列模型，并将潜变量与动作融合为单一token，从而兼得Transformer的序列建模/生成能力与VAE的随机性。实验主张是：该结构在Atari 100k上取得126.7%平均人类归一化分数，在不使用lookahead search的方法中创下纪录，同时在单张RTX 3090上以4.3小时完成对应1.85小时实时交互经验的智能体训练。论文还主张注意力机制对奖励相关大目标或多目标场景（如Amidar、MsPacman、Chopper Command、Gopher）有利，而自编码器与采样随机性对Pong、Breakout这类单一小移动目标不利。
- 研究方法
- 采用基于模型DRL的三步迭代流程：执行当前策略收集真实环境数据入回放缓冲区，从缓冲区采样轨迹更新世界模型，以缓冲区采样点为起点用世界模型想象经验训练策略。世界模型以CNN实现分类VAE编码器/解码器，潜变量为32个类别、每类32类的随机分布，经straight-through梯度采样；z_t与a_t经MLP和拼接融合为单token，输入带后续掩码的GPT类Transformer，隐藏状态经MLP分别预测奖励、延续标志和下一潜分布。训练损失由图像重构、symlog two-hot奖励预测、延续标志二元交叉熵、动力学KL与表示KL组成，β1=0.5、β2=0.1并使用stop-gradient。策略学习仅基于想象：从回放缓冲区取短上下文计算初始后验，推理时从先验采样并用KV cache加速，状态为[z_t,h_t]，actor-critic沿用DreamerV3设置，含λ-return、百分位归一化和critic EMA。评测在Atari 100k的26个游戏上进行，100k样本约合400k帧、1.85小时实时游玩；用5个随机种子、每2500样本步保存检查点、每检查点20个评估回合并取最终检查点均值，与SimPLe、TWM、IRIS、DreamerV3比较且不直接比较MuZero/EfficientZero等lookahead方法；另做模块重排与Transformer层数消融，材料还提及用单条示范轨迹经世界模型想象提升探索效率但未展开细节。
- 学术贡献
- 贡献首先在方法/架构层面：STORM把GPT类Transformer序列模型与分类VAE的随机潜变量结合，并用单一随机潜变量表示图像、将观测与动作融合为单token，区别于IRIS的多token、TWM的Transformer-XL及观测/动作/奖励三token设计，也区别于Dreamer/TransDreamer依赖隐藏状态的重构。其次在基准评测层面，它在Atari 100k上报告126.7%平均人类归一化分数和58.4%中位数，在不使用lookahead search的同类方法中刷新纪录，并给出单卡RTX 3090上4.3小时训练的效率数据，对样本高效基于模型RL的世界模型子领域有直接比较价值。第三，消融研究给出可复用的设计结论：重构损失应直接施加于编码器输出而非序列模型输出，奖励与延续预测应使用隐藏状态而非单帧潜变量，2层Transformer已足够且加深层数未提升最终性能。第四，论文公开代码，材料中未明确提及新的理论定理或材料数据集发掘。
- 存在局限
- 实验范围限于Atari 100k的26个图像游戏与离散动作，材料未报告连续控制、真实机器人或材料引言所述无人系统/制造系统中的验证，因此架构的外推边界未闭合。评测刻意不直接比较MuZero、EfficientZero等lookahead search方法，126.7%的纪录仅在“不使用lookahead search”的方法集合内成立，不能据此推断其优于带搜索的方法。论文承认自回归想象会累积预测误差，且Pong、Breakout等单一小移动目标上因自编码器特性和采样随机性影响注意力权重而表现不佳，说明该结构对目标尺度与随机噪声的鲁棒性并不均匀。消融主要围绕模块重排和2/4/6层Transformer，在Ms Pacman、Pong等有限任务上展示，材料未给出多种子置信区间或统计检验；对加深层数无收益给出的三个原因（相邻帧差异小、数据不足、表示损失影响编码器）属于解释性假设，未通过控制变量实验逐一验证。材料还明确未深入比较带自监督辅助损失的model-free方法，因而世界模型相对这些方法的增益归因未被隔离。
- 进一步研究方向
- 其一，将STORM与MCTS等lookahead search结合，在Atari 100k同一100k样本预算下比较搜索深度/模拟次数与分数的关系，并报告相对纯STORM的额外计算开销，以检验材料所述“未来可组合”的增益。其二，针对Pong、Breakout的失败，设计受控消融比较单随机潜变量与多token表示、不同KL权重对注意力权重和小目标重建的影响，验证小目标信息保留是否改善。其三，利用材料提及的DeepMind Control、Minecraft等环境检验Transformer层数、数据量与表示损失权重三者对性能的交互，直接检验“加深层数无益源于数据不足/表示损失影响编码器”的解释。其四，把材料第5.3节仅提及的单条示范轨迹探索方案在26个游戏上系统评估，变化示范轨迹长度与质量，测量探索效率与最终人类归一化分数；同时与带自监督辅助损失的model-free基线并列，隔离世界模型贡献。
本综述由 AI（deepseek-flash）阅读开放获取全文自动生成，仅供参考，不代表原文作者观点，请以原文为准。
参考文献
相关论文
话题：Reinforcement Learning in Robotics