策略梯度这玩意儿，我最早接触的时候踩过一个特别典型的坑：把REINFORCE的损失函数写反了符号，训练了三百多轮，策略熵一路涨到天上，智能体在环境里疯狂转圈，动作概率几乎均匀分布。当时盯着TensorBoard上那条平得像尺子画出来的回报曲线，一度怀疑是环境本身有问题。后来把梯度公式从头推了一遍才反应过来——策略梯度不是监督学习，它的"标签"是采样出来的回报，符号错一位，整个优化方向就反了。这件事让我意识到，策略梯度算法看起来公式简洁，但真正落地到PyTorch代码里，从采样、算回报、构造损失到更新参数，每一步都有容易翻车的地方。
这篇内容想聊的就是PyTorch强化学习中的策略梯度算法，重点不是把教科书上的公式再抄一遍，而是把从数学到代码的那段"翻译"过程讲透。适合已经了解一点强化学习基本概念、能看懂PyTorch基础代码、但真正动手写策略梯度时总觉得哪里不对劲的读者。我会从策略梯度的核心思想切入，拆解REINFORCE、带基线的REINFORCE、以及Actor-Critic这条演进路线，每一块都配上可直接跑的PyTorch实现思路，再把我自己踩过的坑和调试经验摊开讲。读完你应该能独立写出一个在CartPole这类经典控制任务上稳定收敛的策略梯度智能体，并且知道它为什么不收敛时该往哪儿查。
1. 策略梯度到底在优化什么：从"选动作"到"调概率"
很多人刚学完Q-learning那一套再转过来看策略梯度，第一反应是困惑：价值方法明明已经能通过Q值选动作了，为什么还要单独学一个策略网络？这个问题不搞清楚，后面写代码就是照葫芦画瓢，出了问题完全无从下手。
1.1 价值方法与策略方法的根本分歧
价值方法（Q-learning、DQN这一类）的核心是学一个价值函数，然后用argmax从价值里"挑"动作。它把决策这件事外包给了价值估计，策略本身是隐式的、确定性的。这套思路在离散动作空间里很自然，但一旦动作空间连续，argmax就没法直接算了——你不可能对一个连续区间做穷举。
策略方法走的是另一条路：直接参数化策略本身，记作π(a|s; θ)，网络输出的就是动作的概率分布（离散）或者分布的参数（连续，比如高斯分布的均值和方差）。决策时从分布里采样，而不是取最大值。这个转变带来的直接好处是：连续动作、随机策略、以及"策略本身有结构"的场景都能自然处理。
我个人的理解是，价值方法像"先给每个选项打分再选最高分"，策略方法像"直接训练一个会做选择的直觉"。后者在动作维度高、或者最优策略本身就需要随机性（比如石头剪刀布这类博弈）时优势明显。
1.2 目标函数：期望回报对策略参数的梯度
策略梯度的目标很直白：最大化期望回报
$$J(\theta) = \mathbb{E}{\tau \sim \pi\theta}[R(\tau)]$$
其中τ是一条完整轨迹，R(τ)是这条轨迹的累计回报。问题在于，期望是对策略π_θ采样出来的轨迹求的，而策略参数θ又决定了采样分布——这层依赖关系让直接对θ求导变得棘手，因为你没法把梯度传过"采样"这个随机操作。
策略梯度定理给出了答案：
$$\nabla_\theta J(\theta) = \mathbb{E}{\tau \sim \pi\theta}\left[\sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t|s_t) \cdot R(\tau)\right]$$
这个式子的直觉是：提高那些带来高回报的动作的对数概率，降低低回报动作的概率。∇log π这一项在概率上升时为正，乘以回报R就变成了"回报越高，往这个方向推得越狠"。这就是策略梯度全部的灵魂，后面所有变体都是在这个骨架上做方差缩减。
注意：这里的R(τ)在实际实现中通常换成从时刻t开始的回报（reward-to-go），而不是整条轨迹的总回报。这个细节后面会专门讲，它对训练稳定性影响很大。
1.3 为什么"采样"能绕过不可导的决策
有个细节值得单独说清楚，因为它经常被跳过。策略网络输出概率，然后我们"采样"一个动作——采样这个操作本身是不可导的，梯度没法穿过它。策略梯度定理的巧妙之处在于，它用对数似然的梯度∇log π(a|s)替代了对采样操作的求导，把"不可导的采样"转化成了"可导的概率对数"。数学上这是通过log-derivative trick（也叫score function estimator）实现的：
$$\nabla_\theta \pi_\theta(a|s) = \pi_\theta(a|s) \nabla_\theta \log \pi_\theta(a|s)$$
这个恒等式让期望里的梯度可以写成对log π求导再乘回报的形式。落到PyTorch里，我们不需要手动实现这个trick——只要用torch.log(prob)配合loss = -(log_prob * return).mean()，反向传播时PyTorch自动就算出了正确的梯度。但理解这层原理很重要，否则你看到loss里那个负号、看到为什么要.detach()回报，就会一头雾水。
2. REINFORCE：最朴素的策略梯度实现
REINFORCE是策略梯度最原始的形态，也是理解后续所有变体的地基。它的流程简单到可以用一句话概括：跑完一整条轨迹，用累计回报给每个动作的对数概率加权，做一次梯度上升。
2.1 完整训练循环的PyTorch骨架
先看策略网络。离散动作空间下，输出层维度等于动作数，接softmax得到概率分布：
import torch import torch.nn as nn import torch.nn.functional as F class PolicyNet(nn.Module): def __init__(self, state_dim, action_dim, hidden=128): super().__init__() self.fc1 = nn.Linear(state_dim, hidden) self.fc2 = nn.Linear(hidden, hidden) self.out = nn.Linear(hidden, action_dim) def forward(self, x): x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) return F.softmax(self.out(x), dim=-1)训练循环的关键在于"先采样一整条轨迹，再统一算回报，最后更新一次"。这和DQN那种边交互边更新的节奏完全不同：
def collect_episode(env, policy): log_probs, rewards = [], [] state, _ = env.reset() done = False while not done: state_t = torch.FloatTensor(state).unsqueeze(0) probs = policy(state_t) dist = torch.distributions.Categorical(probs) action = dist.sample() log_probs.append(dist.log_prob(action)) state, reward, terminated, truncated, _ = env.step(action.item()) rewards.append(reward) done = terminated or truncated return log_probs, rewards这里用torch.distributions.Categorical而不是手动torch.multinomial，好处是它直接提供log_prob方法，省去自己算对数的麻烦，也避免了数值稳定性问题。
2.2 回报计算：reward-to-go与折扣因子
拿到奖励序列后，要把它转成每个时刻的回报。这里有个新手常犯的错误：直接用整条轨迹的总回报给每个动作加权。正确做法是用从当前时刻往后的累计折扣回报（reward-to-go）：
def compute_returns(rewards, gamma=0.99): returns = [] R = 0 for r in reversed(rewards): R = r + gamma * R returns.insert(0, R) return returns为什么必须用reward-to-go？因为时刻t的动作只影响t之后的回报，用整条轨迹的总回报会把t之前的奖励也算进来，引入额外的方差。这个改动看起来小，实测对收敛速度的提升相当明显。折扣因子γ一般取0.99，任务越长、越需要"看得远"，γ就越接近1；但γ太接近1会让回报的方差变大，收敛变慢，这是个需要权衡的参数。
2.3 损失函数与一次梯度上升
把log概率和回报乘起来，取负号做梯度下降（等价于对目标做梯度上升）：
def update_policy(optimizer, log_probs, returns): returns_t = torch.FloatTensor(returns) # 标准化回报，降低方差 returns_t = (returns_t - returns_t.mean()) / (returns_t.std() + 1e-8) loss = 0 for log_prob, R in zip(log_probs, returns_t): loss += -log_prob * R optimizer.zero_grad() loss.backward() optimizer.step()注意回报标准化这一步。REINFORCE的方差大是出了名的，把回报减去均值除以标准差，相当于给梯度做了一个缩放，让不同回合之间的梯度尺度更一致。这不是理论必需的，但实测下来几乎是标配操作，不加的话训练曲线会抖得没法看。
提示：
returns_t在构造loss时不需要.detach()，因为它本身就是从奖励算出来的常数，不参与梯度计算。但如果你用的是价值网络估计的回报，那就必须.detach()，否则梯度会错误地流进价值网络。
3. 方差缩减：基线、优势函数与Actor-Critic
REINFORCE能跑，但跑得慢、抖得厉害。原因在于它用采样回报直接估计梯度，而回报的方差很大。策略梯度这条线上几乎所有重要的改进，本质上都在做同一件事：在保持梯度期望不变的前提下，把方差压下去。
3.1 基线为什么能降方差却不引入偏差
引入基线b(s)的核心公式是：
$$\nabla_\theta J(\theta) = \mathbb{E}\left[\nabla_\theta \log \pi_\theta(a|s) \cdot (R - b(s))\right]$$
关键结论是：减去一个只依赖状态、不依赖动作的基线b(s)，梯度的期望不变。证明依赖于一个事实——对任意状态s，Σ_a π(a|s) ∇log π(a|s) = ∇Σ_a π(a|s) = ∇1 = 0。也就是说，基线项乘上∇log π在动作上求和恰好为零，所以它不改变期望，只改变方差。
直觉上，基线的作用是"给回报定一个参照系"。如果某个动作的回报比平均水平高，就提高它的概率；比平均低，就降低。用平均回报做基线，相当于把"绝对好坏"变成了"相对好坏"，梯度的尺度更稳定。实践中基线通常用状态价值函数V(s)来估计，这就自然过渡到了Actor-Critic。
3.2 用价值网络做基线的双网络结构
带基线的REINFORCE需要两个网络：策略网络（Actor）和价值网络（Critic）。价值网络学的是V(s)，用来当基线：
class ValueNet(nn.Module): def __init__(self, state_dim, hidden=128): super().__init__() self.fc1 = nn.Linear(state_dim, hidden) self.fc2 = nn.Linear(hidden, hidden) self.out = nn.Linear(hidden, 1) def forward(self, x): x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) return self.out(x).squeeze(-1)更新时，Actor用(R - V(s))作为权重，Critic用均方误差拟合V(s) → R：
def update_actor_critic(actor_opt, critic_opt, actor, critic, states, log_probs, returns): states_t = torch.FloatTensor(states) returns_t = torch.FloatTensor(returns) values = critic(states_t).detach() advantages = returns_t - values # Actor损失 actor_loss = 0 for log_prob, adv in zip(log_probs, advantages): actor_loss += -log_prob * adv actor_opt.zero_grad() actor_loss.backward() actor_opt.step() # Critic损失 values_pred = critic(states_t) critic_loss = F.mse_loss(values_pred, returns_t) critic_opt.zero_grad() critic_loss.backward() critic_opt.step()这里values = critic(states_t).detach()那个.detach()是必须的。Actor的梯度不应该流进Critic，否则两个网络会互相干扰，训练容易崩。我见过不少人在这里漏了detach，结果Critic被Actor的梯度带偏，价值估计越来越离谱。
3.3 优势函数A(s,a)的统一视角
把R - V(s)这个量正式命名为优势函数A(s,a)，它衡量的是"在状态s下选动作a，比平均水平好多少"。策略梯度的通用形式就变成了：
$$\nabla_\theta J(\theta) = \mathbb{E}\left[\nabla_\theta \log \pi_\theta(a|s) \cdot A(s,a)\right]$$
从这个视角看，REINFORCE是A(s,a)用蒙特卡洛回报估计（无偏但高方差），带基线的版本是用R - V(s)估计（方差降低），而Actor-Critic则是用一步TD误差r + γV(s') - V(s)来估计优势（偏差引入但方差更低）。这三者是同一条谱系上的不同取舍，理解了这个谱系，你就不会把它们当成三个孤立的算法来记。
| 方法 | 优势估计方式 | 偏差 | 方差 | 样本效率 |
|---|---|---|---|---|
| REINFORCE | 蒙特卡洛回报R | 无偏 | 高 | 低 |
| 带基线REINFORCE | R - V(s) | 无偏 | 中 | 中 |
| Actor-Critic | TD误差 | 有偏 | 低 | 高 |
这张表是我自己在选型时最常回看的东西。任务回合短、算力够，REINFORCE加基线就够用；任务回合长、样本宝贵，就得上Actor-Critic甚至更复杂的优势估计（比如GAE）。
4. 实操中的坑：从能跑到跑得好之间隔着什么
代码写出来能跑，和训练真的收敛，中间差着十万八千里。这一节把我调试策略梯度时反复遇到的几个问题摊开讲，每个都给出排查思路。
4.1 梯度爆炸与策略熵坍塌
策略梯度训练中最常见的两种病态，方向正好相反。一种是策略熵坍塌：softmax输出越来越尖锐，几乎总是选同一个动作，探索彻底停止，回报卡在某个局部最优上不去。另一种是梯度爆炸：某次更新步子太大，策略参数被推到极端值，之后再也恢复不过来。
熵坍塌的排查方法是监控策略熵-(probs * log(probs)).sum()。如果它快速掉到接近0，说明探索不足。缓解手段是在损失里加熵正则项：
entropy = dist.entropy().mean() loss = -log_prob * advantage - entropy_coef * entropyentropy_coef一般取0.01左右，太大策略会一直随机、学不到东西，太小起不到正则作用。梯度爆炸则靠梯度裁剪兜底：
torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=0.5)max_norm取0.5到1.0之间比较稳妥。这两个手段我基本是默认加上的，属于"加了不亏"的保险。
4.2 回报标准化与折扣因子的联动
前面提过回报标准化，但这里有个联动细节值得单独说。当你用reward-to-go加折扣因子算回报时，不同时刻的回报尺度差异很大（越靠前的时刻累计的奖励越多），直接标准化整条序列会让靠后的时刻回报被压得很小。我的做法是按整条轨迹的回报做标准化，而不是逐时刻标准化，这样保持了时刻之间的相对关系。另外γ的选择要和回合长度匹配：CartPole这种几百步的回合，γ=0.99合适；如果回合只有十几步，γ可以取0.95甚至更低，避免回报被过度放大。
4.3 采样效率与并行环境的取舍
REINFORCE是on-policy的，采一批数据只能用一次，更新完就得扔掉重采。这意味着样本效率天然很低。在小任务上这不是问题，但任务一复杂，单环境串行采样就成了瓶颈。可行的加速手段是开多个环境并行采样，把多个回合的数据拼成一个大batch再更新。PyTorch本身不负责环境并行，通常借助向量化环境封装来实现。我实测下来，8个环境并行能把wall-clock训练时间压到原来的三分之一左右，但要注意不同环境的回合长度不一致，拼接数据时得按回合边界正确切分，否则reward-to-go会算错。
注意：并行环境拼接数据时，千万不要把不同回合的奖励序列首尾相连后再统一算回报，那样会把上一个回合的回报错误地累加到下一个回合。正确做法是每个回合单独算reward-to-go，再拼起来。
4.4 一个具体的排查案例
有次我在一个连续动作任务上训练，回报曲线前期涨得挺好，到中途突然断崖式下跌再也没起来。查了半天发现是动作分布的方差参数没有做约束——高斯策略的std直接作为网络输出，训练中一度变成负数，torch.distributions.Normal对负std的处理导致采样异常，进而让回报估计彻底失真。修复方法是用softplus或者exp把std约束为正：
mean = self.mean_head(x) std = F.softplus(self.std_head(x)) + 1e-5 dist = torch.distributions.Normal(mean, std)这个坑的教训是：策略网络的输出如果有物理约束（概率非负、方差为正），一定要在结构上保证，而不是指望训练过程"自己学会"。类似的还有离散动作的softmax数值稳定性，PyTorch的F.softmax内部已经做了max减法，但如果你手动实现就得自己注意。
5. 从策略梯度到PPO：为什么工业界更偏爱后者
纯策略梯度（REINFORCE、基础Actor-Critic）在玩具任务上够用，但真正上强度时，几乎所有人都会转向PPO这类算法。理解这个转向的原因，能帮你判断什么时候该停手、什么时候该升级。
5.1 步长敏感：策略梯度最致命的软肋
策略梯度对学习率极其敏感。学习率稍大，一次更新就可能把策略推到很差的区域，而on-policy的特性决定了你没法用旧数据补救，只能眼睁睁看着它崩。这个问题的根源在于：策略梯度的更新没有"信任域"约束，参数空间里走一小步，策略空间里可能走了一大步。
PPO的核心改进就是给更新加了约束——用重要性采样比率π_new/π_old衡量新旧策略的差异，并通过clip把比率限制在[1-ε, 1+ε]区间内，防止单次更新把策略改得太狠。这个clip操作实现起来就几行代码，但效果立竿见影，训练稳定性提升一个档次。
5.2 重要性采样与clip的直觉
重要性采样的作用是让同一批数据能被多次利用（off-policy化），从而提高样本效率。比率r_t = π_new(a|s) / π_old(a|s)衡量新策略相对旧策略对这个动作的偏好变化。PPO的clip目标：
$$L^{clip} = \mathbb{E}\left[\min(r_t A_t, \text{clip}(r_t, 1-\epsilon, 1+\epsilon) A_t)\right]$$
min加clip的组合保证了：当优势为正时，比率超过1+ε就不再给额外奖励；当优势为负时，比率低于1-ε也不再额外惩罚。这样策略不会因为某个动作优势大就疯狂提高它的概率，更新被限制在一个温和的范围内。ε通常取0.2。
5.3 什么时候该从策略梯度升级到PPO
我的经验判断标准是：如果任务回合短、状态维度低、训练能在几分钟内跑完，纯策略梯度加基线完全够用，没必要上PPO增加复杂度。但如果出现以下信号，就该考虑升级了——训练曲线对学习率极度敏感、换个随机种子结果天差地别、单次更新后回报骤降、或者样本采集成本很高（比如需要真实设备交互）。这些信号本质上都指向"更新太激进"，而PPO正是为这个问题设计的。
从策略梯度到PPO的代码迁移其实不难，核心改动就是把loss = -log_prob * advantage换成clip后的重要性采样目标，再加上对旧策略log概率的缓存。理解了策略梯度的骨架，PPO只是在这个骨架上加了一层保险。
6. 写在最后的一点个人体会
策略梯度这条线我从REINFORCE一路写到PPO，最大的感受是：公式的简洁程度和实现的坑数量成反比。策略梯度定理那个期望式子漂亮得不行，但真落到代码里，符号、detach、标准化、熵正则、梯度裁剪，每一个细节都能让训练结果天差地别。
如果让我给刚上手的人一句建议，那就是先把REINFORCE在CartPole上跑到稳定收敛，把回报曲线、策略熵、梯度范数这几个指标都盯熟，再去碰带基线的版本和Actor-Critic。很多人一上来就抄PPO的代码，跑通了也不知道为什么通，出了问题更不知道怎么查。把最朴素的版本吃透，后面每加一个组件你都能说清楚它解决了什么问题、代价是什么，这才是真正把策略梯度学明白了。
另外提一个容易被忽略的点：策略梯度类算法的随机性很大，同一个代码换个种子结果可能差很多。评估时至少跑三到五个种子取平均，别拿单次跑出来的漂亮曲线当结论。我自己就吃过这个亏，某次单种子跑出个特别好的结果，兴冲冲去复现，换种子直接崩，白高兴一场。
