STORM: Transformer-based World Model for Planning
一句话定位
STORM是一种基于Transformer的世界模型，利用自注意力机制进行时序建模和规划，在高维观察空间中实现高效的世界模型学习和多步规划能力。
前置依赖
- World Models (2019): Haarnoja et al. 的VAE + MDN-RNN分离式世界模型
- Dreamer (2020): RSSM架构下的世界模型学习
- MuZero (2020): 隐式世界模型 + MCTS规划
- Transformer (2017): Vaswani et al. 的自注意力机制
- Decision Transformer (2021): 将Transformer应用于序列决策
- RSSM (Recurrent State Space Model): Dreamer中的循环状态空间模型
核心思想
STORM的核心洞察是：Transformer的自注意力机制天然适合建模长程依赖和复杂的状态交互。
传统RSSM-based世界模型的局限：
- RNN的局部性: LSTM/GRU的隐藏状态是压缩的全局向量，长程依赖建模困难
- 固定容量: RNN隐藏维度固定，无法根据任务动态调整
- 顺序处理: 无法并行化时序模拟
STORM使用Transformer替代RNN作为世界模型核心：
- 并行时序模拟: 一次前向计算H步想象轨迹
- 自适应感受野: 注意力机制自动关注相关状态
- 长程建模能力: 适合长视界规划任务
模型结构图
┌─────────────────────────────────────────────────────┐
│ STORM 整体架构 │
└─────────────────────────────────────────────────────┘
输入观察序列:
o_0, o_1, o_2, ..., o_t
│ │ │ │
▼ ▼ ▼ ▼
┌───────────────────────────────────────────────────────────────────────┐
│ 编码器 (Encoder) f │
│ 将每个观察编码为潜在状态 z_t = f(o_t) │
└───────────────────────────────────────────────────────────────────────┘
│
▼
┌───────────────────────────────────────────────────────────────────────┐
│ Transformer世界模型 g (核心创新) │
│ │
│ ┌─────────────────────────────────────────────────────────────┐ │
│ │ Transformer Block │ │
│ │ LayerNorm ──▶ Multi-Head Self-Attention ──▶ MLP │ │
│ └─────────────────────────────────────────────────────────────┘ │
│ │
│ 输入: [z_t, a_{t-1}, r_{t-1}, h_{t-1}] │
│ 输出: 下一状态预测和奖励预测 │
│ │
│ (使用GPT风格的causal attention, 只Attend到过去状态) │
└───────────────────────────────────────────────────────────────────────┘
│
▼
┌───────────────────────────────────────────────────────────────────────┐
│ 预测头 (Prediction Heads) │
│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │
│ │ 策略头 π │ │ 价值头 V │ │ 奖励头 h │ │
│ │ π(a_t|z_t) │ │ V(z_t) │ │ r(z_t,a_t) │ │
│ └──────────────┘ └──────────────┘ └──────────────┘ │
└───────────────────────────────────────────────────────────────────────┘
STORM的Transformer架构详解
时间步: t=0 t=1 t=2 ... t=H
│ │ │ │
▼ ▼ ▼ ▼
┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐
│ z_0 │ │ z_1 │ │ z_2 │ ... │ z_H │
└─────┘ └─────┘ └─────┘ └─────┘
│ │ │ │
▼ ▼ ▼ │
┌─────┐ ┌─────┐ ┌─────┐ │
│ a_0 │ │ a_1 │ │ a_2 │ │
└─────┘ └─────┘ └─────┘ │
│ │ │ │
▼ ▼ ▼ │
┌─────┐ ┌─────┐ ┌─────┐ │
│ r_0 │ │ r_1 │ │ r_2 │ │
└─────┘ └─────┘ └─────┘ │
│ │ │ │
▼ ▼ ▼ ▼
══════════════════════════════════════════════════════════════
║ Transformer Block (causal) ║
║ ║
║ Cross-Attention: Query=z_t, Key/Value=[z_{<t}, a_{<t}] ║
║ ║
══════════════════════════════════════════════════════════════
│ │ │ │
▼ ▼ ▼ ▼
┌─────┐ ┌─────┐ ┌─────┐ ┌─────┐
│ z_1'│ │ z_2'│ │ z_3'│ ... │ V_H │ (价值)
│ r_0 │ │ r_1 │ │ r_2 │ └─────┘
└─────┘ └─────┘ └─────┘
数学推导
1. 编码器
编码器
STORM使用轻量级CNN编码器（区别于MuZero的深度ResNet）：
2. Transformer世界模型
设世界模型
输出：
2.1 自注意力机制
对输入序列中的位置
其中
2.2 因果注意力(Causal Attention)
为确保只使用过去信息构建未来：
这确保
2.3 多头注意力
通常
3. 状态转移建模
STORM使用Transformer建模完整的状态-动作-奖励序列：
关键优势：
- 任意时间步的状态可以通过Attention直接获取，无需RNN的隐藏传播
- 支持跳跃连接: 直接从
和 建立依赖关系（ 可以很大）
4. 规划与想象轨迹
STORM的规划通过并行想象轨迹实现：
给定初始状态 z_0:
# 并行生成H步想象
Z = [z_0] # 状态序列
A = [] # 动作序列
R = [] # 奖励序列
for h in range(H):
# 使用Transformer预测下一状态和奖励
(z_pred, r_pred) = transformer(Z, A, R)
# 采样动作
a_h = policy_head(z_pred)
Z.append(z_pred)
A.append(a_h)
R.append(r_pred)
# 计算累积回报
G = Σ γ^i R[i] + γ^H V(Z[H])
5. 与RSSM的比较
RSSM使用RNN进行状态转移：
STORM使用Transformer：
关键差异：
| 特性 | RSSM (Dreamer) | STORM |
|---|---|---|
| 状态传递 | 隐藏向量 | 完整序列 |
| 依赖建模 | 局部(RNN hidden) | 全局(Attention) |
| 并行性 | 顺序 | 完全并行 |
| 容量 | 固定 | 动态(序列长度×维度) |
| 长程依赖 | 困难 | 自适应 |
6. 训练目标
STORM的联合训练目标：
6.1 价值损失
其中
6.2 策略损失
使用策略梯度估计
6.3 奖励损失
6.4 动力学重建损失
7. Transformer与规划的数学洞见
Transformer的核心优势可以从数学上解释：
设序列长度
RNN的表示容量:
Transformer的表示容量:
这意味着：
- 长视界任务: Transformer可以精确保持长期记忆
- 跳步依赖: 允许
直接依赖 ，不经过中间传播 - 条件生成: 给定任意状态子集，可以精确推断其他状态
推理过程
Planning循环
def plan_with_storm(model, z_init, horizon=20, num_actions=1):
"""
STORM的规划过程
model: STORM Transformer世界模型
z_init: 初始编码状态
horizon: 想象视界
"""
# 初始化序列
Z = [z_init] # 状态序列
A = [] # 动作序列
R = [] # 奖励序列
for h in range(horizon):
# 使用因果Transformer预测
# 输入: [z_0, a_0, r_0, z_1, a_1, r_1, ..., z_h]
context = build_sequence(Z, A, R)
# Transformer前向
z_pred, r_pred = model(context)
# 策略头采样动作
a = model.policy_head(z_pred)
# 追加到序列
Z.append(z_pred)
A.append(a)
R.append(r_pred)
# 返回第一个动作
return A[0]
树搜索变体
STORM也可以结合树搜索：
- 从根节点
开始 - 使用Transformer估计先验
- 对每个动作，Transformer计算
- 评估
作为backup值 - UCB选择下一动作
优点与局限
优点
- 长程建模能力强: 自注意力直接建模任意时间步依赖
- 并行计算效率高: Transformer可并行处理整个序列
- 表示容量大:
vs RNN的 - 可解释性强: Attention weights显示状态间相关性
- 迁移学习友好: 预训练的GPT/PViT可直接使用
- 条件生成: 可在任意历史节点注入观测
局限
- 计算复杂度高: Self-Attention是
- 内存随序列增长: 长视界规划需要大量内存
- 位置编码问题: 时序位置需要精心设计
- 训练不稳定: 大型Transformer训练需要高级优化技巧
- 局部交互建模弱: 可能过度关注全局而忽略局部结构
与前后内容的衔接
前置依赖
STORM继承了多代世界模型的发展：
- World Models: 分离式世界模型架构(encoder/dynamics/reward)
- Dreamer RSSM: 状态空间模型的建模思想
- MuZero: 隐式世界模型+规划框架
- Decision Transformer: 将Transformer引入序列决策
后置发展
STORM的思想影响了后续工作：
- TransDreamer (2022): 在Dreamer架构中用Transformer替代RNN
- Video Transformer World Models (2024): 处理视频输入的Transformer世界模型
- Genie (2024): 基于Transformer的通用世界模型
可复现实现要点
核心Transformer实现
class TransformerWorldModel(nn.Module):
def __init__(self, latent_dim, action_dim, d_model=256, nhead=8, n_layers=4):
super().__init__()
self.d_model = d_model
# 嵌入层
self.state_embed = nn.Linear(latent_dim, d_model)
self.action_embed = nn.Linear(action_dim, d_model)
self.reward_embed = nn.Linear(1, d_model)
# Transformer
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=4*d_model,
dropout=0.1
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
# 预测头
self.state_head = nn.Linear(d_model, latent_dim)
self.reward_head = nn.Linear(d_model, 1)
self.policy_head = nn.Linear(d_model, action_dim)
self.value_head = nn.Linear(d_model, 1)
def forward(self, states, actions, rewards):
# states: [B, L, D]
# actions: [B, L, A]
# rewards: [B, L, 1]
# 嵌入
s_emb = self.state_embed(states)
a_emb = self.action_embed(actions)
r_emb = self.reward_embed(rewards)
# 组合输入 (causal mask自动应用)
x = s_emb + a_emb + r_emb
# Transformer前向
h = self.transformer(x)
# 预测下一步
z_next = self.state_head(h[:, :-1]) # 预测状态
r_pred = self.reward_head(h[:, :-1]) # 预测奖励
pi = self.policy_head(h[:, -1:]) # 策略（最后位置）
V = self.value_head(h[:, -1:]) # 价值（最后位置）
return z_next, r_pred, pi, VCausal Mask实现
def create_causal_mask(seq_len, device):
"""创建因果掩码，确保只Attend到过去"""
mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1)
mask = mask.masked_fill(mask == 1, float('-inf'))
return mask.to(device)Planning实现
def storm_plan(model, z_init, horizon=20):
"""STORM的多步规划"""
B = z_init.size(0)
# 初始化序列
states = [z_init]
actions = []
rewards = []
for h in range(horizon):
# 打包序列
seq_len = len(states)
state_seq = torch.stack(states, dim=1) # [B, H, D]
action_seq = torch.stack(actions, dim=1) if actions else torch.zeros(B, 0, model.action_dim)
reward_seq = torch.stack(rewards, dim=1) if rewards else torch.zeros(B, 0, 1)
# 获取下一步预测
with torch.no_grad():
_, r_pred, a_pred, V_pred = model(state_seq, action_seq, reward_seq)
actions.append(a_pred[:, -1])
rewards.append(r_pred[:, -1])
# 使用dynamics预测下一状态
z_next = model.state_predict(states[-1], actions[-1])
states.append(z_next)
return actions[0] # 返回第一个动作章节摘要
STORM是将Transformer引入世界模型学习的重要尝试，核心创新在于用自注意力替代RNN进行状态转移建模。相比Dreamer的RSSM，STORM具有更强的长程建模能力和并行计算优势；相比MuZero的MCTS，STORM的想象轨迹更加高效且可并行生成。Transformer的全局感受野使其在需要长期记忆的任务中表现出色，但也带来了
关键词
- STORM, Transformer世界模型, 自注意力, 时序建模
- 因果注意力, 并行想象轨迹, 规划
- RSSM对比, 长程依赖, 状态空间模型