2025_NIPS_STORM: Efficient Stochastic Transformer based World Models for Reinforcement Learning

STORM论文总结与核心部分翻译

一、文章主要内容

本文提出了一种基于随机Transformer的世界模型（STORM），用于模型基强化学习（model-based RL）。该模型旨在解决深度强化学习（DRL）中样本效率低、复杂环境建模难以及模型与现实差异导致的性能下降等问题。

核心背景与挑战

1. 传统模型无关RL（如Rainbow、PPO）样本效率低，需大量环境交互，且值函数估计困难、稀疏奖励场景下有效信息不足；
2. 现有模型基RL虽通过自监督学习构建环境仿真模型、利用想象经验优化策略，但存在自回归过程中预测误差累积、RNN类序列模型训练速度慢等问题；
3. 引入随机噪声可缓解模型与现实的差异，但现有Transformer-based世界模型未充分发挥其并行计算优势，且训练耗时较长。

模型结构与方法

1. 世界模型设计：采用分类变分自编码器（categorical VAE）作为图像编码器，将高维图像转换为低维随机潜变量，增强鲁棒性并减少预测误差；融合Transformer作为序列模型，通过掩码自注意力机制捕捉长序列依赖，支持并行计算；将潜变量与动作融合为单个token输入Transformer，提升建模效率。
2. 损失函数：采用端到端自监督训练，总损失包含图像重建损失、奖励预测损失、持续标志预测损失、动力学损失和表示损失，通过KL散度约束潜变量分布与预测分布的一致性。