TRPO 深度解析：从数学原理到代码实现

一、问题背景：为什么需要 TRPO？

强化学习的目标是不断更新策略参数 $\theta$，使累计奖励最大化。最朴素的做法是直接做梯度上升：

$$\theta \leftarrow \theta +\alpha {\nabla }_{\theta }J(\theta )$$

但这有一个致命问题：步子太大会摔跤。一次更新过大，新策略可能远比旧策略差，而且强化学习没有”后悔药”——坏策略采集的数据质量更差，训练可能一路崩塌。

TRPO 的解决思路是：在每次更新时，显式限制新旧策略的差距不超过 $ϵ$：

$$\underset{\theta }{\max }{\mathbb{E}}_t\left[\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t|s_t)}{A}_t\right]\text{s.t.}{\mathbb{E}}_s\left[D_{\text{KL}}\left({\pi }_{{\theta }_{\text{old}}}(\cdot |s)\Vert {\pi }_{\theta }(\cdot |s)\right)\right]\le ϵ$$

注意约束中有 ${\mathbb{E}}_s[\cdot ]$：KL 散度对所有可能出现的状态 $s$ 取期望，保证在整个状态分布上策略变化都不过大。

二、目标函数拆解

策略参数 $\theta$

$\theta$ 是神经网络（策略网络）的全部权重和偏置。策略网络输入状态 $s$，输出动作的概率分布 ${\pi }_{\theta }(\cdot |s)$。

在语言模型场景中，$s_t$ = “原始问题 + 已生成的前 $t$ 个 token”，$a_t$ = “第 $t+1$ 个 token 的选择”。

重要性采样比率

$$r_t(\theta )=\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t|s_t)}$$

我们用旧策略 ${\pi }_{{\theta }_{\text{old}}}$ 采集数据，但希望估计新策略 ${\pi }_{\theta }$ 的表现。重要性采样的核心等式：

$${\mathbb{E}}_{a\sim {\pi }_{{\theta }_{\text{old}}}}\left[r_t(\theta )\cdot f(a)\right]={\mathbb{E}}_{a\sim {\pi }_{\theta }}\left[f(a)\right]$$

• $r_t>1$：新策略比旧策略更倾向选这个动作
• $r_t<1$：新策略比旧策略更不倾向选这个动作
• $r_t=1$：两个策略对此动作概率相同

优势函数 $A_t$

$$A_t=Q(s_t,a_t)-V(s_t)$$

符号	含义
$Q(s_t,a_t)$	在状态 $s_t$ 执行动作 $a_t$ 后的期望累计奖励
$V(s_t)$	在状态 $s_t$ 下，按当前策略执行的期望累计奖励（基线）
$A_t>0$	这个动作比平均水平好，应该强化
$A_t<0$	这个动作比平均水平差，应该抑制

减去基线 $V(s_t)$ 不改变梯度的期望值（因为 ${\mathbb{E}}_{a\sim {\pi }_{\theta }}[V(s_t)]=V(s_t)$ 与动作选择无关），但能大幅降低方差，让训练更稳定。

对时间步 $t$ 取期望

$${\mathbb{E}}_t[\cdots ]\approx \frac{1}{T}\sum _{t=1}^T(\cdots )$$

等价于对一条（或一批）轨迹中所有时间步求平均，用来估计期望值。

三、KL 散度约束与二阶泰勒展开

KL 散度的定义

对于离散分布：

$$D_{\text{KL}}(p\Vert q)=\sum _{a}p(a)\log \frac{p(a)}{q(a)}$$

对于连续分布：

$$D_{\text{KL}}(p\Vert q)=\int p(x)\log \frac{p(x)}{q(x)}dx$$

性质：

• $D_{\text{KL}}(p\Vert p)=0$（自己与自己的散度为零）
• $D_{\text{KL}}(p\Vert q)\ge 0$ 恒成立（由 Jensen 不等式保证）
• 不对称：$D_{\text{KL}}(p\Vert q)\ne D_{\text{KL}}(q\Vert p)$（一般情况下）

对 KL 散度做二阶泰勒展开

设 $\delta =\theta -{\theta }_{\text{old}}$，在 $\delta =0$ 处展开 $D_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}\Vert {\pi }_{{\theta }_{\text{old}}+\delta })$：

