src/plot_curves.py

# -*- coding: utf-8 -*-
"""
绘图脚本 —— 用 PIL 手绘学习曲线 / 损失曲线 / 评估分布 / 基线对比
=================================================================
（工程环境里没有 matplotlib，所以用 PIL 直接画图，零额外依赖）

用法::

    python plot_curves.py --run runs/main_2.5m --baseline runs/baseline_50k --out artifacts

产物（artifacts/ 下）：
    learning_curve.png   训练回报 + 评估回报随步数变化（含 200 分"解决"线）
    losses.png           policy loss / value loss / entropy / KL
    eval_progress.png    每次评估的均值 ± 标准差
    eval_histogram.png   100 局评估的回报分布（对比随机策略）
    compare_baseline.png 课件 5 万步命令 vs 正式训练的对比
"""

import argparse
import csv
import json
import os

import numpy as np
from PIL import Image, ImageDraw, ImageFont

C = {
    "blue": (86, 156, 255), "green": (80, 200, 120), "orange": (255, 165, 60),
    "red": (240, 90, 90), "purple": (170, 120, 240), "gray": (150, 150, 160),
    "cyan": (90, 210, 220), "yellow": (240, 210, 80),
}
BG = (18, 20, 25)
PANEL_BG = (30, 33, 40)
GRID = (55, 60, 70)
TXT = (235, 235, 240)
TXT_DIM = (190, 190, 195)


# --------------------------------------------------------------------------- #
# 基础绘图工具
# --------------------------------------------------------------------------- #
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 _rot_text(img, xy, text, font, fill):
    """把文字旋转 90° 画到 img 上（用于 y 轴标题）。"""
    tmp = Image.new("RGBA", (240, 30), (0, 0, 0, 0))
    ImageDraw.Draw(tmp).text((0, 0), text, font=font, fill=fill)
    tmp = tmp.crop(tmp.getbbox() or (0, 0, 1, 1)).rotate(90, expand=True)
    img.paste(tmp, (int(xy[0]), int(xy[1])), tmp)


def _nan_safe(vals):
    return [v if np.isfinite(v) else 0.0 for v in vals]


