plot.py

"""
plot.py  —  Complete plotting suite with all metrics.

NEW plots (vs old version)
--------------------------
  intrinsic_fraction.png  Curiosity vs task-reward share over training
  icm_cutoff_marker       Vertical line on all plots at ICM cutoff step
  maze_eval_bar.png       Post-training deterministic Maze win-rate bar chart
  seed error bands        All learning curves show mean ± std across seeds

Updated plots
-------------
  win_rate.png          Most interpretable: % episodes solved (Maze: return>0)
  reward_per_step.png   Efficiency: reward / action
  steps_to_solve.png    Winning episodes only — lower = faster maze solves
  kl_clipfrac.png       PPO stability: KL divergence + clip fraction
  grad_norm.png         L2 gradient norm
  icm_loss_split.png    Forward vs inverse loss shown separately
  summary.png           8-panel paper-ready overview figure

Multi-seed aggregation
----------------------
  plot.py auto-detects all CSVs for each condition across seeds.
  All curves show mean ± std across seeds.

Multi-eta comparison mode
-------------------------
  Pass --multi_dir <log_dir>:<eta_label> (repeatable) to overlay all etas on
  shared axes for direct comparison. Each eta gets its own colour; conditions
  (Vanilla / ICM / ICM-cutoff) are distinguished by linestyle.
  Example:
    python plot.py \\
      --multi_dir logs/704750_eta0.001:0.001 \\
      --multi_dir logs/707948_eta0.005:0.005 \\
      --multi_dir logs/704749_eta0.01:0.01 \\
      --out_dir plots/multi_eta

Usage
-----
  python plot.py                                        # auto-detect latest CSVs
  python plot.py --vanilla logs/X.csv --icm logs/Y.csv
  python plot.py ... --transfer logs/transfer_results.csv
"""

import os, glob, argparse
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

# ── Colour palette ────────────────────────────────────────────────────────────
C_VAN = "#4C72B0"   # blue  — vanilla PPO
C_ICM = "#DD8452"   # orange — PPO + ICM
C_CUT = "#55A868"   # green  — ICM cutoff
C_FWD = "#55A868"   # green  — forward model
C_INV = "#C44E52"   # red    — inverse model
ALPHA = 0.20
W     = 20          # smoothing window

# Colours for multi-eta comparison (one colour per eta value)
ETA_COLORS = ["#E377C2", "#17BECF", "#BCBD22", "#9467BD", "#8C564B"]
# Linestyles for condition within each eta
LS_VAN = "-"
LS_ICM = "--"
LS_CUT = ":"

plt.rcParams.update({"font.size": 11, "axes.titlesize": 12,
                      "axes.labelsize": 11, "legend.fontsize": 10,
                      "figure.dpi": 150})


# ── Data loading helpers ───────────────────────────────────────────────────────

def load(p):
    return pd.read_csv(p, na_values=[""])


def find_latest(d, tag):
    """Find latest CSV matching tag, preferring maze source logs over finetune/heist."""
    files = glob.glob(os.path.join(d, f"*{tag}*.csv"))
    if not files:
        return None
    source = [f for f in files if "finetune" not in f and "heist" not in f
              and "transfer" not in f and "zero_shot" not in f
              and "maze_eval" not in f and "master" not in f]
    pool = source if source else files
    return max(pool, key=os.path.getmtime)


def find_all_seed_csvs(log_dir, tag):
    """
    Return all source-training CSVs for a condition tag across all seeds.
    Excludes finetune, heist, transfer, zero_shot, maze_eval, and master CSVs.
    """
    files = glob.glob(os.path.join(log_dir, f"*{tag}*.csv"))
    return [f for f in files
            if "finetune" not in f and "heist" not in f
            and "transfer" not in f and "zero_shot" not in f
            and "maze_eval" not in f and "master" not in f]


def load_all_seeds(log_dir, tag):
    """Load all seed CSVs for a condition tag. Returns list of DataFrames."""
    paths = find_all_seed_csvs(log_dir, tag)
    return [load(p) for p in paths] if paths else []


# ── Smoothing helpers ─────────────────────────────────────────────────────────

def sm(s, w=W):
    return s.rolling(window=w, min_periods=1).mean()


def aggregate_seeds(dfs, x_col, y_col, w=W):
    """
    Aggregate y_col across multiple seed DataFrames aligned on x_col.
    Returns (xs, smoothed_mean, smoothed_std) as numpy arrays.
    Uses pandas merge on x_col — works when all seeds have identical step counts
    (same batch_size × num_iterations). Falls back gracefully to single-seed.
    """
    valid_dfs = [df[[x_col, y_col]].dropna()
                   .set_index(x_col)
                   .loc[lambda d: ~d.index.duplicated(keep="last")]
                 for df in dfs if x_col in df.columns and y_col in df.columns]
    if not valid_dfs:
        return np.array([]), np.array([]), np.array([])

    merged = pd.concat(valid_dfs, axis=1, join="outer").sort_index()
    xs    = merged.index.to_numpy(dtype=np.float64)
    means = merged.mean(axis=1).to_numpy(dtype=np.float64)
    stds  = merged.std(axis=1).fillna(0).to_numpy(dtype=np.float64)

    # Smooth both mean and std
    smooth_m = pd.Series(means).rolling(w, min_periods=1).mean().to_numpy()
    smooth_s = pd.Series(stds).rolling(w, min_periods=1).mean().to_numpy()
    return xs, smooth_m, smooth_s


def shade_seeds(ax, dfs, x_col, y_col, color, label, w=W):
    """Plot mean ± std across seeds with shaded band."""
    xs, means, stds = aggregate_seeds(dfs, x_col, y_col, w)
    if xs.size == 0:
        return
    ax.plot(xs, means, color=color, label=label, linewidth=2)
    ax.fill_between(xs, means - stds, means + stds, color=color, alpha=ALPHA)


def shade(ax, df, x, y, color, label, w=W):
    """Single-dataframe shade (used for ICM-specific plots)."""
    if y not in df.columns:
        return
    sub = df[[x, y]].dropna()
    if sub.empty:
        return
    xs, ys = sub[x], sub[y]
    s_mean = sm(ys, w)
    s_std  = sm(ys, w).rolling(w, min_periods=1).std().fillna(0)  # smoothed std of smoothed series
    ax.plot(xs, s_mean, color=color, label=label, linewidth=2)
    ax.fill_between(xs, s_mean - s_std, s_mean + s_std, color=color, alpha=ALPHA)


def save(fig, path):
    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
    fig.savefig(path, bbox_inches="tight")
    plt.close(fig)
    print(f"Saved -> {path}")


# ── ICM cutoff detection ──────────────────────────────────────────────────────

