repro_bundle/scripts/run_sweep.py

# /// script
# requires-python = ">=3.10"
# dependencies = [
#   "torch",
#   "numpy",
#   "gymnasium",
#   "scipy",
#   "pandas",
#   "huggingface_hub",
# ]
# ///
"""Fresh independent sweep runner for the Rationality paper reproduction.
Runs all variable x seed combos for one env (taxi|cliffwalking) in parallel,
saves CSVs in the repo layout, and uploads to a HF dataset repo.
"""
import argparse, os, subprocess, sys, time
from concurrent.futures import ThreadPoolExecutor, as_completed

ENV = os.environ.get("RATIONALITY_ENV", "taxi")
SEEDS = [int(x) for x in os.environ.get("SEEDS", "1,2,3,4,5").split(",")]
HF_REPO = os.environ.get("HF_REPO", "alphaXiv/rationality-repro-logs")
REPO_DIR = "/workspace/Rationality"

CONFIGS = {
    # exp_reg
    ("taxi", "reg"): dict(experiment="exp_reg", episodes=2000, eps_train=0.25, horizon=500, variables={
        "baseline": {},
        "ln_train": {"layernorm": True},
        "l2_train": {"l2_coef": 1e-5},
        "wn_train": {"weightnorm": True},
    }),
    ("cliffwalking", "reg"): dict(experiment="exp_reg", episodes=1000, eps_train=0.25, horizon=500, variables={
        "baseline": {},
        "ln_train": {"layernorm": True},
        "l2_train": {"l2_coef": 1e-7},
        "wn_train": {"weightnorm": True},
    }),
    # exp_domain_rand
    ("taxi", "domain_rand"): dict(experiment="exp_domain_rand", episodes=2000, eps_train=0.25, horizon=500, variables={
        "baseline": {},
        "envrnd_train_25": {"env_randomization": True, "mix_kernels": 3, "mix_eps_low": 0.05, "mix_eps_high": 0.3},
    }),
    ("cliffwalking", "domain_rand"): dict(experiment="exp_domain_rand", episodes=2000, eps_train=0.25, horizon=500, variables={
        "baseline": {},
        "envrnd_train_25": {"env_randomization": True, "mix_kernels": 5, "mix_eps_low": 0.1, "mix_eps_high": 0.3},
    }),
    # exp_environment_level
    ("taxi", "environment_level"): dict(experiment="exp_environment_level", episodes=2000, horizon=500, variables={
        "default": {"eps_train": 0.0},
        "eps_train_01": {"eps_train": 0.1},
        "eps_train_03": {"eps_train": 0.3},
        "eps_train_05": {"eps_train": 0.5},
        "eps_train_07": {"eps_train": 0.7},
    }),
    ("cliffwalking", "environment_level"): dict(experiment="exp_environment_level", episodes=2000, horizon=500, variables={
        "default": {"eps_train": 0.0},
        "eps_train_01": {"eps_train": 0.1},
        "eps_train_03": {"eps_train": 0.3},
        "eps_train_05": {"eps_train": 0.5},
        "eps_train_07": {"eps_train": 0.7},
    }),
}


def setup_repo():
    if os.path.isdir(os.path.join(REPO_DIR, ".git")):
        print("Repo already present")
        return
    os.makedirs("/workspace", exist_ok=True)
    subprocess.run(["git", "clone", "https://github.com/EVIEHub/Rationality", REPO_DIR], check=True)
    # patch Taxi-v3 -> v4 fallback (v4 has identical transition dynamics to v3)
    p = os.path.join(REPO_DIR, "src/env/taxi.py")
    s = open(p).read()
    s = s.replace('def get_base_P(env_id="Taxi-v3"):\n    env = gym.make(env_id)\n    P = env.unwrapped.P',
                 'def get_base_P(env_id="Taxi-v3"):\n    try:\n        env = gym.make(env_id)\n    except Exception:\n        env = gym.make("Taxi-v4")\n    P = env.unwrapped.P')
    s = s.replace('    env0 = gym.make("Taxi-v3")\n    d0_train = make_d0_gym_valid_uniform(env0, nS)',
                 '    try:\n        env0 = gym.make("Taxi-v3")\n    except Exception:\n        env0 = gym.make("Taxi-v4")\n    d0_train = make_d0_gym_valid_uniform(env0, nS)')
    open(p, "w").write(s)
    # patch cliffwalking.py: CliffWalking-v0 -> v1 fallback (identical transition dynamics)
    cp = os.path.join(REPO_DIR, "src/env/cliffwalking.py")
    cs = open(cp).read()
    cs = cs.replace('def get_base_P(env_id="CliffWalking-v0"):\n    env = gym.make(env_id)\n    P = env.unwrapped.P',
                    'def get_base_P(env_id="CliffWalking-v0"):\n    try:\n        env = gym.make(env_id)\n    except Exception:\n        env = gym.make("CliffWalking-v1")\n    P = env.unwrapped.P')
    # also patch the second gym.make(env_id) call inside build_experiment_finite_horizon
    cs = cs.replace('    env0 = gym.make(env_id)\n    d0_train = make_d0_cliff_start(env0, nS)',
                    '    try:\n        env0 = gym.make(env_id)\n    except Exception:\n        env0 = gym.make("CliffWalking-v1")\n    d0_train = make_d0_cliff_start(env0, nS)')
    open(cp, "w").write(cs)
    # patch train.py to allow configurable LOG_DIR via env var
    tp = os.path.join(REPO_DIR, "train.py")
    ts = open(tp).read()
    ts = ts.replace('LOG_DIR = Path("~/rational_exp/logs").expanduser()',
                    'LOG_DIR = Path(os.environ.get("RATIONALITY_LOG_DIR", "~/rational_exp/logs")).expanduser()')
    if "import os" not in ts:
        ts = ts.replace("import argparse", "import argparse\nimport os", 1)
    open(tp, "w").write(ts)
    print("Patched train.py LOG_DIR")


