MBPO：基于模型的策略优化

一句话总结 · 要点速览

• MBPO 学习一个概率神经网络集成作为环境模型，从真实数据中的状态出发，用当前策略在模型里展开很短的 k 步，生成大量虚拟转移来训练 SAC。
• 短视界展开的理由：模型误差随展开步数累积，而从真实状态“分支”出去的短展开既能利用模型泛化，又不会偏离真实分布太远。
• MBPO 每个真实步做很多次（本站为 10 次）SAC 更新，主要用模型数据；高更新比（UTD）是它样本效率高的重要原因之一。
• 本站 Pendulum 实测（2 个种子）：2000 个真实步时，MBPO 评估回报约 −100，标准 SAC 约 −1300；但只用真实数据、同样每步 10 次更新的 SAC 也达到约 −98，说明在这个简单任务上收益主要来自高更新比。

术语不熟悉？查看 强化学习术语表。

MPC 每一步都要在线规划，决策开销大；另一条路线是把模型当作“数据生成器”，像 Dyna 一样用模拟经验训练一个策略网络，执行时只需一次前向传播。MBPO（Model-Based Policy Optimization，Janner 等，2019）是这条路线在深度强化学习中的代表：它把概率集成模型和 SAC 结合起来，并回答了一个关键问题——模型生成的轨迹应该展开多长？

模型误差与展开长度

用学到的模型 $\hat{p}$ 生成数据训练策略，最担心的是模型误差：策略会利用模型中“回报虚高”的错误，在真实环境中表现很差。误差还会随展开步数累积——第 1 步的预测误差会改变第 2 步的输入，越往后偏得越远。

MBPO 论文给出了一个界：设模型的单步误差为 ${ϵ}_m$、新旧策略的差异为 ${ϵ}_{\pi }$，用模型从初始状态完整展开时，真实回报与模型回报之差大致随 $\frac{\gamma }{(1-\gamma {)}^2}{ϵ}_m$ 增长——对长视界任务非常大。而如果改为从真实数据中的状态出发、只展开 $k$ 步（分支展开，branched rollout），误差项中与模型有关的部分只随 $k$ 增长；论文据此论证了“短展开 + 从真实状态分支”的合理性。直观理解：

• 起点来自真实数据分布，不会因为从头展开而一路漂移；
• 只展开几步，模型误差来不及累积；
• 每个起点都可以采样新的动作，生成真实数据中没有的转移，相当于在真实状态附近做“数据增强”。

算法

初始化策略 π（SAC）、模型集成 {p̂_1..p̂_B}、真实数据 D_env、模型数据 D_model
用随机动作收集 D_env 的初始数据
对每个真实步：
    每隔 E 步：用 D_env 训练模型集成（选出验证误差最小的若干“精英”成员）
               从 D_env 中随机取 M 个状态，用 π 在模型中展开 k 步，把转移加入 D_model
    用 π 与真实环境交互一步，加入 D_env
    做 G 次 SAC 更新，每个批次中真实数据占比很小（如 5%），其余来自 D_model

几个关键设计：

• 概率集成模型：与 PETS 相同，每个成员输出 $(\mathrm{\Delta }s,r)$ 的高斯分布，用负对数似然训练；展开时每条样本随机选一个精英成员并从其分布中采样，从而把模型不确定性反映在生成的数据里。
• 奖励也由模型预测：生成的虚拟转移需要奖励，因此模型同时预测 $r$。
• 高更新比（update-to-data ratio，UTD）：每个真实步做 $G$ 次（论文中 20–40 次）SAC 更新。模型数据量大，才能支撑这么多次更新而不过拟合真实数据。
• 只保留较新的模型数据：策略和模型都在变，旧的模型数据很快过时，因此 $D_{model}$ 只保留最近若干批。

实验：Pendulum-v1

为了在 CPU 上较快跑完，我们把规模缩小：

