用PyTorch手把手复现ChatGPT背后的PPO算法（附完整代码与CartPole实战）

瑶瑶宝

313人浏览 · 2026-03-26 11:02:26

瑶瑶宝 · 2026-03-26 11:02:26 发布

从零实现PPO算法：ChatGPT背后的强化学习核心技术解析

在OpenAI的ChatGPT训练流程中，强化学习阶段使用的近端策略优化（PPO）算法是关键所在。本文将带您深入理解这一算法，并用PyTorch从零开始实现一个完整的PPO模型，最后在经典的CartPole环境中进行实战验证。

1. PPO算法核心原理剖析

PPO作为当前最先进的策略梯度算法，其核心在于平衡策略更新的"步长"——既要有足够的探索来提升性能，又要避免过大的更新导致训练不稳定。这主要通过三个关键设计实现：

比率裁剪（Clipping）机制

ratio = new_probs / old_probs
clipped_ratio = torch.clamp(ratio, 1-ε, 1+ε)
loss = -torch.min(ratio * advantage, clipped_ratio * advantage).mean()

这段代码展示了PPO最核心的clip操作，其中ε通常取0.1-0.3。当新旧策略差异过大时，裁剪机制会限制更新幅度，确保训练稳定性。

优势估计（Advantage Estimation） PPO使用广义优势估计（GAE）来更准确地评估动作价值：

A_t = δ_t + (γλ)δ_{t+1} + (γλ)^2δ_{t+2} + ... 
其中 δ_t = r_t + γV(s_{t+1}) - V(s_t)

GAE通过参数λ（通常0.9-0.95）在偏差和方差间取得平衡。

多轮小批量更新 与传统策略梯度算法不同，PPO会：

收集一批经验数据
对这批数据执行K次（通常3-10次）小批量更新
每次更新使用不同的数据划分

这种设计显著提高了数据利用率。下表对比了PPO与经典策略梯度算法的差异：

特性	传统策略梯度	PPO
更新稳定性	低	高
数据利用率	低	高
超参数敏感性	高	中
并行化潜力	低	高

2. 网络架构设计与实现

PPO采用Actor-Critic架构，包含两个关键组件：

策略网络（Actor）

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, 64)
        self.fc3 = nn.Linear(64, action_dim)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return F.softmax(self.fc3(x), dim=-1)

输出动作的概率分布，在离散动作空间中使用softmax，连续空间则输出高斯分布的均值和方差。

价值网络（Critic）

class ValueNetwork(nn.Module):
    def __init__(self, state_dim):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, 64)
        self.fc2 = nn.Linear(64, 64)
        self.fc3 = nn.Linear(64, 1)
        
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)

输出状态价值估计，用于计算优势函数。实际实现时，两个网络可以共享底层特征提取层。

提示：对于较复杂的环境，可以考虑使用更大的网络或引入注意力机制等现代架构。

3. 完整训练流程实现

PPO的训练过程可分为数据收集、优势计算和参数更新三个阶段：

1. 并行数据收集

def collect_rollouts(env, policy, n_steps):
    states, actions, rewards = [], [], []
    state = env.reset()
    for _ in range(n_steps):
        with torch.no_grad():
            action_probs = policy(torch.FloatTensor(state))
            action = Categorical(action_probs).sample().item()
        
        next_state, reward, done, _ = env.step(action)
        
        states.append(state)
        actions.append(action)
        rewards.append(reward)
        
        state = next_state if not done else env.reset()
    
    return np.array(states), np.array(actions), np.array(rewards)

2. 优势计算与标准化

def compute_advantages(rewards, values, gamma=0.99, lam=0.95):
    deltas = rewards[:-1] + gamma * values[1:] - values[:-1]
    advantages = []
    adv = 0
    for delta in reversed(deltas):
        adv = delta + gamma * lam * adv
        advantages.append(adv)
    advantages = np.array(advantages[::-1])
    return (advantages - advantages.mean()) / (advantages.std() + 1e-8)

3. 核心训练循环

def train_step(policy, optimizer, states, actions, advantages, old_probs, clip_param=0.2):
    # 计算新策略概率
    new_probs = policy(states).gather(1, actions.unsqueeze(1))
    ratios = new_probs / old_probs
    
    # 裁剪目标函数
    surr1 = ratios * advantages
    surr2 = torch.clamp(ratios, 1-clip_param, 1+clip_param) * advantages
    policy_loss = -torch.min(surr1, surr2).mean()
    
    # 价值函数损失
    value_loss = F.mse_loss(policy.value(states), returns)
    
    # 熵奖励
    entropy = -(new_probs * torch.log(new_probs + 1e-5)).mean()
    
    # 总损失
    loss = policy_loss + 0.5 * value_loss - 0.01 * entropy
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

4. CartPole环境实战与调优

在CartPole-v1环境中，我们使用以下配置进行训练：

config = {
    "n_steps": 2048,        # 每次迭代收集的步数
    "batch_size": 64,       # 小批量大小
    "n_epochs": 10,         # 每次迭代的更新轮数
    "gamma": 0.99,          # 折扣因子
    "lam": 0.95,            # GAE参数
    "clip_param": 0.2,      # 裁剪参数ε
    "lr": 3e-4,             # 学习率
    "entropy_coef": 0.01,   # 熵奖励系数
}

训练过程中有几个关键观察点：