用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个离散动作（向左或向右施力）。

完整的策略网络实现需要考虑以下要素：

1. 网络架构：3层全连接网络（输入层、隐藏层、输出层）
2. 激活函数：ReLU用于隐藏层，Softmax用于输出层
3. 梯度计算：PyTorch自动微分机制
4. 动作选择：按概率分布采样而非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算法的训练流程包含几个关键步骤：

1. 轨迹收集：使用当前策略与环境交互生成完整episode
2. 回报计算：计算每个时间步的累积折扣回报
3. 策略更新：执行梯度上升更新网络参数
4. 重复迭代：直到策略收敛或达到最大训练轮次

回报计算需要特别注意折扣因子γ的选择（通常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算法存在高方差问题，以下进阶技术可提升性能：

1. 基线减法：引入状态值函数作为基线减小方差

   # 修改回报计算
   advantage = returns - values
   

2. Actor-Critic架构：同时学习策略和价值函数
3. 并行采样：多个环境并行收集轨迹加速训练
4. 信任域策略优化：约束策略更新幅度

一个简单的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）对训练效果影响显著。同时，定期测试策略性能而非仅依赖训练回报评估，能更准确反映策略的真实质量。