TRPO：信赖域策略优化

一句话总结 · 要点速览

• TRPO 在每次更新时最大化代理目标 E[(π_θ/π_old)·A]，同时约束新旧策略的平均 KL 散度不超过 δ，以保证单调改进、避免策略崩溃。
• 求解时对目标做一阶近似、对 KL 约束做二阶近似，得到自然梯度方向，用共轭梯度法求解并配合回溯线搜索。
• TRPO 稳定但实现复杂、计算开销大，PPO 用裁剪目标以一阶优化近似了 TRPO 的效果。

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

TRPO（Trust Region Policy Optimization，Schulman et al., 2015）通过限制每次更新前后策略之间的 KL 散度，保证策略更新不会“走得太远”，从而获得近似单调的性能提升。它用共轭梯度法求解自然梯度方向，再用回溯线搜索保证约束满足且目标确实提升。TRPO 理论扎实、训练稳定，但实现较复杂、计算开销较大；后来的 PPO 用一阶优化和裁剪近似了它的核心思想，成为更常用的选择。理解 TRPO 有助于真正理解 PPO 为什么要那样设计。

为什么需要信赖域

步长难题

在 策略梯度 和 A2C 中，每一步沿梯度方向更新参数，步长由学习率决定。问题在于：

• 参数空间的距离 ≠ 策略空间的距离：参数的微小变化可能让策略剧烈改变（例如 softmax 接近饱和时），反之亦然。
• 一次坏的更新代价很大：在监督学习中，一步走错下一步还能纠正；在强化学习中，策略变差后采集的数据也会变差，可能陷入恶性循环，性能崩溃后难以恢复。

我们希望每次更新都在策略空间中控制步长，并且尽量保证性能不下降。

性能提升下界

Kakade & Langford (2002) 给出了一个关键恒等式：新策略 $\tilde{\pi }$ 与旧策略 $\pi$ 的性能差可以用旧策略的优势函数表示：

$$J(\tilde{\pi })=J(\pi )+{\mathbb{E}}_{\tau \sim \tilde{\pi }}[\sum _{t=0}^{\mathrm{\infty }}{\gamma }^t{A}_{\pi }(S_t,A_t)]$$

麻烦在于期望是在新策略的状态分布下取的。TRPO 用旧策略的状态分布 ${\rho }_{\pi }$ 代替，得到代理目标（surrogate objective）：

$$L_{\pi }(\tilde{\pi })=J(\pi )+\sum _s{\rho }_{\pi }(s)\sum _a\tilde{\pi }(a\mid s)A_{\pi }(s,a)$$

Schulman 等人证明了：

$$J(\tilde{\pi })\ge L_{\pi }(\tilde{\pi })-C\cdot D_{\text{KL}}^{max}(\pi ,\tilde{\pi }),C=\frac{4ϵ\gamma }{(1-\gamma {)}^2}$$

其中 $ϵ=\underset{s,a}{max}|A_{\pi }(s,a)|$。这意味着：只要每一步都最大化右侧的下界，性能就单调不降（Minorization-Maximization 思想）。

带 KL 约束的优化目标

理论中的惩罚系数 $C$ 太大，会导致步长极小。TRPO 实际采用约束形式，并用平均 KL 代替最大 KL：

$$\begin{array}{rl}\underset{\theta }{max}& {\mathbb{E}}_{s\sim {\rho }_{{\theta }_{\text{old}}},a\sim {\pi }_{{\theta }_{\text{old}}}}[\frac{{\pi }_{\theta }(a\mid s)}{{\pi }_{{\theta }_{\text{old}}}(a\mid s)}{A}_{{\theta }_{\text{old}}}(s,a)]\\ \text{s.t.}& {\mathbb{E}}_{s\sim {\rho }_{{\theta }_{\text{old}}}}\left[D_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}(\cdot \mid s)\parallel {\pi }_{\theta }(\cdot \mid s))\right]\le \delta \end{array}$$

目标中的比值 $\frac{{\pi }_{\theta }}{{\pi }_{{\theta }_{\text{old}}}}$ 是重要性采样权重，使我们能用旧策略的数据评估新策略。$\delta$ 是信赖域半径，常用 0.01。

求解方法

自然梯度

在 ${\theta }_{\text{old}}$ 附近对目标做一阶展开、对约束做二阶展开（KL 在 $\theta ={\theta }_{\text{old}}$ 处的值和梯度都为 0）：

