PPO算法实战：600步搞定Pendulum-v1平衡（附完整PyTorch代码）

PPO算法实战：600步搞定Pendulum-v1平衡（附完整PyTorch代码）

如果你刚开始接触强化学习，面对一堆理论公式和算法变体感到无从下手，那么从一个具体的、可运行的环境开始“动手”可能是最好的解药。Pendulum-v1，这个经典的倒立摆控制问题，就像一个绝佳的“新手村”。它状态和动作空间连续，但又不像一些复杂环境那样需要海量计算资源，非常适合用来理解像PPO（近端策略优化）这类现代策略梯度算法的核心思想与实现细节。很多教程会告诉你PPO有多好，但很少会手把手地带你走一遍，如何用大约600个训练回合（episode）就让智能体学会稳定地平衡摆杆。这篇文章的目的就是填补这个空白：我们将从零开始，构建一个完整的PyTorch实现，并深入探讨那些让训练快速收敛的关键技巧，比如奖励缩放（Reward Scaling）和状态归一化，而不仅仅是把代码扔给你。

1. 理解我们的“训练场”：Pendulum-v1环境剖析

在开始写代码之前，我们必须先彻底了解智能体将要面对的环境。Pendulum-v1模拟的是一个单摆系统：摆杆的一端被固定，另一端可以自由摆动。你的任务是通过施加一个扭矩（力矩），让摆杆从自然下垂的状态摆动并稳定在竖直向上（倒立）的位置。

1.1 状态与动作：智能体感知与交互的维度

这个环境的信息交换遵循标准的“观察-行动-奖励”循环。智能体每一步都需要处理以下信息：

• 状态空间 (State Space): 一个包含3个元素的数组 [cos(theta), sin(theta), theta_dot]。

  • theta 是摆杆与竖直向上方向的夹角（弧度）。环境并不直接给出 theta，而是给出其余弦和正弦值，这避免了角度从 π 跳变到 -π 的不连续性。
  • theta_dot 是摆杆的角速度（弧度/秒）。
  • 注意：有些资料会提到4个状态变量，这通常包含了位置和线速度，但在Gym的Pendulum-v1标准实现中，状态向量就是上述3维。

• 动作空间 (Action Space): 一个连续的标量值，范围在 [-2.0, 2.0] 之间。这个值代表施加在摆杆铰链上的扭矩大小。正负号代表方向。

• 奖励函数 (Reward Function): 奖励的计算公式为 -(theta^2 + 0.1*theta_dot^2 + 0.001*action^2)。

  • 这个设计非常直观：theta 越接近0（竖直），theta_dot 越小（越稳定），action 越小（控制越平滑），得到的负值就越小，即惩罚越小，等价于奖励越高。
  • 理想状态（完全竖直且静止）下的奖励是 0，这是理论上的最大值。因此，训练的目标是让累计奖励（通常是负值）尽可能接近0。

理解这个奖励结构至关重要。它不是一个稀疏奖励（只有成功才给正奖励），而是一个稠密奖励，每一步都根据当前“不完美”的程度给予反馈。这大大降低了学习难度。

1.2 环境的挑战与PPO的适应性

为什么PPO适合这个任务？Pendulum-v1虽然基础，但具备连续控制问题的典型挑战：

1. 连续动作空间：无法像离散动作那样简单地选择最大Q值的动作。PPO通过输出动作分布（如高斯分布）的参数来自然地处理连续动作。
2. 奖励尺度问题：原始奖励始终为负，且绝对值可能很小（比如-0.05）。如果不做处理，计算出的优势函数和梯度也会非常小，导致学习缓慢甚至不稳定。这就是我们引入奖励缩放的核心原因。
3. 探索与利用的平衡：智能体需要尝试不同的扭矩来发现如何向上摆动，同时也要学会精细调节以保持平衡。PPO通过其裁剪机制和策略熵正则化，能较好地维持这一平衡。

下面的表格对比了处理Pendulum-v1时，关键技巧的引入与否对训练可能产生的影响：

