scripts/claim1_reduction.py

"""CLAIM 1 -- PABLO reduces unconstrained Bandit Linear Optimization to standard Online
Linear Optimization (Section 3 / Algorithm 1 / Prop. 2.1, Cor. 2.2, Prop. 2.3).

Four independent tests:
 T1  EXHAUSTIVE check of Proposition 2.1.  A PABLO round has exactly 2d equally likely
     outcomes, so E[ltilde|F], E[||ltilde||^2|F] are computed EXACTLY by enumeration and
     compared with the paper's closed forms.
 T2  Corollary 2.2 constants (a.s. ||ltilde||^2 <= 4 d^2 ||l||^2 and E||ltilde||^2 <= 2d||l||^2),
     plus a BOUNDARY AUDIT: scale H_t by c > 1 so Eq. (4) is violated and locate the exact
     c at which the a.s. bound breaks.
 T3  The reduction identity itself: E[R_T^{uBLO}(u)] = E[R_T^{OLO}(u; ltilde)] for an
     F_0-measurable comparator -- i.e. the bandit problem IS the OLO problem in expectation.
 T4  Black-box modularity: three different OLO subroutines plugged into the same PABLO
     wrapper; the uBLO regret tracks each subroutine's OLO regret on the estimated losses.
"""

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 pablo import (
    Alg6,
    OGD,
    env_rademacher,
    env_stochastic,
    pablo_all_outcomes,
    run_pablo,
)

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


def t1_prop21():
    rng = np.random.default_rng(SEED)
    err_mean, err_2nd, n = 0.0, 0.0, 0
    for d in range(1, 13):
        for _ in range(200):
            w = rng.standard_normal(d) * 10 ** rng.uniform(-2, 2)
            ell = rng.standard_normal(d)
            ell = ell / np.linalg.norm(ell) * rng.uniform(0.1, 3)
            eps = float(10 ** rng.uniform(-3, 0))
            _, LT = pablo_all_outcomes(w, ell, d, eps)
            m = max(float(np.linalg.norm(w)), eps)
            exact_mean = LT.mean(axis=0)
            exact_2nd = float((LT * LT).sum(axis=1).mean())
            # paper's closed forms (Prop 2.1); Tr(H_t) = 1/m^2 for the isotropic choice
            pred_2nd = (
                d * float(np.dot(ell, ell)) + d * float(np.dot(ell, w)) ** 2 / m**2
            )
            scale = max(np.linalg.norm(ell), 1e-12)
            err_mean = max(err_mean, float(np.max(np.abs(exact_mean - ell))) / scale)
            err_2nd = max(err_2nd, abs(exact_2nd - pred_2nd) / max(pred_2nd, 1e-12))
            n += 1
    return dict(
        configs=n,
        d_range=[1, 12],
        max_rel_err_unbiasedness=err_mean,
        max_rel_err_second_moment=err_2nd,
    )


def t2_cor22_and_boundary():
    rng = np.random.default_rng(SEED + 1)
    worst_as, worst_exp, n = 0.0, 0.0, 0
    for d in range(1, 13):
        for _ in range(300):
            w = rng.standard_normal(d) * 10 ** rng.uniform(-2, 2)
            ell = rng.standard_normal(d)
            ell = ell / np.linalg.norm(ell) * rng.uniform(0.1, 3)
            eps = float(10 ** rng.uniform(-3, 0))
            _, LT = pablo_all_outcomes(w, ell, d, eps)
            g2 = float(np.dot(ell, ell))
            worst_as = max(
                worst_as, float(np.max((LT * LT).sum(axis=1))) / (4 * d * d * g2)
            )
            worst_exp = max(
                worst_exp, float((LT * LT).sum(axis=1).mean()) / (2 * d * g2)
            )
            n += 1
    # boundary audit: H_t = c I/(d(||w||^2 v eps^2)); Eq. (4) demands c <= 1.
    # a.s. bound is ||ltilde||^2 <= d ||l||^2 (||w||/m_c + sqrt(d))^2 with m_c = m/sqrt(c),
    # so the ratio to 4d^2||l||^2 grows like (sqrt(c)+sqrt(d))^2/(4d) when ||w||>>eps.
    audit = []
    for c in [1.0, 1.5, 2.0, 4.0, 10.0, 100.0, 1e4]:
        worst = 0.0
        for d in [1, 2, 4, 8, 16]:
            for _ in range(200):
                w = rng.standard_normal(d)
                w = w / np.linalg.norm(w) * 10 ** rng.uniform(-1, 2)
                ell = rng.standard_normal(d)
                ell = ell / np.linalg.norm(ell)
                _, LT = pablo_all_outcomes(w, ell, d, 1e-3, c=c)
                worst = max(worst, float(np.max((LT * LT).sum(axis=1))) / (4 * d * d))
        audit.append(
            dict(c=c, max_ratio_to_Cor22_as_bound=worst, holds=bool(worst <= 1.0))
        )
    return dict(
        configs=n,
        max_ratio_as_bound=worst_as,
        max_ratio_expectation_bound=worst_exp,
        boundary_audit=audit,
    )


