""" ============================================================================= 第 7 讲 · 代码 4：用 PyTorch 实现 PPO（裁剪目标 + GAE） ============================================================================= 对应讲义：lectures/lecture07_actor_critic_ppo.md 第 7、8 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- PPO（Proximal Policy Optimization，Schulman et al., 2017）是目前最常用的策略梯度算法， 也是 RLHF（第 9 讲）的核心。它把本讲的所有思想组合在一起： 1. Actor-Critic：策略网络 + 价值网络（01 文件） 2. GAE：用 λ 平衡优势估计的偏差与方差（02 文件） 3. 替代目标 + 裁剪：用重要性比 r = π_new/π_old 复用同一批数据做**多轮**更新， 同时用裁剪防止策略一次走得太远（03 文件） 与第 6 讲的 REINFORCE 对比：REINFORCE 每批数据只做一次梯度更新就扔掉； PPO 每批数据做 K 轮（epochs）小批量更新，样本效率高得多，而裁剪保证了这样做仍然安全。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- 需要：pip install torch gymnasium python 04_ppo_cartpole.py 标准 PPO（约 40 秒） python 04_ppo_cartpole.py --no-clip 去掉裁剪（同一批数据多轮更新，没有保护） python 04_ppo_cartpole.py --epochs 1 每批数据只更新 1 轮 python 04_ppo_cartpole.py --gae-lambda 1.0 λ = 1（MC 优势） python 04_ppo_cartpole.py --env Acrobot-v1 换一个更难的环境（奖励为负，越接近 0 越好） """ import argparse import time import numpy as np import torch import torch.nn as nn import gymnasium as gym CONFIG = { "env_id": "CartPole-v1", "total_steps": 100_000, "n_steps": 2048, # 每次更新前收集多少步数据（一个"批次"） "epochs": 10, # 每批数据重复训练几轮 "minibatch_size": 64, "gamma": 0.99, "gae_lambda": 0.95, "clip_eps": 0.2, "lr": 3e-4, "vf_coef": 0.5, # 价值损失的权重 "ent_coef": 0.0, # 熵奖励的权重（鼓励探索；CartPole 不需要） "max_grad_norm": 0.5, "hidden": 64, } # ============================================================================= # 第 1 部分：网络 # ============================================================================= def layer_init(layer, std=np.sqrt(2), bias=0.0): """ 正交初始化（orthogonal initialization）：PPO 实现中的一个常见细节。 【为什么要关心初始化？】 策略网络最后一层用很小的标准差（0.01）初始化，让初始策略接近均匀随机—— 一开始"什么都不知道"，就应该充分探索。价值网络最后一层用标准差 1。 这类细节单独看都不起眼，但合在一起会显著影响深度 RL 的稳定性 （Huang et al., "The 37 Implementation Details of PPO"）。 """ nn.init.orthogonal_(layer.weight, std) nn.init.constant_(layer.bias, bias) return layer class ActorCritic(nn.Module): """策略网络（Actor）和价值网络（Critic），两者不共享参数。""" def __init__(self, obs_dim, n_actions, hidden): super().__init__() self.actor = nn.Sequential( layer_init(nn.Linear(obs_dim, hidden)), nn.Tanh(), layer_init(nn.Linear(hidden, hidden)), nn.Tanh(), layer_init(nn.Linear(hidden, n_actions), std=0.01), ) self.critic = nn.Sequential( layer_init(nn.Linear(obs_dim, hidden)), nn.Tanh(), layer_init(nn.Linear(hidden, hidden)), nn.Tanh(), layer_init(nn.Linear(hidden, 1), std=1.0), ) def value(self, obs): return self.critic(obs).squeeze(-1) def dist(self, obs): return torch.distributions.Categorical(logits=self.actor(obs)) # ============================================================================= # 第 2 部分：收集一批数据 # ============================================================================= def collect_rollout(env, model, obs, n_steps, episode_log): """ 用当前策略与环境交互 n_steps 步（可以跨越多个回合）。 【存储的内容】 obs、actions、logp_old（采样时的 log π_old(a|s)，计算重要性比要用）、rewards、values（V(s_t)）、 next_values（V(s_{t+1})，真正终止时为 0）、dones（回合是否在这一步结束，用于 GAE 的截断） 【终止 vs 截断（第 5 讲 7.6 节）】 - 真正终止（杆子倒了）：没有未来，next_value = 0 - 时间截断（到了 500 步）：还有未来，next_value = V(最后的观测)——继续自举 两种情况下，GAE 的递推都要在这里"断开"（不能把下一个回合的 TD 误差加进来）。 【返回】(一个包含上述张量的字典, 下一步的观测) """ buf = {k: [] for k in ["obs", "actions", "logp", "rewards", "values", "next_values", "dones"]} for _ in range(n_steps): obs_t = torch.as_tensor(obs, dtype=torch.float32) with torch.no_grad(): dist = model.dist(obs_t) action = dist.sample() logp = dist.log_prob(action) value = model.value(obs_t) next_obs, reward, terminated, truncated, _ = env.step(action.item()) episode_log["return"] += reward if terminated: next_value = 0.0 else: with torch.no_grad(): next_value = model.value(torch.as_tensor(next_obs, dtype=torch.float32)).item() for k, v in zip(buf, [obs, action.item(), logp.item(), reward, value.item(), next_value, float(terminated or truncated)]): buf[k].append(v) if terminated or truncated: episode_log["returns"].append(episode_log["return"]) episode_log["return"] = 0.0 next_obs, _ = env.reset() obs = next_obs data = {k: torch.as_tensor(np.array(v), dtype=torch.float32) for k, v in buf.items()} data["actions"] = data["actions"].long() return data, obs def compute_gae(data, gamma, lam): """ GAE（02 文件）：从后往前递推 Â_t = δ_t + γλ(1 - done_t) Â_{t+1}。 δ_t = r_t + γ · next_value_t - V(s_t)。 回报目标（用来训练价值网络）：R_t = Â_t + V(s_t)，也就是 λ-回报。 """ T = len(data["rewards"]) adv = torch.zeros(T) running = 0.0 for t in reversed(range(T)): delta = data["rewards"][t] + gamma * data["next_values"][t] - data["values"][t] running = delta + gamma * lam * (1.0 - data["dones"][t]) * running adv[t] = running data["advantages"] = adv data["returns"] = adv + data["values"] # ============================================================================= # 第 3 部分：PPO 更新 # ============================================================================= def ppo_update(model, optimizer, data, cfg, use_clip=True): """ 对一批数据做 epochs 轮、每轮若干个小批量的更新。 【损失函数】（讲义 7.2 节） r_t(θ) = exp(log π_θ(a_t|s_t) - log π_old(a_t|s_t)) （重要性比） L_clip = -mean( min( r_t Â_t, clip(r_t, 1-ε, 1+ε) Â_t ) ) （裁剪的替代目标，取负号变成最小化） L_value = mean( (V_θ(s_t) - R_t)² ) L_entropy= -mean( H[π_θ(·|s_t)] ) 总损失 = L_clip + c_v L_value + c_e L_entropy use_clip=False 时：L = -mean(r_t Â_t)，即未裁剪的替代目标——同样的数据多轮更新，没有任何保护。 【监控指标】 approx_kl ：新旧策略的近似 KL，≈ mean(log π_old - log π_new)。PPO 希望它保持很小（~0.01）。 clip_frac ：概率比被裁剪的样本比例。 【返回】(平均 approx_kl, 平均 clip_frac) """ T = len(data["rewards"]) kls, clipfracs = [], [] for _ in range(cfg["epochs"]): perm = torch.randperm(T) # 每轮打乱顺序 for start in range(0, T, cfg["minibatch_size"]): idx = perm[start:start + cfg["minibatch_size"]] dist = model.dist(data["obs"][idx]) logp_new = dist.log_prob(data["actions"][idx]) ratio = torch.exp(logp_new - data["logp"][idx]) adv = data["advantages"][idx] adv = (adv - adv.mean()) / (adv.std() + 1e-8) # 小批量内标准化优势 if use_clip: surr = torch.min(ratio * adv, torch.clamp(ratio, 1 - cfg["clip_eps"], 1 + cfg["clip_eps"]) * adv) else: surr = ratio * adv loss_pi = -surr.mean() loss_v = ((model.value(data["obs"][idx]) - data["returns"][idx]) ** 2).mean() loss_ent = -dist.entropy().mean() loss = loss_pi + cfg["vf_coef"] * loss_v + cfg["ent_coef"] * loss_ent optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(model.parameters(), cfg["max_grad_norm"]) optimizer.step() with torch.no_grad(): kls.append((data["logp"][idx] - logp_new).mean().item()) clipfracs.append(((ratio - 1).abs() > cfg["clip_eps"]).float().mean().item()) return float(np.mean(kls)), float(np.mean(clipfracs)) def evaluate(model, env_id, n_episodes=10, seed=12345): """用最可能的动作评估策略。""" env = gym.make(env_id) totals = [] for i in range(n_episodes): obs, _ = env.reset(seed=seed + i) done, total = False, 0.0 while not done: with torch.no_grad(): a = int(model.actor(torch.as_tensor(obs, dtype=torch.float32)).argmax().item()) obs, r, terminated, truncated, _ = env.step(a) total += r done = terminated or truncated totals.append(total) env.close() return float(np.mean(totals)) def train(cfg, seed=0, use_clip=True, verbose=True): """PPO 主循环。返回 [(已交互步数, 最近 10 个回合的平均回报, 评估回报, approx_kl, clip_frac), ...]。""" torch.manual_seed(seed) env = gym.make(cfg["env_id"]) obs, _ = env.reset(seed=seed) env.action_space.seed(seed) model = ActorCritic(env.observation_space.shape[0], env.action_space.n, cfg["hidden"]) optimizer = torch.optim.Adam(model.parameters(), lr=cfg["lr"], eps=1e-5) episode_log = {"return": 0.0, "returns": []} log, start = [], time.time() n_updates = cfg["total_steps"] // cfg["n_steps"] if verbose: print(f" {'步数':>7s} {'最近10回合回报':>12s} {'评估回报':>8s} {'approx_KL':>10s} {'裁剪比例':>8s} {'用时':>6s}") for u in range(1, n_updates + 1): data, obs = collect_rollout(env, model, obs, cfg["n_steps"], episode_log) compute_gae(data, cfg["gamma"], cfg["gae_lambda"]) kl, clipfrac = ppo_update(model, optimizer, data, cfg, use_clip) recent = float(np.mean(episode_log["returns"][-10:])) if episode_log["returns"] else float("nan") eval_ret = evaluate(model, cfg["env_id"]) if (u % 5 == 0 or u == n_updates) else None log.append((u * cfg["n_steps"], recent, eval_ret, kl, clipfrac)) if verbose and eval_ret is not None: print(f" {u * cfg['n_steps']:7d} {recent:12.1f} {eval_ret:8.1f} {kl:10.4f} {clipfrac:8.3f} " f"{time.time() - start:5.0f}s") env.close() return log def main(): parser = argparse.ArgumentParser(description="PPO") parser.add_argument("--no-clip", action="store_true", help="去掉裁剪（未裁剪的替代目标）") parser.add_argument("--epochs", type=int, default=CONFIG["epochs"]) parser.add_argument("--gae-lambda", type=float, default=CONFIG["gae_lambda"]) parser.add_argument("--env", type=str, default=CONFIG["env_id"]) parser.add_argument("--steps", type=int, default=CONFIG["total_steps"]) parser.add_argument("--seed", type=int, default=0) args = parser.parse_args() torch.set_num_threads(1) cfg = dict(CONFIG, epochs=args.epochs, gae_lambda=args.gae_lambda, env_id=args.env, total_steps=args.steps) print("=" * 78) print(f"PPO（{'无裁剪' if args.no_clip else '裁剪 ε=0.2'}，{cfg['epochs']} 轮，λ={cfg['gae_lambda']}），" f"{cfg['env_id']}，种子 {args.seed}") print("=" * 78) log = train(cfg, seed=args.seed, use_clip=not args.no_clip) evals = [e for _, _, e, _, _ in log if e is not None] print(f"\n 最后一次评估：{evals[-1]:.1f}；最近 10 个训练回合的平均回报：{log[-1][1]:.1f}") print(f" 整个训练过程中 approx_KL 的最大值：{max(k for _, _, _, k, _ in log):.4f}") if __name__ == "__main__": main()