告别表格！用PyTorch实战策略梯度（Policy Gradient），5步搞定REINFORCE算法

告别表格！用PyTorch实战策略梯度（Policy Gradient），5步搞定REINFORCE算法

在强化学习领域，策略梯度（Policy Gradient）方法正逐渐成为解决复杂决策问题的主流选择。与传统的表格型方法相比，策略梯度通过神经网络直接参数化策略，能够有效应对高维状态空间和连续动作空间的挑战。本文将带您从零开始，使用PyTorch框架实现经典的REINFORCE算法，让您快速掌握这一强大工具的核心实现技巧。

1. 环境准备与策略网络设计

在开始编码之前，我们需要搭建实验环境。CartPole是OpenAI Gym中经典的强化学习测试环境，非常适合初学者理解策略梯度的基本原理。这个环境中，智能体需要控制小车左右移动以保持杆子竖直。

首先安装必要的依赖库：

pip install gym torch numpy matplotlib

策略网络是REINFORCE算法的核心组件，它接收环境状态作为输入，输出各个动作的概率分布。我们使用一个简单的三层全连接神经网络来实现：

import torch import torch.nn as nn import torch.nn.functional as F class PolicyNetwork(nn.Module): def __init__(self, state_dim, action_dim, hidden_size=128): super(PolicyNetwork, self).__init__() self.fc1 = nn.Linear(state_dim, hidden_size) self.fc2 = nn.Linear(hidden_size, action_dim) def forward(self, x): x = F.relu(self.fc1(x)) x = F.softmax(self.fc2(x), dim=-1) return x

这个网络的关键设计要点包括：

• 输入层维度与状态空间维度匹配
• 输出层使用softmax激活函数确保动作概率归一化
• 隐藏层使用ReLU激活函数增强非线性表达能力

实际调试中发现：过大的隐藏层可能导致训练不稳定，而太小的网络又难以捕捉复杂策略。经过多次实验，128个隐藏单元在CartPole环境中表现出较好的平衡。

2. 理解REINFORCE算法的数学基础

REINFORCE算法建立在策略梯度定理的基础上，其核心思想是通过蒙特卡洛采样来估计梯度，然后使用梯度上升法优化策略。策略梯度定理给出了目标函数J(θ)关于参数θ的梯度：

∇θ J(θ) = E[∇θ log πθ(a|s) * Qπ(s,a)]

其中：

• πθ(a|s)是策略网络给出的动作概率
• Qπ(s,a)是状态-动作价值函数
• 期望E[...]表示在策略πθ下轨迹的期望

在实际实现中，我们使用轨迹的回报Gt来近似Qπ(s,a)，这使得算法成为蒙特卡洛方法。这种近似带来了两个重要特性：

1. 无偏性：Gt是Qπ(s,a)的无偏估计
2. 高方差：单条轨迹的回报可能波动很大

关键技巧：为减少方差，通常会引入基线（baseline），最常见的是使用状态价值函数V(s)作为基线。但在基础REINFORCE实现中，我们暂时不使用基线，以保持算法简洁。

3. 完整实现REINFORCE算法

现在我们将上述理论转化为具体的PyTorch实现。完整的REINFORCE算法包含以下几个关键步骤：

def reinforce(env, policy, optimizer, num_episodes=1000, gamma=0.99): for episode in range(num_episodes): state = env.reset() log_probs = [] rewards = [] # 收集轨迹数据 while True: state = torch.FloatTensor(state).unsqueeze(0) action_probs = policy(state) action = torch.multinomial(action_probs, 1).item() next_state, reward, done, _ = env.step(action) log_prob = torch.log(action_probs.squeeze(0)[action]) log_probs.append(log_prob) rewards.append(reward) state = next_state if done: break # 计算折扣回报 returns = [] G = 0 for r in reversed(rewards): G = r + gamma * G returns.insert(0, G) # 归一化回报（重要技巧！） returns = torch.FloatTensor(returns) returns = (returns - returns.mean()) / (returns.std() + 1e-9) # 计算策略梯度损失 policy_loss = [] for log_prob, G in zip(log_probs, returns): policy_loss.append(-log_prob * G) # 更新策略网络 optimizer.zero_grad() policy_loss = torch.stack(policy_loss).sum() policy_loss.backward() optimizer.step()

