Actor-critic with Imagination

一句话定位

Actor-critic with Imagination 结合了 actor-critic 框架与世界模型的想象力 rollouts,通过在潜在空间中进行 imagined trajectories 来计算策略梯度和价值估计。

前前置依赖

  • Actor-critic 基础(policy gradient、value function)
  • 世界模型(RSSM、VAE、dreamer)
  • GAE (Generalized Advantage Estimation)
  • 4-value-based planning - 价值规划基础

核心思想

Actor-critic with Imagination 的核心思想是利用世界模型在潜在空间中进行 imagined rollouts,以此产生”假想”的经验来训练 actor 和 critic。这使得智能体可以像 model-free RL 一样使用 policy gradient 方法,但使用model 生成的经验来避免与真实环境的高昂交互。

与标准 Actor-critic 的对比:

组件标准 Actor-critic+ Imagination
Critic 输入真实回报 想象回报
Rollout 来源真实环境世界模型
数据效率
误差来源策略方差模型误差

模型结构图

┌─────────────────────────────────────────────────────────────┐
│            Actor-critic with Imagination 框架              │
├─────────────────────────────────────────────────────────────┤
│                                                             │
│   ┌──────────────────────────────────────────────────┐     │
│   │                   World Model                      │     │
│   │   p(s'|s,a), p(r|s,a), q(s|o)                    │     │
│   └──────────────────────────────────────────────────┘     │
│                          ▲                                   │
│                          │ imagine                           │
│   ┌──────────────────────────────────────────────────┐     │
│   │              Imagination Rollout                   │     │
│   │  s_0 ──a_0──▶ s_1 ──a_1──▶ s_2 ──...──▶ s_H     │     │
│   │  o_0       r_0       r_1       r_H               │     │
│   │  V(s_0)   A(s_0,a_0) ...                         │     │
│   └──────────────────────────────────────────────────┘     │
│                          │                                   │
│                          ▼                                   │
│   ┌──────────────────────────────────────────────────┐     │
│   │                 Actor (Policy)                    │     │
│   │   π(a|s),  ∇_θ J = E[∇logπ(a|s) · A]            │     │
│   └──────────────────────────────────────────────────┘     │
│                                                             │
│   ┌──────────────────────────────────────────────────┐     │
│   │                 Critic (Value)                    │     │
│   │   V(s),  TD target = r + γV(s')                  │     │
│   └──────────────────────────────────────────────────┘     │
└─────────────────────────────────────────────────────────────┘

数学推导

1. 想象力 Rollout 如何输入 Actor-critic

设世界模型参数为 ,当前 policy 为 ,value function 为

Imagination Rollout 过程:

对于初始状态 (从真实观测编码或起始状态):

生成长度为 的想象轨迹:

2. GAE (Generalized Advantage Estimation)

GAE 是一种低方差的 advantage 估计方法,结合了 TD() 的优点:

K-step advantage 估计:

其中 是 TD error。

GAE():

等价形式:

性质:

  • (TD(0) advantage)
  • (Monte Carlo)
  • :平衡 bias-variance trade-off

在 imagination 中的使用:

对于想象轨迹,用 计算 GAE:

3. 从想象回报计算策略梯度

Policy Gradient 定理(stochastic 策略):

其中 可以是任意 advantage 估计。

使用想象轨迹的梯度估计:

其中 是 imagined trajectories 的数量。

4. 潜在空间中的价值目标计算

Value Target 的计算(潜在空间):

给定想象轨迹

n-step return:

Value loss:

在 latent space 的特殊处理:

由于 是潜在状态,value function 直接在潜在空间操作:

其中 是从 imagined rollout 计算的 target。

5. Dreamer 方法的核心公式

Dreamer 是 Actor-critic with Imagination 的代表性方法:

演员(Actor)Loss:

评论家(Critic)Loss:

模型 Loss:

三者的交替优化:

  1. 固定 ,更新模型
  2. 固定 ,更新
  3. 固定 ,更新

训练细节

Dreamer 训练流程

# Dreamer 伪代码
for iteration in range(num_iterations):
    # 1. 收集真实经验
    for step in range(env_steps):
        a = policy.act(o)
        o_next, r = env.step(a)
        buffer.add((o, a, r, o_next))
        o = o_next
 
    # 2. 学习世界模型
    for _ in range(model_steps):
        batch = buffer.sample()
        # 学习 p(s'|s,a), p(r|s,a), q(s|o)
        model.fit(batch)
 
    # 3. Imagination Rollout
    imagined_trajectories = []
    for _ in range(imagined_rollouts):
        s = encoder(o)  # 编码初始观测
        for h in range(horizon):
            a = policy.act(s)  # 采样动作
            s_next = model.predict(s, a)
            r = model.reward(s, a)
            imagined_trajectories.append((s, a, r, s_next))
            s = s_next
 
    # 4. 计算 advantage (GAE)
    for trajectory in imagined_trajectories:
        compute_gae(trajectory, lambda_, gamma)
 
    # 5. 更新 Actor
    for _ in range(actor_steps):
        policy.update(imagined_trajectories)
 
    # 6. 更新 Critic
    for _ in range(critic_steps):
        value.update(imagined_trajectories)

GAE 参数选择

BiasVariance适用场景
0模型准确、horizon 短
0.9-0.95通用场景
0.99-1.0长 horizon、稀疏奖励

训练中的关键挑战

挑战描述解决思路
复合误差模型误差累积在长 rollout限制 horizon、early stopping
梯度弥散远视距的优势估计不准使用 ,GAE 截断
过估计max Q 导致过估计twin Q-learning, target network
想象-真实分布偏移想象与真实不同混合真实与想象数据

推理/rollout/planning 过程

推理流程

Online Planning with Actor-critic:

与 MPC/MCTS 不同,actor-critic 在推理时直接输出动作,不需要显式规划:

def act(obs, policy, encoder, model):
    # 编码观测
    z = encoder(obs)
 
    # 直接用 actor 输出动作(无需规划)
    a = policy(z)
 
    # 可选:imagination 辅助验证
    z_next = model.predict(z, a)
    value = critic(z_next)
 
    return a

Imagination 的在线使用

Planner vs Actor:

模式规划方式计算成本适用场景
Planner(MPC/MCTS)在线优化低延迟要求低
Actor(Dreamer)直接输出实时控制

混合模式:

部分方法(如 DreamerPro)结合两者:

  1. 用 actor 作为初始策略
  2. 用 imagination rollouts 评估动作
  3. 如有需要,用 CEM refine

优点与局限

优点

  1. 样本效率高:imagination rollouts 生成大量”假想”经验
  2. 无需在线优化:actor 直接输出动作,延迟低
  3. 可微分化:梯度可以从 critic 流回 actor
  4. 兼容性强:可与任何 model-based / model-free 方法结合

局限

  1. 模型误差:imagination 质量受限于世界模型精度
  2. 误差累积:长 horizon rollout 误差指数增长
  3. 多模态困难:单一 policy 难以表达多模态策略
  4. 探索不足:actor-critic 的 exploration 是独立问题

与前后内容的衔接

可复现实现要点

import torch
import torch.nn as nn
from torch.distributions import Normal
 
class Actor(nn.Module):
    def __init__(self, z_dim, action_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(z_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, action_dim * 2)  # mean, log_std
        )
        self.z_dim = z_dim
        self.action_dim = action_dim
 
    def forward(self, z):
        out = self.net(z)
        mean, log_std = out.chunk(2, dim=-1)
        std = torch.exp(log_std)
        dist = Normal(mean, std)
        return dist
 
class Critic(nn.Module):
    def __init__(self, z_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(z_dim, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 1)
        )
 
    def forward(self, z):
        return self.net(z)
 
class Dreamer:
    def __init__(self, encoder, world_model, actor, critic):
        self.encoder = encoder
        self.model = world_model
        self.actor = actor
        self.critic = critic
        self.gamma = 0.99
        self.lambda_ = 0.95
 
    def imagination_rollout(self, z0, horizon=15):
        """生成 imagined trajectories"""
        traj = []
        z = z0
        for h in range(horizon):
            dist = self.actor(z)
            a = dist.sample()
            z_next = self.model.predict(z, a)
            r = self.model.reward(z, a)
            traj.append((z, a, r, z_next))
            z = z_next
        return traj
 
    def compute_gae(self, traj):
        """计算 GAE advantage"""
        values = [self.critic(s) for s, _, _, _ in traj]
        advantages = []
        gae = 0
        for t in reversed(range(len(traj))):
            _, a, r, s_next = traj[t]
            V = values[t]
            V_next = values[t+1] if t < len(traj) - 1 else 0
            delta = r + self.gamma * V_next - V
            gae = delta + self.gamma * self.lambda_ * gae
            advantages.insert(0, gae)
        return advantages
 
    def update(self, obs):
        """Dreamer 更新步骤"""
        z0 = self.encoder(obs)
 
        # Imagination rollout
        traj = self.imagination_rollout(z0)
 
        # 计算 advantage
        advantages = self.compute_gae(traj)
 
        # 更新 actor
        dist = self.actor(z0)
        # ... policy gradient step
 
        # 更新 critic
        for s, _, r, s_next in traj:
            # ... value regression step
            pass

章节摘要

Actor-critic with Imagination 将世界模型的想象能力与 actor-critic 框架结合,通过在潜在空间中进行 imagined rollouts 来生成训练数据。这使得 policy gradient 方法可以在不与真实环境交互的情况下进行训练,极大提升了 sample efficiency。Dreamer 是这一范式的代表性方法,通过 GAE 计算 advantage,利用 imagined trajectories 更新 actor 和 critic。核心挑战在于模型误差累积和想象-真实分布偏移。后续章节将介绍这一方法在连续控制中的特殊应用。

关键词

actor-critic, imagination, imagination rollout, GAE, Generalized Advantage Estimation, Dreamer, policy gradient, value function, latent space, world model, policy optimization, advantage estimation, imagined trajectories