def draw_panel(img, box, series, title, xlabel, ylabel,
               hline=None, legend=True, xticklabels=None, ylim=None):
    """在 box=(x0,y0,x1,y1) 里画坐标系。

    series 元素：{"name","xs","ys","color","kind":"line"/"scatter"/"bar","width"}
    xticklabels: 给定则 x 轴用分类标签（xs 为 0..n-1）。
    ylim: (ymin, ymax) 强制 y 轴范围（超出范围的点不画，避免极端离群点压扁曲线）。
    """
    d = ImageDraw.Draw(img)
    x0, y0, x1, y1 = box
    pad_l, pad_r, pad_t, pad_b = 78, 18, 48, 48
    px0, py0 = x0 + pad_l, y0 + pad_t
    px1, py1 = x1 - pad_r, y1 - pad_b
    d.rectangle([px0, py0, px1, py1], fill=PANEL_BG)

    xs_all, ys_all = [], []
    for s in series:
        xs_all.extend(s["xs"])
        ys_all.extend(s["ys"])
    if hline is not None:
        ys_all.append(hline[0])
    if not xs_all:
        return
    xmin, xmax = float(min(xs_all)), float(max(xs_all))
    if ylim is not None:
        ymin, ymax = float(ylim[0]), float(ylim[1])
    else:
        ymin, ymax = float(min(ys_all)), float(max(ys_all))
        pad_y = (ymax - ymin) * 0.10
        ymin, ymax = ymin - pad_y, ymax + pad_y
    if xmax == xmin:
        xmax = xmin + 1
    if ymax == ymin:
        ymax = ymin + 1

    def px(xv):
        return px0 + (xv - xmin) / (xmax - xmin) * (px1 - px0)

    def py(yv):
        return py1 - (yv - ymin) / (ymax - ymin) * (py1 - py0)

    # 网格
    for i in range(5):
        yv = ymin + (ymax - ymin) * i / 4
        y = py(yv)
        d.line([px0, y, px1, y], fill=GRID, width=1)
        label = f"{yv:,.0f}" if abs(yv) >= 1 else f"{yv:.2f}"
        d.text((x0 + 30, y - 8), label, font=_font(13), fill=TXT_DIM)
        if xticklabels is None:
            xv = xmin + (xmax - xmin) * i / 4
            x = px(xv)
            d.line([x, py0, x, py1], fill=GRID, width=1)
            lab = f"{xv/1e6:.1f}M" if abs(xv) >= 1e6 else (
                f"{xv/1e3:.0f}k" if abs(xv) >= 1e3 else (
                    f"{xv:.0f}" if abs(xv) >= 10 else f"{xv:.1f}"))
            d.text((x - 14, py1 + 10), lab, font=_font(13), fill=TXT_DIM)

    # 分类 x 轴
    if xticklabels is not None:
        n = len(xticklabels)
        width = (px1 - px0) / max(1, n)
        for i, lab in enumerate(xticklabels):
            x = px0 + width * (i + 0.5)
            d.text((x - d.textlength(lab, font=_font(13)) / 2, py1 + 10),
                   lab, font=_font(13), fill=TXT_DIM)

    # 参考线
    if hline is not None:
        yv, label = hline[0], hline[1]
        color = hline[2] if len(hline) > 2 else C["red"]
        y = py(yv)
        for xx in range(int(px0), int(px1), 14):
            d.line([xx, y, min(xx + 7, px1), y], fill=color, width=2)
        d.text((px1 - d.textlength(label, font=_font(13)) - 6, y - 19),
               label, font=_font(13), fill=color)

    # 数据（超出 ylim 的点直接不画，避免极端离群点拖垮坐标轴）
    for s in series:
        color = C.get(s.get("color", "blue"), C["blue"])
        kind, width = s.get("kind", "line"), s.get("width", 2)
        raw = [(xv, yv) for xv, yv in zip(s["xs"], s["ys"]) if ymin <= yv <= ymax]
        pts = [(px(xv), py(yv)) for xv, yv in raw]
        if kind == "bar":
            if xticklabels is not None:
                slot = (px1 - px0) / len(xticklabels)
                bw = slot * 0.5
                centers = [px0 + slot * (i + 0.5) for i in range(len(s["ys"]))]
            else:
                bw = max(4, (px1 - px0) / max(1, len(pts)) * 0.7)
                centers = [x for x, _ in pts]
            for cx, (_, cy) in zip(centers, pts):
                d.rectangle([cx - bw / 2, min(cy, py1), cx + bw / 2, max(cy, py1)], fill=color)
        elif kind == "scatter":
            for (cx, cy) in pts:
                d.ellipse([cx - 3, cy - 3, cx + 3, cy + 3], fill=color)
        else:
            if len(pts) > 1:
                d.line(pts, fill=color, width=width, joint="curve")

    # 标题 / 轴标题
    d.text((x0 + pad_l, y0 + 10), title, font=_font(19), fill=TXT)
    d.text((px0 + (px1 - px0) / 2 - d.textlength(xlabel, font=_font(13)) / 2, y1 - 20),
           xlabel, font=_font(13), fill=TXT_DIM)
    _rot_text(img, (x0 + 3, (py0 + py1) / 2 - 60), ylabel, _font(13), TXT_DIM)

    # 图例
    if legend:
        entries = [(s["name"], s.get("color", "blue")) for s in series]
        lh = 22 * len(entries) + 10
        lw = 16 + max(d.textlength(n, font=_font(13)) for n, _ in entries) + 34
        d.rectangle([px0 + 8, py0 + 8, px0 + 8 + lw, py0 + 8 + lh], fill=(24, 27, 33))
        ly = py0 + 14
        for name, color in entries:
            d.line([px0 + 16, ly + 7, px0 + 38, ly + 7], fill=C[color], width=3)
            d.text((px0 + 46, ly), name, font=_font(13), fill=TXT)
            ly += 22


def moving_average(xs, ys, window):
    if len(ys) <= window:
        return xs, ys
    ys = _nan_safe(ys)
    out = [float(np.mean(ys[max(0, i - window + 1):i + 1])) for i in range(len(ys))]
    return xs, out


# --------------------------------------------------------------------------- #
# 数据
# --------------------------------------------------------------------------- #
def read_csv_rows(path):
    if not os.path.exists(path):
        return []
    with open(path, "r", encoding="utf-8-sig") as f:
        return list(csv.DictReader(f))


def load_run(run_dir):
    train = read_csv_rows(os.path.join(run_dir, "train_log.csv"))
    evals = read_csv_rows(os.path.join(run_dir, "eval_log.csv"))
    summary = {}
    sp = os.path.join(run_dir, "summary.json")
    if os.path.exists(sp):
        with open(sp, "r", encoding="utf-8") as f:
            summary = json.load(f)
    return train, evals, summary


