近端策略优化算法 PPO 详解与 PyTorch 实现

import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import numpy as np
import gym

# 设备配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Actor-Critic 神经网络定义
class ActorCritic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super(ActorCritic, self).__init__()
        # 共享层用于特征提取
        self.shared_layer = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ReLU()
        )
        # Actor: 输出动作概率
        self.actor = nn.Sequential(
            nn.Linear(128, action_dim),
            nn.Softmax(dim=-1)
        )
        # Critic: 输出状态值
        self.critic = nn.Linear(128, 1)

    def forward(self, state):
        shared = self.shared_layer(state)
        action_probs = self.actor(shared)
        state_value = self.critic(shared)
        return action_probs, state_value

# 经验存储类
class Memory:
    def __init__(self):
        self.states = []
        self.actions = []
        self.logprobs = []
        self.rewards = []
        self.is_terminals = []

    def clear(self):
        self.states = []
        self.actions = []
        self.logprobs = []
        self.rewards = []
        self.is_terminals = []

# PPO Agent 类
class PPO:
    def __init__(self, state_dim, action_dim, lr=0.002, gamma=0.99, eps_clip=0.2, K_epochs=4):
        self.policy = ActorCritic(state_dim, action_dim).to(device)
        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)
        self.policy_old = ActorCritic(state_dim, action_dim).to(device)
        self.policy_old.load_state_dict(self.policy.state_dict())
        self.MseLoss = nn.MSELoss()
        self.gamma = gamma
        self.eps_clip = eps_clip
        self.K_epochs = K_epochs

    def select_action(self, state, memory):
        state = torch.FloatTensor(state).to(device)
        action_probs, _ = self.policy_old(state)
        dist = Categorical(action_probs)
        action = dist.sample()
        memory.states.append(state)
        memory.actions.append(action)
        memory.logprobs.append(dist.log_prob(action))
        return action.item()

    def update(self, memory):
        old_states = torch.stack(memory.states).to(device).detach()
        old_actions = torch.stack(memory.actions).to(device).detach()
        old_logprobs = torch.stack(memory.logprobs).to(device).detach()

        # 计算折扣奖励
        rewards = []
        discounted_reward = 0
        for reward, is_terminal in zip(reversed(memory.rewards), reversed(memory.is_terminals)):
            if is_terminal:
                discounted_reward = 0
            discounted_reward = reward + (self.gamma * discounted_reward)
            rewards.insert(0, discounted_reward)
        rewards = torch.tensor(rewards, dtype=torch.float32).to(device)
        rewards = (rewards - rewards.mean()) / (rewards.std() + 1e-7)

        # 多轮更新
        for _ in range(self.K_epochs):
            action_probs, state_values = self.policy(old_states)
            dist = Categorical(action_probs)
            new_logprobs = dist.log_prob(old_actions)
            entropy = dist.entropy()

            ratios = torch.exp(new_logprobs - old_logprobs.detach())
            advantages = rewards - state_values.detach().squeeze()

            surr1 = ratios * advantages
            surr2 = torch.clamp(ratios, 1 - self.eps_clip, 1 + self.eps_clip) * advantages
            loss_actor = -torch.min(surr1, surr2).mean()
            loss_critic = self.MseLoss(state_values.squeeze(), rewards)
            loss = loss_actor + 0.5 * loss_critic - 0.01 * entropy.mean()

            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

        self.policy_old.load_state_dict(self.policy.state_dict())

# 主程序
if __name__ == "__main__":
    env = gym.make("CartPole-v1")
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n
    
    lr = 0.002
    gamma = 0.99
    eps_clip = 0.2
    K_epochs = 4
    max_episodes = 1000
    max_timesteps = 300

    ppo = PPO(state_dim, action_dim, lr, gamma, eps_clip, K_epochs)
    memory = Memory()

    for episode in range(1, max_episodes + 1):
        state = env.reset()
        total_reward = 0
        for t in range(max_timesteps):
            action = ppo.select_action(state, memory)
            state, reward, done, _ = env.step(action)
            memory.rewards.append(reward)
            memory.is_terminals.append(done)
            total_reward += reward
            if done:
                break
        ppo.update(memory)
        memory.clear()
        print(f"Episode {episode}, Total Reward: {total_reward}")

    env.close()

关键实现点说明

1. Actor-Critic 网络结构：ActorCritic 模型通过共享层生成动作概率和状态值。Critic 提供状态价值估计，辅助 Actor 进行策略更新。
2. PPO 的优化目标：使用裁剪的目标函数限制策略更新幅度，避免过度更新导致发散。
3. 内存管理：Memory 类用于存储每一轮的状态、动作、奖励和终止标志，支持批量更新。
4. 奖励归一化：在 update 方法中对奖励进行标准化处理，以加快收敛速度并减少数值不稳定。
5. 训练循环：PPO 从环境中采样，更新策略，并打印每一集的总奖励以便监控训练进度。

注意：上述代码主要用于理解和学习算法原理。若应用于实际项目，可能需要根据具体环境调整超参数（如学习率、折扣因子等）并进行进一步优化。

4. 算法对比总结

总结

1. PPO 是 TRPO 的改进版：PPO 使用简单的裁剪机制代替了 TRPO 的二次优化，显著降低了实现复杂度，同时保持了良好的稳定性和效率。
2. A3C 的并行化设计：A3C 的核心是通过多线程异步更新提升效率，但其稳定性略低于 PPO 和 TRPO。
3. 实用性：PPO 因其简单、稳定、高效的特点，已成为强化学习领域的主流算法；TRPO 更适合需要极高策略稳定性的任务；A3C 在资源受限的场景下表现优异。

PPO 的关键是通过限制策略的变化范围（剪辑），让优化更加稳定，同时通过优势函数引导策略改进，充分利用采样数据。这种平衡使得 PPO 成为许多强化学习任务的默认选择。