scripts/claim2_static.py

"""CLAIM 2 -- Theorem 3.1: expected static regret
      E[R_T(u)] = O~( G eps + (d/kappa) E[ ||u|| sqrt(V_T log_+(d||u||Lambda_T/(G eps))) ] )
   with kappa = sqrt(d) in the norm-oblivious setting and kappa = 1 in the norm-adaptive
   setting, i.e. an effective dimension factor sqrt(d) resp. d (the paper's "sqrt(d)
   separation"; note the paper writes d/kappa, so the SMALLER rate is the oblivious one).

Tests
 A  batched-vs-reference implementation consistency.
 B  scaling of the measured E[R_T(u)] in T, d and ||u|| on the hardest known loss family
    (hypercube bias, Delta = G/sqrt(T), noise variance G^2/(2d)); fitted log-log exponents
    vs the theorem's 1/2, 1/2, 1.
 C  the bound itself, evaluated with explicit constants dropped, for kappa=sqrt(d) and
    kappa=1, on four environments including an adversarial sign-flipping one; also the
    risk-control term R_T(0) = O(G eps).
 D  where kappa comes from: the two moment bounds of Corollary 2.2 measured along real
    trajectories (a.s. 4d^2||l||^2 -> factor d;  conditional expectation 2d||l||^2 ->
    factor sqrt(d)), and the fitted d-exponent of the realised sum_t ||ltilde_t||^2.
 E  a norm-adaptive adversary that fixes ||u|| only after seeing the realised trajectory:
    does the sqrt(d)-worse kappa=1 branch actually get realised?
"""

from __future__ import annotations

import json
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from batch import BatchAlg6, fit_exponent, run_pablo_batch
from pablo import Alg6, env_flip, env_hard, env_rademacher, run_pablo

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
SEED = 20260725
G = 1.0
EPSILON = 1.0  # user-specified risk parameter (the paper's eps; R_T(0) <= O~(G eps))
EPS_PERT = 1e-3  # the varepsilon of Eq. (4)


def make_env(kind, T, d, seed):
    rng = np.random.default_rng(seed)
    if kind == "hard":
        return env_hard(T, d, G, rng)
    if kind == "rademacher":
        return env_rademacher(T, d, G, rng)
    if kind == "flip":
        return env_flip(T, d, G, rng)
    raise ValueError(kind)


def measure(kind, T, d, S, seed):
    losses = make_env(kind, T, d, seed)
    L = 2 * d * G  # a.s. bound on ||ltilde|| (Cor. 2.2)
    olo = BatchAlg6(S, d, T, L, EPSILON / d)  # "tuned with eps/d" (Theorem 3.1)
    r = run_pablo_batch(losses, olo, d, EPS_PERT, np.random.default_rng(seed + 999), S)
    r["clipped"] = olo.clipped
    r["V_T"] = float(np.sum(losses * losses))
    r["finite"] = bool(np.all(np.isfinite(r["cum_play"])))
    return r


def regret_for(r, u):
    return r["cum_play"] - float(np.dot(r["sum_losses"], u))


def worst_direction(r):
    s = r["sum_losses"]
    n = float(np.linalg.norm(s))
    return -s / n if n > 0 else np.zeros_like(s)


def bound_thm31(u_norm, V_T, T, d, kappa):
    """Theorem 3.1 shape with unspecified constants set to 1:
      G*eps + (d/kappa) ||u|| sqrt( V_T log_+( d ||u|| Lambda_T/(G eps) ) ),
    Lambda_T = G sqrt(T) log^2(1+T)."""
    Lam = G * np.sqrt(T) * np.log(1 + T) ** 2
    lg = max(np.log(max(d * u_norm * Lam / (G * EPSILON), 1.0)), 1.0)
    return G * EPSILON + (d / kappa) * u_norm * np.sqrt(V_T * lg)