$$\underset{\mathrm{\Delta }\theta }{max}{g}^{\mathrm{\top }}\mathrm{\Delta }\theta \text{s.t.}\frac{1}{2}\mathrm{\Delta }{\theta }^{\mathrm{\top }}F\mathrm{\Delta }\theta \le \delta$$

其中 $g$ 是代理目标的梯度，$F$ 是 KL 的 Hessian，即 Fisher 信息矩阵。用拉格朗日乘子法得到解析解：

$$\mathrm{\Delta }\theta =\sqrt{\frac{2\delta }{g^{\mathrm{\top }}{F}^{-1}g}}{F}^{-1}g$$

$F^{-1}g$ 就是自然梯度（Kakade, 2001）方向：它在“策略分布的几何”下最陡，与参数化方式无关。

共轭梯度

神经网络参数动辄上万维，显式构造并求逆 $F$ 不可行。TRPO 用共轭梯度法近似求解 $Fx=g$，只需要计算 Fisher-向量积 $Fv$：

$$Fv={\mathrm{\nabla }}_{\theta }\left({\left({\mathrm{\nabla }}_{\theta }{\bar{D}}_{\text{KL}}\right)}^{\mathrm{\top }}v\right)$$

即对 KL 求一次梯度，与 $v$ 做内积，再求一次梯度——两次反向传播，无需构造矩阵。通常 10 次共轭梯度迭代就足够。实践中还会加一个阻尼项 $F+\lambda I$ 以提高数值稳定性。

线搜索

二阶近似可能不准确，直接走满步长可能违反 KL 约束或使目标变差。TRPO 进行回溯线搜索：依次尝试步长 $\mathrm{\Delta }\theta ,\frac{1}{2}\mathrm{\Delta }\theta ,\frac{1}{4}\mathrm{\Delta }\theta ,\dots$，直到同时满足 (1) 实际 KL $\le \delta$；(2) 代理目标确实提升。都不满足则放弃本次更新。

算法流程：

对每次迭代：
    1. 用 π_θold 采集一批轨迹，用 GAE 估计优势 Â
    2. 计算代理目标梯度 g
    3. 用共轭梯度求 x ≈ F⁻¹ g（只用 Fisher-向量积）
    4. 最大步长 β = sqrt(2δ / (xᵀ F x))，候选更新 Δθ = β x
    5. 回溯线搜索：θ ← θold + 0.5^j Δθ，直到 KL ≤ δ 且目标提升
    6. 用回归更新价值网络 V

代码实现

下面是一个完整的 TRPO 实现（CartPole-v1，8 个并行环境，每轮 2048 步，共 50 轮约 10 万步）。策略更新使用共轭梯度与线搜索，价值网络用 Adam 回归。

"""TRPO：CartPole-v1（共轭梯度 + 回溯线搜索 + KL 约束）。"""
import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
from torch.distributions import Categorical, kl_divergence

GAMMA, LAM, MAX_KL, CG_ITERS, DAMPING = 0.99, 0.97, 0.01, 10, 0.1
N_ENVS, N_STEPS, ITERS = 8, 256, 50          # 每轮 2048 步，共约 10 万步
torch.manual_seed(0)
envs = gym.vector.SyncVectorEnv([lambda: gym.make("CartPole-v1") for _ in range(N_ENVS)],
                                autoreset_mode=gym.vector.AutoresetMode.SAME_STEP)
envs = gym.wrappers.vector.RecordEpisodeStatistics(envs)
obs_dim, n_actions = envs.single_observation_space.shape[0], envs.single_action_space.n

policy = nn.Sequential(nn.Linear(obs_dim, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, n_actions))
value = nn.Sequential(nn.Linear(obs_dim, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), nn.Linear(64, 1))
v_opt = torch.optim.Adam(value.parameters(), lr=1e-3)


def flat_grad(y, params, **kw):
    return torch.cat([g.reshape(-1) for g in torch.autograd.grad(y, params, **kw)])


def set_flat(params, flat):
    i = 0
    for p in params:
        n = p.numel(); p.data.copy_(flat[i:i + n].view_as(p)); i += n


