想象与未来预测:基本思想
一句话定位:利用训练好的世界模型生成虚拟轨迹(想象 rollout),使智能体能够在无需真实环境交互的情况下进行规划、策略学习和决策。
前置依赖
- RSSM 基础(后验 vs 先验、确定性 hidden state、随机 latent variable)
- 变分推断(ELBO、重参数化)
- 强化学习基础(策略、奖励、价值函数)
- 模型预测控制(MPC)基本概念
核心思想
Imagination(想象)是基于世界模型的规划 Paradigm 的核心概念。其核心思想是:训练一个世界模型来学习环境动态,然后利用这个模型生成虚拟轨迹(“想象轨迹”),在这些虚拟轨迹上进行策略优化和规划。
核心问题:真实环境的交互成本高昂(样本效率低、危险场景不可行),能否通过学习一个环境模型来模拟与真实环境的交互?
解决方案:
真实环境交互(昂贵) 想象 rollout(廉价)
┌──────────────┐ ┌──────────────────┐
│ o_t ──────► a_t │ h_t ─────► z_t │
│ r_t, o_{t+1}│ │ 先验采样 │
│ │ │ 确定性转移 │
└──────────────┘ │ 策略输出 a_t │
vs │ 奖励预测 r_hat_t │
└──────────────────┘
在真实环境中交互需要每一步都从真实环境获取反馈,而在想象中交互只需要一次前向传播。
想象 rollout 的关键特性:
- 样本效率高:不需要真实环境交互,可以无限次生成虚拟轨迹
- 风险可控:可以在虚拟环境中尝试危险动作而无需承担后果
- 并行化:可以同时生成多条轨迹进行批量规划
- 可微分化:基于梯度的规划可以直接优化动作序列
与无模型的对比:
| 方面 | 无模型强化学习 | 基于模型的强化学习(世界模型) |
|---|---|---|
| 样本效率 | 低(需要大量交互) | 高(模型可生成无限数据) |
| 泛化能力 | 受限于真实交互分布 | 受限于模型质量 |
| 规划能力 | 无(仅依赖策略网络) | 可利用模型进行 MPC/梯度规划 |
| 探索成本 | 高 | 低(可在想象中探索) |
| 模型误差 | 无 | 有(误差累积问题) |
模型结构图
A. 完整 Dreamer 架构
┌─────────────────────────────────────────────────────────────────┐
│ Dreamer Pipeline │
│ │
│ ┌─────────────┐ │
│ │ 环境交互 │ │
│ │ (真实轨迹) │ │
│ └──────┬──────┘ │
│ │ o_t, a_t, r_t │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ 世界模型 (RSSM) │ │
│ │ │ │
│ │ Encoder: o_t ──► h_t │ │
│ │ │ │ │
│ │ 先验: p(z_t|h_{t-1},a_{t-1}) │ │
│ │ 后验: q(z_t|h_t) │ │
│ │ │ │ │
│ │ 转移: h_t = f(h_{t-1}, z_t, a_{t-1}) │ │
│ │ │ │ │
│ │ Decoder: h_t ──► o_hat_t │ │
│ │ Reward: h_t ──► r_hat_t │ │
│ └──────────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ 想象 rollout (基于先验采样) │ │
│ │ │ │
│ │ h_t ──► 策略网络 π(a_t|h_t) ──► a_t │ │
│ │ │ │ │
│ │ ▼ │ │
│ │ z_t ~ p(z_t|h_{t-1}, a_{t-1}) ◄── 先验采样 │ │
│ │ │ │ │
│ │ ▼ │ │
│ │ h_{t+1} = f(h_t, z_t, a_t) │ │
│ │ │ │ │
│ │ ▼ │ │
│ │ 累积奖励计算: Σ γ^t r_hat_t │ │
│ └──────────────────────────────────────────────────────────┘ │
│ │ │
│ ▼ │
│ ┌──────────────────────────────────────────────────────────┐ │
│ │ 策略学习 (Actor-Critic) │ │
│ │ │ │
│ │ 想象轨迹上的价值估计 V^π(h_t) │ │
│ │ 策略梯度更新: ∇_θ π_θ(a_t|h_t) │ │
│ └──────────────────────────────────────────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
B. 想象 rollout 的信息流
时间步 1 时间步 2 时间步 3 时间步 T
│ │ │ │
▼ ▼ ▼ ▼
┌──────┐ ┌──────┐ ┌──────┐ ┌──────┐
│ h_1 │──────►│ h_2 │──────►│ h_3 │──────►│ h_T │
└──┬───┘ └──┬───┘ └──┬───┘ └──┬───┘
│ │ │ │
▼ ▼ ▼ ▼
┌──────┐ ┌──────┐ ┌──────┐ ┌──────┐
│ z_1~ │ │ z_2~ │ │ z_3~ │ │ z_T~ │
│p(z_1│ │p(z_2│ │p(z_3│ │p(z_T│
└──────┘ └──────┘ └──────┘ └──────┘
│ │ │ │
▼ ▼ ▼ ▼
┌──────┐ ┌──────┐ ┌──────┐ ┌──────┐
│ a_1= │ │ a_2= │ │ a_3= │ │ a_T= │
│π(h_1)│ │π(h_2)│ │π(h_3)│ │π(h_T)│
└──────┘ └──────┘ └──────┘ └──────┘
│ │ │ │
▼ ▼ ▼ ▼
┌──────┐ ┌──────┐ ┌──────┐ ┌──────┐
│r_hat_1│ │r_hat_2│ │r_hat_3│ │r_hat_T│
└──────┘ └──────┘ └──────┘ └──────┘
│ │ │ │
└──────────────┴──────────────┴──────────────┘
│
▼
累积奖励 J = Σ γ^t r_hat_t
数学推导
A. 想象 rollout 的生成过程
给定初始隐状态
对于
-
策略采样(或贪婪选择):
-
先验采样:
-
确定性转移:
-
奖励预测:
注意:上述过程中,
B. 累积奖励的计算
想象 rollout 的总回报:
其中
奖励预测网络
C. 基于梯度的规划
在确定性近似下(忽略
问题定义:
确定性近似:令
此时累积奖励对动作的梯度为:
其中
梯度上升更新动作序列:
D. MPC(模型预测控制)
MPC 在每个时间步执行以下步骤:
- 预测阶段:使用世界模型预测
步内的未来轨迹 - 优化阶段:优化动作序列
以最大化累积奖励 - 执行阶段:仅执行第一个动作
,然后重复
数学形式:
MPC 的一个关键优势是闭环反馈:即使模型预测有误,执行真实动作后重新规划可以校正误差。
训练细节
A. 世界模型训练
世界模型(RSSM)的训练已经在上一节讨论。想象 rollout 使用训练好的先验网络和转移网络。
训练数据收集:
# 初始随机策略收集数据
for episode in range(n_episodes):
obs = env.reset()
for step in range(max_steps):
action = random_action()
next_obs, reward, done = env.step(action)
buffer.append(obs, action, reward, next_obs, done)
obs = next_obs
if done:
break世界模型训练:
for epoch in range(n_epochs):
batch = buffer.sample(batch_size)
loss, recon_loss, kl_loss = train_step(model, encoder, optimizer, batch)B. 策略学习(Dreamer)
Dreamer 使用 Dreamer Actor-Critic 算法进行策略学习:
价值估计:
使用 TD 学习近似:
策略梯度:
C. 想象 rollout 的实现
def imagination_rollout(model, policy_net, value_net, h_0, horizon):
"""
生成想象 rollout
model: RSSM 世界模型
policy_net: 策略网络 π(a|h)
value_net: 价值网络 V(h)
h_0: 初始 hidden state
horizon: rollout 长度
"""
h_seq = [h_0]
a_seq = []
r_seq = []
v_seq = []
h_prev = h_0
for t in range(horizon):
# 1. 策略网络输出动作分布
action_logits = policy_net(h_prev) # (batch, action_dim * 2)
action_mu, action_logsigma = action_logits.chunk(2, dim=-1)
action_sigma = torch.exp(action_logsigma)
# 2. 采样动作(重参数化)
action_eps = torch.randn_like(action_mu)
a_t = action_mu + action_sigma * action_eps
# 3. 先验采样 z_t
p_mu, p_sigma = model.prior(h_prev, a_t)
z_t = model.reparameterize(p_mu, p_sigma)
# 4. 确定性转移
h_t = model.forward(h_prev, z_t, a_t)
# 5. 奖励预测
r_hat_t = model.reward_net(h_t)
# 6. 价值估计
v_hat_t = value_net(h_t)
# 记录
h_seq.append(h_t)
a_seq.append(a_t)
r_seq.append(r_hat_t)
v_seq.append(v_hat_t)
h_prev = h_t
return h_seq, a_seq, r_seq, v_seq推理/rollout/规划过程
A. 推理模式(先验 rollout)
想象 rollout 的推理模式:
def plan_with_imagination(model, policy_net, h_0, n_rollouts=10, horizon=15):
"""
通过想象 rollout 进行规划
"""
best_actions = None
best_value = -float('inf')
for _ in range(n_rollouts):
# 生成一条想象轨迹
h_seq, a_seq, r_seq, v_seq = imagination_rollout(
model, policy_net, None, h_0, horizon
)
# 计算累积价值
cumulative_value = sum([gamma**i * v_seq[i] for i in range(len(v_seq))])
if cumulative_value > best_value:
best_value = cumulative_value
best_actions = a_seq
# 返回价值最高的动作序列的第一个动作
return best_actions[0] if best_actions else NoneB. MPC 规划
def mpc_planning(model, h_t, horizon=15, n_candidates=100):
"""
模型预测控制规划
"""
best_action = None
best_value = -float('inf')
for _ in range(n_candidates):
# 1. 随机采样动作序列
action_seq = torch.randn(horizon, action_dim, device=h_t.device)
# 2. 想象 rollout
h_prev = h_t
total_reward = 0
for step in range(horizon):
# 先验采样
p_mu, p_sigma = model.prior(h_prev, action_seq[step])
z_t = model.reparameterize(p_mu, p_sigma)
# 转移
h_t = model.forward(h_prev, z_t, action_seq[step])
# 奖励
r_hat = model.reward_net(h_t)
total_reward += r_hat.item()
h_prev = h_t
# 3. 选择最优动作序列的第一个动作
if total_reward > best_value:
best_value = total_reward
best_action = action_seq[0]
return best_actionC. 与真实 rollout 的对比
| 方面 | 真实 rollout | 想象 rollout |
|---|---|---|
| 数据来源 | 真实环境交互 | 世界模型生成 |
| 成本 | 高(环境交互) | 低(模型前向传播) |
| 随机性 | 真实噪声 | 模型预测的随机性 |
| 误差 | 无 | 有(模型误差累积) |
| 可并行性 | 受限 | 高(可批量生成) |
优点与局限
优点
- 样本效率极高:可以无限次生成虚拟轨迹,无需真实环境交互
- 风险可控:可以在虚拟环境中尝试危险动作
- 可梯度优化:动作序列可以通过梯度下降直接优化
- 与 MPC 结合:每步重新规划,利用真实反馈校正误差
- 支持探索:可以在想象中探索未见过的状态区域
局限
- 模型误差:世界模型不是完美逼近真实动态,想象 rollout 存在误差
- 误差累积:多步 rollout 时误差会累积,导致最终状态严重偏离
- 随机性建模不足:先验分布假设可能无法捕获复杂的多模态动态
- 训练-推理差异:训练时使用后验推断,推理时使用先验采样,这种差异可能导致性能下降
- 复合误差:模型误差、策略误差、价值估计误差会相互放大
与前后内容的衔接
- 前置内容:RSSM 提供了想象 rollout 所需的先验采样和确定性转移机制
- 后续发展:误差累积问题催生了更复杂的架构(ensemble models、dropout 等);时间一致性是多步预测的关键挑战
- 核心连接:想象 rollout 是基于世界模型的规划的核心,也是理解 Dreamer 算法的基础
可复现实现要点
A. 核心想象 rollout 实现
import torch
import torch.nn as nn
import torch.nn.functional as F
class ImaginationRollout:
"""想象 rollout 管理器"""
def __init__(self, rssm_model, policy_net, reward_net):
self.model = rssm_model
self.policy = policy_net
self.reward = reward_net
def forward(self, h_0, action_seq):
"""
给定初始状态和动作序列,生成想象轨迹
h_0: (batch, hidden_dim) - 初始 hidden state
action_seq: (horizon, batch, action_dim) - 动作序列
"""
horizon = action_seq.size(0)
batch_size = h_0.size(0)
h_seq = [h_0]
z_seq = []
r_seq = []
h_prev = h_0
for t in range(horizon):
# 1. 先验采样
p_mu, p_sigma = self.model.prior(h_prev, action_seq[t])
z_t = self.model.reparameterize(p_mu, p_sigma)
# 2. 确定性转移
h_t = self.model.forward(h_prev, z_t, action_seq[t])
# 3. 奖励预测
r_hat_t = self.reward(h_t)
h_seq.append(h_t)
z_seq.append(z_t)
r_seq.append(r_hat_t)
h_prev = h_t
return h_seq, z_seq, r_seq
def compute_returns(self, r_seq, gamma=0.99):
"""
计算累积回报(折扣奖励和)
r_seq: (horizon, batch, 1) - 预测奖励序列
"""
horizon = len(r_seq)
returns = []
R = torch.zeros_like(r_seq[0])
for t in reversed(range(horizon)):
R = r_seq[t] + gamma * R
returns.insert(0, R)
return returns
def policy_rollout(self, h_0, horizon):
"""
使用策略网络生成动作的想象 rollout
h_0: (batch, hidden_dim)
horizon: int
"""
h_seq = [h_0]
a_seq = []
r_seq = []
h_prev = h_0
for t in range(horizon):
# 1. 策略网络采样动作
action_dist = self.policy(h_prev)
a_t = action_dist.rsample() # 重参数化采样
# 2. 先验采样 z_t
p_mu, p_sigma = self.model.prior(h_prev, a_t)
z_t = self.model.reparameterize(p_mu, p_sigma)
# 3. 确定性转移
h_t = self.model.forward(h_prev, z_t, a_t)
# 4. 奖励预测
r_hat_t = self.reward(h_t)
h_seq.append(h_t)
a_seq.append(a_t)
r_seq.append(r_hat_t)
h_prev = h_t
return h_seq, a_seq, r_seqB. Dreamer 策略训练
def train_policy_actor_critic(model, policy_net, value_net, optimizer,
h_0, horizon, gamma=0.99, lambda_=0.95):
"""
Dreamer 的 Actor-Critic 训练步骤
h_0: 初始 hidden state
horizon: 想象 rollout 长度
lambda_: TD(lambda) 的 lambda 参数
"""
# 1. 生成想象 rollout
h_seq, a_seq, r_seq = imagination_rollout(model, policy_net, h_0, horizon)
# 2. 计算回报
returns = compute_returns(r_seq, gamma)
# 3. 价值估计(使用 TD)
values = []
for t in range(horizon):
if t == horizon - 1:
v_t = returns[t] # 最后一层直接用回报
else:
v_t = r_seq[t] + gamma * returns[t + 1]
values.append(v_t)
# 4. 策略梯度(使用 REINFORCE 或价值作为 baseline)
policy_loss = 0
for t in range(horizon):
# log_prob of action under current policy
action_dist = policy_net(h_seq[t])
log_prob = action_dist.log_prob(a_seq[t])
# 优势函数 A_t = Q_t - V_t = returns[t] - values[t]
advantage = returns[t] - values[t].detach()
policy_loss -= (log_prob * advantage).mean()
# 5. 价值损失
value_loss = sum([F.mse_loss(values[t], returns[t].detach())
for t in range(horizon)])
# 6. 更新策略网络
optimizer.zero_grad()
(policy_loss + value_loss).backward()
optimizer.step()
return policy_loss.item(), value_loss.item()C. 关键实现细节
先验采样 vs 后验采样:
- 训练时:使用后验分布
- 推理时:使用先验分布
重参数化技巧:
eps = torch.randn_like(mu)
z = mu + sigma * eps折扣累积回报:
R = 0
for t in reversed(range(horizon)):
R = r_t + gamma * R
returns[t] = R章节摘要
想象 rollout 是基于世界模型的规划 Paradigm 的核心。通过 RSSM 的先验采样和确定性转移机制,我们可以生成虚拟轨迹,在无需真实环境交互的情况下进行策略优化和规划。
核心流程:
:策略网络输出动作 :先验采样 latent variable :确定性转移 :奖励预测 - 累积回报
核心优势:样本效率高、风险可控、支持梯度规划和 MPC
核心问题:误差累积导致多步 rollout 质量下降,这是下一节讨论的核心内容。
关键词
想象 rollout、先验采样、模型预测控制、规划、策略学习、Dreamer、Actor-Critic、累积回报、梯度规划、无模型 vs 基于模型、样本效率、重参数化技巧、确定性转移、折扣因子 TD 学习