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

告别表格用PyTorch实战策略梯度Policy Gradient5步搞定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_size128): 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)这使得算法成为蒙特卡洛方法。这种近似带来了两个重要特性无偏性Gt是Qπ(s,a)的无偏估计高方差单条轨迹的回报可能波动很大关键技巧为减少方差通常会引入基线baseline最常见的是使用状态价值函数V(s)作为基线。但在基础REINFORCE实现中我们暂时不使用基线以保持算法简洁。3. 完整实现REINFORCE算法现在我们将上述理论转化为具体的PyTorch实现。完整的REINFORCE算法包含以下几个关键步骤def reinforce(env, policy, optimizer, num_episodes1000, gamma0.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()这段代码有几个值得注意的实现细节使用torch.multinomial根据概率分布采样动作这是策略梯度方法的关键反向计算折扣回报时gamma参数控制未来奖励的衰减程度回报归一化是稳定训练的重要技巧避免梯度幅度波动过大损失函数中的负号是因为PyTorch默认做梯度下降而我们需要梯度上升性能优化提示在实际应用中可以考虑使用批量更新收集多个episode后再更新来减少梯度更新的方差。4. 训练技巧与常见问题解决即使实现了正确的算法训练过程中仍可能遇到各种问题。以下是我们在CartPole环境中实践总结的经验4.1 学习率选择学习率对REINFORCE的训练效果影响极大。经过多次实验我们发现学习率训练效果0.01训练不稳定容易发散0.001收敛速度适中最终性能好0.0001收敛过慢需要更多episode# 推荐的学习率设置 optimizer torch.optim.Adam(policy.parameters(), lr0.001)4.2 奖励缩放CartPole环境中每一步的奖励都是1这可能导致回报Gt随着episode长度线性增长。我们采用两种方法解决回报归一化如前面代码所示引入折扣因子gamma控制长期奖励的影响实际测试发现gamma0.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))在实现这些进阶技巧时建议逐步添加每次只引入一个改进并验证其效果这样才能准确理解每个组件的作用。