技巧/组件	未采用时的潜在问题	采用后的主要作用
奖励缩放 (Reward Scaling)	梯度微小，学习速度极慢，训练曲线波动大。	将奖励调整到合适的尺度，稳定梯度，加速收敛。
状态归一化 (State Normalization)	不同状态维度量纲和范围不同，可能干扰网络学习。	使输入数据分布稳定，加速网络收敛，提高泛化性。
广义优势估计 (GAE)	优势估计方差高，策略更新噪声大。	在偏差和方差间取得平衡，提供更平滑、更可靠的优势信号。
策略熵正则化	策略过早收敛到次优解，探索不足。	鼓励探索，防止策略过早确定性化，有助于找到更优策略。

提示：在强化学习中，“技巧”往往比算法选择本身更能决定实验的成败。尤其是在像PPO这样已经相对鲁棒的算法上，正确的工程实现细节是快速复现论文结果的关键。

2. PPO核心思想与我们的实现策略

PPO之所以成为深度强化学习领域的“默认选择”之一，源于它在理论严谨性和实现简便性之间取得的出色平衡。它本质上是信任域策略优化（TRPO） 的一个近似，但通过一个简单的裁剪目标函数，避免了TRPO中复杂的二阶优化计算。

2.1 从策略梯度到PPO裁剪目标

策略梯度算法的核心是沿着策略性能梯度上升的方向更新参数。其目标函数是期望优势函数： J(θ) = E[ (π_θ(a|s) / π_θ_old(a|s)) * A(s, a) ] 其中 (π_θ(a|s) / π_θ_old(a|s)) 是新旧策略的概率比，A(s, a) 是优势函数，衡量当前动作比平均动作好多少。

直接优化这个目标可能导致更新步长过大，策略剧烈变化，进而导致性能崩溃。PPO的解决方案是修改目标函数，限制新旧策略的差异：

# PPO 裁剪目标函数的核心代码逻辑
ratio = torch.exp(log_prob_new - log_prob_old) # 概率比
surr1 = ratio * advantages
surr2 = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon) * advantages
policy_loss = -torch.min(surr1, surr2).mean() # 取最小值进行裁剪

• ratio: 新旧策略对数概率的指数，即概率比。
• advantages: 估计的优势函数。
• clip_epsilon: 一个超参数（如0.2），定义了裁剪的范围。
• torch.min(surr1, surr2): 这是PPO的精华。当 ratio 偏离1太远（即策略变化太大）时，surr2 会被裁剪，其值将小于 surr1，因此 torch.min 会选择被裁剪后的目标值，从而限制了策略更新的幅度。

2.2 网络架构设计：Actor与Critic

我们采用在连续控制任务中非常经典的“Actor-Critic”架构，两个网络共享部分特征提取层或完全独立。在我们的实现中，为了清晰起见，让它们完全独立。

• Actor网络 (策略网络): 输入状态，输出动作的概率分布参数。对于连续动作空间，我们通常假设动作服从高斯分布，因此网络输出分布的均值(mu)和对数标准差(log_std)。

  class Actor(nn.Module):
      def __init__(self, state_dim, action_dim):
          super().__init__()
          self.fc1 = nn.Linear(state_dim, 128)
          self.fc2 = nn.Linear(128, 128)
          self.mu_head = nn.Linear(128, action_dim)
          # 将 log_std 作为可训练参数，而非网络输出，更稳定
          self.log_std = nn.Parameter(torch.zeros(action_dim))
  
      def forward(self, x):
          x = torch.relu(self.fc1(x))
          x = torch.relu(self.fc2(x))
          mu = 2.0 * torch.tanh(self.mu_head(x)) # 将均值限制在[-2,2]内，与动作范围匹配
          std = torch.exp(self.log_std) # 确保标准差为正
          return mu, std
  

  关键点：使用 torch.tanh 将均值输出限制在 [-2, 2]，与环境动作范围一致。将 log_std 作为独立参数而非网络输出，训练更稳定。

