Update Down/down/exp/plot_paradox_results.py

Down/down/exp/plot_paradox_results.py

@@ -1,15 +1,23 @@
-"""Plot results produced by exp/run_paradox_tabular.py.
+"""
+Plot results produced by exp/run_paradox_tabular.py.
 
-Creates PNG plots under down/graphics/.
+Creates PNG plots under down/graphics/ by default.
+Now supports:
+  --csv    path to results csv
+  --out    output directory for pngs
+  --prefix filename prefix for plots
 """
 
 from __future__ import annotations
 
+import argparse
 import csv
 import os
 from pathlib import Path
 from typing import Dict, List
 
+import matplotlib
+matplotlib.use("Agg")
 import matplotlib.pyplot as plt
 
 
@@ -18,13 +26,51 @@ def _load_csv(path: Path) -> List[Dict[str, str]]:
         return list(csv.DictReader(f))
 
 
+def _to_float(x: str) -> float:
+    try:
+        return float(x)
+    except Exception:
+        return float("nan")
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description="Plot DOWN paradox tabular results.")
+    p.add_argument(
+        "--csv",
+        type=str,
+        default="",
+        help="Path to results CSV. If omitted, uses TWOQUARKS_RESULTS_DIR/down_paradox_tabular_results.csv or down/results/...",
+    )
+    p.add_argument(
+        "--out",
+        type=str,
+        default="",
+        help="Output directory for PNGs. If omitted, uses TWOQUARKS_GRAPHICS_DIR or down/graphics.",
+    )
+    p.add_argument("--prefix", type=str, default="Down", help="Prefix for output filenames.")
+    return p
+
+
 def main() -> None:
-    Down_dir = Path(__file__).resolve().parents[1]  # .../down
-    results_dir = Path(os.environ.get("DOWN_RESULTS_DIR", (Down_dir / "results").as_posix()))
-    graphics_dir = Path(os.environ.get("DOWN_GRAPHICS_DIR", (Down_dir / "graphics").as_posix()))
-    results_dir.mkdir(parents=True, exist_ok=True)
+    args = _build_parser().parse_args()
+
+    quark_dir = Path(__file__).resolve().parents[1]  # .../down
+
+    # Default dirs (env overrides)
+    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()))
+
+    # Respect CLI overrides
+    if args.out.strip():
+        graphics_dir = Path(args.out)
+
     graphics_dir.mkdir(parents=True, exist_ok=True)
-    results_csv = results_dir / "down_paradox_tabular_results.csv"
+
+    if args.csv.strip():
+        results_csv = Path(args.csv)
+    else:
+        results_dir.mkdir(parents=True, exist_ok=True)
+        results_csv = results_dir / "down_paradox_tabular_results.csv"
 
     if not results_csv.exists():
         raise FileNotFoundError(f"Missing results CSV: {results_csv}")
@@ -33,46 +79,41 @@ def main() -> None:
     if not rows:
         raise RuntimeError(f"CSV is empty: {results_csv}")
 
-    # Normalize columns (older versions used episode_reward/rho_state)
-    def to_float(x: str) -> float:
-        try:
-            return float(x)
-        except Exception:
-            return float("nan")
+    prefix = args.prefix.strip() or "Down"
 
-    # Build per-agent series
+    # Agents present in CSV
     agents = sorted({r.get("agent", "unknown") for r in rows})
-    xs_all = [int(r.get("global_episode", r.get("episode", 0)) or 0) for r in rows]
 
-    # Episode reward plot
+    # Reward plot
     plt.figure()
     for a in agents:
         xs = [int(r.get("global_episode", r.get("episode", 0)) or 0) for r in rows if r.get("agent") == a]
-        ys = [to_float(r.get("episode_reward", r.get("return", "nan"))) for r in rows if r.get("agent") == a]
+        ys = [_to_float(r.get("episode_reward", r.get("return", "nan"))) for r in rows if r.get("agent") == a]
         if xs:
             plt.plot(xs, ys, label=a)
     plt.xlabel("Global episode")
     plt.ylabel("Episode reward")
-    plt.title("DOWN: reward over training")
+    plt.title(f"{prefix}: reward over training")
     plt.legend()
-    out1 = graphics_dir / "down_paradox_tabular_reward.png"
-    plt.savefig(out1, dpi=120, bbox_inches="tight")
+    out1 = graphics_dir / f"{prefix}_paradox_tabular_reward.png"
+    plt.savefig(out1, dpi=160, bbox_inches="tight")
     plt.close()
 
     # Rho plot (if present)
-    if any(("rho_state" in r) for r in rows) or any(("rho" in r) for r in rows):
+    has_rho = any(("rho_state" in r) for r in rows) or any(("rho" in r) for r in rows)
+    if has_rho:
         plt.figure()
         for a in agents:
             xs = [int(r.get("global_episode", r.get("episode", 0)) or 0) for r in rows if r.get("agent") == a]
-            ys = [to_float(r.get("rho_state", r.get("rho", "nan"))) for r in rows if r.get("agent") == a]
+            ys = [_to_float(r.get("rho_state", r.get("rho", "nan"))) for r in rows if r.get("agent") == a]
             if xs:
                 plt.plot(xs, ys, label=a)
         plt.xlabel("Global episode")
         plt.ylabel("rho")
-        plt.title("DOWN: rho over training")
+        plt.title(f"{prefix}: rho over training")
         plt.legend()
-        out2 = graphics_dir / "down_paradox_tabular_rho.png"
-        plt.savefig(out2, dpi=120, bbox_inches="tight")
+        out2 = graphics_dir / f"{prefix}_paradox_tabular_rho.png"
+        plt.savefig(out2, dpi=160, bbox_inches="tight")
         plt.close()
 
     print(f"Saved plots to {graphics_dir}")