code/build_hub_package.py

# -*- coding: utf-8 -*-
"""
把训练好的智能体 + 代码 + 视频 + 评估分数打包成"可以直接上传 Hub"的文件夹。
================================================================
对应 PPT 第 28 页学习目标第 4 条：
  "能够将训练好的智能体和代码上传到 Hub，并附上一个漂亮的视频回放和评估分数"

产物 hub_package/ 结构：
  README.md                    自动填好分数 + 内嵌视频的模型卡
  metrics.json                 三/四个智能体的评估分数汇总
  requirements.txt
  models/qtable_*.pkl          训练出的模型文件（Q 表 + 超参 + 指标）
  videos/replay_*.mp4          回放视频（含总回放 replay_value_based_all.mp4）
  eval_results/eval_*.json     评估分数（逐局明细）
  logs/<run>/train_console.log 训练日志
  logs/<run>/train_curve.png   训练曲线
  assets/qtable_policy_*.png   学到的策略可视化
  code/                        训练 + 评估 + 可视化 + 上传脚本（可复现）
  run_all.ps1                  一键复现脚本

用法：
  D:\\study\\robot\\pyenv\\DRL\\python.exe build_hub_package.py
  D:\\study\\robot\\pyenv\\DRL\\python.exe build_hub_package.py --repo-id 你的用户名/value-based-qlearning
"""

import argparse
import json
import os
import shutil
from datetime import datetime

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# 要打包的智能体（按 PPT 顺序）
AGENTS = [
    dict(run="frozenlake4x4", title="FrozenLake-v1 · 4x4 不滑",
         desc="PPT 第一部分主实验：确定性冰面，从零手写 Q 表，10,000 回合"),
    dict(run="frozenlake4x4_slippery", title="FrozenLake-v1 · 4x4 滑冰版",
         desc="对照实验：只改 is_slippery=True，冰面会打滑，20,000 回合"),
    dict(run="frozenlake8x8", title="FrozenLake-v1 · 8x8",
         desc="进阶实验：地图 8x8 + 拉长探索期，60,000 回合"),
    dict(run="taxi", title="Taxi-v3",
         desc="PPT 第二部分主实验：500 状态 × 6 动作，25,000 回合，官方 100 个评估种子"),
]

CODE_FILES = [
    "train_qlearning.py",
    "evaluate_and_record.py",
    "visualize_qtable.py",
    "make_combined_video.py",
    "build_hub_package.py",
    "upload_to_hub.py",
    "run_all.ps1",
    "01_实践3说明.md",
]

REQUIREMENTS = """gymnasium>=1.0
numpy>=1.24
imageio>=2.31
imageio-ffmpeg>=0.4
Pillow>=9.0
matplotlib>=3.7
tqdm>=4.65
"""


def read_json(path):
    if not os.path.exists(path):
        return None
    with open(path, encoding="utf-8") as f:
        return json.load(f)