def find_icm_cutoff_step(dfs):
    """
    Return the global_step where icm_active first becomes 0, averaged across
    seeds. Returns None if no cutoff is detected.
    """
    cutoff_steps = []
    for df in dfs:
        if "icm_active" not in df.columns:
            continue
        active = df[["global_step", "icm_active"]].dropna()
        off    = active[active["icm_active"] == 0]
        if not off.empty:
            cutoff_steps.append(off["global_step"].min())
    return int(np.mean(cutoff_steps)) if cutoff_steps else None


def add_cutoff_vline(ax, cutoff_step, label="ICM cutoff"):
    """Add a vertical dashed line marking the ICM→PPO transition."""
    if cutoff_step is not None:
        ax.axvline(x=cutoff_step, color="gray", linestyle="--",
                   linewidth=1.2, alpha=0.8, label=label)


# ── Plot functions ────────────────────────────────────────────────────────────

def plot_episodic_return(v_dfs, i_dfs, c_dfs, cutoff, out):
    fig, ax = plt.subplots(figsize=(10, 4))
    shade_seeds(ax, v_dfs, "global_step", "episodic_return", C_VAN, "Vanilla PPO")
    shade_seeds(ax, i_dfs, "global_step", "episodic_return", C_ICM, "PPO + ICM")
    shade_seeds(ax, c_dfs, "global_step", "episodic_return", C_CUT, "PPO + ICM cutoff")
    add_cutoff_vline(ax, cutoff)
    ax.set_title("Episodic Return — Maze (Source Task)")
    ax.set_xlabel("Environment Steps"); ax.set_ylabel("Episodic Return")
    ax.legend(); ax.grid(alpha=0.3); fig.tight_layout(); save(fig, out)


def plot_win_rate(v_dfs, i_dfs, c_dfs, cutoff, out):
    fig, ax = plt.subplots(figsize=(10, 4))
    shade_seeds(ax, v_dfs, "global_step", "win_rate", C_VAN, "Vanilla PPO")
    shade_seeds(ax, i_dfs, "global_step", "win_rate", C_ICM, "PPO + ICM")
    shade_seeds(ax, c_dfs, "global_step", "win_rate", C_CUT, "PPO + ICM cutoff")
    add_cutoff_vline(ax, cutoff)
    ax.set_title("Win Rate — Rolling 100 Episodes (Maze)\n"
                 "Binary: return = 10 → solved, return = 0 → timeout")
    ax.set_xlabel("Environment Steps"); ax.set_ylabel("Win Rate (0 – 1)")
    ax.set_ylim(-0.05, 1.05)
    ax.legend(); ax.grid(alpha=0.3); fig.tight_layout(); save(fig, out)


def plot_reward_per_step(v_dfs, i_dfs, c_dfs, cutoff, out):
    fig, ax = plt.subplots(figsize=(10, 4))
    shade_seeds(ax, v_dfs, "global_step", "reward_per_step", C_VAN, "Vanilla PPO")
    shade_seeds(ax, i_dfs, "global_step", "reward_per_step", C_ICM, "PPO + ICM")
    shade_seeds(ax, c_dfs, "global_step", "reward_per_step", C_CUT, "PPO + ICM cutoff")
    add_cutoff_vline(ax, cutoff)
    ax.set_title("Reward per Step — Efficiency Metric (Maze)\n"
                 "Higher = agent solves maze with fewer actions")
    ax.set_xlabel("Environment Steps"); ax.set_ylabel("Reward / Step")
    ax.legend(); ax.grid(alpha=0.3); fig.tight_layout(); save(fig, out)


def plot_steps_to_solve(v_dfs, i_dfs, c_dfs, cutoff, out):
    fig, ax = plt.subplots(figsize=(10, 4))
    shade_seeds(ax, v_dfs, "global_step", "steps_to_solve", C_VAN, "Vanilla PPO")
    shade_seeds(ax, i_dfs, "global_step", "steps_to_solve", C_ICM, "PPO + ICM")
    shade_seeds(ax, c_dfs, "global_step", "steps_to_solve", C_CUT, "PPO + ICM cutoff")
    add_cutoff_vline(ax, cutoff)
    ax.set_title("Steps to Solve — Winning Episodes Only (Maze)\n"
                 "Lower = more direct path. Blank on timeouts = auto-dropped")
    ax.set_xlabel("Environment Steps"); ax.set_ylabel("Episode Length (wins only)")
    ax.legend(); ax.grid(alpha=0.3); fig.tight_layout(); save(fig, out)


def plot_losses(v_dfs, i_dfs, c_dfs, cutoff, out):
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    for ax, col, title in zip(axes,
            ["policy_loss", "value_loss", "entropy_loss"],
            ["Policy Loss", "Value Loss", "Entropy"]):
        shade_seeds(ax, v_dfs, "global_step", col, C_VAN, "Vanilla PPO")
        shade_seeds(ax, i_dfs, "global_step", col, C_ICM, "PPO + ICM")
        shade_seeds(ax, c_dfs, "global_step", col, C_CUT, "ICM cutoff")
        add_cutoff_vline(ax, cutoff)
        ax.set_title(title); ax.set_xlabel("Steps")
        ax.legend(); ax.grid(alpha=0.3)
    fig.suptitle("PPO Losses — Maze"); fig.tight_layout(); save(fig, out)


def plot_kl_clipfrac(v_dfs, i_dfs, c_dfs, cutoff, out):
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))
    shade_seeds(axes[0], v_dfs, "global_step", "approx_kl", C_VAN, "Vanilla PPO")
    shade_seeds(axes[0], i_dfs, "global_step", "approx_kl", C_ICM, "PPO + ICM")
    shade_seeds(axes[0], c_dfs, "global_step", "approx_kl", C_CUT, "ICM cutoff")
    add_cutoff_vline(axes[0], cutoff)
    axes[0].set_title("Approx KL Divergence")
    axes[0].set_xlabel("Steps"); axes[0].legend(); axes[0].grid(alpha=0.3)

    shade_seeds(axes[1], v_dfs, "global_step", "clipfrac", C_VAN, "Vanilla PPO")
    shade_seeds(axes[1], i_dfs, "global_step", "clipfrac", C_ICM, "PPO + ICM")
    shade_seeds(axes[1], c_dfs, "global_step", "clipfrac", C_CUT, "ICM cutoff")
    add_cutoff_vline(axes[1], cutoff)
    axes[1].set_title("Clip Fraction\n(ideal range: 0.05 – 0.20)")
    axes[1].set_xlabel("Steps"); axes[1].legend(); axes[1].grid(alpha=0.3)
    fig.suptitle("PPO Stability Metrics — Maze"); fig.tight_layout(); save(fig, out)


def plot_grad_norm(v_dfs, i_dfs, c_dfs, cutoff, out):
    fig, ax = plt.subplots(figsize=(10, 4))
    shade_seeds(ax, v_dfs, "global_step", "grad_norm", C_VAN, "Vanilla PPO")
    shade_seeds(ax, i_dfs, "global_step", "grad_norm", C_ICM, "PPO + ICM")
    shade_seeds(ax, c_dfs, "global_step", "grad_norm", C_CUT, "ICM cutoff")
    add_cutoff_vline(ax, cutoff)
    ax.set_title("Policy Gradient Norm — Maze")
    ax.set_xlabel("Steps"); ax.set_ylabel("L2 Gradient Norm")
    ax.legend(); ax.grid(alpha=0.3); fig.tight_layout(); save(fig, out)


