策略梯度定理（Policy Gradient Theorem）是强化学习中极为优雅且核心的基石。对于需要处理连续动作空间、复杂动力学系统的任务（例如自主导航系统中的力矩控制或机器人的运动规划）来说，理解这个定理是设计高效RL智能体的前提。本文将详细解释这个定理的数学原理和具体实现。

轨迹概率计算

在一场游戏或者仿真环境中，我们把环境输出的s 与演员（Actor）输出的动作a 全部组合起来，就是一个轨迹，

给定演员的参数$θ$，我们可以计算某个轨迹$τ$ 发生的概率为：

我们先计算环境输出$s_1$ 的概率$p(s_1)$, 在根据$s_1$ 执行$a_1$的概率$p_θ (a_1|s_1)$ 。在某一场游戏的某一个回合里面，我们会得到$R(τ )$。我们要做的就是调整演员内部的参数$θ$，使得$R(τ )$ 的值越大越好。

所以$R(τ )$ 是一个随机变量。我们能够计算的是$R(τ )$ 的期望值。 给定某一组参数$θ$，我们可计算$r_θ$ 的期望值为:

我们要最大化期望奖励。因为我们要让奖励越大越好，所以可以使用梯度上升（gradient ascent）来最大化期望奖励。

数学前提知识

现在我们补充一点点数学基础知识： 对于自然对数它的导数是：

结合链式法则，如果我们对一个函数取$f(x)$对数然后再求导，结果是:

把分母 $f(x)$乘到左边去，把等式左右互换，就得到了大名鼎鼎的 Log-Derivative Trick 恒等式：

除此之外，我们还需要知道： 对于一个连续型随机变量 $X$，其概率密度函数为 $p(x)$。如果我们想求某个关于 $X$的函数$f(X)$的数学期望，其期望与积分的转换公式为：

• $\mathbb{E}_{X \sim p ( x )} [ \cdot]$ : 表示在概率分布 $p(x)$ 下求期望。

公式推导

现在把期望奖励公式和 Log_derivative trick 公式结合进行推导：

实际上奖励期望值无法直接计算，所以我们用采样的方式采样$N$个$τ$并计算每一个的值，把每一个的值加起来，就可以得到梯度，即：

这里的 $\operatorname{l o g} p_{\theta} \left( a_{t}^{n} \mid s_{t}^{n} \right)$ 可以通过演员网络根据权重参数$\theta$ 计算出动作概率对数。

这个公式推导有个细节，需要对轨迹概率计算公式进行求导：

因为公式其它几项不带$\theta$, 所以求导为0.实际上要计算梯度，首先我们要收集很多s 与a的对（pair），还要知道这些s 与a 在与环境交互的时候，会得到多少奖励。

现在已经知道奖励期望梯度如何计算，下面给出一份代码：

python
import gymnasium as gym
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical

# ==========================================
# 1. 定义策略网络 (Policy Network)
# 对应数学公式中的 π_θ(a|s)
# ==========================================
class PolicyNetwork(nn.Module):
    def __init__(self, obs_dim, act_dim):
        super(PolicyNetwork, self).__init__()
        # 简单的多层感知机 (MLP)
        self.net = nn.Sequential(
            nn.Linear(obs_dim, 64),
            nn.ReLU(),
            nn.Linear(64, act_dim)
        )

    def forward(self, state):
        # 输出每个动作的未归一化概率 (Logits)
        logits = self.net(state)
        # 封装为离散的类别分布 (Categorical Distribution)
        # 这在底层会自动帮我们处理 Softmax 和对数概率的计算
        return Categorical(logits=logits)

# ==========================================
# 2. 训练主循环
# ==========================================
def train_policy_gradient():
    # 创建环境
    env = gym.make('CartPole-v1')
    obs_dim = env.observation_space.shape[0] # 状态维度：4 (位置, 速度, 角度, 角速度)
    act_dim = env.action_space.n             # 动作维度：2 (向左, 向右)

    # 初始化策略网络和优化器
    policy = PolicyNetwork(obs_dim, act_dim)
    optimizer = optim.Adam(policy.parameters(), lr=1e-2)
    
    gamma = 0.99 # 折扣因子 (Discount factor)

    for episode in range(500): # 训练 500 个回合 (Episodes)
        state, _ = env.reset()
        
        # 用于记录这一条轨迹 (Trajectory) 的数据
        log_probs = []
        rewards = []

        # ------------------------------------------
        # 阶段 A：与环境交互，采样一条完整的轨迹 τ
        # ------------------------------------------
        while True:
            # 状态转为 Tensor
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            
            # 策略网络输出动作分布
            action_dist = policy(state_tensor)
            
            # 根据概率分布采样一个动作 a_t
            action = action_dist.sample()
            
            # 记录该动作的对数概率 log p_θ(a_t|s_t)
            log_probs.append(action_dist.log_prob(action))
            
            # 环境执行动作，返回下一步状态和奖励
            next_state, reward, terminated, truncated, _ = env.step(action.item())
            rewards.append(reward)
            
            state = next_state
            
            # 回合结束（杆子倒了，或达到最大步数）
            if terminated or truncated:
                break

        # ------------------------------------------
        # 阶段 B：计算累积折扣回报 (Discounted Reward-to-go)
        # 对应数学公式中的 R(τ)
        # ------------------------------------------
        returns = []
        G = 0
        # 从最后一步往前推算，计算每个时间步 t 之后的总回报
        for r in reversed(rewards):
            G = r + gamma * G
            returns.insert(0, G)
        returns = torch.tensor(returns)
        
        # 【工程Trick】：回报归一化 (Baseline)
        # 减去均值除以标准差，降低方差，防止梯度震荡，加快网络收敛
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)

        # ------------------------------------------
        # 阶段 C：计算策略梯度损失并反向传播
        # 对应公式：∇J(θ) = E [ ∇log p_θ(a_t|s_t) * R ]
        # ------------------------------------------
        policy_loss = []
        for log_prob, G_t in zip(log_probs, returns):
            # 因为 PyTorch 是做梯度下降，我们要最大化目标，所以加负号
            policy_loss.append(-log_prob * G_t)
            
        # 将一个回合中所有时间步的损失求和
        loss = torch.cat(policy_loss).sum()

        # 梯度清零，反向传播，更新网络参数 θ
        optimizer.zero_grad()
        loss.backward()  # 这里触发 Autograd 机制，计算 ∇_θ
        optimizer.step()

        # ------------------------------------------
        # 打印训练进度
        # ------------------------------------------
        if episode % 50 == 0:
            print(f"Episode {episode} | Total Reward: {sum(rewards)} | Loss: {loss.item():.4f}")

    env.close()
    print("Training finished!")

if __name__ == '__main__':
    train_policy_gradient()