def build_readme(repo_id=None, scores=None, now=None):
    raw = f"https://huggingface.co/{repo_id}/resolve/main/" if repo_id else ""
    link = (lambda p: raw + p) if repo_id else (lambda p: p)

    lines = []
    lines.append("# Value-based 实践 · Q-Learning（FrozenLake-v1 + Taxi-v3）")
    lines.append("")
    lines.append("> 《第3周-深度强化学习原理实践》PPT 第 27–40 页 value-based 实践产物："
                 "**从零手写 Q-Learning**（Q 表 + epsilon-greedy + Bellman 更新），"
                 "训练好的智能体 + 训练代码 + 回放视频 + 评估分数。")
    lines.append("")
    lines.append("## 1. 评估分数")
    lines.append("")
    lines.append("| 智能体 | 环境 | 平均回报（100 局） | 成功率 | 平均步数 | 训练回合 |")
    lines.append("|---|---|---|---|---|---|")
    for agent in AGENTS:
        s = (scores or {}).get(agent["run"])
        if s:
            lines.append(
                f"| {agent['title']} | `{s['env_id']}` | **{s['mean_reward']:.2f} ± "
                f"{s['std_reward']:.2f}** | {s['success_rate'] * 100:.0f}% | "
                f"{s['mean_steps']:.1f} | {s.get('n_training_episodes', '—')} |"
            )
        else:
            lines.append(f"| {agent['title']} | — | 待运行 | — | — | — |")
    lines.append("")
    lines.append("> 评估方式：独立环境、贪心策略 `argmax_a Q(s,a)`、每局最多 99 步，跑 100 局；"
                 "FrozenLake 的回报是「是否到达目标」（0/1），Taxi 的回报是"
                 "「每步 −1，成功送达 +20」。")
    lines.append("")

    lines.append("## 2. 回放视频")
    lines.append("")
    if os.path.exists(os.path.join(BASE_DIR, "videos", "replay_value_based_all.mp4")):
        lines.append("**总回放（三个智能体合集）**")
        lines.append("")
        lines.append(f'<video controls width="100%" src="{link("videos/replay_value_based_all.mp4")}"></video>')
        lines.append("")
    for agent in AGENTS:
        video = f"videos/replay_{agent['run']}.mp4"
        if os.path.exists(os.path.join(BASE_DIR, video)):
            lines.append(f"**{agent['title']}**")
            lines.append("")
            lines.append(f'<video controls width="100%" src="{link(video)}"></video>')
            lines.append("")

    lines.append("## 3. 学到的策略（Q 表可视化）")
    lines.append("")
    for agent in AGENTS:
        fig = f"assets/qtable_policy_{agent['run']}.png"
        if os.path.exists(os.path.join(BASE_DIR, fig)):
            lines.append(f"![{agent['title']}]({link(fig)})")
            lines.append("")

    lines.append("## 4. 训练曲线")
    lines.append("")
    for agent in AGENTS:
        fig = f"logs/{agent['run']}/train_curve.png"
        if os.path.exists(os.path.join(BASE_DIR, fig)):
            lines.append(f"![{agent['title']} 训练曲线]({link(fig)})")
            lines.append("")

    lines.append("## 5. 怎么复现")
    lines.append("")
    lines.append("```powershell")
    lines.append("# 环境（本机已有）：conda activate DRL")
    lines.append("conda activate DRL")
    lines.append("")
    lines.append("# 训练（每个实验 1~30 秒）")
    lines.append("python train_qlearning.py --env frozenlake4x4")
    lines.append("python train_qlearning.py --env frozenlake4x4_slippery")
    lines.append("python train_qlearning.py --env frozenlake8x8")
    lines.append("python train_qlearning.py --env taxi")
    lines.append("")
    lines.append("# 评估 + 录视频（生成 eval_results/*.json 和 videos/*.mp4）")
    lines.append("python evaluate_and_record.py --model models/qtable_frozenlake4x4.pkl")
    lines.append("python evaluate_and_record.py --model models/qtable_frozenlake4x4_slippery.pkl --video-episodes 5 --slowmo 4")
    lines.append("python evaluate_and_record.py --model models/qtable_frozenlake8x8.pkl")
    lines.append("python evaluate_and_record.py --model models/qtable_taxi.pkl")
    lines.append("")
    lines.append("# Q 表可视化 + 总回放 + 打包")
    lines.append("python visualize_qtable.py --all")
    lines.append("python make_combined_video.py")
    lines.append("python build_hub_package.py")
    lines.append("```")
    lines.append("")
    lines.append("## 6. 文件说明")
    lines.append("")
    lines.append("- `models/qtable_*.pkl`：模型文件。pickle 格式，内含 Q 表、超参数、训练种子和评估指标；")
    lines.append("  `evaluate_and_record.py` 可以直接加载（不需要重新训练）。")
    lines.append("- `videos/*.mp4`：回放，每帧标注步数/动作/奖励/Q 值条形图，片头片尾是标题卡和分数卡。")
    lines.append("- `eval_results/*.json`：评估分数 + 逐局明细（起点种子、回报、步数、是否成功）。")
    lines.append("- `logs/<run>/train_console.log`：完整训练日志（环境自检、超参数、每 N 回合进度、每次评估）。")
    lines.append("- `code/`：全部代码（训练 / 评估录像 / 可视化 / 打包 / 上传 / 一键脚本）。")
    lines.append("")
    lines.append("## 7. 对应 PPT 页码")
    lines.append("")
    lines.append("| PPT 页 | 内容 | 对应实现 |")
    lines.append("|---|---|---|")
    lines.append("| 29 | Q 表 / Q 函数 / 最优策略 | `train_qlearning.py` 的 `initialize_q_table` / `greedy_policy` |")
    lines.append("| 31–32 | FrozenLake 观测/动作/奖励 | 训练日志里的「环境自检」 |")
    lines.append("| 33 | Q 表初始化 + 贪心策略 | `initialize_q_table` / `greedy_policy` |")
    lines.append("| 34 | epsilon-greedy + 超参数 | `epsilon_greedy_policy` + `EXPERIMENTS` 配置 |")
    lines.append("| 35 | 训练循环（Bellman 更新） | `train_qlearning.py::train` 里的 `Q[s][a] += lr*(...)` |")
    lines.append("| 36–37 | 训练、Q 表可视化、评估、保存模型 | 训练日志 / `visualize_qtable.py` / `evaluate_and_record.py` / `models/*.pkl` |")
    lines.append("| 38–40 | Taxi-v3（500 状态 × 6 动作） | `--env taxi`（含 PPT 的 100 个官方评估种子） |")
    lines.append("| 28 | 上传 Hub + 视频 + 分数 | `build_hub_package.py` + `upload_to_hub.py` + README |")
    lines.append("")
    lines.append(f"- 生成时间：{now:%Y-%m-%d %H:%M}")
    lines.append("")
    return "\n".join(lines)