def plot_icm_rewards(i_dfs, c_dfs, cutoff, out):
    """Intrinsic vs extrinsic reward breakdown for ICM and ICM-cutoff runs."""
    fig, axes = plt.subplots(1, 2, figsize=(12, 4))

    # Left: intrinsic reward over time
    for dfs, color, label in [(i_dfs, C_ICM, "ICM"), (c_dfs, C_CUT, "ICM cutoff")]:
        shade_seeds(axes[0], dfs, "global_step", "intrinsic_reward_mean", color, label)
    add_cutoff_vline(axes[0], cutoff)
    axes[0].set_title("Intrinsic Reward — Curiosity Signal\n"
                       "Decaying = ICM learning to predict next state ✓\n"
                       "Drops to 0 after cutoff for ICM cutoff condition")
    axes[0].set_xlabel("Steps"); axes[0].set_ylabel("Intrinsic Reward")
    axes[0].legend(); axes[0].grid(alpha=0.3)

    # Right: extrinsic reward
    for dfs, color, label in [(i_dfs, C_ICM, "PPO+ICM"), (c_dfs, C_CUT, "ICM cutoff")]:
        shade_seeds(axes[1], dfs, "global_step", "extrinsic_reward_mean", color, label)
    add_cutoff_vline(axes[1], cutoff)
    axes[1].set_title("Extrinsic Reward — Raw Game Signal\n"
                       "Should rise as agent learns to solve the maze")
    axes[1].set_xlabel("Steps"); axes[1].set_ylabel("Normalised Game Reward / Step")
    axes[1].legend(); axes[1].grid(alpha=0.3)

    fig.suptitle("ICM Reward Breakdown — Maze")
    fig.tight_layout(); save(fig, out)


def plot_icm_loss_split(i_dfs, out):
    """Forward vs inverse ICM loss."""
    # Use first seed only for ICM-specific loss curves (same model, not much seed variance)
    if not i_dfs:
        return
    i = i_dfs[0]
    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
    shade(axes[0], i, "global_step", "icm_loss",         C_ICM, "Total")
    shade(axes[1], i, "global_step", "icm_forward_loss", C_FWD, "Forward Loss")
    shade(axes[2], i, "global_step", "icm_inverse_loss", C_INV, "Inverse Loss")
    axes[0].set_title("Total ICM Loss\n(β·fwd + (1-β)·inv, β=0.2)")
    axes[1].set_title("Forward Model Loss\n(predicts next feature vector; should decay slowly)")
    axes[2].set_title("Inverse Model Loss\n(predicts action; baseline=ln(15)≈2.71)")
    for ax in axes:
        ax.set_xlabel("Steps"); ax.legend(); ax.grid(alpha=0.3)
    fig.suptitle("ICM Loss Breakdown — Maze"); fig.tight_layout(); save(fig, out)


def plot_intrinsic_fraction(i_dfs, c_dfs, cutoff, out):
    """
    Intrinsic reward fraction = mean_i / (|mean_i| + |mean_e|).
    Shows how much the agent is driven by curiosity vs task reward over time.
    Fraction → 0 for ICM cutoff runs once cutoff is reached.
    """
    fig, ax = plt.subplots(figsize=(10, 4))
    shade_seeds(ax, i_dfs, "global_step", "intrinsic_reward_fraction",
                C_ICM, "PPO + ICM")
    shade_seeds(ax, c_dfs, "global_step", "intrinsic_reward_fraction",
                C_CUT, "PPO + ICM cutoff")
    add_cutoff_vline(ax, cutoff, label=f"ICM cutoff (~{cutoff//1_000_000}M steps)" if cutoff else "ICM cutoff")
    ax.set_title("Intrinsic Reward Fraction — Maze\n"
                 "mean_intrinsic / (|mean_intrinsic| + |mean_extrinsic|)\n"
                 "Shows curiosity vs task-reward balance; drops to 0 after cutoff")
    ax.set_xlabel("Environment Steps")
    ax.set_ylabel("Intrinsic Reward Fraction (0 – 1)")
    ax.set_ylim(-0.05, 1.05)
    ax.legend(); ax.grid(alpha=0.3); fig.tight_layout(); save(fig, out)


def plot_maze_eval_bar(v_dfs, i_dfs, c_dfs, out):
    """
    Side-by-side bar chart of post-training deterministic Maze eval win-rate.
    Bars show mean across seeds; error bars show std.
    """
    def get_eval_stats(dfs, col):
        vals = []
        for df in dfs:
            if col in df.columns:
                v = df[col].dropna()
                if not v.empty:
                    vals.append(v.iloc[-1])
        if not vals:
            return None, None
        return float(np.mean(vals)), float(np.std(vals)) if len(vals) > 1 else 0.0

    conditions = []
    for label, dfs, color in [("Vanilla", v_dfs, C_VAN),
                               ("ICM", i_dfs, C_ICM),
                               ("ICM cutoff", c_dfs, C_CUT)]:
        mean_wr, std_wr = get_eval_stats(dfs, "maze_eval_win_rate")
        mean_r,  std_r  = get_eval_stats(dfs, "maze_eval_mean_reward")
        if mean_wr is not None:
            conditions.append((label, mean_wr, std_wr, mean_r, std_r, color))

    if not conditions:
        print("  No maze_eval_win_rate data found — skipping maze eval bar chart")
        return

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    labels = [c[0] for c in conditions]
    x      = np.arange(len(labels))

    # Win rate
    ax = axes[0]
    wrs   = [c[1] for c in conditions]
    wrstd = [c[2] for c in conditions]
    colors = [c[5] for c in conditions]
    bars = ax.bar(x, wrs, yerr=wrstd, color=colors, capsize=6,
                  edgecolor="black", linewidth=0.8)
    for bar, w in zip(bars, wrs):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.005,
                f"{w:.1%}", ha="center", va="bottom", fontsize=10)
    ax.set_xticks(x); ax.set_xticklabels(labels)
    ax.set_title("Maze Eval Win Rate\n(deterministic, 200 episodes, post-training)")
    ax.set_ylabel("Win Rate"); ax.set_ylim(0, min(max(wrs)*1.4 + 0.05, 1.05))
    ax.grid(axis="y", alpha=0.3)

    # Mean reward
    ax = axes[1]
    mrs   = [c[3] for c in conditions]
    mrstd = [c[4] for c in conditions]
    bars = ax.bar(x, mrs, yerr=mrstd, color=colors, capsize=6,
                  edgecolor="black", linewidth=0.8)
    for bar, m in zip(bars, mrs):
        ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.01,
                f"{m:.3f}", ha="center", va="bottom", fontsize=10)
    ax.set_xticks(x); ax.set_xticklabels(labels)
    ax.set_title("Maze Eval Mean Reward\n(deterministic, 200 episodes, post-training)")
    ax.set_ylabel("Mean Reward"); ax.grid(axis="y", alpha=0.3)

    fig.suptitle("Post-Training Deterministic Maze Evaluation — All Conditions",
                 fontsize=13)
    fig.tight_layout(); save(fig, out)


