Update Down/AntiDown/exp/plot_corrupted_valley_results.py

Down/AntiDown/exp/plot_corrupted_valley_results.py

@@ -1,10 +1,17 @@
 """
-Plot results for the Corrupted Valley tabular experiment.
+Plot results for the Corrupted Valley tabular experiment (AntiDown).
 
-Expects a CSV file created by `exp/run_corrupted_valley_tabular.py` with columns:
+Expects a CSV created by exp/run_corrupted_valley_tabular.py with columns:
     phase, phase_tag, episode, agent, total_reward, valley_visits
+
+Supports:
+    --csv     path to CSV file
+    --out     output directory for PNGs
+    --prefix  filename prefix
 """
 
+from __future__ import annotations
+
 # --- twoquarks bootstrap (path-stable imports) ---
 import sys
 from pathlib import Path
@@ -13,32 +20,48 @@ if str(_ROOT) not in sys.path:
     sys.path.insert(0, str(_ROOT))
 # -------------------------------------------------
 
-
+import argparse
 import csv
 import os
 from collections import defaultdict
 from typing import Dict, List, Tuple
 
+import matplotlib
+matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
 
-RESULTS_DIR = Path(os.environ.get("TWOQUARKS_RESULTS_DIR", (Path(__file__).resolve().parent.parent / "results").as_posix()))
-RESULTS_CSV = str(RESULTS_DIR / "antidown_corrupted_valley_tabular.csv")
-OUT_DIR = Path(os.environ.get("TWOQUARKS_GRAPHICS_DIR", (Path(__file__).resolve().parent.parent / "graphics").as_posix()))
-OUT_DIR.mkdir(parents=True, exist_ok=True)
-OUT_PREFIX = str(OUT_DIR / "antidown_corrupted_valley_")
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description="Plot AntiDown corrupted valley tabular results.")
+    p.add_argument(
+        "--csv",
+        type=str,
+        default="",
+        help="Path to results CSV. If omitted, uses TWOQUARKS_RESULTS_DIR/antidown_corrupted_valley_tabular.csv",
+    )
+    p.add_argument(
+        "--out",
+        type=str,
+        default="",
+        help="Output directory for PNGs. If omitted, uses TWOQUARKS_GRAPHICS_DIR",
+    )
+    p.add_argument(
+        "--prefix",
+        type=str,
+        default="AntiDown",
+        help="Prefix for output filenames.",
+    )
+    return p
 
 
-def load_results(path: str) -> List[Dict[str, str]]:
-    rows: List[Dict[str, str]] = []
-    with open(path, mode="r", newline="") as f:
-        reader = csv.DictReader(f)
-        for row in reader:
-            rows.append(row)
-    return rows
+def load_results(path: Path) -> List[Dict[str, str]]:
+    with path.open("r", newline="") as f:
+        return list(csv.DictReader(f))
 
 
-def build_global_index(rows: List[Dict[str, str]]) -> Tuple[Dict[int, int], Dict[int, Dict[str, List[Tuple[int, float, float]]]]]:
+def build_global_index(
+    rows: List[Dict[str, str]],
+) -> Tuple[Dict[int, int], Dict[int, Dict[str, List[Tuple[int, float, float]]]]]:
     """
     Returns:
         phase_to_max_ep: phase -> max episode index
@@ -64,8 +87,10 @@ def build_global_index(rows: List[Dict[str, str]]) -> Tuple[Dict[int, int], Dict
         ep = int(r["episode"])
         total_reward = float(r["total_reward"])
         valley_visits = float(r["valley_visits"])
+
         max_ep = phase_to_max_ep[phase]
         global_ep = (phase - 1) * max_ep + ep
+
         data[phase][agent].append((global_ep, total_reward, valley_visits))
 
     return phase_to_max_ep, data
@@ -75,14 +100,14 @@ def plot_metric(
     data: Dict[int, Dict[str, List[Tuple[int, float, float]]]],
     metric_index: int,
     ylabel: str,
-    out_path: str,
+    out_path: Path,
 ) -> None:
     """
-    metric_index: 1 for total_reward, 2 for valley_visits.
+    metric_index: 1 -> total_reward
+                  2 -> valley_visits
     """
     plt.figure()
 
-    # merge across phases, keeping breaks in the global index
     agent_to_xy: Dict[str, Tuple[List[int], List[float]]] = {}
 
     for phase, agents in sorted(data.items()):
@@ -95,7 +120,6 @@ def plot_metric(
             agent_to_xy[agent][1].extend(ys)
 
     for agent, (xs, ys) in agent_to_xy.items():
-        # sort by x
         pairs = sorted(zip(xs, ys), key=lambda p: p[0])
         xs_sorted = [p[0] for p in pairs]
         ys_sorted = [p[1] for p in pairs]
@@ -105,31 +129,56 @@ def plot_metric(
     plt.ylabel(ylabel)
     plt.legend()
     plt.tight_layout()
-    plt.savefig(out_path)
+    plt.savefig(out_path, dpi=160)
     plt.close()
 
 
 def main() -> None:
-    if not os.path.exists(RESULTS_CSV):
-        raise FileNotFoundError(f"Results file not found: {RESULTS_CSV}")
+    args = _build_parser().parse_args()
+
+    quark_dir = Path(__file__).resolve().parents[1]  # .../AntiDown
+
+    # Defaults via env
+    results_dir = Path(os.environ.get("TWOQUARKS_RESULTS_DIR", (quark_dir / "results").as_posix()))
+    graphics_dir = Path(os.environ.get("TWOQUARKS_GRAPHICS_DIR", (quark_dir / "graphics").as_posix()))
+
+    if args.out.strip():
+        graphics_dir = Path(args.out)
+    graphics_dir.mkdir(parents=True, exist_ok=True)
+
+    if args.csv.strip():
+        results_csv = Path(args.csv)
+    else:
+        results_dir.mkdir(parents=True, exist_ok=True)
+        results_csv = results_dir / "antidown_corrupted_valley_tabular.csv"
+
+    if not results_csv.exists():
+        raise FileNotFoundError(f"Results file not found: {results_csv}")
+
+    rows = load_results(results_csv)
+    if not rows:
+        raise RuntimeError(f"CSV is empty: {results_csv}")
 
-    rows = load_results(RESULTS_CSV)
     _, data = build_global_index(rows)
 
+    prefix = args.prefix.strip() or "AntiDown"
+
     plot_metric(
         data,
         metric_index=1,
         ylabel="total reward per episode",
-        out_path=OUT_PREFIX + "reward.png",
+        out_path=graphics_dir / f"{prefix}_corrupted_valley_reward.png",
     )
 
     plot_metric(
         data,
         metric_index=2,
         ylabel="valley visits per episode",
-        out_path=OUT_PREFIX + "valley_visits.png",
+        out_path=graphics_dir / f"{prefix}_corrupted_valley_valley_visits.png",
     )
 
+    print(f"Saved plots to {graphics_dir}")
+
 
 if __name__ == "__main__":
-    main()
+    main()