这段代码有几个值得注意的实现细节：

1. 使用torch.multinomial根据概率分布采样动作，这是策略梯度方法的关键
2. 反向计算折扣回报时，gamma参数控制未来奖励的衰减程度
3. 回报归一化是稳定训练的重要技巧，避免梯度幅度波动过大
4. 损失函数中的负号是因为PyTorch默认做梯度下降，而我们需要梯度上升

性能优化提示：在实际应用中，可以考虑使用批量更新（收集多个episode后再更新）来减少梯度更新的方差。

4. 训练技巧与常见问题解决

即使实现了正确的算法，训练过程中仍可能遇到各种问题。以下是我们在CartPole环境中实践总结的经验：

4.1 学习率选择

学习率对REINFORCE的训练效果影响极大。经过多次实验，我们发现：

# 推荐的学习率设置 optimizer = torch.optim.Adam(policy.parameters(), lr=0.001)

4.2 奖励缩放

CartPole环境中，每一步的奖励都是+1，这可能导致回报Gt随着episode长度线性增长。我们采用两种方法解决：

1. 回报归一化（如前面代码所示）
2. 引入折扣因子gamma控制长期奖励的影响

实际测试发现：gamma=0.99在CartPole中效果良好，既能考虑长期回报，又不至于使Gt过大。

4.3 梯度消失问题

当策略接近最优时，某些动作的概率可能变得非常小，导致对数概率趋近于负无穷，引发梯度消失。我们通过以下方法缓解：

# 在策略网络输出层添加小常数避免零概率 action_probs = policy(state) + 1e-8 action_probs = action_probs / action_probs.sum()

4.4 训练监控

可视化训练过程有助于及时发现问题和调整超参数：

import matplotlib.pyplot as plt def plot_rewards(rewards): plt.plot(rewards) plt.xlabel('Episode') plt.ylabel('Total Reward') plt.title('Training Progress') plt.show()

5. 进阶优化与扩展思路

掌握了基础REINFORCE实现后，我们可以考虑以下几个方向的改进：

5.1 引入基线减少方差

如前所述，引入基线可以显著降低梯度估计的方差。最简单的基线是回报的移动平均：

baseline = 0 beta = 0.9 # 基线更新系数 for log_prob, G in zip(log_probs, returns): baseline = beta * baseline + (1 - beta) * G.item() policy_loss.append(-log_prob * (G - baseline))

5.2 使用Advantage函数

Advantage函数A(s,a) = Q(s,a) - V(s)能更准确地评估动作的相对优势。可以使用n-step TD估计：

# n-step TD估计 n_step = 5 advantages = [] for t in range(len(rewards)): end = min(t + n_step, len(rewards)) G = sum([gamma**(i-t) * rewards[i] for i in range(t, end)]) if end < len(rewards): G += gamma**n_step * torch.max(policy(torch.FloatTensor(next_state))) advantages.append(G - baseline)

5.3 扩展到连续动作空间

REINFORCE也可以应用于连续动作空间，只需修改策略网络输出高斯分布的参数：

class ContinuousPolicy(nn.Module): def __init__(self, state_dim, action_dim): super().__init__() self.fc_mean = nn.Linear(state_dim, action_dim) self.fc_std = nn.Linear(state_dim, action_dim) def forward(self, x): mean = self.fc_mean(x) std = F.softplus(self.fc_std(x)) + 1e-5 return torch.distributions.Normal(mean, std)

5.4 并行采样加速训练

使用多个环境并行采集轨迹可以显著提高数据效率：

from multiprocessing import Process, Queue def worker(env_fn, policy, queue, num_episodes): env = env_fn() for _ in range(num_episodes): # 采集轨迹并放入队列 queue.put(collect_trajectory(env, policy))

在实现这些进阶技巧时，建议逐步添加，每次只引入一个改进并验证其效果，这样才能准确理解每个组件的作用。