# --------------------------------------------------------------------------- #
# 图 1：学习曲线
# --------------------------------------------------------------------------- #
def plot_learning_curve(run_dir, out_path, title):
    train, evals, _ = load_run(run_dir)
    if not train:
        print(f"[跳过] 缺 {run_dir}/train_log.csv")
        return
    steps = [int(r["global_step"]) for r in train]
    rets = [float(r["mean_return"]) for r in train]
    series = []
    if any(np.isfinite(rets)):
        series.append({"name": "train return (raw)", "xs": steps, "ys": rets,
                       "color": "gray", "kind": "line", "width": 1})
        mx, my = moving_average(steps, rets, max(5, len(steps) // 60))
        series.append({"name": "train return (moving avg)", "xs": mx, "ys": my,
                       "color": "blue", "kind": "line", "width": 3})
    if evals:
        ex = [int(r["global_step"]) for r in evals]
        ey = [float(r["eval_mean"]) for r in evals]
        series.append({"name": "eval (argmax policy)", "xs": ex, "ys": ey,
                       "color": "green", "kind": "line", "width": 3})
        series.append({"name": "eval points", "xs": ex, "ys": ey,
                       "color": "green", "kind": "scatter"})
    W, H = 1150, 580
    img = Image.new("RGB", (W, H), BG)
    draw_panel(img, (10, 10, W - 10, H - 10), series, title,
               "environment steps", "episode return",
               hline=(200.0, "solved threshold = 200", C["red"]),
               ylim=(-320, 330))
    d = ImageDraw.Draw(img)
    d.text((100, 44),
           "note: the first 2 argmax evals (-1983 / -142) are off-scale — early policy not yet deterministic, see README",
           font=_font(12), fill=(150, 150, 160))
    img.save(out_path)
    print(f"[图] {out_path}")


# --------------------------------------------------------------------------- #
# 图 2：损失曲线（2×2）
# --------------------------------------------------------------------------- #
def plot_losses(run_dir, out_path):
    train, _, _ = load_run(run_dir)
    if not train:
        return
    steps = [int(r["global_step"]) for r in train]
    W, H = 1150, 800
    img = Image.new("RGB", (W, H), BG)
    panels = [("policy_loss", "1) policy loss (PPO-Clip)", "orange"),
              ("value_loss", "2) value loss", "purple"),
              ("entropy", "3) policy entropy", "cyan"),
              ("approx_kl", "4) approximate KL", "green")]
    for i, (key, title, color) in enumerate(panels):
        ys = [float(r[key]) for r in train]
        r, c = divmod(i, 2)
        box = (10 + c * (W // 2), 10 + r * (H // 2), 10 + (c + 1) * (W // 2) - 10,
               10 + (r + 1) * (H // 2) - 10)
        mx, my = moving_average(steps, ys, max(3, len(ys) // 80))
        draw_panel(img, box,
                   [{"name": "raw", "xs": steps, "ys": _nan_safe(ys), "color": "gray", "width": 1},
                    {"name": "moving avg", "xs": mx, "ys": my, "color": color, "width": 2}],
                   title, "environment steps", key)
    img.save(out_path)
    print(f"[图] {out_path}")


# --------------------------------------------------------------------------- #
# 图 3：训练中的评估进度（均值 ± 标准差）
# --------------------------------------------------------------------------- #
def plot_eval_progress(run_dir, out_path):
    _, evals, _ = load_run(run_dir)
    if not evals:
        return
    steps = [int(r["global_step"]) for r in evals]
    means = [float(r["eval_mean"]) for r in evals]
    stds = [float(r["eval_std"]) for r in evals]
    W, H = 1150, 540
    img = Image.new("RGB", (W, H), BG)
    d = ImageDraw.Draw(img)
    ylo, yhi = -260.0, 340.0
    draw_panel(img, (10, 10, W - 10, H - 10),
               [{"name": "eval mean", "xs": steps, "ys": means, "color": "green", "width": 3}],
               "evaluation during training (argmax policy)", "environment steps", "eval return",
               hline=(200.0, "solved threshold = 200", C["red"]), ylim=(ylo, yhi))
    # 手动画误差棒（与坐标轴同一范围）
    px0, py0, px1, py1 = 10 + 78, 10 + 48, W - 10 - 18, H - 10 - 48
    xmin, xmax = min(steps), max(steps)
    for xv, m, s in zip(steps, means, stds):
        x = px0 + (xv - xmin) / max(1, xmax - xmin) * (px1 - px0)
        hi_, lo_ = m + s, m - s
        y_hi = py1 - (min(max(hi_, ylo), yhi) - ylo) / (yhi - ylo) * (py1 - py0)
        y_lo = py1 - (min(max(lo_, ylo), yhi) - ylo) / (yhi - ylo) * (py1 - py0)
        d.line([x, y_hi, x, y_lo], fill=C["blue"], width=2)
        d.line([x - 4, y_hi, x + 4, y_hi], fill=C["blue"], width=2)
        d.line([x - 4, y_lo, x + 4, y_lo], fill=C["blue"], width=2)
    img.save(out_path)
    print(f"[图] {out_path}")


# --------------------------------------------------------------------------- #
# 图 4：评估回报分布
# --------------------------------------------------------------------------- #
def plot_eval_histogram(eval_json, out_path):
    if not os.path.exists(eval_json):
        return
    with open(eval_json, "r", encoding="utf-8") as f:
        data = json.load(f)
    returns = np.asarray(data["ppo"]["returns"], dtype=float)
    rand = data.get("random_baseline")
    lo, hi = float(returns.min()), float(returns.max())
    pad = max(5.0, (hi - lo) * 0.05)
    bins = np.linspace(lo - pad, hi + pad, 15)
    hist, edges = np.histogram(returns, bins=bins)
    centers = (edges[:-1] + edges[1:]) / 2
    W, H = 1150, 540
    img = Image.new("RGB", (W, H), BG)
    draw_panel(img, (10, 10, W - 10, H - 10),
               [{"name": f"PPO agent ({len(returns)} episodes, argmax)",
                 "xs": centers.tolist(), "ys": hist.tolist(), "color": "blue", "kind": "bar"}],
               "evaluation return distribution (argmax policy)", "episode return",
               "number of episodes",
               hline=(200.0, "solved threshold = 200", C["red"]),
               ylim=(0, max(1, int(hist.max()) * 1.25)))
    if rand:
        d = ImageDraw.Draw(img)
        txt = (f"random policy baseline: mean {rand['mean']:.0f} ± {rand['std']:.0f}, "
               f"landing rate {rand['landing_rate']*100:.0f}%  (far to the left, off-scale)")
        d.text((100, 84), txt, font=_font(13), fill=C["orange"])
    img.save(out_path)
    print(f"[图] {out_path}")


# --------------------------------------------------------------------------- #
# 图 5：课件命令基线 vs 正式训练（横向柱状图）
# --------------------------------------------------------------------------- #
def plot_compare(run_dir, baseline_dir, out_path):
    _, evals, summary = load_run(run_dir)
    _, base_evals, base_summary = load_run(baseline_dir)
    if not evals:
        return
    items = []
    if base_evals:
        items.append((f"course command\n{int(base_summary.get('total_timesteps', 0))/1000:.0f}k steps",
                      float(base_evals[-1]["eval_mean"])))
    items.append((f"this run\n{summary.get('total_timesteps', 0)/1e6:.1f}M steps",
                  float(evals[-1]["eval_mean"])))

    W, H = 980, 460
    img = Image.new("RGB", (W, H), BG)
    d = ImageDraw.Draw(img)
    d.text((24, 16), "course command baseline vs. full PPO training", font=_font(19), fill=TXT)
    left, right, top = 180, W - 150, 80
    vmin = min(v for _, v in items + [("", 0.0)])
    vmax = max(v for _, v in items + [("", 220.0)])
    span = vmax - vmin

    def vx(v):
        return left + (v - vmin) / span * (right - left)

    rowh = 78
    for i, (label, value) in enumerate(items):
        y = top + i * rowh
        d.text((20, y + 10), label, font=_font(14), fill=TXT_DIM)
        color = C["gray"] if i == 0 and len(items) > 1 else C["green"]
        x0, x1 = sorted([vx(0), vx(value)])
        d.rectangle([x0, y, x1, y + 40], fill=color)
        tx = max(vx(value), vx(0.0)) + 8
        d.text((tx, y + 10), f"{value:+.1f}", font=_font(16), fill=TXT)
    # 0 线与 200 线
    for val, color, lab in [(0.0, C["gray"], "0"), (200.0, C["red"], "solved = 200")]:
        x = vx(val)
        d.line([x, top - 12, x, top + len(items) * rowh], fill=color, width=2)
        d.text((x - 20, top - 34), lab, font=_font(13), fill=color)
    img.save(out_path)
    print(f"[图] {out_path}")


def main():
    p = argparse.ArgumentParser(description="画学习曲线与评估图")
    p.add_argument("--run", type=str, default="runs/main_2.5m")
    p.add_argument("--baseline", type=str, default="runs/baseline_50k")
    p.add_argument("--out", type=str, default="artifacts")
    p.add_argument("--title", type=str, default="PPO on LunarLander — implemented from scratch")
    p.add_argument("--eval-json", type=str, default="")
    args = p.parse_args()
    os.makedirs(args.out, exist_ok=True)

    plot_learning_curve(args.run, os.path.join(args.out, "learning_curve.png"), args.title)
    plot_losses(args.run, os.path.join(args.out, "losses.png"))
    plot_eval_progress(args.run, os.path.join(args.out, "eval_progress.png"))
    plot_compare(args.run, args.baseline, os.path.join(args.out, "compare_baseline.png"))
    plot_eval_histogram(args.eval_json or os.path.join(args.out, "eval_results.json"),
                        os.path.join(args.out, "eval_histogram.png"))


if __name__ == "__main__":
    main()