def plot_transfer(t, out):
    """Bar chart: mean reward (left) and win rate (right) for all conditions."""
    if t is None:
        return
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    labels = t["run_tag"].tolist()
    colors = [C_CUT if "cutoff" in l
              else (C_ICM if "icm" in l
              else ("#888888" if "scratch" in l
              else C_VAN)) for l in labels]
    x = np.arange(len(labels))

    # Left: mean reward
    ax = axes[0]
    means = t["mean_reward"].tolist()
    stds  = t["std_reward"].tolist()
    bars  = ax.bar(x, means, yerr=stds, color=colors,
                   capsize=5, edgecolor="black", linewidth=0.8)
    for bar, m in zip(bars, means):
        ax.text(bar.get_x()+bar.get_width()/2, bar.get_height()+0.05,
                f"{m:.3f}", ha="center", va="bottom", fontsize=9)
    ax.set_xticks(x); ax.set_xticklabels(labels, rotation=25, ha="right")
    ax.set_title("Mean Episode Reward")
    ax.set_ylabel("Mean Reward"); ax.grid(axis="y", alpha=0.3)

    # Right: win rate
    ax = axes[1]
    wrs = t["win_rate"].tolist() if "win_rate" in t.columns else [0]*len(labels)
    bars = ax.bar(x, wrs, color=colors, edgecolor="black", linewidth=0.8)
    for bar, w in zip(bars, wrs):
        ax.text(bar.get_x()+bar.get_width()/2, bar.get_height()+0.005,
                f"{w:.1%}", ha="center", va="bottom", fontsize=9)
    ax.set_xticks(x); ax.set_xticklabels(labels, rotation=25, ha="right")
    ax.set_title("Win Rate (episodes with reward > 0)")
    ax.set_ylabel("Win Rate"); ax.set_ylim(0, max(wrs)*1.3 + 0.01)
    ax.grid(axis="y", alpha=0.3)

    fig.suptitle("Transfer to Heist — All Conditions")
    fig.tight_layout(); save(fig, out)


def plot_explained_variance(v_dfs, i_dfs, c_dfs, cutoff, out):
    fig, ax = plt.subplots(figsize=(10, 4))
    shade_seeds(ax, v_dfs, "global_step", "explained_variance", C_VAN, "Vanilla PPO")
    shade_seeds(ax, i_dfs, "global_step", "explained_variance", C_ICM, "PPO + ICM")
    shade_seeds(ax, c_dfs, "global_step", "explained_variance", C_CUT, "ICM cutoff")
    add_cutoff_vline(ax, cutoff)
    ax.set_title("Explained Variance — Maze\n"
                 "1.0 = critic perfectly predicts returns; <0 = worse than mean")
    ax.set_xlabel("Environment Steps"); ax.set_ylabel("Explained Variance")
    ax.axhline(y=0, color="gray", linestyle="--", alpha=0.5)
    ax.legend(); ax.grid(alpha=0.3); fig.tight_layout(); save(fig, out)


def plot_summary(v_dfs, i_dfs, c_dfs, cutoff, log_dir, out):
    """Paper-ready 8-panel overview: source training + transfer results."""
    fig, axes = plt.subplots(2, 4, figsize=(22, 8))
    ax = axes.flatten()

    # Row 1: Source task metrics
    shade_seeds(ax[0], v_dfs, "global_step", "episodic_return", C_VAN, "Vanilla")
    shade_seeds(ax[0], i_dfs, "global_step", "episodic_return", C_ICM, "ICM")
    shade_seeds(ax[0], c_dfs, "global_step", "episodic_return", C_CUT, "ICM cutoff")
    add_cutoff_vline(ax[0], cutoff)
    ax[0].set_title("Episodic Return (Source: Maze)")
    ax[0].set_xlabel("Steps"); ax[0].legend(); ax[0].grid(alpha=0.3)

    shade_seeds(ax[1], v_dfs, "global_step", "win_rate", C_VAN, "Vanilla")
    shade_seeds(ax[1], i_dfs, "global_step", "win_rate", C_ICM, "ICM")
    shade_seeds(ax[1], c_dfs, "global_step", "win_rate", C_CUT, "ICM cutoff")
    add_cutoff_vline(ax[1], cutoff)
    ax[1].set_title("Win Rate (Source)"); ax[1].set_xlabel("Steps")
    ax[1].set_ylim(-0.05, 1.05); ax[1].legend(); ax[1].grid(alpha=0.3)

    shade_seeds(ax[2], i_dfs, "global_step", "intrinsic_reward_mean", C_ICM, "Intrinsic ICM")
    shade_seeds(ax[2], c_dfs, "global_step", "intrinsic_reward_mean", C_CUT, "Intrinsic cutoff")
    shade_seeds(ax[2], i_dfs, "global_step", "extrinsic_reward_mean", C_VAN, "Extrinsic (raw game)")
    add_cutoff_vline(ax[2], cutoff)
    ax[2].set_title("ICM: Intrinsic vs Extrinsic")
    ax[2].set_xlabel("Steps"); ax[2].legend(); ax[2].grid(alpha=0.3)

    shade_seeds(ax[3], i_dfs, "global_step", "intrinsic_reward_fraction", C_ICM, "ICM")
    shade_seeds(ax[3], c_dfs, "global_step", "intrinsic_reward_fraction", C_CUT, "ICM cutoff")
    add_cutoff_vline(ax[3], cutoff)
    ax[3].set_title("Intrinsic Reward Fraction")
    ax[3].set_xlabel("Steps"); ax[3].set_ylim(-0.05, 1.05)
    ax[3].legend(); ax[3].grid(alpha=0.3)

    # Row 2: Losses + Transfer
    shade_seeds(ax[4], v_dfs, "global_step", "value_loss", C_VAN, "Vanilla")
    shade_seeds(ax[4], i_dfs, "global_step", "value_loss", C_ICM, "ICM")
    shade_seeds(ax[4], c_dfs, "global_step", "value_loss", C_CUT, "ICM cutoff")
    add_cutoff_vline(ax[4], cutoff)
    ax[4].set_title("Value Loss")
    ax[4].set_xlabel("Steps"); ax[4].legend(); ax[4].grid(alpha=0.3)

    if i_dfs:
        shade(ax[5], i_dfs[0], "global_step", "icm_forward_loss", C_FWD, "Forward Model")
        shade(ax[5], i_dfs[0], "global_step", "icm_inverse_loss", C_INV, "Inverse Model")
    ax[5].set_title("ICM Loss Split")
    ax[5].set_xlabel("Steps"); ax[5].legend(); ax[5].grid(alpha=0.3)

    # Finetune curves in summary panels 6-7
    ft_files = sorted(glob.glob(os.path.join(log_dir, "heist__finetune*__*.csv")))
    for fp in ft_files:
        name  = os.path.basename(fp).replace(".csv", "")
        parts = name.split("__")
        nl    = parts[1].replace("finetune", "FT") if len(parts) > 1 else "?"
        ag    = parts[2].upper() if len(parts) > 2 else "?"
        label = f"{ag} {nl}"
        color = C_CUT if "cutoff" in name else (C_ICM if "icm" in name else C_VAN)
        ls    = "-" if "1L" in name else "--"
        df    = load(fp)
        n6 = len(ax[6].lines)
        shade(ax[6], df, "global_step", "episodic_return", color, label, w=50)
        if len(ax[6].lines) > n6:
            ax[6].lines[-1].set_linestyle(ls)
        n7 = len(ax[7].lines)
        shade(ax[7], df, "global_step", "win_rate", color, label, w=50)
        if len(ax[7].lines) > n7:
            ax[7].lines[-1].set_linestyle(ls)

    ax[6].set_title("Finetune: Episodic Return")
    ax[6].set_xlabel("Steps"); ax[6].legend(fontsize=8); ax[6].grid(alpha=0.3)
    ax[7].set_title("Finetune: Win Rate")
    ax[7].set_xlabel("Steps"); ax[7].set_ylim(-0.05, 1.05)
    ax[7].legend(fontsize=8); ax[7].grid(alpha=0.3)

    fig.suptitle("Training Summary — Maze Source Task + Heist Transfer  |  ICM Study",
                 fontsize=14)
    fig.tight_layout(); save(fig, out)


