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

2. 经验回放类 (Memory)

PPO 需要收集一批轨迹数据后再进行多次更新，因此需要一个临时存储结构。

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 = []

3. PPO Agent 初始化与动作选择

这里定义了超参数，如折扣因子 gamma、裁剪阈值 eps_clip 以及更新轮数 K_epochs。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

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)
        with torch.no_grad():
            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()

4. 策略更新逻辑

这是 PPO 的灵魂所在。我们在每个 epoch 中计算概率比率，应用裁剪，并联合优化 Actor 和 Critic。

    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())

5. 主训练循环

最后，我们搭建环境并开始训练。这里以 CartPole 为例，这是一个经典的控制任务。

def main():
    env = gym.make("CartPole-v1")
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    ppo = PPO(state_dim, action_dim)
    memory = Memory()

    max_episodes = 1000
    max_timesteps = 300

    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}")

if __name__ == "__main__":
    main()

算法对比：PPO vs TRPO vs A3C

在实际应用中，选择合适的算法取决于具体场景。以下是三种主流策略优化算法的对比：

总结来说：PPO 因其简单、稳定且高效的特性，已成为目前强化学习领域的事实标准。TRPO 虽然理论更严谨，但实现复杂；A3C 则胜在并行速度，但在单卡环境下不如 PPO 稳定。

结语

PPO 的成功在于它在理论严谨性和工程实用性之间找到了完美的平衡点。通过裁剪机制，我们无需复杂的二阶优化即可实现稳定的策略更新。希望这篇详解能帮助你更好地理解其原理，并在实际项目中顺利落地。