# -*- coding: utf-8 -*- """ PPO 训练脚本 —— LunarLander（第 4 周 · 第一个智能体） ===================================================== 对应课件《第4周-强化学习进阶实践内容》第 2~14 页：用 PPO 训练一个月球着陆器智能体。 与第一周（Stable-Baselines3 现成库）不同，本脚本**从零手写 PPO 的全部核心组件**， 每个组件都能在代码里直接看到： 1. Actor-Critic 网络 —— get_action_and_value() 2. rollout 采样 —— 与环境交互收集经验 3. GAE 广义优势估计 —— compute_gae() 4. PPO-Clip 损失 —— 策略/价值/熵三项加权 5. 评估、存盘、视频录制、日志 快速开始 -------- 冒烟测试（约 20 秒）:: python ppo.py --total-timesteps 16384 --num-envs 8 --eval-interval 5 \ --eval-episodes 4 --out-dir runs --run-name smoke 正式训练（CPU 约 15~20 分钟，可达到官方 "solved" 标准 200 分）:: python ppo.py --total-timesteps 2500000 --num-envs 16 --num-steps 256 \ --gamma 0.999 --gae-lambda 0.98 --num-minibatches 8 --update-epochs 10 \ --eval-interval 15 --eval-episodes 16 --capture-video \ --out-dir runs --run-name main_2.5m 复现课件命令（5 万步，会失败——这正是课件想演示的"步数不够"）:: python ppo.py --total-timesteps 50000 --num-envs 4 --num-steps 128 \ --out-dir runs --run-name baseline_50k """ import argparse import csv import json import os import random import time from collections import deque import numpy as np import torch import torch.nn as nn import torch.optim as optim from torch.distributions.categorical import Categorical import gymnasium as gym # =========================================================================== # # 1. 基础工具 # =========================================================================== # def layer_init(layer, std=np.sqrt(2), bias_const=0.0): """正交初始化（PPO 的常规做法，训练更稳定）。""" nn.init.orthogonal_(layer.weight, std) nn.init.constant_(layer.bias, bias_const) return layer def make_env(env_id, seed, idx, render_mode=None): """创建一个环境；用 seed+idx 保证每个并行环境轨迹不同但可复现。""" return gym.make(env_id, render_mode=render_mode) # =========================================================================== # # 2. 向量化环境（手写，行为透明） # =========================================================================== # class VecEnv: """N 个环境并排前进（gymnasium API：reset -> (obs, info), step -> 5 元组）。""" def __init__(self, envs): self.envs = envs self.num_envs = len(envs) self.episode_returns = deque(maxlen=200) self.episode_lengths = deque(maxlen=200) self._cur_ret = np.zeros(self.num_envs, dtype=np.float64) self._cur_len = np.zeros(self.num_envs, dtype=np.int64) self.obs_dim = int(np.prod(envs[0].observation_space.shape)) self.action_dim = int(envs[0].action_space.n) def reset(self, seeds=None): obs = np.zeros((self.num_envs, self.obs_dim), dtype=np.float32) for i, env in enumerate(self.envs): seed = None if seeds is None else int(seeds[i]) o, _ = env.reset(seed=seed) obs[i] = np.asarray(o, dtype=np.float32).ravel() self._cur_ret[i] = 0.0 self._cur_len[i] = 0 return obs def step(self, actions): """同时推进一步；回合结束后自动 reset（不重新播种，保留随机性）。""" obs = np.zeros((self.num_envs, self.obs_dim), dtype=np.float32) rewards = np.zeros(self.num_envs, dtype=np.float32) dones = np.zeros(self.num_envs, dtype=np.float32) infos = [] for i, env in enumerate(self.envs): o, r, terminated, truncated, info = env.step(int(actions[i])) done = bool(terminated or truncated) self._cur_ret[i] += r self._cur_len[i] += 1 info = dict(info) if isinstance(info, dict) else {} if done: info["episode_return"] = float(self._cur_ret[i]) info["episode_length"] = int(self._cur_len[i]) self.episode_returns.append(info["episode_return"]) self.episode_lengths.append(info["episode_length"]) self._cur_ret[i] = 0.0 self._cur_len[i] = 0 o, _ = env.reset() obs[i] = np.asarray(o, dtype=np.float32).ravel() rewards[i] = r dones[i] = float(done) infos.append(info) return obs, rewards, dones, infos def close(self): for env in self.envs: try: env.close() except Exception: pass # =========================================================================== # # 3. Video：用 imageio + imageio-ffmpeg 直接写 mp4（不依赖 moviepy） # =========================================================================== # class VideoRecorder: def __init__(self, path, fps=30): self.path = path self.fps = fps self.frames = [] def add(self, frame): if frame is not None: self.frames.append(np.asarray(frame, dtype=np.uint8)) def save(self): if not self.frames: return None import imageio os.makedirs(os.path.dirname(os.path.abspath(self.path)), exist_ok=True) writer = imageio.get_writer(self.path, fps=self.fps, codec="libx264", quality=8, macro_block_size=1) for frame in self.frames: writer.append_data(frame) writer.close() return self.path # =========================================================================== # # 4. Actor-Critic 网络（课件第 12 页 get_action_and_value 的完整实现） # =========================================================================== # def build_agent(obs_dim, action_dim, hidden=64): class Agent(nn.Module): def __init__(self): super().__init__() # Critic：估计状态价值 V(s) 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), ) # Actor：输出每个动作的 logits 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, action_dim), std=0.01), ) def get_value(self, x): return self.critic(x) def get_action_and_value(self, x, action=None): """给定状态 -> (动作, 对数概率, 熵, 状态价值)。""" logits = self.actor(x) probs = Categorical(logits=logits) if action is None: action = probs.sample() # 训练采样 value = self.critic(x).view(-1) return action, probs.log_prob(action), probs.entropy(), value return Agent() # =========================================================================== # # 5. GAE（课件第 13 页公式的完整实现） # =========================================================================== # def compute_gae(rewards, values, dones, next_done, next_value, gamma, gae_lambda, num_steps): advantages = torch.zeros_like(rewards) lastgaelam = 0 for t in reversed(range(num_steps)): if t == num_steps - 1: nextnonterminal = 1.0 - next_done nextvalues = next_value else: nextnonterminal = 1.0 - dones[t + 1] nextvalues = values[t + 1] # TD 残差 delta = rewards[t] + gamma * nextvalues * nextnonterminal - values[t] # 指数加权平均 advantages[t] = lastgaelam = delta + gamma * gae_lambda * nextnonterminal * lastgaelam returns = advantages + values return advantages, returns # =========================================================================== # # 6. 评估与录像 # =========================================================================== # @torch.no_grad() def evaluate(agent, env_id, device, episodes=10, deterministic=False, seed=12345): """跑若干完整回合，返回 (回报列表, 回合长度列表)。""" env = gym.make(env_id) returns, lengths = [], [] for ep in range(episodes): obs, _ = env.reset(seed=seed + ep) done, ret, ln = False, 0.0, 0 while not done: obs_t = torch.as_tensor(obs, dtype=torch.float32, device=device).unsqueeze(0) if deterministic: action = int(torch.argmax(agent.actor(obs_t), dim=-1).item()) else: action = int(agent.get_action_and_value(obs_t)[0].item()) obs, r, terminated, truncated, _ = env.step(action) ret += float(r) ln += 1 done = bool(terminated or truncated) returns.append(ret) lengths.append(ln) env.close() return returns, lengths @torch.no_grad() def record_episode(agent, env_id, device, path, seed=999, deterministic=True, fps=30): """录制一个完整回合为 mp4，返回 {episode_return, episode_length, frames, video}。""" env = gym.make(env_id, render_mode="rgb_array") obs, _ = env.reset(seed=seed) rec = VideoRecorder(path, fps=fps) done, ret, ln = False, 0.0, 0 while not done and ln < 4000: frame = env.render() if frame is not None: rec.add(frame) obs_t = torch.as_tensor(obs, dtype=torch.float32, device=device).unsqueeze(0) if deterministic: action = int(torch.argmax(agent.actor(obs_t), dim=-1).item()) else: action = int(agent.get_action_and_value(obs_t)[0].item()) obs, r, terminated, truncated, _ = env.step(action) ret += float(r) ln += 1 done = bool(terminated or truncated) env.close() rec.save() return {"episode_return": ret, "episode_length": ln, "frames": len(rec.frames), "video": path} def summarize(returns): arr = np.asarray(returns, dtype=np.float64) return { "episodes": int(arr.size), "mean": float(arr.mean()), "std": float(arr.std()), "min": float(arr.min()), "max": float(arr.max()), "median": float(np.median(arr)), "landing_rate": float(np.mean(arr >= 200.0)), # 官方 solved 标准：平均 >= 200 "p25": float(np.percentile(arr, 25)), "p75": float(np.percentile(arr, 75)), } # =========================================================================== # # 7. 参数 # =========================================================================== # def parse_args(): p = argparse.ArgumentParser(description="从零手写 PPO 训练 LunarLander") # ---- 实验管理 ---- p.add_argument("--env-id", type=str, default="LunarLander-v3", help="环境 ID（gymnasium 中 LunarLander 只有 v3）") p.add_argument("--seed", type=int, default=1, help="随机种子（默认 1）") p.add_argument("--total-timesteps", type=int, default=1_000_000, help="总训练步数") p.add_argument("--out-dir", type=str, default="runs", help="结果根目录") p.add_argument("--run-name", type=str, default="", help="本次实验名（默认自动生成）") p.add_argument("--no-upload", action="store_true", default=False, help="不上传 HuggingFace（上传用 upload_to_hf.py）") p.add_argument("--repo-id", type=str, default=None, help="HuggingFace 仓库 ID（可选）") # ---- 环境 ---- p.add_argument("--num-envs", type=int, default=8, help="并行环境数（默认 8）") p.add_argument("--num-steps", type=int, default=128, help="每个环境每次采样多少步") p.add_argument("--capture-video", action="store_true", help="录制训练过程中的学习视频") # ---- PPO 超参数 ---- p.add_argument("--learning-rate", type=float, default=2.5e-4, help="学习率（默认 2.5e-4）") p.add_argument("--anneal-lr", action="store_true", default=True, help="学习率线性衰减") p.add_argument("--no-anneal-lr", dest="anneal_lr", action="store_false") p.add_argument("--gamma", type=float, default=0.99, help="折扣因子") p.add_argument("--gae-lambda", type=float, default=0.95, help="GAE lambda") p.add_argument("--num-minibatches", type=int, default=4, help="mini-batch 数量") p.add_argument("--update-epochs", type=int, default=4, help="每个 rollout 更新几次") p.add_argument("--norm-adv", action="store_true", default=True, help="优势标准化") p.add_argument("--no-norm-adv", dest="norm_adv", action="store_false") p.add_argument("--clip-coef", type=float, default=0.2, help="PPO 裁剪系数") p.add_argument("--clip-vloss", action="store_true", default=True, help="价值函数裁剪") p.add_argument("--no-clip-vloss", dest="clip_vloss", action="store_false") p.add_argument("--ent-coef", type=float, default=0.01, help="熵奖励系数") p.add_argument("--vf-coef", type=float, default=0.5, help="价值损失系数") p.add_argument("--max-grad-norm", type=float, default=0.5, help="梯度裁剪范数") p.add_argument("--hidden", type=int, default=64, help="隐藏层宽度") # ---- 评估与日志 ---- p.add_argument("--eval-interval", type=int, default=20, help="每多少次 update 评估一次") p.add_argument("--eval-episodes", type=int, default=10, help="每次评估跑多少回合") p.add_argument("--eval-greedy", action="store_true", default=True, help="评估用 argmax（部署策略）；训练采样仍用随机策略") p.add_argument("--eval-stochastic", dest="eval_greedy", action="store_false") p.add_argument("--save-interval", type=int, default=50, help="每多少次 update 存一次盘") p.add_argument("--log-interval", type=int, default=10, help="控制台打印间隔") p.add_argument("--torch-threads", type=int, default=8, help="torch 线程数") return p.parse_args() # =========================================================================== # # 8. 主训练循环 # =========================================================================== # def main(): args = parse_args() torch.set_num_threads(args.torch_threads) random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") run_name = args.run_name or f"PPO-{args.env_id}-{int(time.time())}" run_dir = os.path.join(args.out_dir, run_name) os.makedirs(run_dir, exist_ok=True) with open(os.path.join(run_dir, "config.json"), "w", encoding="utf-8") as f: json.dump(vars(args), f, indent=2, ensure_ascii=False) # ---- 环境与网络 ---- envs = [make_env(args.env_id, args.seed, i) for i in range(args.num_envs)] vec = VecEnv(envs) obs_dim, action_dim = vec.obs_dim, vec.action_dim agent = build_agent(obs_dim, action_dim, args.hidden).to(device) optimizer = optim.Adam(agent.parameters(), lr=args.learning_rate, eps=1e-5) batch_size = args.num_envs * args.num_steps minibatch_size = batch_size // args.num_minibatches num_iterations = max(1, args.total_timesteps // batch_size) print(f"[配置] 环境={args.env_id} 设备={device} 并行环境={args.num_envs} " f"每轮步数={batch_size} 迭代={num_iterations} 总步数={args.total_timesteps}") print(f"[配置] obs_dim={obs_dim} action_dim={action_dim} minibatch={minibatch_size}") # ---- 缓冲区 ---- obs_buf = torch.zeros((args.num_steps, args.num_envs, obs_dim), dtype=torch.float32, device=device) act_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) logp_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) rew_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) done_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) val_buf = torch.zeros((args.num_steps, args.num_envs), dtype=torch.float32, device=device) next_obs = torch.as_tensor(vec.reset(seeds=[args.seed + i for i in range(args.num_envs)]), dtype=torch.float32, device=device) next_done = torch.zeros(args.num_envs, dtype=torch.float32, device=device) # ---- 日志文件 ---- train_csv = open(os.path.join(run_dir, "train_log.csv"), "w", newline="", encoding="utf-8") train_writer = csv.writer(train_csv) train_writer.writerow(["iteration", "global_step", "sps", "episodes_in_window", "mean_return", "std_return", "min_return", "max_return", "mean_ep_length", "policy_loss", "value_loss", "entropy", "approx_kl", "clipfrac", "learning_rate"]) eval_csv = open(os.path.join(run_dir, "eval_log.csv"), "w", newline="", encoding="utf-8") eval_writer = csv.writer(eval_csv) eval_writer.writerow(["iteration", "global_step", "eval_mean", "eval_std", "eval_min", "eval_max", "landing_rate"]) video_points = set() if args.capture_video: video_points = {max(1, int(num_iterations * p)) for p in (0.1, 0.25, 0.4, 0.55, 0.7, 0.85, 1.0)} global_step = 0 t0 = time.time() best_eval = -float("inf") recent_returns = deque(maxlen=100) for iteration in range(1, num_iterations + 1): # ---------- 学习率衰减 ---------- if args.anneal_lr: frac = 1.0 - (iteration - 1.0) / num_iterations lr_now = frac * args.learning_rate optimizer.param_groups[0]["lr"] = lr_now else: lr_now = args.learning_rate # ---------- 采样 ---------- step_returns = [] for t in range(args.num_steps): global_step += args.num_envs obs_buf[t] = next_obs done_buf[t] = next_done with torch.no_grad(): action, logprob, _, value = agent.get_action_and_value(next_obs) val_buf[t] = value.flatten() act_buf[t] = action.float() logp_buf[t] = logprob next_obs_np, reward, done, infos = vec.step(action.cpu().numpy()) rew_buf[t] = torch.as_tensor(reward, dtype=torch.float32, device=device) next_obs = torch.as_tensor(next_obs_np, dtype=torch.float32, device=device) next_done = torch.as_tensor(done, dtype=torch.float32, device=device) for info in infos: if "episode_return" in info: step_returns.append(info["episode_return"]) recent_returns.append(info["episode_return"]) # ---------- GAE ---------- with torch.no_grad(): next_value = agent.get_value(next_obs).reshape(1, -1) advantages, returns = compute_gae( rew_buf, val_buf, done_buf, next_done, next_value, args.gamma, args.gae_lambda, args.num_steps, ) b_obs = obs_buf.reshape((-1, obs_dim)) b_logprobs = logp_buf.reshape(-1) b_actions = act_buf.reshape(-1) b_advantages = advantages.reshape(-1) b_returns = returns.reshape(-1) b_values = val_buf.reshape(-1) # ---------- PPO 更新 ---------- b_inds = np.arange(batch_size) clipfracs = [] pg_loss = torch.tensor(0.0) v_loss = torch.tensor(0.0) entropy = torch.tensor(0.0) approx_kl = torch.tensor(0.0) for _ in range(args.update_epochs): np.random.shuffle(b_inds) for start in range(0, batch_size, minibatch_size): mb = b_inds[start:start + minibatch_size] _, newlogprob, newentropy, newvalue = agent.get_action_and_value( b_obs[mb], b_actions.long()[mb]) logratio = newlogprob - b_logprobs[mb] ratio = logratio.exp() with torch.no_grad(): approx_kl = ((ratio - 1) - logratio).mean() clipfracs.append(((ratio - 1.0).abs() > args.clip_coef).float().mean().item()) mb_adv = b_advantages[mb] if args.norm_adv: mb_adv = (mb_adv - mb_adv.mean()) / (mb_adv.std() + 1e-8) # ① 策略损失（PPO-Clip） pg1 = -mb_adv * ratio pg2 = -mb_adv * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef) pg_loss = torch.max(pg1, pg2).mean() # ② 价值损失（可选裁剪） newvalue = newvalue.view(-1) if args.clip_vloss: v_unclipped = (newvalue - b_returns[mb]) ** 2 v_clipped = b_values[mb] + torch.clamp( newvalue - b_values[mb], -args.clip_coef, args.clip_coef) v_loss = 0.5 * torch.max(v_unclipped, (v_clipped - b_returns[mb]) ** 2).mean() else: v_loss = 0.5 * ((newvalue - b_returns[mb]) ** 2).mean() # ③ 熵奖励 entropy = newentropy.mean() # ④ 总损失 loss = pg_loss - args.ent_coef * entropy + args.vf_coef * v_loss optimizer.zero_grad() loss.backward() nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm) optimizer.step() # ---------- 日志 ---------- sps = int(global_step / max(1e-9, time.time() - t0)) mean_ret = float(np.mean(step_returns)) if step_returns else float("nan") std_ret = float(np.std(step_returns)) if step_returns else float("nan") min_ret = float(np.min(step_returns)) if step_returns else float("nan") max_ret = float(np.max(step_returns)) if step_returns else float("nan") mean_len = float(np.mean(vec.episode_lengths)) if len(vec.episode_lengths) else float("nan") train_writer.writerow([ iteration, global_step, sps, len(step_returns), f"{mean_ret:.3f}", f"{std_ret:.3f}", f"{min_ret:.3f}", f"{max_ret:.3f}", f"{mean_len:.1f}", f"{float(pg_loss.detach()):.5f}", f"{float(v_loss.detach()):.5f}", f"{float(entropy.detach()):.5f}", f"{float(approx_kl.detach()):.5f}", f"{float(np.mean(clipfracs)):.4f}", f"{lr_now:.6g}", ]) train_csv.flush() # ---------- 定期评估 ---------- eval_stats = None if iteration % args.eval_interval == 0 or iteration == num_iterations: rets, _ = evaluate(agent, args.env_id, device, episodes=args.eval_episodes, deterministic=args.eval_greedy, seed=int(args.seed * 1000 + iteration)) eval_stats = summarize(rets) eval_writer.writerow([ iteration, global_step, f"{eval_stats['mean']:.2f}", f"{eval_stats['std']:.2f}", f"{eval_stats['min']:.2f}", f"{eval_stats['max']:.2f}", f"{eval_stats['landing_rate']:.3f}", ]) eval_csv.flush() if eval_stats["mean"] > best_eval: best_eval = eval_stats["mean"] torch.save({ "state_dict": agent.state_dict(), "obs_dim": obs_dim, "action_dim": action_dim, "hidden": args.hidden, "iteration": iteration, "global_step": global_step, "eval_mean": eval_stats["mean"], "config": vars(args), "eval_stats": eval_stats, }, os.path.join(run_dir, "model_best.pth")) if iteration % args.log_interval == 0 or iteration == 1 or iteration == num_iterations: ev = "" if eval_stats is None else f" | eval={eval_stats['mean']:.1f}" print(f"[{iteration:4d}/{num_iterations}] step={global_step:8d} " f"train_ret={mean_ret:8.2f} ep={len(step_returns):3d} " f"len={mean_len:6.1f} pg={float(pg_loss.detach()):7.4f} " f"vl={float(v_loss.detach()):7.4f} " f"ent={float(entropy.detach()):6.3f} sps={sps:5d}{ev}", flush=True) # ---------- 存盘 ---------- if iteration % args.save_interval == 0 or iteration == num_iterations: torch.save({ "state_dict": agent.state_dict(), "obs_dim": obs_dim, "action_dim": action_dim, "hidden": args.hidden, "iteration": iteration, "global_step": global_step, "config": vars(args), }, os.path.join(run_dir, "model_ckpt.pth")) # ---------- 学习过程视频 ---------- if args.capture_video and iteration in video_points: path = os.path.join(run_dir, "videos", f"learning_{global_step:08d}.mp4") try: info = record_episode(agent, args.env_id, device, path, seed=args.seed * 7 + iteration, deterministic=True) print(f" [视频] step={global_step} 回报={info['episode_return']:.1f} -> {path}", flush=True) except Exception as e: print(f" [视频] 失败: {e}", flush=True) # ---------- 收尾 ---------- torch.save({ "state_dict": agent.state_dict(), "obs_dim": obs_dim, "action_dim": action_dim, "hidden": args.hidden, "iteration": num_iterations, "global_step": global_step, "config": vars(args), }, os.path.join(run_dir, "model_final.pth")) final_rets, _ = evaluate(agent, args.env_id, device, episodes=min(50, max(10, args.eval_episodes)), deterministic=True, seed=args.seed * 999 + 7) summary = { "run_name": run_name, "run_dir": run_dir, "env_id": args.env_id, "seed": args.seed, "total_timesteps": args.total_timesteps, "num_envs": args.num_envs, "num_steps": args.num_steps, "gamma": args.gamma, "gae_lambda": args.gae_lambda, "num_minibatches": args.num_minibatches, "update_epochs": args.update_epochs, "learning_rate": args.learning_rate, "device": str(device), "train_seconds": round(time.time() - t0, 1), "final_eval": summarize(final_rets), "best_eval_mean": best_eval, "train_mean_return_last100": float(np.mean(recent_returns)) if recent_returns else None, } with open(os.path.join(run_dir, "summary.json"), "w", encoding="utf-8") as f: json.dump(summary, f, indent=2, ensure_ascii=False) train_csv.close() eval_csv.close() vec.close() print(f"\n[完成] 用时 {summary['train_seconds']/60:.1f} 分钟 | " f"最终评估({summary['final_eval']['episodes']}局,argmax) = " f"{summary['final_eval']['mean']:.2f} ± {summary['final_eval']['std']:.2f} | " f"着陆率 {summary['final_eval']['landing_rate']*100:.0f}%") print(f"[完成] 结果目录: {run_dir}") if __name__ == "__main__": main()