def plot_finetune_curves(log_dir, out):
    """Plot episodic return and win rate for all finetune runs."""
    ft_files = sorted(glob.glob(os.path.join(log_dir, "heist__finetune*__*.csv")))
    if not ft_files:
        return

    fig, axes = plt.subplots(1, 2, figsize=(14, 5))

    for fp in ft_files:
        name     = os.path.basename(fp).replace(".csv", "")
        parts    = name.split("__")
        n_layers = parts[1].replace("finetune", "FT") if len(parts) > 1 else "?"
        agent    = parts[2].upper() if len(parts) > 2 else "?"
        label    = f"{agent} {n_layers}"
        color    = C_CUT if "cutoff" in name else (C_ICM if "icm" in name else C_VAN)
        ls       = "-" if "1L" in name else "--"

        df = load(fp)
        n0 = len(axes[0].lines)
        shade(axes[0], df, "global_step", "episodic_return", color, label, w=50)
        if len(axes[0].lines) > n0:
            axes[0].lines[-1].set_linestyle(ls)
        n1 = len(axes[1].lines)
        shade(axes[1], df, "global_step", "win_rate", color, label, w=50)
        if len(axes[1].lines) > n1:
            axes[1].lines[-1].set_linestyle(ls)

    axes[0].set_title("Episodic Return during Finetune")
    axes[0].set_xlabel("Steps"); axes[0].set_ylabel("Episodic Return")
    axes[0].legend(); axes[0].grid(alpha=0.3)

    axes[1].set_title("Win Rate during Finetune")
    axes[1].set_xlabel("Steps"); axes[1].set_ylabel("Win Rate")
    axes[1].set_ylim(-0.05, 1.05)
    axes[1].legend(); axes[1].grid(alpha=0.3)

    fig.suptitle("Finetune Transfer to Heist — Learning Curves\n"
                 "Solid = 1 layer unfrozen, Dashed = 2 layers unfrozen")
    fig.tight_layout(); save(fig, out)


def plot_zero_shot_histogram(log_dir, out):
    """Per-episode reward distribution from zero_shot_episodes.csv."""
    zs_path = os.path.join(log_dir, "zero_shot_episodes.csv")
    if not os.path.exists(zs_path):
        return

    zs   = pd.read_csv(zs_path)
    tags = zs["run_tag"].unique()
    fig, axes = plt.subplots(1, len(tags), figsize=(6*len(tags), 4), squeeze=False)

    for ax, tag in zip(axes[0], tags):
        sub   = zs[zs["run_tag"] == tag]
        color = C_CUT if "cutoff" in tag else (C_ICM if "icm" in tag else C_VAN)
        ax.hist(sub["reward"], bins=30, color=color, edgecolor="black",
                alpha=0.75, linewidth=0.8)
        wr = sub["win"].mean()
        ax.set_title(f"{tag}\nmean={sub['reward'].mean():.3f}  "
                     f"wr={wr:.1%}  n={len(sub)}")
        ax.set_xlabel("Episode Reward"); ax.set_ylabel("Count")
        ax.grid(axis="y", alpha=0.3)

    fig.suptitle("Zero-Shot Transfer — Episode Reward Distribution")
    fig.tight_layout(); save(fig, out)


# ── Multi-eta comparison plots ────────────────────────────────────────────────

def _multi_eta_entries(multi_dirs):
    """
    Parse list of 'log_dir:eta_label' strings.
    For each entry load vanilla / icm / icm_cutoff DataFrames.
    Returns list of dicts: {eta, log_dir, v_dfs, i_dfs, c_dfs, color, cutoff}.
    """
    entries = []
    for idx, spec in enumerate(multi_dirs):
        if ":" not in spec:
            print(f"WARNING: --multi_dir '{spec}' has no ':eta' suffix — skipping")
            continue
        log_dir, eta = spec.rsplit(":", 1)
        color = ETA_COLORS[idx % len(ETA_COLORS)]
        v_dfs = load_all_seeds(log_dir, "vanilla")
        i_dfs = load_all_seeds(log_dir, f"icm_eta{eta}")
        c_dfs = load_all_seeds(log_dir, f"icm_cutoff_eta{eta}")
        cutoff = find_icm_cutoff_step(c_dfs)
        entries.append(dict(eta=eta, log_dir=log_dir, v_dfs=v_dfs, i_dfs=i_dfs,
                            c_dfs=c_dfs, color=color, cutoff=cutoff))
        print(f"  η={eta}: vanilla={len(v_dfs)} ICM={len(i_dfs)} cutoff={len(c_dfs)} "
              f"cutoff_step={cutoff}")
    return entries


def _shade_multi(ax, dfs, x_col, y_col, color, label, ls):
    """shade_seeds wrapper that applies a specific linestyle."""
    xs, means, stds = aggregate_seeds(dfs, x_col, y_col)
    if xs.size == 0:
        return
    line, = ax.plot(xs, means, color=color, label=label, linewidth=2, linestyle=ls)
    ax.fill_between(xs, means - stds, means + stds, color=color, alpha=ALPHA)


