PPO算法实战：用Python从零实现倒立摆控制（附完整代码）

在强化学习领域，PPO（Proximal Policy Optimization）算法因其出色的稳定性和样本效率，成为近年来最受欢迎的算法之一。不同于传统策略梯度方法容易出现的训练不稳定问题，PPO通过引入"邻近策略优化"机制，在保证性能的同时大幅降低了调参难度。本文将带您从零开始，用Python和PyTorch实现一个完整的PPO智能体，并在经典的倒立摆控制任务上进行实战演练。

1. 环境搭建与核心概念

倒立摆（Pendulum）是一个经典的强化学习基准环境，其目标是控制摆杆使其保持直立状态。这个看似简单的任务实际上包含了连续动作空间控制的核心挑战——需要智能体学会精细调节力矩大小。

我们先安装必要的依赖库：

pip install gymnasium torch numpy matplotlib

PPO算法的核心思想可以概括为三点：

• Actor-Critic架构：同时学习策略函数（Actor）和价值函数（Critic）
• 重要性采样：利用旧策略收集的数据来估计新策略的预期回报
• 裁剪机制：限制策略更新的幅度，防止训练不稳定

import gymnasium as gym
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from collections import deque
import random

2. 神经网络设计

PPO需要两个神经网络：策略网络（Actor）和价值网络（Critic）。我们将使用PyTorch实现这两个网络。

2.1 策略网络实现

策略网络输出动作的概率分布。对于连续动作空间，我们通常输出高斯分布的均值和标准差：

class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=64):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.mean = nn.Linear(hidden_dim, action_dim)
        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))
        mean = torch.tanh(self.mean(x))  # 输出在[-1,1]范围内
        std = torch.exp(self.log_std)
        return torch.distributions.Normal(mean, std)

2.2 价值网络实现

价值网络评估当前状态的价值，用于计算优势函数：

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

3. PPO算法实现

PPO的核心在于其特殊的损失函数设计。我们实现一个完整的PPO智能体类：

class PPOAgent:
    def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99, 
                 gae_lambda=0.95, clip_ratio=0.2, train_iters=4):
        self.policy = PolicyNetwork(state_dim, action_dim)
        self.value_net = ValueNetwork(state_dim)
        self.optimizer = optim.Adam([
            {'params': self.policy.parameters()},
            {'params': self.value_net.parameters()}
        ], lr=lr)
        
        self.gamma = gamma
        self.gae_lambda = gae_lambda
        self.clip_ratio = clip_ratio
        self.train_iters = train_iters
    
    def compute_gae(self, rewards, values, dones):
        advantages = np.zeros_like(rewards)
        last_advantage = 0
        
        for t in reversed(range(len(rewards))):
            delta = rewards[t] + self.gamma * values[t+1] * (1-dones[t]) - values[t]
            advantages[t] = delta + self.gamma * self.gae_lambda * (1-dones[t]) * last_advantage
            last_advantage = advantages[t]
        
        returns = advantages + values[:-1]
        return advantages, returns
    
    def update(self, states, actions, old_log_probs, rewards, dones):
        states = torch.FloatTensor(states)
        actions = torch.FloatTensor(actions)
        old_log_probs = torch.FloatTensor(old_log_probs)
        
        # 计算GAE和回报
        with torch.no_grad():
            values = self.value_net(states).squeeze().numpy()
        
        advantages, returns = self.compute_gae(rewards, values, dones)
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
        advantages = torch.FloatTensor(advantages)
        returns = torch.FloatTensor(returns)
        
        # 多轮训练
        for _ in range(self.train_iters):
            # 计算新策略的概率
            dist = self.policy(states)
            new_log_probs = dist.log_prob(actions).sum(-1)
            entropy = dist.entropy().mean()
            
            # 计算比率和裁剪目标
            ratios = torch.exp(new_log_probs - old_log_probs)
            surr1 = ratios * advantages
            surr2 = torch.clamp(ratios, 1-self.clip_ratio, 1+self.clip_ratio) * advantages
            policy_loss = -torch.min(surr1, surr2).mean()
            
            # 价值函数损失
            value_pred = self.value_net(states).squeeze()
            value_loss = 0.5 * (value_pred - returns).pow(2).mean()
            
            # 总损失
            loss = policy_loss + 0.5 * value_loss - 0.01 * entropy
            
            # 更新参数
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

4. 训练流程与技巧

训练PPO算法需要特别注意数据收集和更新的方式。以下是完整的训练循环：