• Critic网络 (价值网络): 输入状态，输出该状态的标量价值估计 V(s)。这个值用于计算优势函数 A(s, a) = Q(s, a) - V(s)，但我们通常通过GAE来估计优势，Critic的输出是GAE计算的基础。

  class Critic(nn.Module):
      def __init__(self, state_dim):
          super().__init__()
          self.fc1 = nn.Linear(state_dim, 128)
          self.fc2 = nn.Linear(128, 128)
          self.v_head = nn.Linear(128, 1)
  
      def forward(self, x):
          x = torch.relu(self.fc1(x))
          x = torch.relu(self.fc2(x))
          return self.v_head(x)
  

2.3 训练流程概览

PPO是一种同策略（on-policy） 算法，意味着用于更新的数据必须由当前策略生成。其训练遵循一个“收集数据-更新策略”的循环：

1. 交互收集轨迹：使用当前策略与环境交互，收集一定数量（如2048步）的状态、动作、奖励、下一个状态、终止标志等数据。
2. 优势估计：利用收集到的整条轨迹数据和Critic网络，使用GAE方法计算每个时间步的优势函数估计值。
3. 多轮小批量更新：将收集到的数据视为一个“批次”，在更新阶段，将这个批次的数据打乱，分成多个小批量（mini-batch），对每个小批量执行PPO的裁剪目标函数优化和Critic的均方误差优化。这个过程通常会重复多个轮次（K_epochs，如10次）。
4. 迭代：清空缓冲区，用更新后的策略回到第1步，继续收集新数据。

注意：PPO的“近端”体现在第3步。它允许使用同一批数据对策略进行多次更新（K_epochs），但通过裁剪目标函数确保每次更新都不会离生成这批数据的旧策略太远，从而在样本效率和稳定性之间取得折衷。

3. 关键实现技巧深度解析

要让PPO在Pendulum-v1上快速收敛，仅仅实现基础算法是不够的。以下几个技巧是工程实现中的“胜负手”。

3.1 奖励缩放：稳定训练的第一道保险

Pendulum-v1的原始奖励始终为负，且绝对值很小。这会导致计算出的优势函数 A(s,a) 和策略梯度非常微小，在浮点数计算中容易受到噪声影响，学习极其缓慢。

奖励缩放的核心思想是动态地调整奖励的尺度，使其均值和方差处在一个对神经网络优化器（如Adam）友好的范围内。我们实现了一个 RewardScaling 类，它不仅仅做简单的归一化，而是结合了折扣累积奖励的在线标准化：

class RewardScaling:
    def __init__(self, shape, gamma):
        self.gamma = gamma
        self.running_ms = RunningMeanStd(shape=shape) # 动态计算均值和标准差
        self.R = np.zeros(shape) # 用于计算折扣累积奖励

    def __call__(self, x):
        # 更新折扣累积奖励：R_t = r_t + γ * R_{t+1}
        self.R = self.gamma * self.R + x
        # 用这个累积奖励的序列来更新运行统计量
        self.running_ms.update(self.R)
        # 将当前奖励x除以其标准差（不加均值中心化，以保留正负号）
        x = x / (self.running_ms.std + 1e-8)
        return x

    def reset(self):
        # 每个episode开始时，重置累积奖励R
        self.R = np.zeros(self.shape)

• 工作原理：我们不是直接缩放单步奖励 r_t，而是先计算一个折扣累积奖励 R_t，然后用大量 R_t 序列的统计量（这里只用标准差）来缩放当前的 r_t。这样缩放后的奖励具有更稳定的分布。
• 为什么有效：它保证了用于计算优势函数的奖励信号具有大致稳定的幅度，避免了梯度爆炸或消失。在实践中，这个技巧常常能直接将训练速度提升一个数量级。

3.2 状态归一化：为网络提供“标准餐”

神经网络，尤其是深度网络，喜欢输入数据是零均值、单位方差的。状态归一化就是为了满足这个偏好。

class Normalization:
    def __init__(self, shape):
        self.running_ms = RunningMeanStd(shape=shape)

    def __call__(self, x, update=True):
        if update:
            self.running_ms.update(x)
        # 执行 (x - mean) / std
        x = (x - self.running_ms.mean) / (self.running_ms.std + 1e-8)
        return x