def run_one(args_tuple):
    env, key, cfg = args_tuple
    exp_name = cfg["experiment"]
    episodes = cfg["episodes"]
    horizon = cfg["horizon"]
    eps_train = cfg.get("eps_train", 0.25)
    var_name, var_args = cfg["_var"]
    log_root = os.path.join("/workspace/logs", env, exp_name, var_name)
    os.makedirs(log_root, exist_ok=True)
    seed = cfg["_seed"]
    out_csv = os.path.join(log_root, f"result_{seed}.csv")
    if os.path.exists(out_csv) and os.path.getsize(out_csv) > 50:
        return (var_name, seed, "skip", 0.0)
    cmd = [sys.executable, "train.py",
           "--algo", "dqn", "--env", env, "--device", "cpu",
           "--num_episodes", str(episodes), "--eval_every", "20",
           "--horizon", str(horizon), "--eps_train", str(eps_train),
           "--seed", str(seed), "--experiment", exp_name, "--variable", var_name]
    for k, v in var_args.items():
        if isinstance(v, bool):
            if v: cmd.append(f"--{k}")
        else:
            cmd.extend([f"--{k}", str(v)])
    t0 = time.time()
    env_copy = dict(os.environ)
    env_copy["OMP_NUM_THREADS"] = "1"
    env_copy["MKL_NUM_THREADS"] = "1"
    env_copy["TORCH_NUM_THREADS"] = "1"
    env_copy["RATIONALITY_LOG_DIR"] = "/workspace/logs"
    with open(out_csv + ".log", "w") as lf:
        r = subprocess.run(cmd, cwd=REPO_DIR, env=env_copy,
                          stdout=lf, stderr=subprocess.STDOUT, timeout=2400)
    dt = time.time() - t0
    ok = r.returncode == 0 and os.path.exists(out_csv) and os.path.getsize(out_csv) > 50
    return (var_name, seed, "ok" if ok else "FAIL", dt)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--env", default=ENV)
    parser.add_argument("--experiments", default="reg,domain_rand,environment_level")
    parser.add_argument("--workers", type=int, default=32)
    args = parser.parse_args()

    setup_repo()

    jobs = []
    for exp in args.experiments.split(","):
        key = (args.env, exp)
        if key not in CONFIGS:
            print(f"Skipping unknown {key}"); continue
        cfg = CONFIGS[key]
        for var_name, var_args in cfg["variables"].items():
            for seed in SEEDS:
                jc = dict(cfg); del jc["variables"]; jc["_var"] = (var_name, var_args); jc["_seed"] = seed
                jobs.append((args.env, key, jc))

    print(f"Submitting {len(jobs)} runs on {args.workers} workers")
    t0 = time.time()
    done = 0
    with ThreadPoolExecutor(max_workers=args.workers) as ex:
        futs = {ex.submit(run_one, j): j for j in jobs}
        for f in as_completed(futs):
            var_name, seed, status, dt = f.result()
            done += 1
            print(f"[{done}/{len(jobs)}] {var_name} seed{seed}: {status} ({dt:.0f}s)", flush=True)
    print(f"All done in {time.time()-t0:.0f}s")

    # Aggregate summary
    print("\n=== SUMMARY (mean rational_risk_gap over last-15 eval episodes, per variable) ===")
    import pandas as pd, glob
    for exp in args.experiments.split(","):
        key = (args.env, exp)
        if key not in CONFIGS: continue
        exp_name = CONFIGS[key]["experiment"]
        base = os.path.join("/workspace/logs", args.env, exp_name)
        for var in sorted(os.listdir(base)) if os.path.isdir(base) else []:
            vdir = os.path.join(base, var)
            csvs = sorted(glob.glob(os.path.join(vdir, "result_*.csv")))
            means = []
            for f in csvs:
                d = pd.read_csv(f)
                eps = sorted(d["episode"].unique())
                d = d[d["episode"].isin(eps[-15:])]
                means.append(d["rational_risk_gap"].mean())
            if means:
                import numpy as np
                print(f"  {args.env}/{exp_name}/{var}: {np.mean(means):.3f} ± {np.std(means,ddof=1):.3f} (n={len(means)})")

    # Upload to HF dataset repo
    print(f"\nUploading logs to {HF_REPO}")
    from huggingface_hub import HfApi
    api = HfApi()
    api.create_repo(repo_id=HF_REPO, repo_type="dataset", exist_ok=True)
    api.upload_folder(folder_path="/workspace/logs", repo_id=HF_REPO,
                      repo_type="dataset", path_in_repo=f"fresh_logs/{args.env}")
    print("Upload complete.")


if __name__ == "__main__":
    main()