$$D_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}\Vert {\pi }_{{\theta }_{\text{old}}+\delta })=\underset{\text{零阶项}=0}{\underbrace{D_{\text{KL}}{|}_{\delta =0}}}+\underset{\text{一阶项}=0}{\underbrace{{\nabla }_{\delta }{D}_{\text{KL}}{|}_{\delta =0}^{\top }\delta }}+\underset{\text{二阶项}}{\underbrace{\frac{1}{2}{\delta }^{\top }F\delta }}+O(\Vert \delta {\Vert }^3)$$

零阶项 = 0

$$D_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}\Vert {\pi }_{{\theta }_{\text{old}}+\delta }){|}_{\delta =0}=D_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}\Vert {\pi }_{{\theta }_{\text{old}}})=\sum _a{\pi }_{{\theta }_{\text{old}}}(a|s)\log \frac{{\pi }_{{\theta }_{\text{old}}}(a|s)}{{\pi }_{{\theta }_{\text{old}}}(a|s)}=0$$

一阶项 = 0（完整推导）

对 $D_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}\Vert {\pi }_{\theta })$ 关于 $\theta$ 求梯度：

$${\nabla }_{\theta }{D}_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}\Vert {\pi }_{\theta })={\nabla }_{\theta }\sum _a{\pi }_{{\theta }_{\text{old}}}(a|s)\log \frac{{\pi }_{{\theta }_{\text{old}}}(a|s)}{{\pi }_{\theta }(a|s)}=-{\nabla }_{\theta }\sum _a{\pi }_{{\theta }_{\text{old}}}(a|s)\log {\pi }_{\theta }(a|s)$$

$$=-\sum _a{\pi }_{{\theta }_{\text{old}}}(a|s)\cdot \frac{{\nabla }_{\theta }{\pi }_{\theta }(a|s)}{{\pi }_{\theta }(a|s)}$$

在 $\theta ={\theta }_{\text{old}}$ 处，${\pi }_{\theta }={\pi }_{{\theta }_{\text{old}}}$，代入：

$${{\nabla }_{\theta }{D}_{\text{KL}}|}_{\theta ={\theta }_{\text{old}}}=-\sum _a{\nabla }_{\theta }{\pi }_{\theta }(a|s){|}_{\theta ={\theta }_{\text{old}}}=-{\nabla }_{\theta }\underset{=1,\forall \theta }{\underbrace{\sum _a{\pi }_{\theta }(a|s)}}{|}_{\theta ={\theta }_{\text{old}}}=0$$

由于概率归一化 ${\sum }_a{\pi }_{\theta }(a|s)=1$ 对任意 $\theta$ 成立，其关于 $\theta$ 的梯度恒为零向量。

二阶项：从一维到多维

一维情形：$f(x_0+\delta )\approx f(x_0)+f^{\prime }(x_0)\delta +\frac{1}{2}{f}^{\prime \prime }(x_0){\delta }^2$

多维推广：

$$f({\theta }_0+\delta )\approx f({\theta }_0)+\nabla f({\theta }_0{)}^{\top }\delta +\frac{1}{2}{\delta }^{\top }H({\theta }_0)\delta$$

其中 Hessian 矩阵的元素为 $H_{ij}=\frac{{\partial }^{2}f}{\partial {\theta }_i\partial {\theta }_j}$，描述各参数方向上的曲率。

对 KL 散度而言，这个 Hessian 恰好等于 Fisher 信息矩阵 $F$（见下节证明）。

三阶项的形状

$$\frac{1}{6}\sum _{i,j,k}{\mathcal{T}}_{ijk}{\delta }_i{\delta }_j{\delta }_k,{\mathcal{T}}_{ijk}=\frac{{\partial }^{3}f}{\partial {\theta }_i\partial {\theta }_j\partial {\theta }_k}$$

$p$ 阶项的一般规律：系数 $\frac{1}{p!}$ 乘以 $p$ 阶张量与 $\delta$ 的 $p$ 次缩并。TRPO 忽略三阶及以上项。

四、Fisher 信息矩阵

定义