def conjugate_gradient(Avp, b, iters):
    """求解 F x = b，只需要 Fisher-向量积 Avp(v) = F v，无需显式构造 F。"""
    x, r, p = torch.zeros_like(b), b.clone(), b.clone()
    rr = r @ r
    for _ in range(iters):
        Ap = Avp(p)
        alpha = rr / (p @ Ap)
        x += alpha * p; r -= alpha * Ap
        rr_new = r @ r
        if rr_new < 1e-10:
            break
        p = r + (rr_new / rr) * p; rr = rr_new
    return x


params = list(policy.parameters())
obs = torch.as_tensor(envs.reset(seed=0)[0], dtype=torch.float32)
ep_returns = []
for it in range(ITERS):
    # ---------- 采样 ----------
    B_obs, B_act, B_rew, B_done, B_val = [], [], [], [], []
    for t in range(N_STEPS):
        with torch.no_grad():
            a = Categorical(logits=policy(obs)).sample()
            B_val.append(value(obs).squeeze(-1))
        nobs, r, term, trunc, info = envs.step(a.numpy())
        r = torch.as_tensor(r, dtype=torch.float32)
        for i in np.flatnonzero(trunc):
            with torch.no_grad():
                r[i] += GAMMA * value(torch.as_tensor(info["final_obs"][i], dtype=torch.float32)).item()
        if "episode" in info:
            ep_returns += list(info["episode"]["r"][info["_episode"]])
        B_obs.append(obs); B_act.append(a); B_rew.append(r)
        B_done.append(torch.as_tensor(term | trunc, dtype=torch.float32))
        obs = torch.as_tensor(nobs, dtype=torch.float32)
    with torch.no_grad():
        next_v = value(obs).squeeze(-1)
    adv, gae = [], 0
    for t in reversed(range(N_STEPS)):
        v_next = next_v if t == N_STEPS - 1 else B_val[t + 1]
        delta = B_rew[t] + GAMMA * v_next * (1 - B_done[t]) - B_val[t]
        gae = delta + GAMMA * LAM * (1 - B_done[t]) * gae
        adv.insert(0, gae)
    adv, vals = torch.stack(adv).reshape(-1), torch.stack(B_val).reshape(-1)
    ret = adv + vals
    adv = (adv - adv.mean()) / (adv.std() + 1e-8)
    O, A = torch.stack(B_obs).reshape(-1, obs_dim), torch.stack(B_act).reshape(-1)

    # ---------- 策略更新：max  g^T Δθ   s.t.  ½ Δθ^T F Δθ ≤ δ ----------
    with torch.no_grad():
        old_dist = Categorical(logits=policy(O))
        old_logp = old_dist.log_prob(A)

    def surrogate():
        return ((policy_logp() - old_logp).exp() * adv).mean()

    def policy_logp():
        return Categorical(logits=policy(O)).log_prob(A)

    def mean_kl():
        return kl_divergence(old_dist, Categorical(logits=policy(O))).mean()

    g = flat_grad(surrogate(), params)

    def fvp(v):                                     # Fisher-向量积（KL 的 Hessian 乘向量）
        kl_g = flat_grad(mean_kl(), params, create_graph=True)
        return flat_grad(kl_g @ v, params) + DAMPING * v

    step_dir = conjugate_gradient(fvp, g, CG_ITERS)
    step_size = torch.sqrt(2 * MAX_KL / (step_dir @ fvp(step_dir)))
    full_step = step_size * step_dir
    old_params = torch.cat([p.data.reshape(-1) for p in params])
    with torch.no_grad():
        old_surr = surrogate()
    for frac in 0.5 ** np.arange(10):               # 回溯线搜索：满足 KL 约束且目标提升
        set_flat(params, old_params + frac * full_step)
        with torch.no_grad():
            if mean_kl() <= MAX_KL and surrogate() > old_surr:
                break
    else:
        set_flat(params, old_params)

    # ---------- 价值函数回归 ----------
    for _ in range(10):
        for idx in torch.randperm(len(O)).split(256):
            v_loss = ((value(O[idx]).squeeze(-1) - ret[idx]) ** 2).mean()
            v_opt.zero_grad(); v_loss.backward(); v_opt.step()

    if (it + 1) % 10 == 0:
        print(f"iter {it + 1:3d} | steps {(it + 1) * N_ENVS * N_STEPS:6d} | 最近20回合平均回报 {np.mean(ep_returns[-20:]):6.1f}")