def test_A_consistency():
    d, T, S = 5, 300, 400
    losses = make_env("hard", T, d, SEED)
    L = 2 * d * G
    rb = run_pablo_batch(
        losses,
        BatchAlg6(S, d, T, L, EPSILON / d),
        d,
        EPS_PERT,
        np.random.default_rng(4242),
        S,
    )
    ref = [
        run_pablo(
            losses,
            Alg6(d, T, L, EPSILON / d),
            d,
            EPS_PERT,
            np.random.default_rng(9000 + s),
        )["regret"]
        for s in range(400)
    ]
    u = np.zeros(d)
    b = regret_for(rb, u)
    return dict(
        note="independent estimates of E[R_T(0)] from the batched and the reference "
        "single-seed code paths on the same loss sequence",
        batch_mean=float(b.mean()),
        batch_stderr=float(b.std(ddof=1) / np.sqrt(S)),
        reference_mean=float(np.mean(ref)),
        reference_stderr=float(np.std(ref, ddof=1) / np.sqrt(len(ref))),
    )


def test_B_scalings():
    out = {}
    d, S = 8, 400
    Ts = [250, 500, 1000, 2000, 4000, 8000]
    means, sems, maxw = [], [], []
    for T in Ts:
        r = measure("hard", T, d, S, SEED + T)
        R = regret_for(r, 1.0 * worst_direction(r))
        means.append(float(R.mean()))
        sems.append(float(R.std(ddof=1) / np.sqrt(S)))
        maxw.append(float(r["max_w"].max()))
    sl, _, se = fit_exponent(Ts, means)
    out["T_sweep"] = dict(
        env="hard",
        d=d,
        seeds=S,
        u_norm=1.0,
        T=Ts,
        mean_regret=means,
        sem=sems,
        fitted_exponent=sl,
        stderr=se,
        predicted=0.5,
        max_iterate_norm=maxw,
    )

    T, S = 2000, 400
    ds = [2, 4, 8, 16, 32, 64]
    means, sems, lt = [], [], []
    for dd in ds:
        r = measure("hard", T, dd, S, SEED + 31 * dd)
        R = regret_for(r, 1.0 * worst_direction(r))
        means.append(float(R.mean()))
        sems.append(float(R.std(ddof=1) / np.sqrt(S)))
        lt.append(float(np.mean(np.sqrt(r["sum_lt_sq"]))))
    sl, _, se = fit_exponent(ds, means)
    sl2, _, se2 = fit_exponent(ds, lt)
    out["d_sweep"] = dict(
        env="hard",
        T=T,
        seeds=S,
        u_norm=1.0,
        d=ds,
        mean_regret=means,
        sem=sems,
        fitted_exponent=sl,
        stderr=se,
        predicted_norm_oblivious=0.5,
        predicted_norm_adaptive=1.0,
        mean_sqrt_sum_ltilde_sq=lt,
        fitted_d_exponent_of_sqrt_sum_ltilde_sq=sl2,
        stderr2=se2,
    )

    d, T, S = 8, 2000, 400
    r = measure("hard", T, d, S, SEED + 5)
    dirn = worst_direction(r)
    norms = [0.25, 0.5, 1, 2, 4, 8, 16, 32, 64]
    means, sems, over_log = [], [], []
    for un in norms:
        R = regret_for(r, un * dirn)
        means.append(float(R.mean()))
        sems.append(float(R.std(ddof=1) / np.sqrt(S)))
        Lam = G * np.sqrt(T) * np.log(1 + T) ** 2
        lg = max(np.log(max(d * un * Lam / (G * EPSILON), 1.0)), 1.0)
        over_log.append(float(R.mean()) / np.sqrt(lg))
    sl, _, se = fit_exponent(norms, means)
    sl3, _, se3 = fit_exponent(norms, over_log)
    out["u_sweep"] = dict(
        env="hard",
        d=d,
        T=T,
        seeds=S,
        u_norm=norms,
        mean_regret=means,
        sem=sems,
        fitted_exponent=sl,
        stderr=se,
        predicted=1.0,
        fitted_exponent_after_dividing_by_sqrt_log=sl3,
        stderr3=se3,
    )
    return out


