src/evaluate.py

# -*- 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()