code/exp_bandit.py

"""Claims 2-5 audit — EntUCB regret vs the paper's bounds.

--claim 2 : entropic regret of fixed-eps EntUCB vs the two-term bound
            sigma sqrt(2T log(2/delta)) + 2 Cbar beta_T(delta) sqrt(T logdet(...))
            [OpenReview Thm 5.1; arXiv v2: width+martingale terms of Thm 4.1]
--claim 3 : Kantorovich regret of EntUCB with eps_t = eta t^-eta (Thm 4.1 /
            OpenReview Thm 5.2): sublinearity + full bound; plus the
            Carlier-Pegon entropic-gap audit  Ent - Kant <= kappa eps log(1/eps).
--claim 4 : parametric sqrt(NT) rate (Prop 5.5 / OpenReview Cor 5.3): sweep N, T.
--claim 5 : basis-decay interpolation (Thm 5.3 / OpenReview Cor 5.4):
            zeta(n) = 1 - n^-q, n_t = ceil(t^{1/(q+1)}), regret exponent vs
            (q+2)/(2q+2); includes a mis-scheduled control.

All runs verify the optimism certificate and confidence-set coverage per step.
"""
import argparse
import json
import os
import sys
import time

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from botlib import parallel_runs, martingale_term

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "results")
os.makedirs(OUT, exist_ok=True)
PROCS = int(os.environ.get("PROCS", "0"))


def summarize(rs):
    agg = {
        "ts": rs[0].ts.tolist(),
        "kant_regret_mean": np.mean([r.kant_regret for r in rs], axis=0).tolist(),
        "kant_regret_max": np.max([r.kant_regret for r in rs], axis=0).tolist(),
        "kant_regret_noisy_max": np.max([r.kant_regret_noisy for r in rs], axis=0).tolist(),
        "ent_regret_mean": np.mean([r.ent_regret for r in rs], axis=0).tolist(),
        "ent_regret_max": np.max([r.ent_regret for r in rs], axis=0).tolist(),
        "ent_regret_noisy_max": np.max([r.ent_regret_noisy for r in rs], axis=0).tolist(),
        "bound_literal_mean": np.mean([r.bound_literal for r in rs], axis=0).tolist(),
        "bound_literal_min": np.min([r.bound_literal for r in rs], axis=0).tolist(),
        "bound_canonical_min": np.min([r.bound_canonical for r in rs], axis=0).tolist(),
        "beta_mean": np.mean([r.beta for r in rs], axis=0).tolist(),
        "logdet_mean": np.mean([r.logdet for r in rs], axis=0).tolist(),
        "cert_viol_frac_max": float(max(r.cert_viol_frac for r in rs)),
        "coverage_rate": float(np.mean([r.covered_all for r in rs])),
        "n_seeds": len(rs),
    }
    return agg


def fit_exponent(ts, curve, lo_frac=0.25):
    ts, curve = np.asarray(ts, float), np.asarray(curve, float)
    m = (ts >= lo_frac * ts[-1]) & (curve > 0)
    if m.sum() < 3:
        return float("nan")
    A = np.vstack([np.log(ts[m]), np.ones(m.sum())]).T
    slope, _ = np.linalg.lstsq(A, np.log(curve[m]), rcond=None)[0]
    return float(slope)


def base_cfg(**kw):
    cfg = dict(K=8, Kp=8, inst_seed=0, cost_kind="smooth", T=4000, sigma=0.1,
               delta=0.1, lam=1.0, Cbar_mult=1.1, run_seed=0, opt_iters=3)
    cfg.update(kw)
    return cfg


def claim2(full: bool):
    K = 8 if full else 4
    T = 4000 if full else 250
    seeds = 16 if full else 3
    out = {"config": dict(K=K, T=T, seeds=seeds, sigma=0.1, delta=0.1, lam=1.0)}
    for eps in [0.05, 0.01]:
        cfgs = [base_cfg(K=K, Kp=K, T=T, inst_seed=500 + s % 6,
                         eps={"kind": "const", "v": eps}, run_seed=6000 + s)
                for s in range(seeds)]
        rs = parallel_runs(cfgs, PROCS)
        agg = summarize(rs)
        ts = np.array(agg["ts"])
        worst = np.array(agg["ent_regret_noisy_max"])
        bnd_lit = np.array(agg["bound_literal_min"])
        bnd_can = np.array(agg["bound_canonical_min"])
        agg["holds_literal"] = bool(np.all(worst <= bnd_lit))
        agg["holds_canonical"] = bool(np.all(worst <= bnd_can))
        agg["min_ratio_literal"] = float(np.min(bnd_lit / np.maximum(worst, 1e-12)))
        agg["ent_regret_exponent"] = fit_exponent(ts, np.maximum(agg["ent_regret_mean"], 1e-12))
        out[f"eps_{eps}"] = agg
        print(f"[claim2 eps={eps}] holds_literal={agg['holds_literal']} "
              f"(min ratio {agg['min_ratio_literal']:.2f}); holds_canonical={agg['holds_canonical']}; "
              f"growth exp {agg['ent_regret_exponent']:.2f}; cert viol {agg['cert_viol_frac_max']:.4f}; "
              f"coverage {agg['coverage_rate']:.2f}", flush=True)
    return out