def plot_multi_eta_curve(entries, x_col, y_col, title, ylabel, out, ylim=None):
    """Generic multi-eta learning curve: one colour per eta, linestyle per condition."""
    fig, ax = plt.subplots(figsize=(12, 5))
    for e in entries:
        eta, color = e["eta"], e["color"]
        _shade_multi(ax, e["v_dfs"], x_col, y_col, color, f"Vanilla η={eta}", LS_VAN)
        _shade_multi(ax, e["i_dfs"], x_col, y_col, color, f"ICM η={eta}",     LS_ICM)
        _shade_multi(ax, e["c_dfs"], x_col, y_col, color, f"ICM-cut η={eta}", LS_CUT)
        if e["cutoff"]:
            ax.axvline(e["cutoff"], color=color, linestyle=":", linewidth=0.8, alpha=0.5)
    ax.set_title(title); ax.set_xlabel("Environment Steps"); ax.set_ylabel(ylabel)
    if ylim:
        ax.set_ylim(*ylim)
    ax.legend(fontsize=8, ncol=len(entries)); ax.grid(alpha=0.3)
    fig.tight_layout(); save(fig, out)


def plot_multi_eta_curves(entries, out_dir):
    """Generate all learning-curve comparison plots across etas."""
    specs = [
        ("win_rate",          "Win Rate — All Etas (Maze)",           "Win Rate",        (-0.05, 1.05)),
        ("episodic_return",   "Episodic Return — All Etas (Maze)",    "Episodic Return", None),
        ("reward_per_step",   "Reward/Step — All Etas (Maze)",        "Reward / Step",   None),
        ("steps_to_solve",    "Steps to Solve — All Etas (Maze)",     "Episode Length",  None),
        ("grad_norm",         "Grad Norm — All Etas (Maze)",          "L2 Grad Norm",    None),
        ("approx_kl",         "Approx KL — All Etas (Maze)",         "KL",              None),
        ("clipfrac",          "Clip Fraction — All Etas (Maze)",      "Clip Frac",       None),
        ("explained_variance","Explained Variance — All Etas (Maze)", "Expl. Variance",  None),
        ("intrinsic_reward_fraction", "Intrinsic Fraction — All Etas (Maze)",
                                      "Intr. Fraction", (-0.05, 1.05)),
    ]
    for col, title, ylabel, ylim in specs:
        plot_multi_eta_curve(entries, "global_step", col, title, ylabel,
                             f"{out_dir}/{col}.png", ylim=ylim)

    # ICM-only: intrinsic reward for ICM and cutoff conditions
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    for e in entries:
        eta, color = e["eta"], e["color"]
        _shade_multi(axes[0], e["i_dfs"], "global_step", "intrinsic_reward_mean",
                     color, f"ICM η={eta}", LS_ICM)
        _shade_multi(axes[0], e["c_dfs"], "global_step", "intrinsic_reward_mean",
                     color, f"ICM-cut η={eta}", LS_CUT)
        _shade_multi(axes[1], e["i_dfs"], "global_step", "extrinsic_reward_mean",
                     color, f"ICM η={eta}", LS_ICM)
        _shade_multi(axes[1], e["c_dfs"], "global_step", "extrinsic_reward_mean",
                     color, f"ICM-cut η={eta}", LS_CUT)
    axes[0].set_title("Intrinsic Reward — All Etas"); axes[0].set_xlabel("Steps")
    axes[0].legend(fontsize=8); axes[0].grid(alpha=0.3)
    axes[1].set_title("Extrinsic Reward — All Etas"); axes[1].set_xlabel("Steps")
    axes[1].legend(fontsize=8); axes[1].grid(alpha=0.3)
    fig.suptitle("ICM Reward Breakdown — All Etas"); fig.tight_layout()
    save(fig, f"{out_dir}/icm_rewards.png")

    # PPO Losses (3-panel)
    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
    for ax, col, title in zip(axes,
            ["policy_loss", "value_loss", "entropy_loss"],
            ["Policy Loss", "Value Loss", "Entropy"]):
        for e in entries:
            eta, color = e["eta"], e["color"]
            _shade_multi(ax, e["v_dfs"], "global_step", col, color, f"Vanilla η={eta}", LS_VAN)
            _shade_multi(ax, e["i_dfs"], "global_step", col, color, f"ICM η={eta}",     LS_ICM)
            _shade_multi(ax, e["c_dfs"], "global_step", col, color, f"ICM-cut η={eta}", LS_CUT)
        ax.set_title(title); ax.set_xlabel("Steps")
        ax.legend(fontsize=7); ax.grid(alpha=0.3)
    fig.suptitle("PPO Losses — All Etas"); fig.tight_layout()
    save(fig, f"{out_dir}/losses.png")

    # ICM loss split (forward vs inverse) — ICM condition only
    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
    for ax, col, title in zip(axes,
            ["icm_loss", "icm_forward_loss", "icm_inverse_loss"],
            ["Total ICM Loss", "Forward Loss", "Inverse Loss"]):
        for e in entries:
            eta, color = e["eta"], e["color"]
            _shade_multi(ax, e["i_dfs"], "global_step", col, color, f"ICM η={eta}", LS_ICM)
        ax.set_title(title); ax.set_xlabel("Steps")
        ax.legend(fontsize=8); ax.grid(alpha=0.3)
    fig.suptitle("ICM Loss Breakdown — All Etas"); fig.tight_layout()
    save(fig, f"{out_dir}/icm_loss_split.png")


def plot_multi_eta_bar(entries, out):
    """
    Grouped bar chart: x-axis = condition (Vanilla / ICM / ICM-cutoff),
    each group has one bar per eta, coloured by eta.
    """
    conditions = ["Vanilla", "ICM", "ICM cutoff"]
    n_eta = len(entries)
    bar_w = 0.8 / n_eta
    x = np.arange(len(conditions))

    fig, axes = plt.subplots(1, 2, figsize=(12, 5))

    for ax_idx, (metric, ylabel, title) in enumerate([
        ("maze_eval_win_rate",    "Win Rate",    "Maze Eval Win Rate"),
        ("maze_eval_mean_reward", "Mean Reward", "Maze Eval Mean Reward"),
    ]):
        ax = axes[ax_idx]
        for i, e in enumerate(entries):
            eta, color = e["eta"], e["color"]
            vals = []
            for dfs_key in ["v_dfs", "i_dfs", "c_dfs"]:
                dfs = e[dfs_key]
                seeds = []
                for df in dfs:
                    if metric in df.columns:
                        v = df[metric].dropna()
                        if not v.empty:
                            seeds.append(v.iloc[-1])
                vals.append((float(np.mean(seeds)) if seeds else 0.0,
                             float(np.std(seeds))  if len(seeds) > 1 else 0.0))
            means = [v[0] for v in vals]
            stds  = [v[1] for v in vals]
            offset = (i - n_eta/2 + 0.5) * bar_w
            bars = ax.bar(x + offset, means, bar_w * 0.9, yerr=stds,
                          label=f"η={eta}", color=color,
                          capsize=4, edgecolor="black", linewidth=0.7)
            for bar, m in zip(bars, means):
                if m > 0:
                    ax.text(bar.get_x() + bar.get_width()/2,
                            bar.get_height() + (max(means)*0.02 if max(means) > 0 else 0.005),
                            f"{m:.2f}" if metric == "maze_eval_mean_reward" else f"{m:.1%}",
                            ha="center", va="bottom", fontsize=7)
        ax.set_xticks(x); ax.set_xticklabels(conditions)
        ax.set_title(title + "\n(deterministic, 200 episodes, post-training)")
        ax.set_ylabel(ylabel); ax.legend(fontsize=9); ax.grid(axis="y", alpha=0.3)

    fig.suptitle("Post-Training Maze Evaluation — Eta Comparison", fontsize=13)
    fig.tight_layout(); save(fig, out)


