近端策略优化(Proximal Policy Optimization, PPO)作为强化学习领域的重要算法，在众多实际应用中展现出卓越的性能。本文将详细介绍PPO算法的核心原理，并提供完整的PyTorch实现方案。
PPO算法在强化学习任务中具有显著优势：即使未经过精细的超参数调优，也能在Atari游戏环境等复杂场景中取得优异表现。该算法不仅在传统强化学习任务中表现出色，还被广泛应用于大语言模型的对齐优化过程。因此掌握PPO算法对于深入理解现代强化学习技术具有重要意义。
本文将通过Lunar Lander环境演示PPO算法的完整实现过程。文章重点阐述算法的核心概念和实现细节，通过适当的修改，本实现方案可扩展至其他强化学习环境。本文专注于高层次的算法理解，为读者提供系统性的技术资源。
PPO算法核心组件
PPO算法由四个核心组件构成：环境交互模块、智能体决策系统、优势函数计算以及策略更新裁剪机制。每个组件在算法整体架构中发挥着关键作用。
环境交互模块
环境是智能体进行学习和决策的载体。这里我们选用Lunar Lander作为测试环境，这是一个二维物理模拟场景，要求着陆器在月球表面的指定区域安全着陆。环境模块负责提供状态观测信息，接收智能体的动作指令，并根据任务完成情况反馈相应的奖励信号。
有效的环境设计和奖励函数是成功训练的基础。智能体需要从环境中获取充分的状态信息以做出合理决策，同时需要通过明确的奖励信号了解其行为的优劣程度。奖励信号的质量直接影响智能体的学习效率和最终性能。
智能体决策系统
PPO采用演员-评论家(Actor-Critic)架构设计智能体决策系统。演员网络负责根据当前状态选择最优动作，而评论家网络则评估当前状态的价值期望。
演员网络的作用类似于决策执行者，根据观测到的环境状态输出动作概率分布。评论家网络则充当价值评估者，预测在当前状态下能够获得的累积奖励期望值。当评论家的价值估计出现偏差时，通常表明智能体的策略仍有改进空间。
优势函数
优势函数用于量化特定动作相对于评论家期望值的优劣程度。正优势值表示该动作的表现优于期望，应当增强此类行为的选择概率；负优势值则表示表现不佳，需要降低此类行为的选择概率。
相比直接使用原始奖励值，优势函数能够提供更稳定的训练信号。智能体仅在实际表现与预期之间存在显著差异时才进行大幅度的策略调整，这种机制有效避免了训练过程中的不必要波动。本实现采用广义优势估计(Generalized Advantage Estimation, GAE)方法计算优势值。
PPO策略更新裁剪机制
PPO算法的核心创新在于引入策略更新的裁剪机制，这是其相对于传统策略梯度方法的关键改进。在强化学习训练过程中，过大的策略更新可能导致训练失稳，使智能体突然丢失已学习的有效策略。
这种现象可以类比为在狭窄山脊上行走的登山者：如果步伐过大或方向偏离，很容易失足跌落深谷，重新攀登将耗费大量时间和精力。PPO通过实施裁剪约束，确保每次策略更新都在安全范围内进行，保持学习过程的稳定性和连续性。
依赖库安装与环境配置
本实现需要安装gymnasium库及其相关依赖来运行Lunar Lander环境，PyTorch用于神经网络的构建和训练，以及tensordict库来管理训练数据。tensordict是一个先进的数据管理工具，允许将PyTorch张量作为字典元素进行操作，支持通过键值索引和检索数据项。这种设计使得数据管理更加灵活高效，同时保持张量和tensordict在GPU上的计算能力，与PyTorch工作流程无缝集成。
!pip install swig gymnasium torch tensordict pyvirtualdisplay !pip install "gymnasium[box2d]"
为确保实验结果的可重现性，我们设置统一的随机种子：
import random import torch import numpy as np seed = 777 random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.backends.cudnn.deterministic = True
需要注意的是，由于向量化环境初始化过程中存在的随机性因素，完全的结果重现仍然面临技术挑战。虽然代码能够正常运行且智能体训练过程稳定，但具体的数值结果在不同运行之间可能存在差异。
Lunar Lander环境配置
Lunar Lander是gymnasium库中的经典强化学习环境，任务目标是控制着陆器在二维空间中导航，最终在月球表面的指定着陆区域安全降落。智能体根据着陆过程中的姿态稳定性、着陆柔和度以及任务完成速度获得相应奖励。着陆器具备四种控制动作，环境在每个时间步提供八维状态向量描述着陆器的当前状态。详细的环境说明可参考官方文档。
为提高训练数据收集效率，我们创建10个并行的Lunar Lander环境实例。同时配置独立的评估环境，用于记录智能体在各个训练阶段的表现视频，便于直观评估训练效果。
import gymnasium as gym from gymnasium.wrappers import RecordVideo # 创建将运行10个模拟的向量化环境env_name = "LunarLander-v3" num_envs = 10 envs = gym.make_vec(env_name, num_envs=num_envs, vectorization_mode="sync") # 创建我们的评估录制环境，用于当我们# 想要测试我们的智能体在训练的各个阶段表现如何时使用env = gym.make(env_name, render_mode="rgb_array") trigger = lambda t: True recording_output_directory = "./checkpoint_videos" recording_env = RecordVideo(env, recording_output_directory, episode_trigger=trigger)
环境交互辅助类设计
为提高代码的模块化程度和可维护性，我们设计了专门的环境交互辅助类，负责处理训练数据收集和评估过程。该类封装了与环境的所有交互操作，包括训练环境、评估环境、智能体以及计算设备的管理。
主要功能包括：数据rollout收集用于获取训练样本，以及评估rollout执行用于性能测试和视频记录。评估过程中的视频录制功能由RecordVideo包装器自动完成。
from tensordict import TensorDict
import torch
# 关于这个类的快速说明# 它假设训练环境在终端状态时自动重置，例如向量化环境# 并且它假设评估环境不会自动重置# 我还将num_steps_per_rollout固定为收集器初始化时的值，但这可以在# get rollout函数中参数化class PPORolloutCollector:
def __init__(self, agent, envs, num_steps_per_rollout, device, eval_env=None):
self.agent = agent
self.envs = envs
self.num_envs = envs.num_envs
self.num_steps_per_rollout = num_steps_per_rollout
self.device = device
self.eval_env = eval_env
self.obs_shape = envs.single_observation_space.shape
self.action_shape = envs.single_action_space.shape
self.initial_buffer_shape = (num_steps_per_rollout, envs.num_envs)
# 继续重置环境并存储观察
# 并将next_done设置为false(我们假设环境不能以终端状态开始)
obs, _ = envs.reset()
self.next_obs = torch.Tensor(obs).to(device)
self.next_done = torch.zeros(num_envs).to(device)
# 创建一个空缓冲区，将保存观察、来自智能体的动作、该动作的对数概率
# 当前观察的评论家估计、从环境中获得的采取动作的实际奖励，以及
# 动作是否导致终端状态
# 我们故意选择不为每个观察记录"下一状态"，这基本上会使
# 缓冲区的大小翻倍，由于我们按顺序收集观察并且在这里不打乱
# 任何下游操作都可以只使用数组中的下一个值作为下一状态
def _create_buffer(self):
return TensorDict({
"obs": torch.zeros(self.initial_buffer_shape + self.obs_shape).to(self.device),
"actions": torch.zeros(self.initial_buffer_shape + self.action_shape).to(self.device),
"log_probs": torch.zeros(self.initial_buffer_shape).to(self.device),
"rewards": torch.zeros(self.initial_buffer_shape).to(self.device),
"dones": torch.zeros(self.initial_buffer_shape).to(self.device),
"critic_values": torch.zeros(self.initial_buffer_shape).to(self.device),
})
# 将收集num_steps_per_rollout个观察对我们的训练环境的函数
def get_next_rollout(self):
buffer = self._create_buffer()
# 获取最后记录的观察以及该观察是否为终端
next_obs = self.next_obs
next_done = self.next_done
# 收集rollout
for t in range(self.num_steps_per_rollout):
# 记录当前观察和终端状态
buffer["obs"][t] = next_obs
buffer["dones"][t] = next_done
# 查询智能体下一个动作、该动作的对数概率和评论家估计
with torch.no_grad():
action, log_prob, entropy = self.agent.get_actor_values(next_obs)
critic_value = self.agent.get_critic_value(next_obs)
# 记录值
buffer["actions"][t] = action
buffer["log_probs"][t] = log_prob
buffer["critic_values"][t] = critic_value.flatten()
# 执行动作
next_obs, reward, terminations, truncations, infos = envs.step(action.cpu().numpy())
# 形状化并存储奖励
reward = torch.tensor(reward).to(self.device).view(-1)
buffer["rewards"][t] = reward
# 一些环境会终止(意味着智能体处于最终状态)，
# 其他会截断(例如达到时间限制但不在终端状态)
# 这些是重要的区别，但对我们的目的来说意味着同样的事情，模拟结束了
# 所以如果任一为真且模拟重置，我们将next done设置为true
next_done = np.logical_or(terminations, truncations)
# 为下一轮存储下一个obs和next done
next_obs, next_done = torch.Tensor(next_obs).to(self.device), torch.Tensor(next_done).to(self.device)
# 在缓冲区中存储下一个obs和next done
# 这是为了在稍后计算优势时处理边缘情况
buffer['next_obs'] = next_obs
buffer['next_done'] = next_done
# 我们还需要这个下一状态的评论家估计，我们将使用它来引导最终状态的奖励
# 当我们计算gae时
with torch.no_grad():
buffer['next_value'] = self.agent.get_critic_value(next_obs).reshape(1, -1)
self.next_obs = next_obs
self.next_done = next_done
return buffer
# 用于在我们的环境上评估智能体，将运行整个模拟直到终止
# 与上面非常相似，唯一的区别是我们将手动检查环境是否终止
# 然后结束循环
# 将返回一个普通的python字典，包含奖励、熵、奖励平均值、我们智能体的平均熵，以及每次运行的总奖励
def run_eval_rollout(self, num_episodes: int = 5):
assert self.eval_env is not None, "No eval_env provided."
rewards_per_timestep = []
entropies_per_timestep = []
final_rewards = []
total_entropies = []
for _ in range(num_episodes):
obs, _ = self.eval_env.reset()
obs = torch.tensor(obs, device=self.device).unsqueeze(0)
done = False
episode_rewards = []
episode_entropies = []
while not done:
with torch.no_grad():
action, _, entropy = self.agent.get_actor_values(obs)
action = action.squeeze()
obs_np, reward, term, trunc, _ = self.eval_env.step(action.cpu().numpy())
done = term or trunc
obs = torch.tensor(obs_np, device=self.device).unsqueeze(0)
episode_rewards.append(float(reward))
episode_entropies.append(entropy.item())
rewards_per_timestep.append(episode_rewards)
entropies_per_timestep.append(episode_entropies)
final_rewards.append(sum(episode_rewards))
total_entropies.append(sum(episode_entropies) / len(episode_entropies))
return {
"rewards_per_timestep": rewards_per_timestep,
"average_reward_per_run": sum(final_rewards) / len(final_rewards),
"average_entropy_per_run": sum(total_entropies) / len(total_entropies),
"entropies_per_timestep": entropies_per_timestep,
"final_rewards": final_rewards,
}
智能体网络架构设计
本实现采用演员-评论家架构构建智能体决策系统。该架构的一个重要特点是演员网络和评论家网络可以完全独立，也可以共享部分网络层。在复杂环境中，通常建议让两个网络共享前期特征提取层，使得双方都能从对方的损失更新中受益。然而，在本示例中，为了清晰展示两个网络的独立性，我们采用完全分离的网络结构。
借助PyTorch模块的参数管理机制，我们无需将两个网络实现为独立的类，可以在单一Agent类中统一管理。
import torch.nn as nn from torch.distributions.categorical import Categorical # 一个初始化辅助函数。在强化学习问题中，正交初始化# 网络的权重可能会有帮助，这意味着每层的输出尽可能# 不相关。这可以通过帮助梯度更好地流动和避免可能# 从朴素初始化中产生的相关特征问题来改善训练稳定性和效率。# 可以把这想象成确保网络不会以交叉的线路开始。# 这个函数取自CleanRL的PPO实现。def layer_init(layer, std=np.sqrt(2), bias_const=0.0): torch.nn.init.orthogonal_(layer.weight, std) torch.nn.init.constant_(layer.bias, bias_const) return layer class Agent(nn.Module): def __init__(self, envs): super().__init__() # 我们的输入数组有多大？ # 注意如果我们做的是图像观察之类的 # 我们需要重新设计这个，但lunar lander只提供 # 代表世界状态的数字数组 input_shape = np.array(envs.single_observation_space.shape).prod() # 我们的智能体将有多少个动作？ action_shape = envs.single_action_space.n # 创建一个3层评论家，它将预测 # 我们的智能体在给定观察下预期收到的总奖励 self.critic = nn.Sequential( layer_init(nn.Linear(input_shape, 64)), nn.Tanh(), layer_init(nn.Linear(64, 64)), nn.Tanh(), layer_init(nn.Linear(64,1), std=1.0) ) # 创建一个3层演员，它将输出一个概率分布 # 表示在给定观察下最好采取哪个动作的概率 self.actor = nn.Sequential( layer_init(nn.Linear(input_shape, 64)), nn.Tanh(), layer_init(nn.Linear(64, 64)), nn.Tanh(), layer_init(nn.Linear(64, action_shape), std=1.0) ) def save(self, path): torch.save(self.state_dict(), path) def load(self, path): self.load_state_dict(torch.load(path)) def get_critic_value(self, x): return self.critic(x) def get_actor_values(self, x, action=None): logits = self.actor(x) probs = Categorical(logits=logits) if action is None: action = probs.sample() return action, probs.log_prob(action), probs.entropy()
网络初始化采用正交初始化策略，这是强化学习领域的常见实践。正交初始化确保各网络层输出之间的去相关性，有助于改善梯度流动并避免特征相关性问题，从而提升训练的稳定性和效率。这种初始化方法可以形象地理解为确保网络不会以"交叉连线"的状态开始训练。
创建智能体实例并配置计算设备：
# 如果可用，使用gpu
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
agent = Agent(envs).to(device)
print("using device: ", device)
配置优化器，选择Adam优化算法并设置学习率为2.5e-4，这是PPO训练的常用起始参数：
import torch.optim as optim learning_rate = 2.5e-4 optimizer = optim.Adam(agent.parameters(), lr=learning_rate, eps=1e-5)
优势函数理论与实现
优势函数用于量化特定动作相对于评论家价值估计的表现差异。当智能体在游戏中的典型得分为15分时，这个分数就成为性能基准。如果智能体采用新策略获得17分，则优势值为2，表明新策略的表现超出预期。
采用优势函数替代原始奖励信号能够稳定训练过程，防止智能体满足于次优策略。当我们使用优势校正值训练评论家网络时：
相比直接预测原始回报，这种方法通常能实现更快的收敛速度和更稳定的评论家更新过程。
广义优势估计(GAE)算法
广义优势估计是计算优势函数的经典方法。该方法不直接比较完整轨迹回报与价值估计(这种方式存在较大噪声)，而是采用时间差分残差的平滑求和。这种设计在偏差和方差之间取得平衡，即使在评论家网络估计不够精确的情况下也能保持实用性。
来自 “ ITPUB博客 ” ，链接：https://blog.itpub.net/70047311/viewspace-3094953/ ，如需转载，请注明出处，否则将追究法律责任。