def t3_reduction_identity():
    """E[sum_t <l_t, wtilde_t - u>]  ==  E[sum_t <ltilde_t, w_t - u>]  (Prop. 2.3, ghost
    term vanishes for an F_0-measurable comparator).  Monte-Carlo with paired seeds."""
    res = []
    for d, T, nseed in [(2, 200, 4000), (5, 300, 4000), (10, 200, 3000)]:
        rng0 = np.random.default_rng(SEED + 7 * d)
        G = 1.0
        losses = env_rademacher(T, d, G, rng0)
        u = rng0.standard_normal(d)
        u = u / np.linalg.norm(u) * 2.0
        U = np.repeat(u[None, :], T, axis=0)
        a, b = [], []
        for s in range(nseed):
            rng = np.random.default_rng(1000 + s)
            out = run_pablo(
                losses, Alg6(d, T, 2 * d * G, 1.0), d, 1e-3, rng, comparators=U
            )
            a.append(out["regret"])
            b.append(out["olo_regret"])
        a, b = np.array(a), np.array(b)
        diff = a - b
        se = float(diff.std(ddof=1) / np.sqrt(nseed))
        res.append(
            dict(
                d=d,
                T=T,
                seeds=nseed,
                mean_uBLO_regret=float(a.mean()),
                mean_OLO_regret_on_estimates=float(b.mean()),
                mean_difference=float(diff.mean()),
                stderr_of_difference=se,
                z_score=float(diff.mean() / se) if se > 0 else 0.0,
            )
        )
    return res


def t4_modularity():
    """Any OLO learner can be dropped in: the uBLO regret in expectation equals that
    learner's OLO regret on the estimated losses, whatever the learner is."""
    d, T, G, nseed = 6, 500, 1.0, 2000
    rng0 = np.random.default_rng(SEED + 3)
    losses = env_stochastic(T, d, G, rng0)
    u = np.zeros(d)
    u[0] = 3.0
    U = np.repeat(u[None, :], T, axis=0)
    L = 2 * d * G
    makers = {
        "Alg6 (paper, parameter-free)": lambda: Alg6(d, T, L, 1.0),
        "OGD lr=1/(L sqrt(T))": lambda: OGD(d, 1.0 / (L * np.sqrt(T))),
        "OGD lr=10/(L sqrt(T)) (mistuned)": lambda: OGD(d, 10.0 / (L * np.sqrt(T))),
    }
    out = {}
    for name, mk in makers.items():
        a, b = [], []
        for s in range(nseed):
            rng = np.random.default_rng(5000 + s)
            r = run_pablo(losses, mk(), d, 1e-3, rng, comparators=U)
            a.append(r["regret"])
            b.append(r["olo_regret"])
        a, b = np.array(a), np.array(b)
        out[name] = dict(
            mean_uBLO_regret=float(a.mean()),
            mean_OLO_regret=float(b.mean()),
            rel_gap=float(abs(a.mean() - b.mean()) / max(abs(b.mean()), 1e-9)),
            stderr=float((a - b).std(ddof=1) / np.sqrt(nseed)),
        )
    return dict(d=d, T=T, seeds=nseed, comparator_norm=3.0, learners=out)


def t5_conditional_exact():
    """Rao-Blackwellised (zero-Monte-Carlo-error) version of T3/T4: along REAL PABLO
    trajectories, enumerate all 2d outcomes of every round and check the two conditional
    identities that make the reduction work,
        E[<l_t, wtilde_t - u>|F_{t-1}] = <l_t, w_t - u> = E[<ltilde_t, w_t - u>|F_{t-1}].
    """
    worst_play, worst_est, rounds = 0.0, 0.0, 0
    for d, T in [(3, 300), (8, 300), (16, 200)]:
        G = 1.0
        rng0 = np.random.default_rng(SEED + 11 * d)
        losses = env_stochastic(T, d, G, rng0)
        u = rng0.standard_normal(d)
        u = u / np.linalg.norm(u) * 2.0
        U = np.repeat(u[None, :], T, axis=0)
        for mk in [
            lambda: Alg6(d, T, 2 * d * G, 1.0),
            lambda: OGD(d, 1.0 / (2 * d * G * np.sqrt(T))),
        ]:
            rng = np.random.default_rng(77)
            r = run_pablo(losses, mk(), d, 1e-3, rng, comparators=U, record=True)
            for t in range(T):
                w = r["w"][t]
                WT, LT = pablo_all_outcomes(w, losses[t], d, 1e-3)
                truth = float(np.dot(losses[t], w - u))
                e_play = float(np.mean(WT @ losses[t])) - float(np.dot(losses[t], u))
                e_est = float(np.mean(LT @ (w - u)))
                sc = max(abs(truth), 1e-9)
                worst_play = max(worst_play, abs(e_play - truth) / sc)
                worst_est = max(worst_est, abs(e_est - truth) / sc)
                rounds += 1
    return dict(
        rounds_checked=rounds,
        max_rel_err_played_point=worst_play,
        max_rel_err_estimated_loss=worst_est,
    )


if __name__ == "__main__":
    res = dict(
        claim="claim-1 PABLO reduces uBLO to OLO",
        seed=SEED,
        t5_conditional_identity_exact=t5_conditional_exact(),
        t1_proposition_2_1_exhaustive=t1_prop21(),
        t2_corollary_2_2_and_boundary=t2_cor22_and_boundary(),
        t3_reduction_identity=t3_reduction_identity(),
        t4_black_box_modularity=t4_modularity(),
    )
    os.makedirs(OUT, exist_ok=True)
    with open(os.path.join(OUT, "claim1_reduction.json"), "w") as f:
        json.dump(res, f, indent=1)
    print(json.dumps(res, indent=1))