scripts/claim5_lowerbound.py

"""CLAIM 5 -- Theorem 5.2: the folklore Omega(sqrt(dT)) minimax lower bound for adversarial
linear bandits on the unit Euclidean ball.

The theorem's construction (part 1) is implemented verbatim: theta in {+-Delta}^d with
Delta = Theta(T^{-1/2}), losses l_t = theta + xi_t with xi_t ~ N(0, (2d)^{-1} I_d), learner
plays z_t in the unit ball B_d and observes only <l_t, z_t>.  The regret measured is the
direction regret of the paper's Eq. (7),
      R^Z_T = sum_t <z_t, theta> + T ||theta||      (comparator u = -theta/||theta||).

Tests
 1  Run several strong bandit algorithms on the construction and check the claimed floor
    R^Z_T >= sqrt(dT)/64 ^ T/(12d) for every one of them, for the adversary's best Delta.
 2  Fit the d- and T-exponents of the smallest regret achieved by any algorithm we ran
    (empirical minimax envelope) -- predicted 1/2 and 1/2.
 3  MECHANISM: the proof sketch attributes the sqrt(dT) rate (rather than d sqrt(T)) to the
    1/d noise-variance constraint.  Sweep the noise variance sigma^2 and check that the
    achievable regret behaves as sigma * d * sqrt(T), so that sigma^2 = 1/(2d) gives
    sqrt(dT/2).
 4  INFORMATION ARGUMENT re-derived numerically: along real trajectories, compute
    sum_i KL(P_theta || P_theta^(i)) for the d sign-flipped neighbours and check it equals
    4 d Delta^2 sum_t E[z_{ti}^2/||z_t||^2] = 4 d Delta^2 T; then check that when this is
    O(1) the algorithms recover no more coordinate signs than chance.
 5  Part 2 of the theorem: the same construction with losses truncated to ||l_t|| <= 1.
"""

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 fit_exponent

OUT = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "outputs"
)
SEED = 20260725


# ---------------------------------------------------------------------------- algorithms
def alg_thompson(T, d, theta, sigma, rng, truncate=False, lam=1.0):
    """Linear Thompson sampling on the unit ball (a strong stochastic linear-bandit
    baseline).  Regressors are normalised so the observation noise is homoscedastic."""
    Vinv = np.eye(d) / lam
    b = np.zeros(d)
    reg_dir = 0.0
    zs = np.zeros((T, d))
    L = np.eye(d) * sigma
    for t in range(T):
        th_hat = Vinv @ b
        if t % 25 == 0:              # posterior square root, refreshed periodically
            try:
                L = np.linalg.cholesky(sigma * sigma * Vinv + 1e-12 * np.eye(d))
            except np.linalg.LinAlgError:
                L = np.eye(d) * sigma
        th_s = th_hat + L @ rng.standard_normal(d)
        n = np.linalg.norm(th_s)
        z = -th_s / n if n > 0 else rng.standard_normal(d) / np.sqrt(d)
        ell = theta + rng.standard_normal(d) * sigma
        if truncate:
            nl = np.linalg.norm(ell)
            ell = ell * min(1.0, 1.0 / nl)
        y = float(np.dot(ell, z))
        reg_dir += float(np.dot(z, theta))
        zs[t] = z
        Vz = Vinv @ z
        Vinv -= np.outer(Vz, Vz) / (1.0 + float(np.dot(z, Vz)))
        b += z * y
    return reg_dir + T * np.linalg.norm(theta), zs


