scripts/pablo.py

"""Independent re-implementation of PABLO (Jacobsen, Baudry, Ito, Cesa-Bianchi, ICML 2026),
"A Perturbation Approach to Unconstrained Linear Bandits" (arXiv:2603.28201).

Everything here is written from the paper's pseudocode only (Algorithm 1, Algorithm 5,
Algorithm 6, Eq. (4)); no author code exists (no repository is linked in the paper, on
OpenReview, or on arXiv).

Conventions
-----------
* Domain W = R^d (unconstrained), losses f_t(w) = <l_t, w>, ||l_t|| <= G.
* uBLO feedback: the learner plays wtilde_t and observes only the scalar <l_t, wtilde_t>.
* Regret is measured on the *played* point:  R_T(u_{1:T}) = sum_t <l_t, wtilde_t - u_t>.
"""

from __future__ import annotations

import numpy as np

# --------------------------------------------------------------------------------------
# Algorithm 1: PABLO  (Section 2 of the paper)
# --------------------------------------------------------------------------------------


def pablo_perturbation(w, d, eps):
    """Isotropic choice of Eq. (4): H_t = I / (d * (||w_t||^2 v eps^2)), so that
    H_t <= I / (d(||w_t||^2 v eps^2)) holds with equality (Corollary 2.2 applies).

    Eigenvectors of H_t are the standard basis; S = {+-e_i}.
    Returns (H_scale, m) where H_t = (1/(d m^2)) I and m = max(||w||, eps).
    """
    m = max(float(np.linalg.norm(w)), eps)
    return 1.0 / (d * m * m), m


def pablo_step(w, ell, d, eps, rng):
    """One round of Algorithm 1 with the isotropic H_t of Eq. (4).

    H_t^{-1/2} = sqrt(d) * m * I ,  H_t^{1/2} = 1/(sqrt(d) m) * I.
      wtilde_t = w_t + H_t^{-1/2} s_t = w_t + sqrt(d) * m * s_t
      ltilde_t = d H_t^{1/2} s_t <wtilde_t, l_t> = (sqrt(d)/m) * s_t * <wtilde_t, l_t>
    Returns (wtilde, ltilde, i, sigma).
    """
    _, m = pablo_perturbation(w, d, eps)
    i = int(rng.integers(d))
    sigma = 1.0 if rng.integers(2) == 0 else -1.0
    s = np.zeros(d)
    s[i] = sigma
    wt = w + np.sqrt(d) * m * s
    obs = float(np.dot(wt, ell))  # <-- the ONLY feedback the learner gets
    lt = (np.sqrt(d) / m) * s * obs
    return wt, lt, i, sigma


def pablo_all_outcomes(w, ell, d, eps, c=1.0):
    """Exhaustive enumeration of all 2d equally-likely perturbation outcomes of a PABLO
    round.  Used for exact (non-Monte-Carlo) verification of Proposition 2.1."""
    m = max(float(np.linalg.norm(w)), eps) / np.sqrt(c)   # H = c I/(d(||w||^2 v eps^2))
    W = np.repeat(w[None, :], 2 * d, axis=0)
    S = np.zeros((2 * d, d))
    for i in range(d):
        S[2 * i, i] = 1.0
        S[2 * i + 1, i] = -1.0
    WT = W + np.sqrt(d) * m * S
    obs = WT @ ell
    LT = (np.sqrt(d) / m) * S * obs[:, None]
    return WT, LT


# --------------------------------------------------------------------------------------
# Algorithm 5: refined dynamic base algorithm for OLO (Appendix E.2)
# --------------------------------------------------------------------------------------
#   psi(w) = (k/eta) int_0^{||w-w1||} log(x/alpha + 1) dx
#   theta_t = (k/eta) log(||w_t-w1||/alpha + 1) (w_t-w1)/||w_t-w1||  -  g_t
#   w_{t+1} = w1 + theta_t/||theta_t|| * alpha [ exp( (eta/k)(||theta_t|| - eta/2 ||g_t||^2
#                                                           - gamma) ) - 1 ]_+
# (the exponent's eta/k factor is the inverse of the k/eta scale of psi; it is the unique
#  parsing that makes the mirror-descent optimality condition and Theorem E.6 consistent,
#  and Theorem E.6's inequality is checked numerically in scripts/claim3_dynamic.py)


class Alg6:
    """Algorithm 6 -- dynamic algorithm for unconstrained OLO.

    Runs one Algorithm-5 instance per step size eta_i = min(2^i/(T L), 1/L) and plays the
    SUM of their iterates.  alpha = epsilon/T, gamma = L/T, k = 4, w1 = 0.
    Requires no knowledge of P_T, of the comparator sequence, or of ||u||.

    Vectorised across the eta-grid: state W has shape (|S|, d).
    """

    EXP_CLIP = 300.0

    def __init__(self, d, T, L, epsilon, k=4.0):
        self.d, self.T, self.L, self.eps, self.k = d, T, L, epsilon, k
        n = int(np.ceil(np.log2(max(T, 2)))) + 1
        etas = np.array([min(2.0**i / (T * L), 1.0 / L) for i in range(n)])
        self.etas = etas
        self.alpha = epsilon / T
        self.gamma = L / T
        self.W = np.zeros((len(etas), d))
        self.clipped = 0

    def play(self):
        return self.W.sum(axis=0)

    def update(self, g):
        k, alpha, gamma = self.k, self.alpha, self.gamma
        etas = self.etas[:, None]
        nw = np.linalg.norm(self.W, axis=1, keepdims=True)
        safe = np.where(nw > 0, nw, 1.0)
        theta = (k / etas) * np.log(nw / alpha + 1.0) * (self.W / safe) - g[None, :]
        nt = np.linalg.norm(theta, axis=1, keepdims=True)
        gg = float(np.dot(g, g))
        x = (etas / k) * (nt - 0.5 * etas * gg - gamma)
        if np.any(x > self.EXP_CLIP):
            self.clipped += int(np.sum(x > self.EXP_CLIP))
        x = np.clip(x, -np.inf, self.EXP_CLIP)
        mag = alpha * np.maximum(np.expm1(x), 0.0)
        safe_t = np.where(nt > 0, nt, 1.0)
        self.W = mag * theta / safe_t