我们使用一个动态计算运行均值和标准差的类 RunningMeanStd。在训练时 (update=True)，我们不断用看到的状态更新这些统计量；在评估时 (update=False)，则使用固定的统计量进行归一化。这确保了训练和评估阶段的数据变换是一致的。

3.3 广义优势估计：更优的价值信号

优势函数 A(s,a) 衡量的是在状态 s 下采取动作 a 相对于平均情况的好坏。直接使用 Q(s,a) - V(s) 需要估计Q值，比较复杂。GAE提供了一种利用多步TD误差来平衡估计偏差和方差的优雅方法：

# GAE 计算片段
deltas = rewards + gamma * (1 - dones) * next_values - values
advantages = []
gae = 0
for delta, done in zip(reversed(deltas), reversed(dones)):
    gae = delta + gamma * lambda_ * gae * (1 - done)
    advantages.insert(0, gae) # 从后往前插入
advantages = torch.tensor(advantages)

• lambda_ 是一个介于0和1之间的超参数。当 lambda_=0 时，GAE 退化为单步TD误差（高偏差，低方差）；当 lambda_=1 时，GAE 等价于蒙特卡洛回报（低偏差，高方差）。通常取 0.95 作为一个不错的折衷。
• 从轨迹末尾反向计算是GAE实现的标准做法，非常高效。

4. 完整训练流程与代码逐行指南

现在，我们将所有组件串联起来，构建完整的训练循环。以下是核心步骤和对应的代码逻辑。

4.1 初始化与环境设置

首先，导入必要的库并创建环境。注意，我们创建了两个环境实例：一个用于训练（env），一个用于阶段性评估（env_evaluate）。评估环境通常设置为可渲染模式，方便我们直观观察智能体的表现。

import gym
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

env = gym.make('Pendulum-v1') # 训练环境
state_dim = env.observation_space.shape[0] # 3
action_dim = env.action_space.shape[0] # 1

# 用于阶段性评估的环境，可以开启渲染
env_evaluate = gym.make('Pendulum-v1', render_mode='human')
agent = PPOAgent(state_dim, action_dim)
state_norm = Normalization(shape=state_dim)
reward_scaling = RewardScaling(shape=1, gamma=0.99)
replay_buffer = ReplayBuffer(buffer_size=2048, state_dim=state_dim, action_dim=action_dim)

4.2 主训练循环

训练循环的核心是交替进行数据收集和策略更新。

max_episodes = 600 # 目标训练回合数
max_steps_per_episode = 200 # 每个episode最大步数

for episode in range(max_episodes):
    state, _ = env.reset()
    episode_reward = 0
    reward_scaling.reset() # 每个episode开始前重置奖励缩放器
    state = state_norm(state) # 归一化初始状态

    for step in range(max_steps_per_episode):
        # 1. 选择动作
        with torch.no_grad():
            action, log_prob = agent.select_action(state)
            # action是从策略分布中采样的，log_prob是该动作的对数概率

        # 2. 与环境交互
        next_state, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated

        # 3. 处理奖励和状态
        scaled_reward = reward_scaling(reward) # 奖励缩放
        episode_reward += reward # 记录原始奖励用于日志
        next_state = state_norm(next_state) # 归一化下一状态

        # 4. 存储转移样本
        # dw (dead/win) 表示是否因终止而结束，用于GAE计算
        dw = True if (terminated and not truncated) else False
        replay_buffer.store(state, action, log_prob, scaled_reward, next_state, dw, done)

        state = next_state

        # 5. 当缓冲区满时，进行策略更新
        if replay_buffer.count == replay_buffer.buffer_size:
            agent.update(replay_buffer)
            replay_buffer.clear() # 清空on-policy缓冲区

        if done:
            break

    # 6. 日志记录与评估
    print(f'Episode {episode:4d}, Total Steps:{step:4d}, Reward: {episode_reward:+.3f}')

    # 每100个episode评估一次策略
    if episode % 100 == 0:
        evaluate_policy(env_evaluate, agent, state_norm)

4.3 策略更新函数详解

agent.update 函数是PPO算法的核心。它从缓冲区获取数据，计算优势，并执行多轮小批量更新。