def train_ppo(env_name="Pendulum-v1", max_episodes=1000, max_steps=200):
    env = gym.make(env_name)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.shape[0]
    
    agent = PPOAgent(state_dim, action_dim)
    episode_rewards = []
    
    for episode in range(max_episodes):
        state, _ = env.reset()
        states, actions, rewards, dones, old_log_probs = [], [], [], [], []
        episode_reward = 0
        
        for step in range(max_steps):
            # 选择动作
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            with torch.no_grad():
                dist = agent.policy(state_tensor)
                action = dist.sample()
                log_prob = dist.log_prob(action).sum(-1)
                value = agent.value_net(state_tensor)
            
            next_state, reward, done, _, _ = env.step(action.numpy()[0])
            
            # 存储经验
            states.append(state)
            actions.append(action.numpy()[0])
            rewards.append(reward)
            dones.append(done)
            old_log_probs.append(log_prob.item())
            
            state = next_state
            episode_reward += reward
            
            if done:
                break
        
        # 更新智能体
        agent.update(states, actions, old_log_probs, rewards, dones)
        episode_rewards.append(episode_reward)
        
        # 打印训练进度
        if episode % 10 == 0:
            avg_reward = np.mean(episode_rewards[-10:])
            print(f"Episode {episode}, Avg Reward: {avg_reward:.1f}")
    
    return episode_rewards

4.1 训练技巧

1. GAE参数调整：

   • λ=0.95通常效果不错，但可以尝试0.9-0.99之间的值
   • 较高的λ值会使优势估计偏差更小但方差更大

2. 裁剪比例选择：

   • 常见值为0.1-0.3
   • 更小的值意味着更保守的更新，训练更稳定但可能更慢

3. 批量大小：

   • 每个episode的步数不宜过少
   • 可以考虑多个episode的数据合并后再更新

5. 结果可视化与分析

训练完成后，我们可以观察智能体的表现并进行可视化：

def evaluate_agent(agent, env, n_episodes=5, render=True):
    for episode in range(n_episodes):
        state, _ = env.reset()
        total_reward = 0
        
        for step in range(200):
            if render:
                env.render()
            
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            with torch.no_grad():
                dist = agent.policy(state_tensor)
                action = dist.mean
            
            next_state, reward, done, _, _ = env.step(action.numpy()[0])
            total_reward += reward
            state = next_state
            
            if done:
                break
        
        print(f"Evaluation Episode {episode}, Reward: {total_reward:.1f}")
    env.close()

# 训练并评估
rewards = train_ppo()
env = gym.make("Pendulum-v1", render_mode="human")
agent = PPOAgent(env.observation_space.shape[0], env.action_space.shape[0])
evaluate_agent(agent, env)

5.1 性能指标分析

成功的PPO实现通常会在倒立摆环境中表现出以下特征：

• 在100-200个episode后开始学习有效策略
• 稳定后的episode奖励通常在-200到-50之间（Pendulum-v1的奖励范围）
• 策略网络输出的动作会随着摆杆角度变化而平滑调整

6. 进阶优化方向

对于希望进一步提升PPO性能的开发者，可以考虑以下优化：

1. 网络架构改进：

   • 增加层数或神经元数量
   • 尝试不同的激活函数（如Swish、Mish）
   • 添加层归一化（LayerNorm）

2. 训练过程优化：

   • 实现并行环境采样
   • 添加学习率调度
   • 引入课程学习（Curriculum Learning）

3. 算法增强：

   • 结合好奇心驱动探索
   • 添加模型不确定性估计
   • 实现分布式PPO变体

# 进阶PPO实现示例 - 并行环境采样
from torch.utils.data import DataLoader, TensorDataset

class ParallelPPOAgent(PPOAgent):
    def __init__(self, state_dim, action_dim, n_envs=4, **kwargs):
        super().__init__(state_dim, action_dim, **kwargs)
        self.n_envs = n_envs
    
    def parallel_collect(self, envs, steps_per_env=64):
        # 在多个环境中并行收集数据
        all_states, all_actions = [], []
        all_rewards, all_dones = [], []
        all_log_probs = []
        
        states = [env.reset()[0] for env in envs]
        
        for _ in range(steps_per_env):
            states_tensor = torch.FloatTensor(np.array(states))
            with torch.no_grad():
                dists = self.policy(states_tensor)
                actions = dists.sample()
                log_probs = dists.log_prob(actions).sum(-1)
                values = self.value_net(states_tensor).squeeze()
            
            next_states, rewards, dones = [], [], []
            for i, env in enumerate(envs):
                next_state, reward, done, _, _ = env.step(actions[i].numpy())
                next_states.append(next_state)
                rewards.append(reward)
                dones.append(done)
            
            # 存储数据
            all_states.extend(states)
            all_actions.extend(actions.numpy())
            all_rewards.extend(rewards)
            all_dones.extend(dones)
            all_log_probs.extend(log_probs.numpy())
            
            states = next_states
        
        return all_states, all_actions, all_rewards, all_dones, all_log_probs

在实际项目中，PPO算法的超参数调优往往需要多次实验。建议使用如Weights & Biases或TensorBoard等工具记录训练过程，方便分析比较不同配置的效果。