repro-bundle/scripts/claim3_gradient_variance.py

"""Claim 3 (Theorem 3): O(sqrt(dW_T)+d) for linear, O(d sqrt(W_T)+d) for
convex, O((d/lam) log(dW_T)) for lambda-strongly-convex, in terms of the
gradient VARIANCE W_T = sup_{x1..xT in X} sum_t ||grad f_t(x_t) - mu_T||^2
(Eq 3.3), rather than gradient variation V_T.

Same algorithm (Algorithm 1) and same step-size forms as Theorem 1
(linear/convex: eta_t = R/sqrt(d^2+Vbar_{t-1}); strongly convex:
eta_t=1/(lambda t)) -- Theorem 3 is a reinterpretation of the same
non-consecutive-gradient-variation analysis in terms of W_T. We compute W_T
directly from the realized gradient sequence and check the bound-ratio stays
bounded across dimension sweeps, for linear, convex, and strongly-convex
function families.

FIXED after review: comparator sequence now uses the "iid" adversarial mode
(bco_core.DriftingQuadratic, mode="iid") instead of a slow random walk, so
regret is a genuine (non-benign, non-trivially-negative) stress test of the
bound -- see bco_core.py's docstring for why the random-walk mode let an
adaptive player trivially beat the fixed offline comparator.
"""
import json
import sys

import numpy as np

sys.path.insert(0, "scripts")
from bco_core import DriftingQuadratic, run_algorithm1  # noqa: E402

R = 1.0
T = 4000
D_VALUES = [2, 4, 8, 16, 32, 64]
N_SEEDS = 20


def gradient_variance(env: DriftingQuadratic, radius: float) -> float:
    # W_T = sup_{x_1..x_T} sum_t ||grad f_t(x_t) - mu_T||^2. Since grad f_t(x)
    # = lam*x + c_t is affine and the sup over the ball is attained at the
    # boundary aligned with (c_t - mu_c), we approximate the sup by evaluating
    # at x_t = R * (c_t - mean(c))/||c_t-mean(c)|| for each t (a valid, if not
    # exactly optimal, feasible point achieving a large deviation) -- this
    # gives a legitimate LOWER bound on the true sup, adequate for the
    # bound-ratio sanity check we run here.
    c = env.c[1:env.T + 1]
    mu_c = c.mean(axis=0)
    grads = []
    for t in range(1, env.T + 1):
        direction = c[t - 1] - mu_c
        norm = np.linalg.norm(direction)
        x = radius * direction / norm if norm > 1e-9 else np.zeros(env.d)
        grads.append(env.grad(t, x))
    grads = np.array(grads)
    mu = grads.mean(axis=0)
    return float(np.sum((grads - mu) ** 2))


def eta_linear_convex(d, T, L=1.0):
    def eta_fn(t, vbar_prev):
        return R / np.sqrt(d ** 2 + vbar_prev + 1e-8)
    return eta_fn


def eta_strongly_convex(lam):
    def eta_fn(t, vbar_prev):
        return 1.0 / (lam * t)
    return eta_fn


results = {"T": T, "d_values": D_VALUES, "families": {}}
for family, lam, bound_fn, label in [
    ("linear", 0.0, lambda d, W: np.sqrt(d * max(W, 1e-9)) + d, "O(sqrt(d W_T) + d)"),
    ("convex", 0.15, lambda d, W: d * np.sqrt(max(W, 1e-9)) + d, "O(d sqrt(W_T) + d)"),
    ("strongly_convex", 0.5, lambda d, W: (d / 0.5) * np.log(max(d * W, 2.0)), "O((d/lam) log(d W_T))"),
]:
    per_d = {}
    for d in D_VALUES:
        regrets, Ws = [], []
        for seed in range(N_SEEDS):
            np.random.seed(3000 * d + seed)
            env = DriftingQuadratic(d=d, T=T, lam=lam, seed=seed, mode="iid", c_scale=1.0)
            eta_fn = eta_strongly_convex(lam) if family == "strongly_convex" else eta_linear_convex(d, T)
            res = run_algorithm1(env, R, eta_fn)
            regrets.append(res["regret"])
            Ws.append(gradient_variance(env, R))
        mean_regret = float(np.mean(regrets))
        mean_W = float(np.mean(Ws))
        bound = bound_fn(d, mean_W)
        per_d[d] = {
            "mean_regret": mean_regret, "mean_W_T": mean_W,
            "bound_shape": bound, "regret_over_bound": mean_regret / bound if bound > 0 else None,
        }
        print(f"[{family}] d={d:3d}: regret={mean_regret:10.3f}  W_T={mean_W:10.3f}  "
              f"bound({label})={bound:10.3f}  ratio={mean_regret/bound:.4f}")
    results["families"][family] = {"label": label, "per_d": per_d}

with open("outputs/claim3_gradient_variance.json", "w") as f:
    json.dump(results, f, indent=2)
print("\nWrote outputs/claim3_gradient_variance.json")