• 共 5000 个真实步，前 500 步随机动作；
• 模型：7 个成员的概率集成（3 个隐藏层 × 200，Swish 激活），选验证误差最小的 5 个为精英；每 250 个真实步重训一次，每次从 2000 个真实状态出发展开 $k=1$ 步；
• MBPO：每个真实步做 10 次 SAC 更新，批次中 5% 为真实数据；
• 对照组：标准 SAC（每个真实步 1 次更新，只用真实数据），以及 SAC×10（每个真实步 10 次更新，只用真实数据）。第二个对照用来区分“模型数据”和“更多的更新次数”各自的贡献。
• 每 1000 个真实步用确定性策略评估 5 个回合，2 个随机种子。

"""MBPO（基于模型的策略优化）与 SAC 的样本效率对比：Pendulum-v1，相同的真实交互步数。
MBPO = SAC + 概率神经网络集成模型 + 从真实状态出发的短视界模型展开（这里 k = 1）。
对照组 sac10：同样每个真实步做 10 次更新，但只用真实数据，用来区分“模型数据”与“更新次数”的作用。
用法：python mbpo_pendulum.py [mbpo|sac|sac10] [种子]；DQN_QUICK=1 时只跑很短的版本用于检查代码。"""
import copy
import os
import sys
import time

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Normal

QUICK = os.environ.get("DQN_QUICK") == "1"
torch.set_num_threads(1)
GAMMA, TAU, LR, BATCH = 0.99, 0.005, 3e-4, 256
TOTAL_STEPS = 800 if QUICK else 5_000
START_STEPS = 500
MODEL_EVERY, ROLLOUT_STARTS, ROLLOUT_LEN = 250, 2_000, 1   # 每 250 步重训模型并生成一批模型数据
UTD, REAL_RATIO = 10, 0.05                                  # MBPO：每个真实步做 10 次 SAC 更新，5% 真实数据
N_ENS, N_ELITE = 7, 5
OBS, ACT, A_MAX = 3, 1, 2.0


def mlp(i, o, h=256):
    return nn.Sequential(nn.Linear(i, h), nn.ReLU(), nn.Linear(h, h), nn.ReLU(), nn.Linear(h, o))