def claim3(full: bool):
    K = 8 if full else 4
    T = 4000 if full else 250
    seeds = 8 if full else 3
    out = {"config": dict(K=K, T=T, seeds=seeds)}
    for eta in ([0.6, 0.75, 0.9] if full else [0.75]):
        cfgs = [base_cfg(K=K, Kp=K, T=T, inst_seed=500 + s % 6,
                         eps={"kind": "power", "eta": eta}, run_seed=7000 + s,
                         kappa=0.0, eta_for_bound=eta)
                for s in range(seeds)]
        rs = parallel_runs(cfgs, PROCS)
        agg = summarize(rs)
        ts = np.array(agg["ts"])
        agg["kant_exponent_mean"] = fit_exponent(ts, np.maximum(agg["kant_regret_mean"], 1e-12))
        agg["holds_literal_noisy"] = bool(np.all(
            np.array(agg["kant_regret_noisy_max"]) <= np.array(agg["bound_literal_min"])))
        out[f"eta_{eta}"] = agg
        print(f"[claim3 eta={eta}] kant exponent {agg['kant_exponent_mean']:.3f}; "
              f"bound holds {agg['holds_literal_noisy']}; cert {agg['cert_viol_frac_max']:.4f}", flush=True)

    # Carlier-Pegon gap audit
    from botlib import sinkhorn_log
    import ot as pot
    gap_res = {}
    for alpha, costf in [(1.0, lambda X, Y: np.abs(X - Y)),
                         (0.5, lambda X, Y: np.sqrt(np.abs(X - Y)))]:
        for m in ([50, 100, 200] if full else [40]):
            x = (np.arange(m) + 0.5) / m
            mu = np.full(m, 1 / m)
            X, Y = np.meshgrid(x, x, indexing="ij")
            cost = costf(X, Y)
            kv = float(np.sum(pot.emd(mu, mu, cost) * cost))
            epss = np.geomspace(0.2, 0.002, 10)
            gaps = []
            for eps in epss:
                P, _, _ = sinkhorn_log(mu, mu, cost, float(eps), n_iter=30000, tol=1e-11)
                ent = float(np.sum(P * cost)) + float(eps) * float(np.sum(P[P > 0] * np.log(P[P > 0] * m * m)))
                gaps.append(ent - kv)
            slope = np.array(gaps) / (epss * np.log(1 / epss))
            gap_res[f"alpha_{alpha}_m_{m}"] = {"eps": epss.tolist(), "gap": gaps,
                                               "gap_over_epsloge": slope.tolist(),
                                               "kappa_bound": 1.0 / alpha}
            print(f"[claim3 gap] alpha={alpha} m={m}: slope in [{slope.min():.3f},{slope.max():.3f}] "
                  f"vs kappa {1 / alpha:.1f}", flush=True)
    out["carlier_pegon"] = gap_res
    return out