def update(self, replay_buffer):
    # 从缓冲区获取数据并转换为张量
    states, actions, old_log_probs, rewards, next_states, dws, dones = replay_buffer.get()

    # 计算优势函数 (GAE) 和目标价值
    with torch.no_grad():
        values = self.critic(states)
        next_values = self.critic(next_states)
        deltas = rewards + self.gamma * (1 - dws) * next_values - values
        advantages = self.compute_gae(deltas, dones)
        targets = advantages + values

    # 将数据打乱，准备进行多轮小批量更新
    indices = np.arange(replay_buffer.buffer_size)
    for _ in range(self.K_epochs):
        np.random.shuffle(indices)
        for start in range(0, replay_buffer.buffer_size, self.mini_batch_size):
            end = start + self.mini_batch_size
            batch_indices = indices[start:end]

            batch_states = states[batch_indices]
            batch_actions = actions[batch_indices]
            batch_old_log_probs = old_log_probs[batch_indices]
            batch_advantages = advantages[batch_indices]
            batch_targets = targets[batch_indices]

            # 计算新的动作对数概率和熵
            dist_now = self.actor.get_dist(batch_states)
            log_probs_now = dist_now.log_prob(batch_actions).sum(dim=-1, keepdim=True)
            entropy = dist_now.entropy().sum(dim=-1, keepdim=True)

            # PPO 裁剪目标函数
            ratios = torch.exp(log_probs_now - batch_old_log_probs)
            surr1 = ratios * batch_advantages
            surr2 = torch.clamp(ratios, 1 - self.clip_epsilon, 1 + self.clip_epsilon) * batch_advantages
            actor_loss = -torch.min(surr1, surr2).mean() - self.entropy_coef * entropy.mean()

            # Critic 损失 (MSE)
            values_now = self.critic(batch_states)
            critic_loss = nn.functional.mse_loss(batch_targets, values_now)

            # 更新网络
            self.actor_optimizer.zero_grad()
            actor_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.actor.parameters(), 0.5) # 梯度裁剪
            self.actor_optimizer.step()

            self.critic_optimizer.zero_grad()
            critic_loss.backward()
            torch.nn.utils.clip_grad_norm_(self.critic.parameters(), 0.5)
            self.critic_optimizer.step()

4.4 训练曲线分析与调试建议

运行上述代码，你可能会观察到训练曲线呈现以下模式：

• 初期（~前50个episode）：奖励波动很大，智能体在随机探索。
• 中期（~50-300个episode）：奖励开始呈现上升趋势，智能体逐渐学会将摆杆向上摆动。
• 后期（~300-600个episode）：奖励趋于稳定，接近0，智能体能够熟练地将摆杆平衡在竖直位置。

如果训练不理想，可以检查以下几点：

1. 奖励是否缩放？ 检查 reward_scaling 输出的值是否在 [-10, 10] 左右的范围。如果还是像 -0.05 这样，说明缩放没起作用。
2. 优势函数是否合理？ 打印出 advantages 的均值和标准差。理想情况下，均值应接近0，标准差在一个合理的量级（如1-5）。如果绝对值非常小（如1e-3），可能是奖励缩放或GAE计算有误。
3. 策略更新幅度是否过大？ 监控 ratios（概率比）的值。如果大量 ratios 超出了裁剪范围 [0.8, 1.2]，说明单次更新步长可能太大，可以尝试减小学习率或增大 clip_epsilon。
4. 熵值是否下降过快？ 策略熵（entropy）衡量策略的随机性。在训练初期，熵应该保持一定水平以鼓励探索；随着训练进行，熵会逐渐下降，策略变得确定。如果熵在最初几十个episode就骤降到接近0，说明探索不足，可以适当增大 entropy_coef。

我在自己的机器上跑这个实现，通常在第400到500个episode左右，智能体就能非常稳定地平衡摆杆了，累计奖励（未缩放的原奖励）能稳定在 -10 以上（距离理论最优值0很近了）。最关键的是，整个训练过程非常平稳，没有出现性能突然崩溃的情况，这正体现了PPO结合这些工程技巧后的鲁棒性。