| |
| """ |
| 评估脚本 —— 加载训练好的模型，输出指标 + 演示视频 + 关键帧 + 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 |
|
|
|
|
| |
| |
| |
| @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 |
| |
| 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]} 步)") |
| |
| 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() |
|
|
