为什么它能成为强化学习的“黄金标准”？深扒 Proximal Policy Optimization (PPO) 的核心奥秘 原创

# --- compatibility shim for NumPy>=2.0 ---
import numpy as np
if not hasattr(np, "bool8"):
    np.bool8 = np.bool_

# --- imports ---
import gymnasium as gym   # use gymnasium; if you must keep gym
import torch
import torch.nn as nn
import torch.optim as optim

# Actor-Critic network definition
class ActorCritic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.fc1 = nn.Linear(state_dim, 64)
        self.fc2 = nn.Linear(64, 64)
        self.policy_logits = nn.Linear(64, action_dim)  # unnormalized action logits
        self.value = nn.Linear(64, 1)

    def forward(self, state):
        # state can be 1D (obs_dim,) or 2D (batch, obs_dim)
        x = torch.relu(self.fc1(state))
        x = torch.relu(self.fc2(x))
        return self.policy_logits(x), self.value(x)

# Initialize environment and model
env = gym.make("CartPole-v1")
obs_space = env.observation_space
act_space = env.action_space

obs_dim = obs_space.shape[0]
act_dim = act_space.n

model = ActorCritic(obs_dim, act_dim)
optimizer = optim.Adam(model.parameters(), lr=3e-4)

# PPO hyperparameters
epochs = 50               # number of training iterations
steps_per_epoch = 1000    # timesteps per epoch (per update batch)
gamma = 0.99              # discount factor
lam = 0.95                # GAE lambda
clip_epsilon = 0.2        # PPO clip parameter
K_epochs = 4              # update epochs per batch
ent_coef = 0.01
vf_coef = 0.5

for epoch in range(epochs):
    # Storage buffers for this epoch
    observations, actions = [], []
    rewards, dones = [], []
    values, log_probs = [], []

    # Reset env (Gymnasium returns (obs, info))
    obs, _ = env.reset()

    for t in range(steps_per_epoch):
        obs_tensor = torch.tensor(obs, dtype=torch.float32)

        with torch.no_grad():
            logits, value = model(obs_tensor)
            dist = torch.distributions.Categorical(logits=logits)
            action = dist.sample()
            log_prob = dist.log_prob(action)

        # Step env (Gymnasium returns 5-tuple)
        next_obs, reward, terminated, truncated, _ = env.step(action.item())
        done = terminated or truncated

        # Store transition
        observations.append(obs_tensor)
        actions.append(action)
        rewards.append(float(reward))
        dones.append(done)
        values.append(float(value.item()))
        log_probs.append(float(log_prob.item()))

        obs = next_obs
        ifdone:
            obs, _ = env.reset()

    # Bootstrap last value for GAE (from final obs of the epoch)
    with torch.no_grad():
        last_v = model(torch.tensor(obs, dtype=torch.float32))[1].item()

    # Compute GAE advantages and returns
    advantages = []
    gae = 0.0
    # Append bootstrap to values (so values[t+1] is valid)
    values_plus = values + [last_v]
    for t in reversed(range(len(rewards))):
        nonterminal = 0.0 if dones[t] else 1.0
        delta = rewards[t] + gamma * values_plus[t + 1] * nonterminal - values_plus[t]
        gae = delta + gamma * lam * nonterminal * gae
        advantages.insert(0, gae)

    returns = [adv + v for adv, v in zip(advantages, values)]

    # Convert buffers to tensors
    obs_tensor = torch.stack(observations)  # (N, obs_dim)
    act_tensor = torch.tensor([a.item() for a in actions], dtype=torch.long)
    adv_tensor = torch.tensor(advantages, dtype=torch.float32)
    ret_tensor = torch.tensor(returns, dtype=torch.float32)
    old_log_probs = torch.tensor(log_probs, dtype=torch.float32)

    # Normalize advantages
    adv_tensor = (adv_tensor - adv_tensor.mean()) / (adv_tensor.std() + 1e-8)

    # PPO policy and value update
    for _ in range(K_epochs):
        logits, value_pred = model(obs_tensor)
        dist = torch.distributions.Categorical(logits=logits)
        new_log_probs = dist.log_prob(act_tensor)
        entropy = dist.entropy().mean()

        # Probability ratio r_t(theta)
        ratio = torch.exp(new_log_probs - old_log_probs)

        # Clipped objective
        surr1 = ratio * adv_tensor
        surr2 = torch.clamp(ratio, 1 - clip_epsilon, 1 + clip_epsilon) * adv_tensor
        policy_loss = -torch.min(surr1, surr2).mean()

        # Value loss
        value_loss = nn.functional.mse_loss(value_pred.squeeze(-1), ret_tensor)

        # Total loss
        loss = policy_loss + vf_coef * value_loss - ent_coef * entropy

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