# -*- coding: utf-8 -*- """ 评估脚本 —— 加载训练好的模型，输出指标 + 演示视频 + 关键帧 + GIF ================================================================= 用法:: # 100 局严格评估（argmax 部署策略），并录 3 个回合的视频 python evaluate.py --model runs/main_2.5m/model_best.pth --episodes 100 \ --video-episodes 3 --out artifacts 产物： artifacts/eval_results.json 所有逐局回报 + 统计指标（含随机策略对照） artifacts/videos/replay_ep*.mp4 artifacts/keyframes.png 关键帧拼图 artifacts/replay.gif 模型卡里可直接显示的动图 """ import argparse import json import os import numpy as np import torch import gymnasium as gym from PIL import Image, ImageDraw, ImageFont from ppo import build_agent, summarize, VideoRecorder # --------------------------------------------------------------------------- # # 模型加载 # --------------------------------------------------------------------------- # def load_agent(path, device): ckpt = torch.load(path, map_location=device, weights_only=False) agent = build_agent(ckpt["obs_dim"], ckpt["action_dim"], ckpt.get("hidden", 64)).to(device) agent.load_state_dict(ckpt["state_dict"]) agent.eval() return agent, ckpt # --------------------------------------------------------------------------- # # 跑回合（agent=None 时是随机策略基线） # --------------------------------------------------------------------------- # @torch.no_grad() def run_episodes(agent, env_id, device, episodes, deterministic=True, seed=1000, collect_frames=0): env = gym.make(env_id, render_mode="rgb_array" if collect_frames else None) returns, lengths, frames_episodes = [], [], [] for ep in range(episodes): obs, _ = env.reset(seed=seed + ep) done, ret, ln = False, 0.0, 0 frames = [] while not done: if collect_frames and ep < collect_frames: f = env.render() if f is not None: frames.append(np.asarray(f, dtype=np.uint8)) t = torch.as_tensor(obs, dtype=torch.float32, device=device).unsqueeze(0) if agent is None: action = int(env.action_space.sample()) elif deterministic: action = int(torch.argmax(agent.actor(t), dim=-1).item()) else: action = int(agent.get_action_and_value(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) if collect_frames and ep < collect_frames: frames_episodes.append(frames) env.close() return returns, lengths, frames_episodes # --------------------------------------------------------------------------- # # 可视化小工具 # --------------------------------------------------------------------------- # def _font(size): for name in ("arial.ttf", "segoeui.ttf", "msyh.ttc"): try: return ImageFont.truetype(name, size) except Exception: continue return ImageFont.load_default() def make_keyframes(frames_list, path, title="LunarLander PPO — evaluation keyframes"): """把若干个回合均匀抽帧，拼成 2×3 关键帧图。""" if not frames_list: return None # 每个回合取中间均匀 3 帧，最多 6 帧 picks = [] for frames in frames_list: if not frames: continue idxs = np.linspace(0, len(frames) - 1, 3).astype(int) picks.extend([frames[i] for i in idxs]) picks = picks[:6] if not picks: return None cell_w = 400 h0, w0 = picks[0].shape[:2] cell_h = int(h0 * cell_w / w0) cols, rows = 3, 2 pad, top = 12, 56 W = cols * cell_w + (cols + 1) * pad H = top + rows * (cell_h + 28) + pad canvas = Image.new("RGB", (W, H), (24, 26, 32)) draw = ImageDraw.Draw(canvas) draw.text((pad, 16), title, font=_font(22), fill=(240, 240, 240)) for i, f in enumerate(picks): img = Image.fromarray(f).resize((cell_w, cell_h), Image.LANCZOS) r, c = divmod(i, cols) x = pad + c * (cell_w + pad) y = top + r * (cell_h + 28) canvas.paste(img, (x, y)) draw.text((x + 4, y + cell_h + 4), f"frame {i + 1}", font=_font(14), fill=(180, 180, 180)) canvas.save(path) return path def make_gif(frames, path, scale=0.5, stride=2, duration=40): """把回合帧序列压成 gif（便于在模型卡里直接显示）。""" if not frames: return None small = [Image.fromarray(f).resize( (int(f.shape[1] * scale), int(f.shape[0] * scale)), Image.LANCZOS) for f in frames[::stride]] if not small: return None small[0].save(path, save_all=True, append_images=small[1:], duration=duration, loop=0, optimize=True) return path # --------------------------------------------------------------------------- # # 主流程 # --------------------------------------------------------------------------- # def main(): p = argparse.ArgumentParser(description="评估 LunarLander PPO 智能体") p.add_argument("--model", type=str, default="runs/main_2.5m/model_best.pth") p.add_argument("--env-id", type=str, default="LunarLander-v3") p.add_argument("--episodes", type=int, default=100, help="评估回合数（官方标准 100 局）") p.add_argument("--seed", type=int, default=2024) p.add_argument("--stochastic", action="store_true", help="用随机策略而不是 argmax") p.add_argument("--video-episodes", type=int, default=3, help="录几个回合的视频") p.add_argument("--out", type=str, default="artifacts", help="输出目录") p.add_argument("--baseline-random", action="store_true", default=True, help="同时评估随机策略作为对照") p.add_argument("--no-baseline", dest="baseline_random", action="store_false") p.add_argument("--torch-threads", type=int, default=8) args = p.parse_args() torch.set_num_threads(args.torch_threads) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") os.makedirs(args.out, exist_ok=True) os.makedirs(os.path.join(args.out, "videos"), exist_ok=True) agent, ckpt = load_agent(args.model, device) print(f"[评估] 模型: {args.model}") print(f"[评估] 训练信息: iteration={ckpt.get('iteration')} " f"step={ckpt.get('global_step')} eval_mean={ckpt.get('eval_mean')}") print(f"[评估] 环境={args.env_id} 回合数={args.episodes} " f"策略={'随机采样' if args.stochastic else 'argmax(确定性)'}") # ---- 主评估 ---- returns, lengths, frames_list = run_episodes( agent, args.env_id, device, args.episodes, deterministic=not args.stochastic, seed=args.seed, collect_frames=max(args.video_episodes, 2), ) ppo_stats = summarize(returns) print(f"\n[结果] PPO 智能体 {args.episodes} 局: " f"平均={ppo_stats['mean']:.2f} ± {ppo_stats['std']:.2f} | " f"中位数={ppo_stats['median']:.1f} | min={ppo_stats['min']:.1f} max={ppo_stats['max']:.1f} | " f"着陆率(>=200)={ppo_stats['landing_rate']*100:.0f}%") # ---- 随机策略对照 ---- random_stats = None if args.baseline_random: r_returns, r_lengths, _ = run_episodes( None, args.env_id, device, args.episodes, deterministic=False, seed=args.seed, collect_frames=0) random_stats = summarize(r_returns) print(f"[对照] 随机策略 {args.episodes} 局: " f"平均={random_stats['mean']:.2f} ± {random_stats['std']:.2f} | " f"着陆率={random_stats['landing_rate']*100:.0f}%") # ---- 录像 ---- videos, gif_path = [], None if args.video_episodes > 0: n_vid = min(args.video_episodes, len(frames_list)) for i in range(n_vid): path = os.path.join(args.out, "videos", f"replay_ep{i + 1}.mp4") rec = VideoRecorder(path, fps=30) for f in frames_list[i]: rec.add(f) rec.save() videos.append({"path": path, "episode_return": returns[i], "episode_length": lengths[i], "frames": len(frames_list[i])}) print(f"[录像] {path} (回报 {returns[i]:.1f}, {lengths[i]} 步)") # 选回报最高的一局做 gif best_i = int(np.argmax(returns[:len(frames_list)])) gif_path = make_gif(frames_list[best_i], os.path.join(args.out, "replay.gif")) print(f"[动图] {gif_path}") keyframes = make_keyframes(frames_list, os.path.join(args.out, "keyframes.png")) # ---- 保存结果 ---- result = { "model": os.path.abspath(args.model), "env_id": args.env_id, "policy": "stochastic" if args.stochastic else "argmax(deterministic)", "episodes": args.episodes, "seed": args.seed, "solved_threshold": 200.0, "solved": bool(ppo_stats["mean"] >= 200.0), "ppo": {**ppo_stats, "returns": [float(x) for x in returns], "lengths": [int(x) for x in lengths]}, "random_baseline": None if random_stats is None else {**random_stats, "returns": [float(x) for x in r_returns]}, "videos": videos, "checkpoint": {k: ckpt.get(k) for k in ("iteration", "global_step", "eval_mean", "hidden")}, } out_json = os.path.join(args.out, "eval_results.json") with open(out_json, "w", encoding="utf-8") as f: json.dump(result, f, indent=2, ensure_ascii=False) print(f"\n[完成] 指标 -> {out_json}") if random_stats is not None: gain = ppo_stats["mean"] - random_stats["mean"] print(f"[对比] 随机 {random_stats['mean']:.1f} -> PPO {ppo_stats['mean']:.1f} " f"(提升 {gain:+.1f} 分)") if __name__ == "__main__": main()