Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| from __future__ import annotations | |
| import argparse | |
| import csv | |
| import sys | |
| from pathlib import Path | |
| def find_latest() -> Path | None: | |
| paths = sorted(Path("outputs").glob("*/reward_log.csv"), key=lambda p: p.stat().st_mtime, reverse=True) | |
| if paths: | |
| return paths[0] | |
| p = Path("docs/driftshield_proof_reward_log.csv") | |
| return p if p.exists() else None | |
| def load(path: Path): | |
| with open(path, newline="") as f: | |
| r = list(csv.DictReader(f)) | |
| if not r: | |
| return None | |
| return r | |
| def main() -> int: | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("csv_path", nargs="?", type=Path) | |
| ap.add_argument("--out", type=Path, default=None) | |
| args = ap.parse_args() | |
| p = args.csv_path or find_latest() | |
| if not p or not p.exists(): | |
| print("No CSV found. Provide path or add docs/driftshield_proof_reward_log.csv") | |
| return 1 | |
| rows = load(p) | |
| if not rows: | |
| return 1 | |
| try: | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| except ImportError: | |
| print("Install matplotlib: pip install matplotlib") | |
| return 1 | |
| eps = [int(x["episode"]) for x in rows if x.get("episode", "").isdigit()] | |
| if not eps: | |
| eps = list(range(1, len(rows) + 1)) | |
| tot = [float(x.get("total_reward", x.get("total", 0))) for x in rows] | |
| inv = [float(x.get("investigation", 0)) for x in rows] | |
| rout = [float(x.get("routing", 0)) for x in rows] | |
| pen = [float(x.get("penalty", x.get("penalty_sum", 0))) for x in rows] | |
| p95 = [1.0 if t >= 0.9 else 0.0 for t in tot] | |
| fig, axes = plt.subplots(2, 2, figsize=(11, 7)) | |
| ax0, ax1, ax2, ax3 = axes[0, 0], axes[0, 1], axes[1, 0], axes[1, 1] | |
| n = list(range(1, len(tot) + 1)) if not eps or len(eps) != len(tot) else eps | |
| if len(n) != len(tot): | |
| n = list(range(1, len(tot) + 1)) | |
| ax0.plot(n, tot, "-o", markersize=2, color="#1f77b4") | |
| ax0.set_ylabel("total_reward") | |
| ax0.set_title("DriftShield — total reward per episode") | |
| ax0.grid(alpha=0.3) | |
| ax1.plot(n, inv, label="investigation", alpha=0.8) | |
| ax1.plot(n, rout, label="routing", alpha=0.8) | |
| ax1.set_ylabel("subscore (0-1)") | |
| ax1.set_title("Base components (sampled)") | |
| ax1.legend(fontsize=8) | |
| ax1.grid(alpha=0.3) | |
| ax2.plot(n, pen, color="darkred", alpha=0.85) | |
| ax2.set_ylabel("penalty_sum") | |
| ax2.set_title("Penalty mass") | |
| ax2.grid(alpha=0.3) | |
| ax3.plot(n, p95, drawstyle="steps-post", color="#2ca02c") | |
| ax3.set_ylabel("1 if total>=0.9 else 0") | |
| ax3.set_xlabel("episode (or index)") | |
| ax3.set_title("Pass-rate proxy") | |
| ax3.set_ylim(-0.1, 1.1) | |
| ax3.grid(alpha=0.3) | |
| fig.suptitle(f"Source: {p.name}", fontsize=9) | |
| fig.tight_layout() | |
| out = args.out or p.with_name("reward_curve.png") | |
| fig.savefig(out, dpi=150) | |
| print(f"Wrote {out.resolve()}") | |
| return 0 | |
| if __name__ == "__main__": | |
| sys.exit(main()) | |