# ---------- 训练结束：用确定性（argmax）策略评估 10 个回合 ----------
eval_env = gym.make("CartPole-v1")
scores = []
for i in range(10):
    o, _ = eval_env.reset(seed=10_000 + i)
    done, total = False, 0.0
    while not done:
        with torch.no_grad():
            act = int(policy(torch.as_tensor(o, dtype=torch.float32)).argmax())
        o, r, te, tr, _ = eval_env.step(act)
        total += r; done = te or tr
    scores.append(total)
print(f"贪心评估（10 回合）平均回报: {np.mean(scores):.1f}")

实际运行输出：

iter  10 | steps  20480 | 最近20回合平均回报  226.2
iter  20 | steps  40960 | 最近20回合平均回报  484.7
iter  30 | steps  61440 | 最近20回合平均回报  490.0
iter  40 | steps  81920 | 最近20回合平均回报  500.0
iter  50 | steps 102400 | 最近20回合平均回报  500.0
贪心评估（10 回合）平均回报: 500.0

TRPO 在约 4 万步时训练回报就接近满分，之后保持稳定，没有出现 A2C 那样的大幅波动——这正是信赖域约束带来的稳定性。（这是单个种子的结果，不同种子会有差异。）

与 PPO 的关系

	TRPO	PPO
约束方式	硬 KL 约束	裁剪概率比（或 KL 惩罚）
优化方法	共轭梯度 + 线搜索（二阶）	小批量 SGD / Adam（一阶）
每批数据的使用	一次更新	多轮（epoch）小批量更新
实现复杂度	高	低
与共享网络、Dropout 等兼容	较差	好

PPO 的作者正是 TRPO 的作者。PPO 保留了“不要让新策略离旧策略太远”的核心思想，但用简单得多的方式实现。

实践技巧

• $\delta$（max_kl）：0.01 是常用值；太大失去约束效果，太小学习缓慢。
• 阻尼系数：0.1 左右，防止 Fisher 矩阵近似病态。
• 批量要大：TRPO 每次迭代只做一次策略更新，需要较大批量（数千步）保证梯度估计准确。
• 价值函数：可以用多个 epoch 的 Adam 回归，也可以像原论文那样用信赖域方法更新。

常见误区

• 把 KL 方向写反：TRPO 约束的是 $D_{\text{KL}}({\pi }_{\text{old}}\parallel {\pi }_{\theta })$；对平均 KL 求 Hessian 时旧分布要 detach。
• 共轭梯度用了 $F$ 的显式矩阵：这在大网络上不可行，应使用 Fisher-向量积。
• 线搜索只检查 KL 不检查目标：可能接受一个让目标变差的更新。
• 对 TRPO 期望过高：单调提升保证建立在精确优势和无穷样本的假设上，实际仍可能波动。

小结

• TRPO 通过最大化代理目标、约束平均 KL 散度来控制策略更新幅度。
• 理论基础是性能差恒等式和由此导出的性能提升下界。
• 自然梯度方向 $F^{-1}g$ 由共轭梯度 + Fisher-向量积高效求得，线搜索保证约束。
• PPO 用一阶方法近似了 TRPO 的思想，是现在更常用的选择。

练习题

1. 把 MAX_KL 改为 0.001 和 0.1，比较学习速度和稳定性。
2. 去掉线搜索（直接走满步长），观察训练是否变得不稳定。
3. 证明对于离散分布，KL 散度在 $\theta ={\theta }_{\text{old}}$ 处的 Hessian 等于 Fisher 信息矩阵。
4. 在代码中记录每次更新后的实际 KL，验证它是否始终不超过 $\delta$。
5. 推导带约束问题 $max{g}^{\mathrm{\top }}x,\text{s.t.}\frac{1}{2}{x}^{\mathrm{\top }}Fx\le \delta$ 的解析解。

参考文献

• Kakade, S. (2001). A Natural Policy Gradient. NeurIPS.
• Kakade, S., & Langford, J. (2002). Approximately Optimal Approximate Reinforcement Learning. ICML.
• Schulman, J., Levine, S., Moritz, P., Jordan, M., & Abbeel, P. (2015). Trust Region Policy Optimization. ICML. arXiv:1502.05477
• Schulman, J. et al. (2016). High-Dimensional Continuous Control Using Generalized Advantage Estimation. ICLR. arXiv:1506.02438

下一步

上一节：Actor-Critic：A2C / A3C　｜　继续阅读：PPO