repro_bundle/scripts/make_figures.py

"""Generate Plotly HTML figures + raw CSV summaries from the bundled authors' logs
(experiment_1 = evaluated-policy state distribution) for Claims 4 & 5.
"""
import os, glob
import numpy as np
import pandas as pd
import plotly.graph_objects as go

BASE = "rational_exp_1/logs"
OUT = "figures_out"
os.makedirs(OUT, exist_ok=True)

VAR_LABEL = {
    "baseline": "baseline DQN", "ln_train": "LayerNorm", "l2_train": "L2 reg",
    "wn_train": "WeightNorm", "envrnd_train_25": "domain rand.",
    "default": "shift 0%", "eps_train_01": "shift 10%", "eps_train_03": "shift 30%",
    "eps_train_05": "shift 50%", "eps_train_07": "shift 70%",
}

def load_stats(env, experiment, y_col="rational_risk_gap", max_episode=None, smooth=3):
    exp_dir = os.path.join(BASE, env, experiment)
    out = {}
    for var in sorted(os.listdir(exp_dir)):
        vdir = os.path.join(exp_dir, var)
        if not os.path.isdir(vdir): continue
        csvs = sorted(glob.glob(os.path.join(vdir, "result_*.csv")))
        if not csvs: continue
        df_all = pd.concat([pd.read_csv(f)[["episode", y_col]] for f in csvs], ignore_index=True)
        if max_episode is not None: df_all = df_all[df_all["episode"] <= max_episode]
        stat = df_all.groupby("episode")[y_col].agg(["mean", "std"]).reset_index().sort_values("episode")
        if smooth and smooth > 1:
            stat["mean"] = stat["mean"].rolling(smooth, min_periods=1, center=True).mean()
            stat["std"] = stat["std"].rolling(smooth, min_periods=1, center=True).mean()
        out[var] = stat
    return out

def make_fig(env, experiment, title, max_episode, ymax=None, fname=None):
    stats = load_stats(env, experiment, max_episode=max_episode)
    fig = go.Figure()
    for var, st in stats.items():
        label = VAR_LABEL.get(var, var)
        x = st["episode"]; y = st["mean"]
        fig.add_trace(go.Scatter(x=x, y=y, name=label, mode="lines",
                                 line=dict(width=2),
                                 hovertemplate=f"{label}<br>ep=%{{x}}<br>gap=%{{y:.1f}}<extra></extra>"))
    fig.update_layout(title=title, xaxis_title="episode", yaxis_title="rational risk gap",
                      template="plotly_white", height=380, legend=dict(font=dict(size=11)),
                      margin=dict(l=50, r=20, t=50, b=40))
    if ymax: fig.update_yaxes(range=[0, ymax])
    fn = fname or f"{env}_{experiment}.html"
    path = os.path.join(OUT, fn)
    fig.write_html(path, include_plotlyjs="cdn", full_html=False)
    # raw csv
    rows = []
    for var, st in stats.items():
        for _, r in st.iterrows():
            rows.append({"variable": var, "episode": int(r["episode"]), "mean": r["mean"], "std": r["std"]})
    pd.DataFrame(rows).to_csv(os.path.join(OUT, fn.replace(".html", ".csv")), index=False)
    return path

# Claim 4: regularisation (Figure 2a)
make_fig("taxi", "exp_reg", "Claim 4 — Taxi: regularisation vs baseline DQN (rational risk gap)", 1400, ymax=400, fname="claim4_taxi_reg.html")
make_fig("cliffwalking", "exp_reg", "Claim 4 — CliffWalking: regularisation vs baseline DQN (rational risk gap)", 1400, ymax=1200, fname="claim4_cliff_reg.html")
# Claim 5a: domain randomisation (Figure 2b)
make_fig("taxi", "exp_domain_rand", "Claim 5 — Taxi: domain randomisation vs baseline (rational risk gap)", 800, ymax=500, fname="claim5a_taxi_dr.html")
make_fig("cliffwalking", "exp_domain_rand", "Claim 5 — CliffWalking: domain randomisation vs baseline (rational risk gap)", 1500, ymax=2000, fname="claim5a_cliff_dr.html")
# Claim 5b: environment shift magnitude (Figure 3)
make_fig("taxi", "exp_environment_level", "Claim 5 — Taxi: environment shift magnitude (rational risk gap)", 900, ymax=700, fname="claim5b_taxi_shift.html")
make_fig("cliffwalking", "exp_environment_level", "Claim 5 — CliffWalking: environment shift magnitude (rational risk gap)", 900, ymax=5500, fname="claim5b_cliff_shift.html")

print("Generated figures in", OUT)
for f in sorted(os.listdir(OUT)): print(" ", f)