QuarksLab / Strange /plot_dual_results.py
TwoQuarks's picture
Upload 100 files
ccba775 verified
Raw
History Blame Contribute Delete
3.06 kB
from __future__ import annotations
import sys
from pathlib import Path
import pandas as pd
import matplotlib.pyplot as plt
ROOT = Path(__file__).resolve().parent
STRANGE_CSV = ROOT / "strange" / "results" / "strange_hypothesis_lab.csv"
ANTISTRANGE_CSV = ROOT / "AntiStrange" / "results" / "antistrange_hypothesis_lab.csv"
OUT_DIR = ROOT / "graphics"
def _episode_return(df: pd.DataFrame) -> pd.DataFrame:
# Prefer explicit episode_return/return, else sum reward per episode.
if {"episode", "episode_return"}.issubset(df.columns):
s = df.groupby("episode")["episode_return"].mean().reset_index()
s = s.rename(columns={"episode_return": "value"})
return s
if {"episode", "return"}.issubset(df.columns):
s = df.groupby("episode")["return"].mean().reset_index()
s = s.rename(columns={"return": "value"})
return s
if {"episode", "reward"}.issubset(df.columns):
s = df.groupby("episode")["reward"].sum().reset_index()
s = s.rename(columns={"reward": "value"})
return s
raise KeyError(f"Need columns episode + (episode_return|return|reward). Found: {list(df.columns)}")
def _episode_mean(df: pd.DataFrame, col: str) -> pd.DataFrame | None:
if {"episode", col}.issubset(df.columns):
s = df.groupby("episode")[col].mean().reset_index()
s = s.rename(columns={col: "value"})
return s
return None
def main() -> int:
OUT_DIR.mkdir(parents=True, exist_ok=True)
if not STRANGE_CSV.exists():
print(f"[dual-plot] Missing {STRANGE_CSV}")
return 2
if not ANTISTRANGE_CSV.exists():
print(f"[dual-plot] Missing {ANTISTRANGE_CSV}")
return 2
df_s = pd.read_csv(STRANGE_CSV)
df_a = pd.read_csv(ANTISTRANGE_CSV)
# --- Dual episode return ---
s_ret = _episode_return(df_s)
a_ret = _episode_return(df_a)
plt.figure()
plt.plot(s_ret["episode"], s_ret["value"], label="Strange")
plt.plot(a_ret["episode"], a_ret["value"], label="AntiStrange")
plt.xlabel("episode")
plt.ylabel("episode_return")
plt.title("Strange vs AntiStrange — Episode Return")
plt.legend()
plt.tight_layout()
plt.savefig(OUT_DIR / "dual_episode_return.png", dpi=180)
plt.close()
# --- Optional dual stability ---
s_stab = _episode_mean(df_s, "stability")
a_stab = _episode_mean(df_a, "stability")
if s_stab is not None or a_stab is not None:
plt.figure()
if s_stab is not None:
plt.plot(s_stab["episode"], s_stab["value"], label="Strange")
if a_stab is not None:
plt.plot(a_stab["episode"], a_stab["value"], label="AntiStrange")
plt.xlabel("episode")
plt.ylabel("stability (mean per episode)")
plt.title("Strange vs AntiStrange — Stability")
plt.legend()
plt.tight_layout()
plt.savefig(OUT_DIR / "dual_stability.png", dpi=180)
plt.close()
print(f"[dual-plot] Saved plots to {OUT_DIR}")
return 0
if __name__ == "__main__":
raise SystemExit(main())