class OGD:
    """Baseline: online gradient descent with a fixed step size (needs oracle tuning)."""

    def __init__(self, d, lr):
        self.w = np.zeros(d)
        self.lr = lr

    def play(self):
        return self.w

    def update(self, g):
        self.w = self.w - self.lr * g


# --------------------------------------------------------------------------------------
# The full uBLO learner: PABLO wrapped around an OLO subroutine
# --------------------------------------------------------------------------------------


def run_pablo(losses, olo, d, eps, rng, comparators=None, record=False):
    """Run PABLO on a loss sequence.

    losses      : (T, d) array of l_t (the environment; never revealed to the learner)
    olo         : OLO subroutine object with .play() / .update(g)
    comparators : (T, d) comparator sequence u_{1:T} (only used for bookkeeping)

    Returns a dict with the realised regret and the diagnostics needed by the claims.
    """
    T = len(losses)
    reg = 0.0
    sum_lt_sq = 0.0  # sum ||ltilde_t||^2  (drives the kappa discussion)
    sum_l_sq = 0.0  # V_T = sum ||l_t||^2
    sum_w_sq = 0.0  # sum ||w_t||^2   (Proposition 4.1)
    olo_reg = 0.0  # regret of the OLO subroutine on the estimated losses
    max_lt = 0.0
    ws, wts = [], []
    for t in range(T):
        w = olo.play()
        ell = losses[t]
        wt, lt, _, _ = pablo_step(w, ell, d, eps, rng)
        u = comparators[t] if comparators is not None else np.zeros(d)
        reg += float(np.dot(ell, wt - u))
        olo_reg += float(np.dot(lt, w - u))
        sum_lt_sq += float(np.dot(lt, lt))
        sum_l_sq += float(np.dot(ell, ell))
        sum_w_sq += float(np.dot(w, w))
        max_lt = max(max_lt, float(np.linalg.norm(lt)))
        if record:
            ws.append(w.copy())
            wts.append(wt.copy())
        olo.update(lt)
    out = dict(
        regret=reg,
        olo_regret=olo_reg,
        sum_ltilde_sq=sum_lt_sq,
        V_T=sum_l_sq,
        sum_w_sq=sum_w_sq,
        max_ltilde=max_lt,
    )
    if record:
        out["w"] = np.array(ws)
        out["wtilde"] = np.array(wts)
    return out


# --------------------------------------------------------------------------------------
# Environments
# --------------------------------------------------------------------------------------


def env_stochastic(T, d, G, rng, bias=None, sigma=None):
    """l_t = theta + noise, projected to ||l_t|| <= G (sub-Gaussian stochastic instance)."""
    if bias is None:
        bias = np.zeros(d)
        bias[0] = G / np.sqrt(T)
    if sigma is None:
        sigma = G / np.sqrt(2 * d)
    L = bias[None, :] + sigma * rng.standard_normal((T, d))
    n = np.linalg.norm(L, axis=1, keepdims=True)
    return L * np.minimum(1.0, G / np.maximum(n, 1e-12))


def env_rademacher(T, d, G, rng):
    """Oblivious adversarial instance: l_t = G * (random +-1 coordinate direction)."""
    L = np.zeros((T, d))
    idx = rng.integers(d, size=T)
    sgn = rng.integers(2, size=T) * 2.0 - 1.0
    L[np.arange(T), idx] = G * sgn
    return L


def env_switching(T, d, G, rng, n_blocks=8):
    """Piecewise-constant adversarial instance (used with moving comparators)."""
    L = np.zeros((T, d))
    bl = np.array_split(np.arange(T), n_blocks)
    for b in bl:
        v = rng.standard_normal(d)
        v = G * v / np.linalg.norm(v)
        L[b] = v
    return L


def env_hard(T, d, G, rng, c=1.0, ret_theta=False):
    """The hypercube instance underlying Theorem 5.2 / the minimax rate: theta in {+-Delta}^d
    with Delta = c*G/sqrt(T), plus Gaussian noise with per-coordinate variance G^2/(2d) so
    that E||l_t||^2 <= G^2.  This is the loss family on which sqrt(dT) is the right rate."""
    theta = (rng.integers(2, size=d) * 2.0 - 1.0) * (c * G / np.sqrt(T))
    L = theta[None, :] + rng.standard_normal((T, d)) * (G / np.sqrt(2 * d))
    n = np.linalg.norm(L, axis=1, keepdims=True)
    L = L * np.minimum(1.0, G / np.maximum(n, 1e-12))
    return (L, theta) if ret_theta else L


def env_flip(T, d, G, rng, amp=0.25, n_flips=3):
    """Adversarial instance that first rewards scaling up and then reverses sign: the
    stress test for the risk-control term R_T(0) <= O(G eps) in Theorem 3.1."""
    v = rng.standard_normal(d)
    v = amp * G * v / np.linalg.norm(v)
    L = np.zeros((T, d))
    sign = -1.0
    for b, idx in enumerate(np.array_split(np.arange(T), n_flips + 1)):
        L[idx] = sign * v
        sign = -sign
    return L