def plot_multi_eta_zero_shot(entries, out):
    """Bar chart of zero-shot win rate on Heist, grouped by condition, coloured by eta."""
    conditions = ["vanilla_zero_shot", "icm_zero_shot", "icm_cutoff_zero_shot"]
    cond_labels = ["Vanilla ZS", "ICM ZS", "ICM-cut ZS"]
    n_eta = len(entries)
    bar_w = 0.8 / n_eta
    x = np.arange(len(conditions))

    fig, ax = plt.subplots(figsize=(10, 5))
    for i, e in enumerate(entries):
        eta, color = e["eta"], e["color"]
        means = []
        zs_path = os.path.join(e["log_dir"], "zero_shot_episodes.csv")
        zs = pd.read_csv(zs_path) if os.path.exists(zs_path) else pd.DataFrame()
        for cond in conditions:
            if not zs.empty and "run_tag" in zs.columns:
                sub = zs[zs["run_tag"] == cond]
                means.append(float(sub["win"].mean()) if not sub.empty and "win" in sub.columns else 0.0)
            else:
                means.append(0.0)
        offset = (i - n_eta/2 + 0.5) * bar_w
        bars = ax.bar(x + offset, means, bar_w * 0.9, label=f"η={eta}", color=color,
                      edgecolor="black", linewidth=0.7)
        for bar, m in zip(bars, means):
            if m > 0:
                ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.0005,
                        f"{m:.1%}", ha="center", va="bottom", fontsize=7)
    ax.set_xticks(x); ax.set_xticklabels(cond_labels)
    ax.set_title("Zero-Shot Transfer to Heist — Win Rate by Eta")
    ax.set_ylabel("Win Rate"); ax.legend(fontsize=9); ax.grid(axis="y", alpha=0.3)
    fig.tight_layout(); save(fig, out)


def plot_multi_eta_transfer(entries, out):
    """
    Grouped bar chart of zero-shot + finetune win rates from transfer_results.csv,
    one colour per eta, grouped by run_tag condition.
    """
    # Collect all transfer results, tagging each row with eta
    frames = []
    for e in entries:
        tp = os.path.join(e["log_dir"], "transfer_results.csv")
        if not os.path.exists(tp):
            continue
        df = pd.read_csv(tp)
        df["_eta"] = e["eta"]
        df["_color"] = e["color"]
        frames.append(df)
    if not frames:
        print("  No transfer_results.csv found — skipping transfer comparison")
        return

    combined = pd.concat(frames, ignore_index=True)
    if "win_rate" not in combined.columns:
        print("  transfer_results.csv has no win_rate column — skipping")
        return

    run_tags = combined["run_tag"].unique()
    n_eta = len(entries)
    bar_w = 0.8 / n_eta
    x = np.arange(len(run_tags))

    fig, ax = plt.subplots(figsize=(max(12, len(run_tags)*2), 5))
    for i, e in enumerate(entries):
        eta, color = e["eta"], e["color"]
        sub = combined[combined["_eta"] == eta]
        means = []
        for tag in run_tags:
            row = sub[sub["run_tag"] == tag]
            means.append(float(row["win_rate"].iloc[0]) if not row.empty else 0.0)
        offset = (i - n_eta/2 + 0.5) * bar_w
        bars = ax.bar(x + offset, means, bar_w * 0.9, label=f"η={eta}", color=color,
                      edgecolor="black", linewidth=0.7)
        for bar, m in zip(bars, means):
            if m > 0:
                ax.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 0.001,
                        f"{m:.1%}", ha="center", va="bottom", fontsize=7)
    ax.set_xticks(x); ax.set_xticklabels(run_tags, rotation=25, ha="right")
    ax.set_title("Transfer to Heist — Win Rate by Eta & Condition")
    ax.set_ylabel("Win Rate"); ax.legend(fontsize=9); ax.grid(axis="y", alpha=0.3)
    fig.tight_layout(); save(fig, out)


def plot_multi_eta_finetune(entries, out):
    """Finetune win-rate curves on Heist, one colour per eta, linestyle per condition."""
    fig, ax = plt.subplots(figsize=(12, 5))
    for e in entries:
        eta, color, log_dir = e["eta"], e["color"], e["log_dir"]
        ft_files = glob.glob(os.path.join(log_dir, "heist__finetune*__icm__*.csv"))
        ft_files += glob.glob(os.path.join(log_dir, "heist__finetune*__vanilla__*.csv"))
        ft_files += glob.glob(os.path.join(log_dir, "heist__finetune*__icm_cutoff__*.csv"))
        for fp in sorted(ft_files):
            name = os.path.basename(fp).replace(".csv", "")
            parts = name.split("__")
            agent = parts[2].upper() if len(parts) > 2 else "?"
            ls = LS_ICM if "icm_cutoff" in name else (LS_CUT if "icm" in name else LS_VAN)
            df = load(fp)
            col = "win_rate" if "win_rate" in df.columns else None
            if col is None:
                continue
            sub = df[["global_step", col]].dropna()
            if sub.empty:
                continue
            label = f"{agent} η={eta}"
            ax.plot(sub["global_step"], sm(sub[col]), color=color,
                    linestyle=ls, linewidth=1.5, label=label)
    ax.set_title("Finetune Transfer to Heist — Win Rate (All Etas)")
    ax.set_xlabel("Steps"); ax.set_ylabel("Win Rate")
    ax.set_ylim(-0.05, 1.05)
    ax.legend(fontsize=7, ncol=max(1, len(entries))); ax.grid(alpha=0.3)
    fig.tight_layout(); save(fig, out)