$$F(\theta )={\mathbb{E}}_{s\sim \rho ,a\sim {\pi }_{\theta }(\cdot |s)}\left[{\nabla }_{\theta }\log {\pi }_{\theta }(a|s){\nabla }_{\theta }\log {\pi }_{\theta }(a|s{)}^{\top }\right]$$

其中 ${\nabla }_{\theta }\log {\pi }_{\theta }(a|s)$ 称为 score 函数（得分函数）。$F$ 是 score 函数外积的期望，是一个 $n\times n$ 的半正定矩阵。

对数导数技巧

$${\nabla }_{\theta }\log {\pi }_{\theta }(a|s)=\frac{{\nabla }_{\theta }{\pi }_{\theta }(a|s)}{{\pi }_{\theta }(a|s)}$$

这是策略梯度定理的核心变换：将”对概率求梯度”转化为”对对数概率求梯度”，数值上更稳定，且便于利用 ${\mathbb{E}}_{a\sim {\pi }_{\theta }}[\cdot ]$ 进行期望计算。

Fisher 矩阵 = KL 散度的 Hessian（证明）

对 $D_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}\Vert {\pi }_{\theta })$ 关于 $\theta$ 求二阶导，在 $\theta ={\theta }_{\text{old}}$ 处：

$$H_{ij}=-\frac{{\partial }^2}{\partial {\theta }_i\partial {\theta }_j}\sum _a{\pi }_{{\theta }_{\text{old}}}(a|s)\log {\pi }_{\theta }(a|s){|}_{\theta ={\theta }_{\text{old}}}$$

$$=\sum _a{\pi }_{{\theta }_{\text{old}}}(a|s)\cdot \frac{\partial \log {\pi }_{\theta }(a|s)}{\partial {\theta }_i}\cdot \frac{\partial \log {\pi }_{\theta }(a|s)}{\partial {\theta }_j}{|}_{\theta ={\theta }_{\text{old}}}=F_{ij}$$

同时有等价形式（利用 score 函数期望为零推导）：

$$F=-{\mathbb{E}}_{a\sim {\pi }_{\theta }}\left[{\nabla }_{\theta }^2\log {\pi }_{\theta }(a|s)\right]$$

自然梯度

在 KL 约束 $\frac{1}{2}{\delta }^{\top }F\delta \le ϵ$ 下最大化线性目标 $g^{\top }\delta$（$g={\nabla }_{\theta }J(\theta )$ 为策略梯度），用拉格朗日乘子法：

$$\mathcal{L}(\delta ,\lambda )=g^{\top }\delta -\lambda \left(\frac{1}{2}{\delta }^{\top }F\delta -ϵ\right)$$

$$\frac{\partial \mathcal{L}}{\partial \delta }=g-\lambda F\delta =0\Rightarrow {\delta }^{\ast }=\frac{1}{\lambda }{F}^{-1}g$$

步长由约束 $\frac{1}{2}{\delta }^{\ast \top }F{\delta }^{\ast }=ϵ$ 确定，解得 $\frac{1}{\lambda }=\sqrt{\frac{2ϵ}{g^{\top }{F}^{-1}g}}$，最终：

$$\boxed{{\delta }^{\ast }=\sqrt{\frac{2ϵ}{g^{\top }{F}^{-1}g}}{F}^{-1}g}$$

这称为自然梯度方向，考虑了参数空间的黎曼曲率，比普通梯度方向更高效。

为什么 Fisher 矩阵很贵

网络规模	$n$	存储 $F$	求逆 $F^{-1}$
小型 MLP	$10^4$	400 MB	可行但慢
中型策略网络	$10^6$	≈ 4 TB	不可行
大语言模型	$10^9$	$\gg$ PB	完全不可行

这正是 TRPO 需要共轭梯度法的根本原因。

五、共轭梯度法求解 $F^{-1}g$

关键洞察

我们要的不是整个逆矩阵，只是一个向量 $x=F^{-1}g$，等价于求解线性方程组：

$$Fx=g$$

共轭梯度法（Conjugate Gradient, CG）专门解 $Ax=b$ 这类对称正定线性方程组（Fisher 矩阵恰好满足这一条件），且有超强性质：

只需要能计算”$F$ 与某向量 $v$ 的乘积 $Fv$“，不需要显式存储 $F$

