告别表格，用PyTorch实战REINFORCE算法：从零搭建你的第一个策略梯度模型
用PyTorch实战REINFORCE算法：从零搭建策略梯度模型
在强化学习领域，策略梯度方法因其直接优化策略的特性而备受关注。不同于基于价值的方法需要先估计价值函数再推导策略，策略梯度算法直接对策略参数进行梯度上升，特别适合处理连续动作空间和高维状态空间问题。本文将聚焦经典的REINFORCE算法，使用PyTorch框架在CartPole环境中实现完整的策略梯度模型。
1. 策略梯度基础与REINFORCE原理
策略梯度方法的核心思想是通过参数化策略函数π(a|s;θ)，直接优化策略参数θ以最大化期望回报。REINFORCE作为最早的策略梯度算法之一，其更新规则可表示为：
θ ← θ + α * G_t * ∇lnπ(a_t|s_t;θ)其中α是学习率，G_t是从时间步t开始的累积折扣回报，∇lnπ(a_t|s_t;θ)是策略对数概率的梯度。这种更新方式具有直观的解释：增加带来高回报动作的概率，减少低回报动作的概率。
策略网络通常设计为：
class PolicyNetwork(nn.Module): def __init__(self, state_dim, action_dim): super().__init__() self.fc1 = nn.Linear(state_dim, 64) self.fc2 = nn.Linear(64, action_dim) def forward(self, state): x = F.relu(self.fc1(state)) return F.softmax(self.fc2(x), dim=-1)注意：输出层使用softmax确保动作概率在(0,1)区间且和为1，这是REINFORCE算法的关键要求
2. 环境准备与策略网络实现
我们选择OpenAI Gym中的CartPole-v1环境作为测试平台。该环境状态空间包含4个连续变量（小车位置、速度、杆角度和角速度），动作空间有2个离散动作（向左或向右施力）。
完整的策略网络实现需要考虑以下要素：
- 网络架构：3层全连接网络（输入层、隐藏层、输出层）
- 激活函数：ReLU用于隐藏层，Softmax用于输出层
- 梯度计算：PyTorch自动微分机制
- 动作选择：按概率分布采样而非argmax
import torch import torch.nn as nn import torch.optim as optim import torch.nn.functional as F class PolicyGradientAgent: def __init__(self, state_dim, action_dim, lr=0.01): self.policy = PolicyNetwork(state_dim, action_dim) self.optimizer = optim.Adam(self.policy.parameters(), lr=lr) self.saved_log_probs = [] self.rewards = []3. 训练流程与回报计算
REINFORCE算法的训练流程包含几个关键步骤：
- 轨迹收集：使用当前策略与环境交互生成完整episode
- 回报计算：计算每个时间步的累积折扣回报
- 策略更新：执行梯度上升更新网络参数
- 重复迭代：直到策略收敛或达到最大训练轮次
回报计算需要特别注意折扣因子γ的选择（通常0.9-0.99）：
def calculate_returns(rewards, gamma=0.99): returns = [] R = 0 for r in reversed(rewards): R = r + gamma * R returns.insert(0, R) returns = torch.tensor(returns) returns = (returns - returns.mean()) / (returns.std() + 1e-9) # 标准化 return returns提示：回报标准化有助于稳定训练，但需保留足够方差促进探索
4. 策略优化与梯度上升
策略优化的核心在于正确计算策略梯度并执行参数更新。PyTorch的自动微分简化了这一过程：
def update_policy(): returns = calculate_returns(agent.rewards) policy_loss = [] for log_prob, R in zip(agent.saved_log_probs, returns): policy_loss.append(-log_prob * R) # 负号因为PyTorch默认最小化 agent.optimizer.zero_grad() policy_loss = torch.cat(policy_loss).sum() policy_loss.backward() agent.optimizer.step() del agent.rewards[:] del agent.saved_log_probs[:]关键点说明：
- 损失函数设计为负的加权对数概率
- 每个episode结束后执行一次更新（蒙特卡洛方法）
- 需要手动清空存储的轨迹数据
5. 训练技巧与常见问题
实际训练中可能遇到的挑战及解决方案：
| 问题 | 原因 | 解决方案 |
|---|---|---|
| 训练不稳定 | 高方差梯度 | 回报标准化，减小学习率 |
| 收敛速度慢 | 探索不足 | 增加初始熵正则化 |
| 过早收敛 | 局部最优 | 保留最小探索概率 |
| 梯度爆炸 | 步长过大 | 梯度裁剪，自适应优化器 |
熵正则化是提升探索的有效技术：
def forward(self, state): x = F.relu(self.fc1(state)) action_probs = F.softmax(self.fc2(x), dim=-1) entropy = -torch.sum(action_probs * torch.log(action_probs)) return action_probs, entropy在损失函数中加入熵项：
policy_loss = (-log_prob * R) - 0.01 * entropy # 熵系数需调优6. 完整训练循环实现
将各组件整合为完整训练流程：
def train(env, agent, episodes=1000, max_steps=1000): for ep in range(episodes): state = env.reset() for t in range(max_steps): state = torch.FloatTensor(state) action_probs = agent.policy(state) action = torch.multinomial(action_probs, 1).item() next_state, reward, done, _ = env.step(action) agent.rewards.append(reward) agent.saved_log_probs.append(torch.log(action_probs[action])) state = next_state if done: break update_policy() if ep % 50 == 0: print(f"Episode {ep}, Reward: {sum(agent.rewards)}")训练过程中可观察到回报的逐步提升，典型的学习曲线会呈现从低回报快速上升到稳定高回报的过程。在CartPole环境中，良好的策略通常能在100-200个episode内达到最大回报500。
7. 高级改进方向
基础REINFORCE算法存在高方差问题，以下进阶技术可提升性能：
- 基线减法：引入状态值函数作为基线减小方差
# 修改回报计算 advantage = returns - values - Actor-Critic架构：同时学习策略和价值函数
- 并行采样：多个环境并行收集轨迹加速训练
- 信任域策略优化：约束策略更新幅度
一个简单的Actor-Critic实现框架：
class ActorCritic(nn.Module): def __init__(self, state_dim, action_dim): super().__init__() self.actor = nn.Sequential( nn.Linear(state_dim, 64), nn.ReLU(), nn.Linear(64, action_dim), nn.Softmax(dim=-1) ) self.critic = nn.Sequential( nn.Linear(state_dim, 64), nn.ReLU(), nn.Linear(64, 1) )在实际项目中，我发现合理设置初始学习率（通常3e-4到1e-3）和折扣因子（0.9-0.99）对训练效果影响显著。同时，定期测试策略性能而非仅依赖训练回报评估，能更准确反映策略的真实质量。