PPO算法实战:用PyTorch从零搭建强化学习智能体(附完整代码)

📅 发布时间:2026/7/4 14:58:14 👁️ 浏览次数:
PPO算法实战:用PyTorch从零搭建强化学习智能体(附完整代码)
PPO算法实战用PyTorch从零搭建强化学习智能体附完整代码如果你已经对强化学习的基本概念有所了解比如知道什么是马尔可夫决策过程也尝试过用Q-learning玩过CartPole那么下一步你很可能想挑战一个更强大、也更实用的算法——近端策略优化也就是PPO。这个由OpenAI在2017年提出的算法几乎成了近年来解决复杂控制问题的“标配”。它不像DQN那样只处理离散动作也不像传统策略梯度那样容易“跑偏”而是在稳定性和实现难度之间找到了一个绝佳的平衡点。但理论归理论真正想把它用起来你会发现网上很多教程要么过于学术化公式满天飞要么就是代码片段零散关键的工程细节语焉不详。比如那个著名的clip参数到底设多少优势函数用GAE估计时λ怎么调网络结构怎么设计才不容易训崩这些问题不亲手写一遍代码、不经历几次训练失败是很难有深刻体会的。这篇文章就是为你准备的。我们将完全从零开始用PyTorch搭建一个完整的PPO智能体。我不会只给你一个“黑箱”代码而是会带你一步步理解每一行代码背后的设计逻辑并分享我在实际项目中调试PPO时踩过的坑和总结出的经验。我们的目标很明确让你不仅能跑通代码更能理解如何根据不同的任务去调整和优化它。1. 环境搭建与核心概念再梳理在动手写代码之前我们需要一个“练兵场”。这里我选择Gymnasium库中的CartPole-v1和Pendulum-v1环境。前者是经典的离散动作控制问题后者则是连续动作空间能更好地展示PPO处理连续控制的能力。确保你的环境已经就绪pip install gymnasium torch numpyPPO的核心思想可以用一个比喻来理解教一个新手学走路。传统的策略梯度方法就像告诉学生“往这个方向走一步”但步子迈多大完全靠感觉容易摔倒策略更新不稳定。TRPO则像给学生脚上绑了根绳子严格限制他每一步的移动范围通过复杂的KL散度约束计算虽然安全但行动不便。PPO则更聪明它给学生定了个规矩“你这一步的移动幅度不能超过上一步的10%到30%”。这个规矩就是剪切Clipping它用非常简单的一阶优化就实现了稳定的策略更新。整个PPO的训练流程是一个“收集-计算-更新”的循环收集轨迹用当前的策略演员网络与环境交互收集一批(状态 动作 奖励 下一个状态)的数据。计算优势利用价值网络评论家网络和收集的数据估算每个(状态 动作)对的优势值即这个动作比平均表现好多少。更新网络用计算出的优势值通过一个特殊的剪切目标函数来更新演员网络使其更倾向于选择优势高的动作同时更新评论家网络让它对状态价值的估计更准。注意PPO的一个关键技巧是数据重用。我们不会用一批数据更新一次就扔掉而是会用它进行多轮例如3-10个epoch的小批量梯度更新这大大提高了样本效率。2. 网络架构设计与实现PPO采用Actor-Critic架构两个网络通常共享一部分底层特征提取层以提升学习效率并减少参数。但对于入门我们先实现一个结构清晰、易于理解的分离式网络。2.1 定义Actor-Critic网络我们的网络需要处理两种输出演员网络输出动作的概率分布离散或分布参数连续评论家网络输出一个标量代表当前状态的价值。import torch import torch.nn as nn import torch.nn.functional as F import numpy as np class ActorCritic(nn.Module): 一个简单的Actor-Critic网络。 对于离散动作空间Actor输出动作logits。 对于连续动作空间Actor输出高斯分布的均值和标准差。 def __init__(self, state_dim, action_dim, is_continuousFalse): super(ActorCritic, self).__init__() self.is_continuous is_continuous # 共享的特征提取层 self.shared_layers nn.Sequential( nn.Linear(state_dim, 64), nn.Tanh(), nn.Linear(64, 64), nn.Tanh(), ) # 演员网络 (策略网络) self.actor nn.Linear(64, action_dim) if self.is_continuous: # 对于连续动作我们还需要一个输出对数标准差的层 self.log_std nn.Parameter(torch.zeros(1, action_dim)) # 评论家网络 (价值网络) self.critic nn.Linear(64, 1) def forward(self, state): 前向传播返回动作分布和对数概率以及状态价值。 在实际训练中我们通常分开调用actor和critic。 features self.shared_layers(state) value self.critic(features) return value def act(self, state): 用于与环境交互时采样动作。 with torch.no_grad(): features self.shared_layers(state) value self.critic(features) if self.is_continuous: mean self.actor(features) std torch.exp(self.log_std) dist torch.distributions.Normal(mean, std) action dist.sample() log_prob dist.log_prob(action).sum(dim-1) action action.clamp(-2.0, 2.0) # 示例对动作进行裁剪 else: logits self.actor(features) dist torch.distributions.Categorical(logitslogits) action dist.sample() log_prob dist.log_prob(action) return action.cpu().numpy(), value.cpu().numpy(), log_prob.cpu().numpy() def evaluate(self, state, action): 用于批量评估计算给定状态-动作对的log prob、价值和熵。 features self.shared_layers(state) value self.critic(features).squeeze() if self.is_continuous: mean self.actor(features) std torch.exp(self.log_std) dist torch.distributions.Normal(mean, std) log_prob dist.log_prob(action).sum(dim-1) entropy dist.entropy().sum(dim-1).mean() else: logits self.actor(features) dist torch.distributions.Categorical(logitslogits) log_prob dist.log_prob(action.squeeze()) entropy dist.entropy().mean() return log_prob, value, entropy这个网络类包含了三个关键方法act: 在环境中探索时使用根据当前策略采样一个动作并返回动作、状态价值和动作的对数概率用于后续计算概率比。evaluate: 在更新阶段使用给定一批状态和动作计算对应的对数概率、状态价值和策略的熵熵用于鼓励探索。2.2 经验回放缓冲区的设计PPO通常使用在线策略学习即用当前策略收集的数据来更新当前策略。我们需要一个缓冲区来存储单次或多次环境交互的轨迹数据。class PPOBuffer: def __init__(self, state_dim, action_dim, buffer_size, gamma0.99, gae_lambda0.95): self.state_dim state_dim self.action_dim action_dim self.buffer_size buffer_size self.gamma gamma self.gae_lambda gae_lambda # 初始化存储空间 self.states np.zeros((buffer_size, state_dim), dtypenp.float32) self.actions np.zeros((buffer_size, action_dim) if action_dim 1 else (buffer_size, 1), dtypenp.float32) self.rewards np.zeros(buffer_size, dtypenp.float32) self.values np.zeros(buffer_size, dtypenp.float32) self.log_probs np.zeros(buffer_size, dtypenp.float32) self.dones np.zeros(buffer_size, dtypenp.float32) self.ptr 0 # 当前指针位置 self.path_start_idx 0 # 当前轨迹的起始索引 def store(self, state, action, reward, value, log_prob, done): 存储单步交互数据。 idx self.ptr self.states[idx] state self.actions[idx] action self.rewards[idx] reward self.values[idx] value self.log_probs[idx] log_prob self.dones[idx] done self.ptr 1 def finish_path(self, last_value0): 当一条轨迹结束时doneTrue调用此函数。 计算从 path_start_idx 到 ptr-1 这段轨迹的GAE优势估计和回报。 path_slice slice(self.path_start_idx, self.ptr) rewards np.append(self.rewards[path_slice], last_value) values np.append(self.values[path_slice], last_value) dones np.append(self.dones[path_slice], 0) # 最后一步不是终止 # GAE 和 回报计算 deltas rewards[:-1] self.gamma * values[1:] * (1 - dones[1:]) - values[:-1] advantages np.zeros_like(deltas, dtypenp.float32) last_gae_lam 0 for t in reversed(range(len(deltas))): last_gae_lam deltas[t] self.gamma * self.gae_lambda * (1 - dones[t1]) * last_gae_lam advantages[t] last_gae_lam returns advantages values[:-1] # 将计算好的优势值和回报存回缓冲区这里简化处理实际可能需要额外数组存储 # 我们这里直接返回计算好的这一批数据清空该段缓冲区 batch_states self.states[path_slice] batch_actions self.actions[path_slice] batch_log_probs_old self.log_probs[path_slice] batch_advantages advantages batch_returns returns # 重置当前轨迹起始点 self.path_start_idx self.ptr return batch_states, batch_actions, batch_log_probs_old, batch_advantages, batch_returns def get(self): 当缓冲区满或需要更新时获取所有数据。 # 这里需要处理最后一段未完成的轨迹如果有 if self.path_start_idx ! self.ptr: print(f警告缓冲区中有未结束的轨迹数据从索引 {self.path_start_idx} 到 {self.ptr-1}。) # 可以选择用最后的状态价值来结束它或者丢弃。这里为了简单我们假设最后一步价值为0。 last_val 0 data self.finish_path(last_val) # 注意这里简化了实际应该将数据整合到总批次中。我们假设每次更新前都会清空缓冲区。 pass # 返回所有数据并重置缓冲区 self.ptr, self.path_start_idx 0, 0 # 实际实现中这里应返回整合后的所有轨迹数据 # 为简化我们假设在训练循环中每收集一条完整轨迹就立即计算并更新这个缓冲区的核心是finish_path方法它使用**广义优势估计GAE**来计算优势函数。GAE是PPO稳定训练的关键它通过参数λ在偏差和方差之间做权衡。提示gae_lambda参数接近1时优势估计方差小但偏差大接近0时偏差小但方差大。通常设置在0.9到0.98之间。3. PPO-Clip损失函数的代码实现与解析这是PPO算法的核心。剪切目标函数看起来有点复杂但拆解后非常直观。其目的是在鼓励策略向优势高的方向更新时防止单次更新步子迈得太大。def compute_ppo_loss(actor_critic, states, actions, old_log_probs, advantages, returns, clip_epsilon0.2, value_coef0.5, entropy_coef0.01): 计算PPO的总损失包含策略损失、价值损失和熵奖励。 参数: actor_critic: 网络模型 states, actions, old_log_probs, advantages, returns: 从缓冲区获取的数据 clip_epsilon: 剪切参数通常为0.1-0.3 value_coef: 价值损失权重 entropy_coef: 熵奖励权重用于鼓励探索 # 将numpy数组转换为PyTorch张量 states torch.FloatTensor(states) actions torch.FloatTensor(actions) if actions.dtype np.float32 else torch.LongTensor(actions) old_log_probs torch.FloatTensor(old_log_probs) advantages torch.FloatTensor(advantages) returns torch.FloatTensor(returns) # 评估当前策略 log_probs, values, entropy actor_critic.evaluate(states, actions) # 1. 计算策略损失 (PPO-Clip) ratios torch.exp(log_probs - old_log_probs.detach()) # 概率比 r_t(θ) surr1 ratios * advantages surr2 torch.clamp(ratios, 1.0 - clip_epsilon, 1.0 clip_epsilon) * advantages policy_loss -torch.min(surr1, surr2).mean() # 取负号是因为我们要最大化目标函数 # 2. 计算价值损失 (MSE) value_loss F.mse_loss(values, returns) # 3. 计算总损失 total_loss policy_loss value_coef * value_loss - entropy_coef * entropy # 额外信息用于监控 clip_fraction torch.mean((torch.abs(ratios - 1.0) clip_epsilon).float()).item() return total_loss, policy_loss.item(), value_loss.item(), entropy.item(), clip_fraction让我们深入看看policy_loss那几行代码ratios: 新策略与旧策略在给定状态下选择同一动作的概率比。ratios 1意味着新策略更倾向于该动作ratios 1则相反。surr1 ratios * advantages: 这是未剪切的原始目标。如果优势advantages为正我们希望ratios增大即更大概率选择该动作反之亦然。surr2 torch.clamp(...) * advantages: 这是剪切后的目标。torch.clamp将ratios限制在[1-ε, 1ε]区间内防止ratios变得过大或过小。torch.min(surr1, surr2):这是PPO的精华所在。我们取surr1和surr2中较小的那个作为最终目标。这意味着当优势为正且ratios 1ε时我们鼓励ratios增长surr1被选中。当优势为正但ratios已经超过1ε时我们不再鼓励它继续增长因为剪切后的目标surr2更小成为了被选中的项。这防止了策略因单次更新而改变过大。优势为负时逻辑相反。最后取负均值因为PyTorch的优化器默认是最小化损失而我们需要最大化这个剪切目标。下表总结了关键超参数的经验取值范围和影响超参数典型取值范围作用与影响Clip Epsilon (ε)0.1 ~ 0.3核心参数。控制策略更新的最大幅度。值越小更新越保守稳定但学习可能变慢值越大更新更激进但可能不稳定。GAE Lambda (λ)0.9 ~ 0.98控制优势估计的偏差-方差权衡。越接近1估计方差越小更平滑但偏差可能越大。折扣因子 (γ)0.99 ~ 0.999衡量未来奖励的重要性。越接近1智能体越有远见但训练可能更不稳定。价值损失系数0.5价值网络损失在总损失中的权重。平衡策略学习和价值估计。熵系数0.01熵奖励的权重。鼓励探索防止策略过早收敛到次优解。太大则过度探索太小则探索不足。学习率3e-4 (Adam常用)优化器的步长。PPO对学习率敏感通常使用较小的固定学习率或衰减学习率。每批数据更新轮数3 ~ 10同一批经验数据重复用于梯度更新的次数。提高数据利用率但过多可能导致过拟合。小批量大小64 ~ 256每次梯度更新使用的样本数。太小噪声大太大计算慢且可能降低泛化性。4. 完整的训练循环与调试技巧现在我们将所有部分组装起来形成一个完整的训练脚本。这里以连续动作的Pendulum-v1环境为例。import gymnasium as gym from torch.optim import Adam def train_ppo(env_namePendulum-v1, total_timesteps100000, max_ep_len200, clip_epsilon0.2, gamma0.99, gae_lambda0.95, policy_lr3e-4, value_lr1e-3, train_epochs10, batch_size64): # 创建环境 env gym.make(env_name) state_dim env.observation_space.shape[0] action_dim env.action_space.shape[0] is_continuous True # 初始化模型和优化器 model ActorCritic(state_dim, action_dim, is_continuous) optimizer Adam([ {params: model.shared_layers.parameters(), lr: policy_lr}, {params: model.actor.parameters(), lr: policy_lr}, {params: model.critic.parameters(), lr: value_lr}, {params: model.log_std, lr: policy_lr} if is_continuous else [] ]) # 训练循环 timestep 0 while timestep total_timesteps: state, _ env.reset() ep_reward 0 ep_len 0 # 收集一条轨迹的数据 states, actions, rewards, values, log_probs, dones [], [], [], [], [], [] for _ in range(max_ep_len): timestep 1 # 交互并存储 action, value, log_prob model.act(torch.FloatTensor(state).unsqueeze(0)) next_state, reward, terminated, truncated, _ env.step(action[0]) done terminated or truncated states.append(state) actions.append(action[0]) rewards.append(reward) values.append(value[0][0]) log_probs.append(log_prob[0]) dones.append(done) state next_state ep_reward reward ep_len 1 if done: break # 轨迹结束计算最后状态的价值用于GAE last_value model(torch.FloatTensor(state).unsqueeze(0)).item() if not done else 0 # 计算优势函数和回报 (这里简化未使用完整的GAE缓冲区仅作演示) returns [] advantages [] R last_value adv 0 for r, v, done_flag in zip(reversed(rewards), reversed(values), reversed(dones)): if done_flag: R 0 adv 0 td_error r gamma * R - v adv td_error gamma * gae_lambda * adv R r gamma * R returns.insert(0, R) advantages.insert(0, adv) # 转换为张量 states_t torch.FloatTensor(states) actions_t torch.FloatTensor(actions) old_log_probs_t torch.FloatTensor(log_probs) returns_t torch.FloatTensor(returns) advantages_t torch.FloatTensor(advantages) # 归一化优势这是一个稳定训练的重要技巧 advantages_t (advantages_t - advantages_t.mean()) / (advantages_t.std() 1e-8) # 使用当前轨迹的数据进行多轮更新 for _ in range(train_epochs): # 随机打乱数据并分成小批量 indices np.arange(len(states)) np.random.shuffle(indices) for start in range(0, len(states), batch_size): end start batch_size batch_indices indices[start:end] batch_states states_t[batch_indices] batch_actions actions_t[batch_indices] batch_old_log_probs old_log_probs_t[batch_indices] batch_advantages advantages_t[batch_indices] batch_returns returns_t[batch_indices] # 计算损失并更新 total_loss, policy_loss, value_loss, entropy, clip_frac compute_ppo_loss( model, batch_states, batch_actions, batch_old_log_probs, batch_advantages, batch_returns, clip_epsilon ) optimizer.zero_grad() total_loss.backward() # 可以添加梯度裁剪防止梯度爆炸 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm0.5) optimizer.step() # 输出本回合信息 print(fTimestep: {timestep}, Episode Reward: {ep_reward:.2f}, Length: {ep_len}, fPolicy Loss: {policy_loss:.4f}, Value Loss: {value_loss:.4f}, Clip Frac: {clip_frac:.3f}) env.close() if __name__ __main__: train_ppo()在实际运行这个脚本时你可能会遇到训练曲线震荡、奖励不增长甚至下降的情况。别担心这很正常。PPO虽然稳定但依然有许多“旋钮”需要调试。下面是我从多次实验中总结的几个关键调试点奖励曲线震荡剧烈首先检查优势归一化是否做了。advantages_t (advantages_t - advantages_t.mean()) / (advantages_t.std() 1e-8)这行代码至关重要它能将优势值缩放到均值为0、标准差为1的分布极大稳定了策略梯度更新。如果还震荡尝试减小clip_epsilon如从0.2调到0.1或降低学习率。价值损失一直很高价值网络评论家学不好会导致优势估计不准进而影响策略更新。可以尝试给价值网络使用单独、更小的学习率如策略网络用3e-4价值网络用1e-3。增加价值网络的更新次数在train_epochs循环内可以单独为价值网络多更新几次。检查回报returns的计算是否正确特别是回合结束时的last_value处理。策略熵降为零智能体停止探索如果熵奖励项entropy很快趋近于0说明策略迅速收敛到一个确定性策略可能陷入局部最优。可以适当增大entropy_coef如从0.01调到0.02或者在损失函数中明确加入熵最大化项我们的代码已经包含。Clip Fraction 持续很高clip_fraction指标反映了有多大比例的概率比ratios被剪切了。如果这个值持续高于0.2或0.3说明clip_epsilon设置可能过小限制了策略的必要更新或者学习率太大导致单步更新幅度总是超标。需要综合调整这两个参数。调试PPO是一个需要耐心和经验的过程。一个有效的方法是保持其他参数不变每次只调整一个参数并观察多个回合的平均奖励和上述监控指标的变化趋势。使用tensorboard或wandb等工具记录训练曲线能帮助你更直观地分析问题。最后当你的智能体在CartPole或Pendulum上表现良好后可以尝试更复杂的环境如MuJoCo中的HalfCheetah或Humanoid。这时你可能需要更深的网络、更精细的超参数调优以及并行环境采样来加速数据收集。PyTorch的torch.multiprocessing或SubprocVecEnv来自stable-baselines3库可以帮助你实现这一点这也是工业级PPO实现的标准配置。