CG 迭代 $k$ 步即可（$k\ll n$），总计算量 $O(kn)$。

Hessian-Vector Product（HVP）

利用自动微分，$Fv$ 可以通过两次反向传播计算：

$$Fv={\nabla }_{\theta }\left[{\left({\nabla }_{\theta }{D}_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}\Vert {\pi }_{\theta })\right)}^{\top }v\right]$$

步骤：

1. 前向传播，计算 $D_{\text{KL}}$
2. 第一次反向传播：得到 $g_{\text{KL}}={\nabla }_{\theta }{D}_{\text{KL}}$（需要 create_graph=True 保留计算图）
3. 计算标量 $s=g_{\text{KL}}^{\top }v$
4. 第二次反向传播：${\nabla }_{\theta }s=Fv$

全程只需 $O(n)$ 内存，完全可行。

方法	存储	计算复杂度
显式存储并求逆 $F^{-1}$	$O(n^2)$	$O(n^3)$
共轭梯度 + HVP	$O(n)$	$O(kn)$，$k\ll n$

六、TRPO 伪代码逐行解析

import torch
import torch.nn as nn
from torch.distributions import Categorical

class TRPO_Concept(nn.Module):
    def __init__(self, policy_net):
        super().__init__()
        self.policy = policy_net

nn.Module 是 PyTorch 所有神经网络的基类，提供参数管理、保存/加载等功能。super().__init__() 必须调用，否则参数追踪机制不会生效。

def get_kl(self, old_dist, new_dist):
    return torch.distributions.kl.kl_divergence(old_dist, new_dist).mean()

计算一批状态上的平均 KL 散度：

$${\bar{D}}_{\text{KL}}=\frac{1}{|B|}\sum _{s\in B}{D}_{\text{KL}}\left({\pi }_{{\theta }_{\text{old}}}(\cdot |s)\Vert {\pi }_{\theta }(\cdot |s)\right)$$

.mean() 对 batch 内所有状态取平均，对应约束中的 ${\mathbb{E}}_s[\cdot ]$。

def trpo_step(self, states, actions, advantages, old_log_probs, max_kl=0.01):

参数	数学符号	形状
states	$\{s_t\}$	[batch, state_dim]
actions	$\{a_t\}$	[batch]
advantages	$\{A_t\}$	[batch]
old_log_probs	${\log\pi_{\theta_\text{old}}(a_t	s_t)}$
max_kl	$ϵ$	标量，默认 0.01

Step 1：计算当前策略分布

logits = self.policy(states)
new_dist = Categorical(logits=logits)
new_log_probs = new_dist.log_prob(actions)

• logits：网络输出原始得分，形状 [batch, action_dim]
• Categorical：内部执行 softmax，将 logits 转成合法概率分布：

$${\pi }_{\theta }(a|s)=\frac{e^{z_a}}{{\sum }_{a^{\prime }}{e}^{z_{a^{\prime }}}}$$

• log_prob(actions)：取出实际执行动作的 $\log {\pi }_{\theta }(a_t|s_t)$，形状 [batch]

用 logits 而非直接输出概率，是因为 softmax 内部做了数值稳定处理（log-sum-exp trick），避免概率连乘下溢到 0。

Step 2：计算重要性采样比率

ratio = torch.exp(new_log_probs - old_log_probs)

对应公式 $r_t(\theta )=\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }_{\text{old}}}(a_t|s_t)}$，利用对数性质：

$$r_t(\theta )=\exp (\log {\pi }_{\theta }(a_t|s_t)-\log {\pi }_{{\theta }_{\text{old}}}(a_t|s_t))$$

先减后 exp，比直接相除数值更稳定（避免两个极小概率相除产生数值爆炸）。

Step 3：计算代理目标函数

loss = -(ratio * advantages).mean()

对应目标函数：

$$L(\theta )={\mathbb{E}}_t\left[r_t(\theta )\cdot A_t\right]\approx \frac{1}{|B|}\sum _t{r}_t(\theta )\cdot A_t$$

为什么取负号？

PyTorch 的优化器只做梯度下降（最小化），而我们想最大化 $L(\theta )$：