def claim4(full: bool):
    Ks = [3, 4, 6, 8] if full else [3, 4]
    Ts = [500, 1000, 2000, 4000] if full else [150, 300]
    seeds = 8 if full else 2
    lam, eta, sigma, delta = 1.0, 0.5, 0.1, 0.1
    out = {"config": dict(Ks=Ks, Ts=Ts, seeds=seeds, eta=eta)}
    table = []
    for K in Ks:
        N = K * K
        cfgs = [base_cfg(K=K, Kp=K, T=max(Ts), inst_seed=900 + s % 6,
                         cost_kind="planted", planted={"kind": "unit", "seed": 800 + s},
                         eps={"kind": "power", "eta": eta}, run_seed=8000 + s,
                         kappa=0.0, eta_for_bound=eta, Cbar_mult=1.1)
                for s in range(seeds)]
        rs = parallel_runs(cfgs, PROCS)
        agg = summarize(rs)
        ts = np.array(agg["ts"])
        for T in Ts:
            i = int(np.argmin(np.abs(ts - T)))
            bnd = (2 * 1.1 * np.sqrt(N * T) * np.log(1 / lam + T * 1.1 ** 2 / N)
                   + martingale_term(sigma, delta, T))
            table.append(dict(N=N, T=int(ts[i]), regret=agg["kant_regret_mean"][i],
                              regret_worst=agg["kant_regret_max"][i], prop55_bound=bnd))
        out[f"N_{N}_curve"] = {"ts": ts.tolist(),
                               "kant_regret_mean": agg["kant_regret_mean"],
                               "kant_regret_max": agg["kant_regret_max"]}
        out[f"N_{N}_audit"] = {k: agg[k] for k in ["cert_viol_frac_max", "coverage_rate", "n_seeds"]}
        print(f"[claim4 N={N}] final regret {agg['kant_regret_mean'][-1]:.2f}; "
              f"cert {agg['cert_viol_frac_max']:.4f}; coverage {agg['coverage_rate']:.2f}", flush=True)
    out["table"] = table
    tb = [r for r in table if r["regret"] > 0]
    Amat = np.array([[np.log(r["N"]), np.log(r["T"]), 1.0] for r in tb])
    y = np.log([r["regret"] for r in tb])
    coef, *_ = np.linalg.lstsq(Amat, y, rcond=None)
    out["fit_exponents"] = {"N_exp": float(coef[0]), "T_exp": float(coef[1])}
    out["bound_holds_all"] = bool(all(r["regret_worst"] <= r["prop55_bound"] for r in table))
    print(f"[claim4] R ~ N^{coef[0]:.2f} T^{coef[1]:.2f} (claim 0.5, 0.5); "
          f"bound holds: {out['bound_holds_all']}", flush=True)
    return out


def claim5(full: bool):
    K = 16 if full else 8
    N = K * K
    T = 4000 if full else 400
    seeds = 4 if full else 2
    qs = [0.5, 1.0, 2.0, 4.0] if full else [0.5, 2.0]
    eta = 0.5
    out = {"config": dict(K=K, N=N, T=T, seeds=seeds, qs=qs, eta=eta)}
    for q in qs:
        cfgs = [base_cfg(K=K, Kp=K, T=T, inst_seed=1300 + s % 4,
                         cost_kind="planted", planted={"kind": "decay", "q": q, "seed": 1200 + s},
                         eps={"kind": "power", "eta": eta},
                         n_sched={"kind": "power", "q": q},
                         run_seed=9000 + s, kappa=0.0, eta_for_bound=eta)
                for s in range(seeds)]
        rs = parallel_runs(cfgs, PROCS)
        agg = summarize(rs)
        agg["exponent"] = fit_exponent(np.array(agg["ts"]), np.maximum(agg["kant_regret_mean"], 1e-12))
        agg["predicted_exponent"] = (q + 2) / (2 * q + 2)
        out[f"q_{q}"] = agg
        print(f"[claim5 q={q}] exponent {agg['exponent']:.3f} vs predicted "
              f"{agg['predicted_exponent']:.3f}; cert {agg['cert_viol_frac_max']:.4f}; "
              f"coverage {agg['coverage_rate']:.2f}", flush=True)

    # control: schedule tuned for q=4 on a q=0.5 instance
    cfgs = [base_cfg(K=K, Kp=K, T=T, inst_seed=1300 + s % 4,
                     cost_kind="planted", planted={"kind": "decay", "q": 0.5, "seed": 1200 + s},
                     eps={"kind": "power", "eta": eta},
                     n_sched={"kind": "power", "q": 4.0},
                     run_seed=9000 + s, kappa=0.0, eta_for_bound=eta)
            for s in range(seeds)]
    rs = parallel_runs(cfgs, PROCS)
    agg = summarize(rs)
    agg["exponent"] = fit_exponent(np.array(agg["ts"]), np.maximum(agg["kant_regret_mean"], 1e-12))
    out["control_missched"] = agg
    print(f"[claim5 control] q_true=0.5, q_sched=4: exponent {agg['exponent']:.3f} "
          f"(well-tuned was {out['q_0.5']['exponent']:.3f})", flush=True)
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--claim", type=int, required=True, choices=[2, 3, 4, 5])
    ap.add_argument("--full", action="store_true")
    ap.add_argument("--tag", default="")
    ap.add_argument("--out", default=OUT)
    args = ap.parse_args()
    t0 = time.time()
    out = {2: claim2, 3: claim3, 4: claim4, 5: claim5}[args.claim](args.full)
    out["wall_seconds"] = time.time() - t0
    tag = args.tag or ("full" if args.full else "smoke")
    os.makedirs(args.out, exist_ok=True)
    path = os.path.join(args.out, f"claim{args.claim}_{tag}.json")
    with open(path, "w") as f:
        json.dump(out, f, indent=2)
    print(f"saved {path}  ({out['wall_seconds']:.1f}s)")


if __name__ == "__main__":
    main()