想象与未来预测:基本思想

一句话定位:利用训练好的世界模型生成虚拟轨迹(想象 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 的关键特性

  1. 样本效率高:不需要真实环境交互,可以无限次生成虚拟轨迹
  2. 风险可控:可以在虚拟环境中尝试危险动作而无需承担后果
  3. 并行化:可以同时生成多条轨迹进行批量规划
  4. 可微分化:基于梯度的规划可以直接优化动作序列

与无模型的对比

方面无模型强化学习基于模型的强化学习(世界模型)
样本效率低(需要大量交互)高(模型可生成无限数据)
泛化能力受限于真实交互分布受限于模型质量
规划能力无(仅依赖策略网络)可利用模型进行 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 的生成过程

给定初始隐状态 和策略 ,想象 rollout 的生成过程为:

对于

  1. 策略采样(或贪婪选择):

  2. 先验采样

  3. 确定性转移

  4. 奖励预测

注意:上述过程中, 的更新使用了先验分布,而不是后验分布。这是因为在想象中我们没有未来观测可供后验校正。

B. 累积奖励的计算

想象 rollout 的总回报:

其中 是折扣因子。

奖励预测网络 通常与观测解码器共享特征,通过均方误差或交叉熵损失训练:

C. 基于梯度的规划

在确定性近似下(忽略 的随机性),我们可以直接对动作序列进行梯度优化。

问题定义

确定性近似:令 (使用先验均值),则:

此时累积奖励对动作的梯度为:

其中 通过 RNN 的反向传播计算(BPTT)。

梯度上升更新动作序列

D. MPC(模型预测控制)

MPC 在每个时间步执行以下步骤:

  1. 预测阶段:使用世界模型预测 步内的未来轨迹
  2. 优化阶段:优化动作序列 以最大化累积奖励
  3. 执行阶段:仅执行第一个动作 ,然后重复

数学形式

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 None

B. 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_action

C. 与真实 rollout 的对比

方面真实 rollout想象 rollout
数据来源真实环境交互世界模型生成
成本高(环境交互)低(模型前向传播)
随机性真实噪声模型预测的随机性
误差有(模型误差累积)
可并行性受限高(可批量生成)

优点与局限

优点

  1. 样本效率极高:可以无限次生成虚拟轨迹,无需真实环境交互
  2. 风险可控:可以在虚拟环境中尝试危险动作
  3. 可梯度优化:动作序列可以通过梯度下降直接优化
  4. 与 MPC 结合:每步重新规划,利用真实反馈校正误差
  5. 支持探索:可以在想象中探索未见过的状态区域

局限

  1. 模型误差:世界模型不是完美逼近真实动态,想象 rollout 存在误差
  2. 误差累积:多步 rollout 时误差会累积,导致最终状态严重偏离
  3. 随机性建模不足:先验分布假设可能无法捕获复杂的多模态动态
  4. 训练-推理差异:训练时使用后验推断,推理时使用先验采样,这种差异可能导致性能下降
  5. 复合误差:模型误差、策略误差、价值估计误差会相互放大

与前后内容的衔接

  • 前置内容: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_seq

B. 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 的先验采样和确定性转移机制,我们可以生成虚拟轨迹,在无需真实环境交互的情况下进行策略优化和规划。

核心流程

  1. :策略网络输出动作
  2. :先验采样 latent variable
  3. :确定性转移
  4. :奖励预测
  5. 累积回报

核心优势:样本效率高、风险可控、支持梯度规划和 MPC

核心问题:误差累积导致多步 rollout 质量下降,这是下一节讨论的核心内容。


关键词

想象 rollout、先验采样、模型预测控制、规划、策略学习、Dreamer、Actor-Critic、累积回报、梯度规划、无模型 vs 基于模型、样本效率、重参数化技巧、确定性转移、折扣因子 TD 学习