# ── CLI ───────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--vanilla",  default=None, help="Vanilla PPO CSV log path (single seed override)")
    parser.add_argument("--icm",      default=None, help="ICM CSV log path (single seed override)")
    parser.add_argument("--transfer", default=None, help="transfer_results.csv path")
    parser.add_argument("--log_dir",  default="logs")
    parser.add_argument("--out_dir",  default="plots")
    parser.add_argument("--icm_tag",  default="__icm__",
                        help="Tag to match ICM run CSVs. Default '__icm__' avoids matching "
                             "'icm_cutoff'. For eta sweeps use e.g. 'icm_eta0.005'.")
    parser.add_argument("--multi_dir", action="append", default=[],
                        metavar="LOG_DIR:ETA",
                        help="Multi-eta comparison mode. Repeatable. "
                             "Format: 'logs/704749_eta0.01:0.01'. "
                             "When provided, generates comparison plots overlaying all etas.")
    args = parser.parse_args()

    # ── Multi-eta comparison mode ──────────────────────────────────────────────
    if args.multi_dir:
        d = args.out_dir
        print(f"Multi-eta mode: {len(args.multi_dir)} eta(s) → {d}/")
        entries = _multi_eta_entries(args.multi_dir)
        if not entries:
            print("No valid --multi_dir entries found.")
            exit(1)
        plot_multi_eta_curves(entries, d)
        plot_multi_eta_bar(entries,        f"{d}/maze_eval_bar.png")
        plot_multi_eta_zero_shot(entries,  f"{d}/zero_shot_bar.png")
        plot_multi_eta_transfer(entries,   f"{d}/transfer_comparison.png")
        plot_multi_eta_finetune(entries,   f"{d}/finetune_curves.png")
        print(f"\nAll multi-eta plots saved to ./{d}/")
        exit(0)

    log_dir = args.log_dir
    d       = args.out_dir

    # Load all seed CSVs for each condition (multi-seed aggregation)
    v_dfs = load_all_seeds(log_dir, "vanilla")
    i_dfs = load_all_seeds(log_dir, args.icm_tag)   # configurable: default avoids 'icm_cutoff'
    c_dfs = load_all_seeds(log_dir, "icm_cutoff")

    # Single-seed overrides from CLI
    if args.vanilla:
        v_dfs = [load(args.vanilla)]
    if args.icm:
        i_dfs = [load(args.icm)]

    # No fallback to find_latest — it doesn't apply source-only filters and can
    # accidentally pick up finetune/heist CSVs (e.g. heist__finetune__icm__*.csv
    # matching an --icm_tag of "__icm__"). If filtered search found nothing, bail.

    if not v_dfs and not i_dfs:
        print("Could not auto-detect logs. Use --vanilla and --icm.")
        exit(1)

    # Ensure empty lists instead of None for safe iteration
    v_dfs = v_dfs or []
    i_dfs = i_dfs or []
    c_dfs = c_dfs or []

    tp = args.transfer or find_latest(log_dir, "transfer_results")
    t  = load(tp) if tp else None

    n_v = len(v_dfs); n_i = len(i_dfs); n_c = len(c_dfs)
    print(f"Vanilla CSVs    : {n_v} seeds")
    print(f"ICM CSVs        : {n_i} seeds")
    print(f"ICM cutoff CSVs : {n_c} seeds")
    print(f"Transfer log    : {tp or 'not found'}")

    # Detect ICM cutoff step from icm_cutoff condition data
    cutoff_step = find_icm_cutoff_step(c_dfs)
    if cutoff_step:
        print(f"ICM cutoff step : {cutoff_step:,}")

    # Dummy single-df for backwards-compat functions that need a df not list
    v = v_dfs[0] if v_dfs else pd.DataFrame()
    i = i_dfs[0] if i_dfs else pd.DataFrame()

    plot_episodic_return(v_dfs, i_dfs, c_dfs, cutoff_step, f"{d}/episodic_return.png")
    plot_win_rate(v_dfs, i_dfs, c_dfs, cutoff_step,        f"{d}/win_rate.png")
    plot_reward_per_step(v_dfs, i_dfs, c_dfs, cutoff_step, f"{d}/reward_per_step.png")
    plot_steps_to_solve(v_dfs, i_dfs, c_dfs, cutoff_step,  f"{d}/steps_to_solve.png")
    plot_losses(v_dfs, i_dfs, c_dfs, cutoff_step,          f"{d}/losses.png")
    plot_kl_clipfrac(v_dfs, i_dfs, c_dfs, cutoff_step,     f"{d}/kl_clipfrac.png")
    plot_grad_norm(v_dfs, i_dfs, c_dfs, cutoff_step,       f"{d}/grad_norm.png")
    plot_icm_rewards(i_dfs, c_dfs, cutoff_step,            f"{d}/icm_rewards.png")
    plot_icm_loss_split(i_dfs,                              f"{d}/icm_loss_split.png")
    plot_intrinsic_fraction(i_dfs, c_dfs, cutoff_step,     f"{d}/intrinsic_fraction.png")
    plot_maze_eval_bar(v_dfs, i_dfs, c_dfs,                f"{d}/maze_eval_bar.png")
    plot_transfer(t,                                        f"{d}/transfer_comparison.png")
    plot_explained_variance(v_dfs, i_dfs, c_dfs, cutoff_step, f"{d}/explained_variance.png")
    plot_summary(v_dfs, i_dfs, c_dfs, cutoff_step, log_dir, f"{d}/summary.png")
    plot_finetune_curves(log_dir,                           f"{d}/finetune_curves.png")
    plot_zero_shot_histogram(log_dir,                       f"{d}/zero_shot_histogram.png")

    print(f"\nAll plots saved to ./{d}/")

    # ── Terminal summary ──────────────────────────────────────────────────────
    print()
    print("╔══════════════════════════════════════════════════════════════╗")
    print("║          TRAINING SUMMARY  (last 100 logged episodes)        ║")
    print("╠══════════════════════════════════════════════════════════════╣")
    for tag, dfs in [("Vanilla", v_dfs), ("ICM", i_dfs), ("ICM-cut", c_dfs)]:
        if not dfs:
            continue
        ret_vals = [df["episodic_return"].dropna().tail(100).mean()
                    for df in dfs if "episodic_return" in df.columns]
        wr_vals  = [df["win_rate"].dropna().tail(100).mean()
                    for df in dfs if "win_rate" in df.columns]
        ret = float(np.mean(ret_vals)) if ret_vals else float("nan")
        wr  = float(np.mean(wr_vals))  if wr_vals  else float("nan")
        n   = len(dfs)
        print(f"║  {tag:<8}  return={ret:6.3f}   win_rate={wr:.3f}   seeds={n}           ║")
    print("╚══════════════════════════════════════════════════════════════╝")

    if t is not None:
        print()
        print("╔══════════════════════════════════════════════════════════════╗")
        print("║              TRANSFER RESULTS  (Heist)                       ║")
        print("╠══════════════════════════════════════════════════════════════╣")
        for _, row in t.iterrows():
            wr_str = f"  win_rate={row['win_rate']:.3f}" if "win_rate" in row else ""
            print(f"║  {row['run_tag']:<28} mean={row['mean_reward']:.4f}  "
                  f"std={row['std_reward']:.4f}{wr_str}  ║")
        print("╚══════════════════════════════════════════════════════════════╝")