def copy_file(src, dst):
    os.makedirs(os.path.dirname(dst), exist_ok=True)
    shutil.copy2(src, dst)


def main():
    parser = argparse.ArgumentParser(description="打包成可上传 Hub 的文件夹")
    parser.add_argument("--repo-id", type=str, default=None,
                        help="你的 HF 用户名/仓库名，例如 zhangsan/value-based-qlearning")
    parser.add_argument("--out", type=str, default=os.path.join(BASE_DIR, "hub_package"))
    args = parser.parse_args()

    out = args.out
    if os.path.exists(out):
        shutil.rmtree(out)
    os.makedirs(out)
    now = datetime.now()

    # 1) 读取所有评估分数
    scores = {}
    for agent in AGENTS:
        run = agent["run"]
        s = read_json(os.path.join(BASE_DIR, "eval_results", f"eval_{run}.json"))
        if s is None:
            print(f"[警告] 缺少 eval_results/eval_{run}.json，"
                  f"请先跑 evaluate_and_record.py --model models/qtable_{run}.pkl")
            continue
        scores[run] = s
    metrics = {}
    for run, s in scores.items():
        metrics[run] = {k: s[k] for k in
                        ("env_id", "n_episodes", "mean_reward", "std_reward",
                         "success_rate", "mean_steps", "qtable_shape",
                         "qtable_nonzero", "evaluated_at")}

    # 2) 拷模型
    for agent in AGENTS:
        src = os.path.join(BASE_DIR, "models", f"qtable_{agent['run']}.pkl")
        if os.path.exists(src):
            copy_file(src, os.path.join(out, "models", os.path.basename(src)))
        else:
            print(f"[警告] 缺少模型 {src}")

    # 3) 拷视频（含总回放）
    video_dir = os.path.join(BASE_DIR, "videos")
    if os.path.isdir(video_dir):
        for name in sorted(os.listdir(video_dir)):
            if name.endswith(".mp4"):
                copy_file(os.path.join(video_dir, name),
                          os.path.join(out, "videos", name))

    # 4) 拷评估 json
    for agent in AGENTS:
        src = os.path.join(BASE_DIR, "eval_results", f"eval_{agent['run']}.json")
        if os.path.exists(src):
            copy_file(src, os.path.join(out, "eval_results", os.path.basename(src)))

    # 5) 拷日志 + 训练曲线
    for agent in AGENTS:
        run_dir = os.path.join(BASE_DIR, "logs", agent["run"])
        for name in ("train_console.log", "train_curve.png", "progress.csv"):
            src = os.path.join(run_dir, name)
            if os.path.exists(src):
                copy_file(src, os.path.join(out, "logs", agent["run"], name))

    # 6) 拷 Q 表可视化
    for agent in AGENTS:
        src = os.path.join(BASE_DIR, "assets", f"qtable_policy_{agent['run']}.png")
        if os.path.exists(src):
            copy_file(src, os.path.join(out, "assets", os.path.basename(src)))

    # 7) 拷代码
    for name in CODE_FILES:
        src = os.path.join(BASE_DIR, name)
        if os.path.exists(src):
            copy_file(src, os.path.join(out, "code", name))

    # 8) 写 README / metrics / requirements
    with open(os.path.join(out, "README.md"), "w", encoding="utf-8") as f:
        f.write(build_readme(repo_id=args.repo_id, scores=scores, now=now))
    with open(os.path.join(out, "metrics.json"), "w", encoding="utf-8") as f:
        json.dump(metrics, f, ensure_ascii=False, indent=2)
    with open(os.path.join(out, "requirements.txt"), "w", encoding="utf-8") as f:
        f.write(REQUIREMENTS)

    n_files = sum(len(files) for _, _, files in os.walk(out))
    print(f"[完成] hub_package/ 共 {n_files} 个文件")
    for root, dirs, files in os.walk(out):
        rel = os.path.relpath(root, out)
        for name in sorted(files):
            path = os.path.join(root, name)
            print(f"  {os.path.join(rel, name) if rel != '.' else name}"
                  f"  ({os.path.getsize(path) / 1024:.0f} KB)")
    print("\n下一步上传：")
    print("  hf auth login                       # 粘贴你的 write token")
    print("  python upload_to_hub.py --repo-id 你的用户名/value-based-qlearning")


if __name__ == "__main__":
    main()