$$\underset{\theta }{\max }L(\theta )⟺\underset{\theta }{\min }-L(\theta )$$

类比：你想找山顶，但手里只有”下坡工具”——把地图翻转，让山顶变谷底，下到谷底就等于爬上了山顶。强化学习里凡是要最大化的目标，在代码中均加负号。

Step 4：计算梯度并展平

grads = torch.autograd.grad(loss, self.policy.parameters(), create_graph=True)
flat_grads = torch.cat([g.view(-1) for g in grads])

• torch.autograd.grad：手动触发反向传播，返回 $g={\nabla }_{\theta }L(\theta )$
• create_graph=True：保留计算图，使后续可对梯度再次求导（HVP 必需）

各层参数形状不同，展平操作将所有梯度拼成一个长向量：

$$g=\left[\underset{\text{第一层权重}}{\underbrace{g_{W_1}^{\text{vec}}}},\underset{\text{第一层偏置}}{\underbrace{g_{b_1}}},\underset{\text{第二层权重}}{\underbrace{g_{W_2}^{\text{vec}}}},\dots \right]\in {\mathbb{R}}^n$$

例如一个小网络：

Linear(64→32): weight [32,64] → 2048 个数，bias [32] → 32 个数
Linear(32→8):  weight [8,32]  →  256 个数，bias  [8] →  8 个数
flat_grads: shape [2344]

Step 5（省略）：共轭梯度求解更新方向

# step_direction = conjugate_gradient(HVP_function, flat_grads)

这里要解：

$$Fd=g⟹d=F^{-1}g$$

其中 $F$ 通过 HVP 隐式使用（见第五节）。

得到更新方向 $d$ 后，TRPO 还需做线性搜索（Line Search）：沿方向 $d$ 找到满足以下两个条件的最大步长 $\alpha$：

1. KL 约束：${\mathbb{E}}_s\left[D_{\text{KL}}({\pi }_{{\theta }_{\text{old}}}\Vert {\pi }_{\theta +\alpha d})\right]\le ϵ$
2. 目标提升：$L(\theta +\alpha d)\ge L(\theta )$

最终更新：$\theta \leftarrow \theta +\alpha d$

七、完整数据流

states  ──→  policy_net  ──→  logits  ──→  Categorical
                                               │
                                               ▼
                                       new_log_probs [batch]
                                               │
                             old_log_probs ────┤
                                               ▼
                                  ratio = exp(new - old)  [batch]
                                               │
                             advantages ───────┤
                                               ▼
                               L = (ratio × adv).mean()
                               loss = -L
                                               │
                                               ▼
                               g = ∂loss/∂θ  (create_graph=True)
                               flat_grads ∈ ℝⁿ
                                               │
                                               ▼
                         共轭梯度法：解 Fd = g
                         （每步用 HVP 算 Fv，无需存 F）
                                               │
                                               ▼
                         线性搜索：找满足 KL ≤ ε 的步长 α
                                               │
                                               ▼
                              θ ← θ + α · d

八、总结与对比

特性	普通策略梯度	TRPO
更新规则	$\theta \leftarrow \theta +\alpha g$	$\theta \leftarrow \theta +\alpha F^{-1}g$
步长控制	手动设置学习率	KL 约束自动确定
稳定性	容易崩塌	有理论保证单调提升
计算代价	低 $O(n)$	高 $O(kn)$，需 CG + HVP
实现难度	简单	复杂

TRPO 的核心贡献是将”策略更新不能太激进”这一直觉严格数学化，并给出可行的计算方案。其后继者 PPO 用更简单的裁剪目标函数近似实现了类似效果，成为目前最主流的策略优化算法：

$$L^{\text{CLIP}}(\theta )={\mathbb{E}}_t\left[\min \left(r_t(\theta )A_t,\text{clip}(r_t(\theta ),1-ϵ,1+ϵ)\cdot A_t\right)\right]$$

参考

• Schulman et al., Trust Region Policy Optimization, ICML 2015. arXiv:1502.05477
• Kakade & Langford, Approximately Optimal Approximate Reinforcement Learning, ICML 2002.
• OpenAI Spinning Up: TRPO
• 配套阅读：MiniMind 源码阅读与架构分析（同属「源码解析」合集）