def alg_etc(T, d, theta, sigma, rng, frac=None, truncate=False):
    """Explore-then-commit: explore coordinate directions, then commit to -theta_hat/||.||.
    frac=None optimises the exploration budget over a grid (an *advantage* given to the
    algorithm, which only strengthens a lower-bound test)."""
    best = None
    fracs = [frac] if frac is not None else [0.05, 0.1, 0.2, 0.4, 0.8, 1.0]
    for f in fracs:
        T0 = max(d, int(f * T) // d * d)
        T0 = min(T0, T)
        est = np.zeros(d)
        reg = 0.0
        for t in range(T0):
            i = t % d
            z = np.zeros(d)
            z[i] = 1.0
            ell = theta + rng.standard_normal(d) * sigma
            if truncate:
                ell = ell * min(1.0, 1.0 / np.linalg.norm(ell))
            est[i] += float(np.dot(ell, z))
            reg += float(np.dot(z, theta))
        est = est / max(T0 // d, 1)
        n = np.linalg.norm(est)
        z = -est / n if n > 0 else np.zeros(d)
        reg += (T - T0) * float(np.dot(z, theta))
        r = reg + T * np.linalg.norm(theta)
        best = r if best is None else min(best, r)
    return best


def alg_osmd_ball(T, d, theta, sigma, rng, eta=None, truncate=False):
    """Bandit gradient descent on the ball with a one-coordinate importance-weighted
    estimator (the SCRiBLe/OSMD-style baseline the paper compares against); the step size is
    optimised over a grid, again in the algorithm's favour."""
    best = None
    etas = [eta] if eta is not None else [0.02, 0.05, 0.1, 0.3, 1.0]
    for e in etas:
        w = np.zeros(d)
        reg = 0.0
        gamma = min(0.5, np.sqrt(d / T))  # exploration rate
        for t in range(T):
            if rng.random() < gamma:
                i = int(rng.integers(d))
                s = 1.0 if rng.integers(2) == 0 else -1.0
                z = np.zeros(d)
                z[i] = s
                explore = True
            else:
                z = w.copy()
                nz = np.linalg.norm(z)
                if nz > 1:
                    z = z / nz
                explore = False
            ell = theta + rng.standard_normal(d) * sigma
            if truncate:
                ell = ell * min(1.0, 1.0 / np.linalg.norm(ell))
            y = float(np.dot(ell, z))
            reg += float(np.dot(z, theta))
            if explore:
                ghat = (d / gamma) * z * y
                w = w - e * ghat / np.sqrt(T)
                nw = np.linalg.norm(w)
                if nw > 1:
                    w = w / nw
        r = reg + T * np.linalg.norm(theta)
        best = r if best is None else min(best, r)
    return best


ALGS = {
    "thompson": lambda *a, **k: alg_thompson(*a, **k)[0],
    "etc": alg_etc,
    "osmd_ball": alg_osmd_ball,
}


def theorem_floor(d, T):
    return min(np.sqrt(d * T) / 64.0, T / (12.0 * d))


def test1_and_2(seeds=15):
    rows = []
    cfgs = [(d, 4000) for d in [4, 8, 16, 32]] + [
        (16, T) for T in [500, 1000, 2000, 8000]
    ]
    deltas = [0.25, 0.5, 1.0, 2.0, 4.0]
    for d, T in cfgs:
        sigma = 1.0 / np.sqrt(2 * d)
        per_alg = {}
        for name, fn in ALGS.items():
            worst = -np.inf
            for c in deltas:
                Delta = c / np.sqrt(T)
                vals = []
                for s in range(seeds):
                    rng = np.random.default_rng(SEED + 1000 * s + d + T + int(100 * c))
                    theta = (rng.integers(2, size=d) * 2.0 - 1.0) * Delta
                    vals.append(fn(T, d, theta, sigma, rng))
                m = float(np.mean(vals))
                if m > worst:
                    worst = m
                    per_alg[name] = dict(
                        regret=m,
                        best_Delta_coefficient=c,
                        sem=float(np.std(vals, ddof=1) / np.sqrt(seeds)),
                    )
        env = min(v["regret"] for v in per_alg.values())
        rows.append(
            dict(
                d=d,
                T=T,
                sqrt_dT=float(np.sqrt(d * T)),
                theorem_floor=float(theorem_floor(d, T)),
                per_algorithm=per_alg,
                empirical_minimax_envelope=env,
                envelope_over_sqrt_dT=env / float(np.sqrt(d * T)),
                above_theorem_floor=bool(env >= theorem_floor(d, T)),
            )
        )
    dsw = [r for r in rows if r["T"] == 4000]
    tsw = [r for r in rows if r["d"] == 16]
    sd, _, sed = fit_exponent(
        [r["d"] for r in dsw], [r["empirical_minimax_envelope"] for r in dsw]
    )
    st, _, set_ = fit_exponent(
        [r["T"] for r in tsw], [r["empirical_minimax_envelope"] for r in tsw]
    )
    return dict(
        seeds=seeds,
        rows=rows,
        fitted_d_exponent=sd,
        stderr_d=sed,
        fitted_T_exponent=st,
        stderr_T=set_,
        predicted=0.5,
    )


def test3_variance_mechanism(seeds=12):
    """regret ~ sigma * d * sqrt(T):  the 1/d variance is what turns d sqrt(T) into
    sqrt(dT)."""
    d, T = 16, 4000
    rows = []
    for s2 in [1.0 / (2 * d), 1.0 / d, 0.25, 0.5, 1.0]:
        sigma = np.sqrt(s2)
        worst = -np.inf
        for c in [0.25, 0.5, 1.0, 2.0, 4.0]:
            Delta = c * sigma * np.sqrt(2 * d) / np.sqrt(T)  # scaled to the noise level
            vals = []
            for s in range(seeds):
                rng = np.random.default_rng(SEED + 77 * s + int(1000 * s2))
                theta = (rng.integers(2, size=d) * 2.0 - 1.0) * Delta
                vals.append(min(fn(T, d, theta, sigma, rng) for fn in ALGS.values()))
            worst = max(worst, float(np.mean(vals)))
        rows.append(
            dict(
                sigma_sq=s2,
                envelope_regret=worst,
                sigma_d_sqrtT=float(sigma * d * np.sqrt(T)),
                ratio=worst / float(sigma * d * np.sqrt(T)),
            )
        )
    sl, _, se = fit_exponent(
        [r["sigma_sq"] for r in rows], [r["envelope_regret"] for r in rows]
    )
    return dict(
        d=d,
        T=T,
        seeds=seeds,
        rows=rows,
        fitted_sigma_sq_exponent=sl,
        stderr=se,
        predicted_exponent=0.5,
        note="regret ~ sigma d sqrt(T) => exponent 1/2 in sigma^2; at sigma^2=1/(2d) "
        "this equals sqrt(dT/2)",
    )


def test4_information(seeds=25):
    """Numerical re-derivation of the KL/Pinsker step and its consequence."""
    rows = []
    d, T = 16, 4000
    sigma = 1.0 / np.sqrt(2 * d)
    for c in [0.25, 0.5, 1.0, 2.0, 8.0]:
        Delta = c / np.sqrt(T)
        kls, correct, regs = [], [], []
        for s in range(seeds):
            rng = np.random.default_rng(SEED + 5 * s + int(10 * c))
            theta = (rng.integers(2, size=d) * 2.0 - 1.0) * Delta
            r, zs = alg_thompson(T, d, theta, sigma, rng)
            nz2 = np.sum(zs * zs, axis=1)
            nz2 = np.where(nz2 > 0, nz2, 1.0)
            # KL per round between theta and its i-th sign flip (Gaussian observations with
            # variance sigma^2 ||z_t||^2):  (2 Delta z_{ti})^2 / (2 sigma^2 ||z_t||^2)
            kl_i = np.sum(
                (2 * Delta) ** 2 * zs**2 / (2 * sigma**2 * nz2[:, None]), axis=0
            )
            kls.append(float(np.sum(kl_i)))
            # how many coordinate signs did the algorithm's final direction get right?
            correct.append(float(np.mean(np.sign(-zs[-1]) == np.sign(theta))))
            regs.append(r)
        rows.append(
            dict(
                Delta_coefficient=c,
                mean_total_KL=float(np.mean(kls)),
                predicted_total_KL_4dDelta2T=float(
                    4 * d * Delta**2 * T / (2 * sigma**2) / d * 1.0
                ),
                closed_form_2Delta2T_over_sigma2=float(2 * Delta**2 * T / sigma**2),
                mean_fraction_of_signs_recovered=float(np.mean(correct)),
                mean_regret=float(np.mean(regs)),
                regret_over_sqrt_dT=float(np.mean(regs) / np.sqrt(d * T)),
            )
        )
    return dict(
        d=d,
        T=T,
        seeds=seeds,
        rows=rows,
        note="sum_i KL = 2 Delta^2 T / sigma^2 exactly (||z_t||=1 for this "
        "algorithm); with sigma^2 = 1/(2d) and Delta = c/sqrt(T) this is "
        "4 c^2 d, i.e. O(1) KL per coordinate -- the signs are unidentifiable "
        "and chance-level recovery is expected",
    )


def test5_truncated(seeds=12):
    rows = []
    for d in [8, 16, 32]:
        T = 4000
        sigma = 1.0 / np.sqrt(2 * d)
        best = np.inf
        for name, fn in ALGS.items():
            worst = -np.inf
            for c in [0.5, 1.0, 2.0]:
                Delta = c / np.sqrt(T)
                vals = []
                for s in range(seeds):
                    rng = np.random.default_rng(SEED + 9 * s + d + int(10 * c))
                    theta = (rng.integers(2, size=d) * 2.0 - 1.0) * Delta
                    vals.append(fn(T, d, theta, sigma, rng, truncate=True))
                worst = max(worst, float(np.mean(vals)))
            best = min(best, worst)
        rows.append(
            dict(
                d=d,
                T=T,
                envelope_regret=best,
                over_sqrt_dT=best / float(np.sqrt(d * T)),
                above_theorem_floor=bool(best >= theorem_floor(d, T)),
            )
        )
    return rows


if __name__ == "__main__":
    res = dict(
        claim="claim-5 Theorem 5.2 Omega(sqrt(dT)) lower bound on the unit ball",
        seed=SEED,
        construction="theta in {+-Delta}^d, Delta = c/sqrt(T), l_t = theta + "
        "N(0,(2d)^{-1} I), action set = unit Euclidean ball",
        T1_T2_rates=test1_and_2(),
        T3_noise_variance_mechanism=test3_variance_mechanism(),
        T4_information_argument=test4_information(),
        T5_truncated_bounded_losses=test5_truncated(),
    )
    os.makedirs(OUT, exist_ok=True)
    with open(os.path.join(OUT, "claim5_lowerbound.json"), "w") as f:
        json.dump(res, f, indent=1)
    print(json.dumps(res, indent=1))