Update Down/plot_dual_down_antidown.py

Down/plot_dual_down_antidown.py

@@ -1,10 +1,8 @@
 """
 Overlay plot: Down vs AntiDown (tabular experiments).
 
-Reads CSVs from TWOQUARKS_RESULTS_DIR (or ./results) and writes PNG into
-TWOQUARKS_GRAPHICS_DIR (or ./graphics).
-
-This is intentionally simple and dependency-light (pandas optional).
+Reads CSVs from a shared results directory and writes a PNG overlay plot.
+This is the canonical comparative result for the Down vs AntiDown experiment.
 """
 
 from __future__ import annotations
@@ -13,6 +11,7 @@ from __future__ import annotations
 import matplotlib
 matplotlib.use("Agg")
 
+import argparse
 import csv
 import os
 from collections import defaultdict
@@ -22,6 +21,10 @@ from typing import Dict, List, Tuple
 import matplotlib.pyplot as plt
 
 
+# ----------------------------
+# Utilities
+# ----------------------------
+
 def _read_csv(path: Path) -> List[Dict[str, str]]:
     with path.open("r", newline="") as f:
         return list(csv.DictReader(f))
@@ -33,35 +36,34 @@ def _mean(xs: List[float]) -> float:
     return sum(xs) / len(xs)
 
 
+# ----------------------------
+# Series builders
+# ----------------------------
+
 def _down_series(rows: List[Dict[str, str]]) -> Tuple[List[int], List[float]]:
     """
     Down CSV fields (expected):
       phase, episode, agent, episode_reward, [global_episode], ...
-    We aggregate by (phase, episode) averaging across agents.
+    Aggregate by (phase, episode), averaging across agents.
     """
     bucket: Dict[Tuple[int, int], List[float]] = defaultdict(list)
+
     for r in rows:
         try:
-            phase = int(r.get("phase", "0"))
-            ep = int(r.get("episode", "0"))
-            rew = float(r.get("episode_reward", "nan"))
+            phase = int(r["phase"])
+            ep = int(r["episode"])
+            rew = float(r["episode_reward"])
         except Exception:
             continue
         bucket[(phase, ep)].append(rew)
 
-    # Sort by phase then episode, make a single x-axis that is "global episode"
-    keys = sorted(bucket.keys())
     xs, ys = [], []
     g = 0
-    last_phase = None
-    for (phase, ep) in keys:
-        if last_phase is None:
-            last_phase = phase
-        if phase != last_phase:
-            last_phase = phase
+    for key in sorted(bucket.keys()):
         xs.append(g)
-        ys.append(_mean(bucket[(phase, ep)]))
+        ys.append(_mean(bucket[key]))
         g += 1
+
     return xs, ys
 
 
@@ -69,55 +71,101 @@ def _antidown_series(rows: List[Dict[str, str]]) -> Tuple[List[int], List[float]
     """
     AntiDown CSV fields (expected):
       phase, episode, agent, total_reward, ...
-    We aggregate by (phase, episode) averaging across agents.
+    Aggregate by (phase, episode), averaging across agents.
     """
     bucket: Dict[Tuple[int, int], List[float]] = defaultdict(list)
+
     for r in rows:
         try:
-            phase = int(r.get("phase", "0"))
-            ep = int(r.get("episode", "0"))
-            rew = float(r.get("total_reward", "nan"))
+            phase = int(r["phase"])
+            ep = int(r["episode"])
+            rew = float(r["total_reward"])
         except Exception:
             continue
         bucket[(phase, ep)].append(rew)
 
-    keys = sorted(bucket.keys())
     xs, ys = [], []
     g = 0
-    last_phase = None
-    for (phase, ep) in keys:
-        if last_phase is None:
-            last_phase = phase
-        if phase != last_phase:
-            last_phase = phase
+    for key in sorted(bucket.keys()):
         xs.append(g)
-        ys.append(_mean(bucket[(phase, ep)]))
+        ys.append(_mean(bucket[key]))
         g += 1
+
     return xs, ys
 
 
+# ----------------------------
+# CLI
+# ----------------------------
+
+def _build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(description="Plot Down vs AntiDown (dual overlay).")
+    p.add_argument(
+        "--results_dir",
+        type=str,
+        default="",
+        help="Directory containing both CSVs. Defaults to TWOQUARKS_RESULTS_DIR or ./results.",
+    )
+    p.add_argument(
+        "--out",
+        type=str,
+        default="",
+        help="Output directory for PNG. Defaults to TWOQUARKS_GRAPHICS_DIR or ./graphics.",
+    )
+    p.add_argument(
+        "--prefix",
+        type=str,
+        default="dual_down_antidown",
+        help="Filename prefix for output image.",
+    )
+    return p
+
+
+# ----------------------------
+# Main
+# ----------------------------
+
 def main() -> None:
+    args = _build_parser().parse_args()
+
     root = Path(__file__).resolve().parent
-    results_dir = Path(os.environ.get("TWOQUARKS_RESULTS_DIR", (root / "results").as_posix()))
-    graphics_dir = Path(os.environ.get("TWOQUARKS_GRAPHICS_DIR", (root / "graphics").as_posix()))
+
+    results_dir = Path(
+        args.results_dir
+        or os.environ.get("TWOQUARKS_RESULTS_DIR", (root / "results").as_posix())
+    )
+
+    graphics_dir = Path(
+        args.out
+        or os.environ.get("TWOQUARKS_GRAPHICS_DIR", (root / "graphics").as_posix())
+    )
     graphics_dir.mkdir(parents=True, exist_ok=True)
 
     down_csv = results_dir / "down_paradox_tabular_results.csv"
     anti_csv = results_dir / "antidown_corrupted_valley_tabular.csv"
 
-    if not down_csv.exists() or not anti_csv.exists():
-        print(f"[dual] missing CSVs. down={down_csv.exists()} anti={anti_csv.exists()}. Skipping.")
-        return
+    if not down_csv.exists():
+        raise FileNotFoundError(f"Missing Down CSV: {down_csv}")
+    if not anti_csv.exists():
+        raise FileNotFoundError(f"Missing AntiDown CSV: {anti_csv}")
 
     down_rows = _read_csv(down_csv)
     anti_rows = _read_csv(anti_csv)
-    if not down_rows or not anti_rows:
-        print("[dual] one CSV is empty. Skipping.")
-        return
+
+    if not down_rows:
+        raise RuntimeError("Down CSV is empty.")
+    if not anti_rows:
+        raise RuntimeError("AntiDown CSV is empty.")
 
     x1, y1 = _down_series(down_rows)
     x2, y2 = _antidown_series(anti_rows)
 
+    if len(x1) != len(x2):
+        print(
+            f"[dual] WARNING: series length mismatch "
+            f"(Down={len(x1)}, AntiDown={len(x2)})"
+        )
+
     plt.figure(figsize=(10, 5))
     plt.plot(x1, y1, label="Down (mean across agents)")
     plt.plot(x2, y2, label="AntiDown (mean across agents)")
@@ -125,10 +173,13 @@ def main() -> None:
     plt.ylabel("Episode return")
     plt.title("Down vs AntiDown — Tabular returns (mean across agents)")
     plt.legend()
-    out = graphics_dir / "dual_down_antidown_episode_return.png"
     plt.tight_layout()
-    plt.savefig(out, dpi=160)
-    print(f"[dual] saved {out}")
+
+    out_path = graphics_dir / f"{args.prefix}_episode_return.png"
+    plt.savefig(out_path, dpi=160)
+    plt.close()
+
+    print(f"[dual] saved {out_path}")
 
 
 if __name__ == "__main__":