class Actor(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = mlp(OBS, 2 * ACT)

    def forward(self, s):
        mu, log_std = self.net(s).chunk(2, -1)
        dist = Normal(mu, log_std.clamp(-20, 2).exp())
        u = dist.rsample(); a = torch.tanh(u)
        logp = dist.log_prob(u).sum(-1) - torch.log(1 - a.pow(2) + 1e-6).sum(-1)
        return a, logp, torch.tanh(mu)


class Ensemble(nn.Module):
    """N 个概率网络并行：输入 (s, a)，输出 (Δs, r) 的高斯分布。"""

    def __init__(self):
        super().__init__()
        def P(*shape):
            return nn.Parameter(torch.randn(N_ENS, *shape) / np.sqrt(shape[0]))
        self.W = nn.ParameterList([P(OBS + ACT, 200), P(200, 200), P(200, 200), P(200, 2 * (OBS + 1))])
        self.b = nn.ParameterList([nn.Parameter(torch.zeros(N_ENS, 1, d)) for d in (200, 200, 200, 2 * (OBS + 1))])
        self.max_lv = nn.Parameter(torch.full((1, 1, OBS + 1), 0.5)); self.min_lv = nn.Parameter(torch.full((1, 1, OBS + 1), -10.0))
        self.elites = list(range(N_ELITE))

    def forward(self, x):
        for i, (W, b) in enumerate(zip(self.W, self.b)):
            x = torch.baddbmm(b, x, W)
            if i < 3:
                x = x * torch.sigmoid(x)
        mean, lv = x.chunk(2, -1)
        lv = self.max_lv - F.softplus(self.max_lv - lv); lv = self.min_lv + F.softplus(lv - self.min_lv)
        return mean, lv

    def train_on(self, X, Y, rng, epochs):
        X, Y = torch.as_tensor(X), torch.as_tensor(Y)
        self.mu, self.sd = X.mean(0), X.std(0) + 1e-6; self.ymu, self.ysd = Y.mean(0), Y.std(0) + 1e-6
        Xn, Yn = (X - self.mu) / self.sd, (Y - self.ymu) / self.ysd
        n_val = min(500, len(X) // 5); perm = torch.as_tensor(rng.permutation(len(X)))
        val, tr = perm[:n_val], perm[n_val:]
        opt = torch.optim.Adam(self.parameters(), lr=1e-3, weight_decay=1e-5)
        boot = tr[torch.as_tensor(rng.integers(0, len(tr), (N_ENS, len(tr))))]
        for _ in range(epochs):
            order = boot[:, torch.randperm(len(tr))]
            for i in range(0, len(tr), 256):
                j = order[:, i:i + 256]
                mean, lv = self(Xn[j])
                loss = (((mean - Yn[j]) ** 2) * torch.exp(-lv) + lv).mean() + 0.01 * (self.max_lv.sum() - self.min_lv.sum())
                opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():                    # 按验证集 MSE 选出最好的 N_ELITE 个成员
            mean, _ = self(Xn[val].expand(N_ENS, -1, -1))
            mse = ((mean - Yn[val]) ** 2).mean((1, 2))
            self.elites = mse.argsort()[:N_ELITE].tolist()
        return mse.sort().values[:N_ELITE].mean().item()

    @torch.no_grad()
    def step(self, s, a, gen):
        x = ((torch.cat([s, a], -1) - self.mu) / self.sd).expand(N_ENS, -1, -1)
        mean, lv = self(x)
        sample = mean + torch.exp(0.5 * lv) * torch.randn(mean.shape, generator=gen)
        pick = torch.as_tensor(self.elites)[torch.randint(0, N_ELITE, (len(s),), generator=gen)]
        out = sample[pick, torch.arange(len(s))] * self.ysd + self.ymu    # 每条样本随机选一个精英成员
        return s + out[:, :OBS], out[:, OBS:]


def run(method, seed):
    rng = np.random.default_rng(seed); torch.manual_seed(seed); gen = torch.Generator().manual_seed(seed)
    env, test_env = gym.make("Pendulum-v1"), gym.make("Pendulum-v1")
    actor = Actor(); q1, q2 = mlp(OBS + ACT, 1), mlp(OBS + ACT, 1)
    q1_t, q2_t = copy.deepcopy(q1), copy.deepcopy(q2)
    opt_a = torch.optim.Adam(actor.parameters(), lr=LR)
    opt_q = torch.optim.Adam(list(q1.parameters()) + list(q2.parameters()), lr=LR)
    log_alpha = torch.zeros(1, requires_grad=True); opt_al = torch.optim.Adam([log_alpha], lr=LR)
    real = {k: torch.zeros(TOTAL_STEPS, d) for k, d in (("s", OBS), ("a", ACT), ("r", 1), ("s2", OBS))}
    cap = ROLLOUT_STARTS * ROLLOUT_LEN * 8                     # 模型数据只保留最近若干批
    fake = {k: torch.zeros(cap, d) for k, d in (("s", OBS), ("a", ACT), ("r", 1), ("s2", OBS))}
    n_fake = ptr = 0
    model = Ensemble() if method == "mbpo" else None

    def sac_update(s, a, r, s2):                # Pendulum 没有终止状态（只有时间截断），所以不需要 done
        alpha = log_alpha.exp().detach()
        with torch.no_grad():
            a2, lp2, _ = actor(s2); sa2 = torch.cat([s2, a2], 1)
            y = r + GAMMA * (torch.min(q1_t(sa2), q2_t(sa2)) - alpha * lp2[:, None])
        sa = torch.cat([s, a], 1)
        lq = F.mse_loss(q1(sa), y) + F.mse_loss(q2(sa), y)
        opt_q.zero_grad(); lq.backward(); opt_q.step()
        an, lp, _ = actor(s); san = torch.cat([s, an], 1)
        la = (alpha * lp - torch.min(q1(san), q2(san)).squeeze(1)).mean()
        opt_a.zero_grad(); la.backward(); opt_a.step()
        lal = -(log_alpha * (lp.detach() - ACT)).mean()
        opt_al.zero_grad(); lal.backward(); opt_al.step()
        with torch.no_grad():
            for n, t in ((q1, q1_t), (q2, q2_t)):
                for p, tp in zip(n.parameters(), t.parameters()):
                    tp.mul_(1 - TAU).add_(TAU * p)

    def evaluate(n=5):
        tot = 0.0
        for k in range(n):
            o, _ = test_env.reset(seed=10_000 + k); done = False
            while not done:
                with torch.no_grad():
                    a = actor(torch.as_tensor(o, dtype=torch.float32))[2].numpy()
                o, r, te, tr, _ = test_env.step(a * A_MAX); tot += r; done = te or tr
        return tot / n

    o, _ = env.reset(seed=seed); env.action_space.seed(seed); t0 = time.time(); log = []
    for step in range(TOTAL_STEPS):
        if step < START_STEPS:
            a = rng.uniform(-1, 1, ACT).astype(np.float32)
        else:
            with torch.no_grad():
                a = actor(torch.as_tensor(o, dtype=torch.float32))[0].numpy()
        o2, r, te, tr, _ = env.step(a * A_MAX)
        for k, v in (("s", o), ("a", a), ("r", [r]), ("s2", o2)):
            real[k][step] = torch.as_tensor(np.asarray(v, np.float32))
        o = o2
        if te or tr:
            o, _ = env.reset()
        n_real = step + 1
        if step + 1 < START_STEPS:
            continue
        if method == "mbpo":
            if (step + 1 - START_STEPS) % MODEL_EVERY == 0:
                X = torch.cat([real["s"][:n_real], real["a"][:n_real]], 1).numpy()
                Y = torch.cat([real["s2"][:n_real] - real["s"][:n_real], real["r"][:n_real]], 1).numpy()
                mse = model.train_on(X, Y, rng, epochs=5 if QUICK else 30)
                s = real["s"][torch.as_tensor(rng.integers(0, n_real, ROLLOUT_STARTS))]
                for _ in range(ROLLOUT_LEN):          # 从真实状态出发，用当前策略在模型里走 k 步
                    with torch.no_grad():
                        a_m = actor(s)[0]
                    s2_m, r_m = model.step(s, a_m, gen)
                    m = len(s); idx = (torch.arange(m) + ptr) % cap
                    for k, v in (("s", s), ("a", a_m), ("r", r_m), ("s2", s2_m)):
                        fake[k][idx] = v
                    ptr = (ptr + m) % cap; n_fake = min(n_fake + m, cap); s = s2_m
            for _ in range(UTD):
                n_r = int(BATCH * REAL_RATIO)
                i = torch.as_tensor(rng.integers(0, n_real, n_r)); j = torch.as_tensor(rng.integers(0, n_fake, BATCH - n_r))
                sac_update(*(torch.cat([real[k][i], fake[k][j]]) for k in ("s", "a", "r", "s2")))
        else:
            for _ in range(UTD if method == "sac10" else 1):
                i = torch.as_tensor(rng.integers(0, n_real, BATCH))
                sac_update(*(real[k][i] for k in ("s", "a", "r", "s2")))
        if (step + 1) % (200 if QUICK else 1000) == 0:
            log.append((step + 1, evaluate()))
            print(f"  {method} 种子{seed} 真实步数 {step + 1:5d}  评估回报 {log[-1][1]:8.1f}  （{time.time() - t0:.0f}s）", flush=True)
    return log


if __name__ == "__main__":
    methods, seeds = ("sac", "sac10", "mbpo"), ([0] if QUICK else [0, 1])
    if len(sys.argv) == 3:
        methods, seeds = (sys.argv[1],), [int(sys.argv[2])]
    for m in methods:
        for sd in seeds:
            log = run(m, sd)
            print(f"{m:4s} 种子{sd}: " + "  ".join(f"{s}:{v:.0f}" for s, v in log))

实际运行结果（评估回报，越接近 0 越好）：

方法	种子	1000 步	2000 步	3000 步	4000 步	5000 步
SAC	0	−1612	−1290	−1330	−129	−123
SAC	1	−1590	−1312	−317	−125	−124
SAC×10（只用真实数据）	0	−1114	−97	−97	−97	−98
SAC×10（只用真实数据）	1	−375	−98	−121	−98	−98
MBPO	0	−1316	−100	−98	−98	−100
MBPO	1	−385	−99	−98	−97	−99

结果解读：

• MBPO 比标准 SAC 的样本效率高得多：MBPO 在 2000 个真实步时已经稳定在 −100 左右，标准 SAC 要到 3000–4000 步才学会，且最终评估回报（约 −123）还略低。
• 但同更新比的 SAC×10 表现几乎一样：只用真实数据、每步更新 10 次的 SAC 在 2000 步时也达到约 −98。也就是说，在 Pendulum 这个简单任务上，MBPO 的优势基本来自高更新比，而不是模型生成的数据。我们如实给出这个对照，避免把收益错误地归因于模型。
• 为什么会这样：Pendulum 状态只有 3 维、动力学平滑，几千条真实数据就覆盖了相关的状态空间，Q 网络在高更新比下也没有出现明显的过拟合或高估。MBPO 论文中模型数据的价值主要体现在 Hopper、Walker2d、Ant、Humanoid 等高维任务上；后续工作（如 REDQ，Chen 等，2021）也表明，在无模型 SAC 中加入集成 Q 等正则后，高更新比本身就能取得与 MBPO 相近的样本效率。
• 计算代价：5000 个真实步，MBPO 用时约 560 秒，SAC×10 约 430 秒，SAC 约 35 秒（三者运行时机器负载不同，只能看量级）。

实践要点

• 模型训练要早停或选精英：用留出的验证集选成员，避免使用过拟合的成员生成数据。
• 展开长度可以随训练逐步增大：模型变准后适当增加 $k$，论文对部分任务使用了随训练线性增长的 $k$。
• 真实数据比例不要为零：少量真实数据能“锚定”Q 函数，防止完全被模型误差带偏。
• 先做对照实验：引入模型之前，先试试只提高 SAC 的更新比（必要时加集成 Q 或层归一化等正则），确认模型数据确实带来了额外收益。

小结

设计	作用
概率集成 + 精英成员	表示不确定性，减少策略对模型误差的利用
从真实状态分支的 $k$ 步短展开	控制模型误差累积，同时生成大量新数据
高更新比的 SAC	充分利用模型数据，提高样本效率
小比例真实数据	让价值函数不完全依赖模型

练习题

1. 把展开长度 $k$ 改为 3 或 5，观察 Pendulum 上的学习曲线；如果变差，检查模型数据的奖励分布与真实数据有何不同。
2. 把 MBPO 的真实数据比例从 5% 改为 0% 和 50%，比较结果。
3. 在 SAC×10 中把更新比继续提高到 20，观察是否开始出现 Q 值高估或训练不稳定。
4. 用 MBPO 论文的误差界解释：为什么从初始状态出发、展开整个回合的“全长展开”在长视界任务中不可取？
5. 修改代码，记录模型在验证集上的预测误差随真实步数的变化，与策略回报的变化放在一起分析。

参考文献

• Sutton, R. S. (1990). Integrated architectures for learning, planning, and reacting based on approximating dynamic programming. ICML.
• Janner, M., Fu, J., Zhang, M., & Levine, S. (2019). When to trust your model: Model-based policy optimization. NeurIPS. arXiv:1906.08253
• Chua, K., Calandra, R., McAllister, R., & Levine, S. (2018). Deep reinforcement learning in a handful of trials using probabilistic dynamics models. NeurIPS.
• Haarnoja, T., Zhou, A., Abbeel, P., & Levine, S. (2018). Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. ICML.
• Chen, X., Wang, C., Zhou, Z., & Ross, K. (2021). Randomized ensembled double Q-learning: Learning fast without a model. ICLR.（REDQ）

下一步

上一节：模型预测控制（MPC）　｜　继续阅读：离线强化学习