def test_C_bound():
    rows = []
    grid = [0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128]
    for kind in ["hard", "rademacher", "flip"]:
        for d in [2, 8, 32]:
            T, S = 2000, 200
            r = measure(kind, T, d, S, SEED + 101 * d + len(kind))
            dirn = worst_direction(r)
            best_ob, best_ad, arg = -np.inf, -np.inf, None
            for un in grid:
                m = float(np.mean(regret_for(r, un * dirn)))
                ro = m / bound_thm31(un, r["V_T"], T, d, np.sqrt(d))
                ra = m / bound_thm31(un, r["V_T"], T, d, 1.0)
                if ro > best_ob:
                    best_ob, arg = ro, un
                best_ad = max(best_ad, ra)
            R0 = regret_for(r, np.zeros(d))
            rows.append(
                dict(
                    env=kind,
                    d=d,
                    T=T,
                    seeds=S,
                    finite=r["finite"],
                    clipped=r["clipped"],
                    max_ratio_to_kappa_sqrt_d_bound=best_ob,
                    argmax_u_norm=arg,
                    max_ratio_to_kappa_1_bound=best_ad,
                    mean_R_T_0=float(R0.mean()),
                    max_R_T_0=float(R0.max()),
                    risk_control_budget_G_eps=G * EPSILON,
                    max_iterate_norm=float(r["max_w"].max()),
                )
            )
    return rows


def test_D_kappa_mechanism():
    rows = []
    for d in [2, 4, 8, 16, 32, 64]:
        T, S = 1000, 200
        r = measure("hard", T, d, S, SEED + 77 * d)
        V = r["V_T"]
        rows.append(
            dict(
                d=d,
                mean_sum_ltilde_sq_over_V_T=float(np.mean(r["sum_lt_sq"])) / V,
                max_seed_sum_ltilde_sq_over_V_T=float(np.max(r["sum_lt_sq"])) / V,
                Cor22_as_bound_4d2=4.0 * d * d,
                Cor22_expectation_bound_2d=2.0 * d,
            )
        )
    sl, _, se = fit_exponent(
        [x["d"] for x in rows],
        [np.sqrt(x["mean_sum_ltilde_sq_over_V_T"]) for x in rows],
    )
    return dict(
        rows=rows,
        fitted_d_exponent_of_realised_dimension_factor=sl,
        stderr=se,
        predicted_if_kappa_sqrt_d=0.5,
        predicted_if_kappa_1=1.0,
    )


def test_E_norm_adaptive():
    rows = []
    grid = np.array([0.25, 0.5, 1, 2, 4, 8, 16, 32, 64, 128])
    for d in [2, 4, 8, 16, 32]:
        T, S = 2000, 400
        r = measure("hard", T, d, S, SEED + 13 * d)
        dirn = worst_direction(r)
        V = r["V_T"]
        bnd = np.array([bound_thm31(un, V, T, d, np.sqrt(d)) for un in grid])
        ob = max(
            float(np.mean(regret_for(r, un * dirn))) / b for un, b in zip(grid, bnd)
        )
        per_seed = np.stack([regret_for(r, un * dirn) for un in grid])  # (grid, S)
        ratios = np.max(per_seed / bnd[:, None], axis=0)
        rows.append(
            dict(
                d=d,
                T=T,
                seeds=S,
                best_oblivious_norm_ratio=ob,
                mean_norm_adaptive_ratio=float(ratios.mean()),
                max_norm_adaptive_ratio=float(ratios.max()),
                adaptive_over_oblivious=float(ratios.mean()) / ob,
                predicted_adaptive_over_oblivious_if_gap_real=float(np.sqrt(d)),
                frac_seeds_exceeding_oblivious_bound=float(np.mean(ratios > 1.0)),
            )
        )
    return rows


if __name__ == "__main__":
    res = dict(
        claim="claim-2 Theorem 3.1 expected static regret",
        seed=SEED,
        settings=dict(
            G=G,
            epsilon=EPSILON,
            varepsilon_eq4=EPS_PERT,
            olo_subroutine="Algorithm 6 of the paper (Appendix E.2); its "
            "static specialisation satisfies exactly the parameter-free OLO "
            "guarantee that Theorem 3.1 requires of JC22 Algorithm 4",
            perturbation="isotropic H_t of Eq. (4), equality case",
        ),
        A_batch_consistency=test_A_consistency(),
        B_scalings=test_B_scalings(),
        C_bound_check=test_C_bound(),
        D_kappa_mechanism=test_D_kappa_mechanism(),
        E_norm_adaptive_adversary=test_E_norm_adaptive(),
    )
    os.makedirs(OUT, exist_ok=True)
    with open(os.path.join(OUT, "claim2_static.json"), "w") as f:
        json.dump(res, f, indent=1)
    print(json.dumps(res, indent=1))