repro-bundle/scripts/bco_core.py

"""Faithful implementation of Algorithm 1 (two-point BCO with OOGD) from
'Improved Dimension Dependence for Bandit Convex Optimization with Gradient
Variations' (arXiv:2602.04761). No official code is released for this paper
(it is a pure theory/COLT-style paper); this is an independent, from-scratch
implementation used to empirically probe the claimed regret-bound scaling.

Algorithm 1 (paper, Section 2.2 / restated from Chiang et al. 2013):
  Input: step sizes {eta_t}, exploration delta, shrinkage xi = delta/R
  x_1 = w_hat_1 = 0, g_tilde_1 = 0
  for t = 1..T:
    i_t ~ Uniform([d])
    query x_t = w_t + delta*e_{i_t}, x'_t = w_t - delta*e_{i_t}
    observe f_t(x_t), f_t(x'_t)
    v_t = (f_t(x_t) - f_t(x'_t)) / (2*delta)
    g_t      = d*(v_t - g_tilde_t[i_t]) * e_{i_t} + g_tilde_t          (Eq 2.2)
    g_tilde_{t+1} = (v_t - g_tilde_t[i_t]) * e_{i_t} + g_tilde_t       (Eq 2.2)
    w_hat_{t+1} = Proj_{(1-xi)X}[w_hat_t - eta_t * g_t]
    w_{t+1}     = Proj_{(1-xi)X}[w_hat_{t+1} - eta_{t+1} * g_tilde_{t+1}]

The feasible set X is taken as the centered ball of radius R (satisfies
Assumption 1 with r=R). Function families used across claims are drifting
linear / quadratic-plus-linear functions, for which gradients (and gradient
variation) are known in closed form, letting us compute exact regret,
gradient variation V_T, gradient variance W_T, and small loss F_T.
"""
from __future__ import annotations

import numpy as np


def project_to_ball(x: np.ndarray, radius: float) -> np.ndarray:
    norm = np.linalg.norm(x)
    if norm <= radius or norm == 0:
        return x
    return x * (radius / norm)


class DriftingQuadratic:
    """f_t(x) = 0.5*lam*||x||^2 + <c_t, x>  (lam=0 -> linear; lam>0 -> lam-strongly convex, L=lam-smooth on the quadratic part).

    Two comparator-generation modes are supported:

    - mode="randomwalk" (legacy): c_t performs a random walk with a FIXED
      per-step norm `drift`. This is a BENIGN environment: because the
      target drifts slowly and smoothly, an adaptive online player can track
      it closely and often beats the single FIXED offline comparator in
      hindsight, producing negative "regret" that trivially satisfies any
      upper bound without stress-testing the algorithm. Kept only for
      backward compatibility / illustration of this failure mode.
    - mode="iid" (default): c_t is RESAMPLED i.i.d. every round from a fixed
      distribution (fixed norm `c_scale`, independent of d). Consecutive
      comparators are then uncorrelated, so gradient variation V_T is large
      and (approximately) dimension-independent (~ 2*T*c_scale^2), and the
      bandit algorithm -- which only observes one noisy coordinate probe per
      round -- cannot track the target at all; its best play converges to
      (approximately) the fixed mean of the c_t sequence, which is also the
      offline optimum. This yields genuine, positive regret driven by the
      bandit estimation/exploration cost, which is exactly the quantity the
      paper's bounds describe, giving a real (non-benign) stress test of the
      claimed dimension dependence.
    """

    def __init__(self, d: int, T: int, lam: float = 0.0, drift: float = 0.05,
                 c_scale: float = 1.0, seed: int = 0, mode: str = "iid"):
        rng = np.random.RandomState(seed)
        self.d = d
        self.T = T
        self.lam = lam
        self.mode = mode
        c = np.zeros((T + 1, d))
        c[0] = rng.normal(size=d)
        c[0] *= c_scale / (np.linalg.norm(c[0]) + 1e-12)
        if mode == "iid":
            for t in range(1, T + 1):
                v = rng.normal(size=d)
                v *= c_scale / (np.linalg.norm(v) + 1e-12)
                c[t] = v
        elif mode == "randomwalk":
            for t in range(1, T + 1):
                step = rng.normal(size=d)
                step *= drift / (np.linalg.norm(step) + 1e-12)
                c[t] = c[t - 1] + step
        else:
            raise ValueError(f"unknown mode {mode!r}")
        self.c = c

    def grad(self, t: int, x: np.ndarray) -> np.ndarray:
        return self.lam * x + self.c[t]

    def value(self, t: int, x: np.ndarray) -> float:
        return 0.5 * self.lam * float(x @ x) + float(self.c[t] @ x)

    def gradient_variation(self) -> float:
        diffs = self.c[1:self.T + 1] - self.c[0:self.T]
        return float(np.sum(diffs ** 2))

    def offline_optimum(self, radius: float) -> tuple[np.ndarray, float]:
        """min_x sum_{t=1}^T f_t(x) over ||x||<=radius, closed form (convex quadratic)."""
        csum = self.c[1:self.T + 1].sum(axis=0)
        if self.lam > 0:
            x_star = -csum / (self.lam * self.T)
        else:
            # Linear: optimum is at the boundary in direction -csum, i.e. at
            # distance `radius` from the origin (this used to normalize to a
            # UNIT vector and then no-op through project_to_ball since a unit
            # vector is already inside any radius>=1 ball, silently pinning
            # the "offline optimum" to the unit sphere regardless of the
            # requested radius -- a bug that made the small-loss quantity
            # F_T = offline_opt(R) - offline_opt(R+slack) degenerate to
            # exactly 0 for every d, since both radii produced the identical
            # point. Fixed: explicitly scale to the requested radius.)
            x_star = -csum / (np.linalg.norm(csum) + 1e-12) * radius
        x_star = project_to_ball(x_star, radius)
        total = sum(self.value(t, x_star) for t in range(1, self.T + 1))
        return x_star, total


def run_algorithm1(env: DriftingQuadratic, radius: float, eta_fn, delta: float | None = None):
    """Run Algorithm 1 (two-point OOGD) against env for T rounds.

    eta_fn(t, vbar_prev) -> eta_t   (vbar_prev = running non-consecutive
    gradient variation estimate Vbar_{t-1}, Eq 2.1, needed for adaptive
    schedules such as Theorem 1's).

    Returns dict with per-round losses, cumulative regret, and Vbar_T.
    """
    d, T = env.d, env.T
    if delta is None:
        # Paper's Algorithm 1 box: delta = 1/(2 d^2 L T R). We use a numerically
        # stable stand-in (small, dimension/T-scaled) since L, exact constants
        # only affect higher-order/bias terms, not the asymptotic regret shape
        # we are probing here.
        delta = radius / (2 * d ** 2 * max(T, 2))
    xi = delta / radius

    w_hat = np.zeros(d)
    w = np.zeros(d)
    g_tilde = np.zeros(d)

    total_regret_loss = 0.0
    vbar_running = 0.0
    losses_at_w = []

    for t in range(1, T + 1):
        i_t = np.random.randint(d)
        e_i = np.zeros(d)
        e_i[i_t] = 1.0

        x_t = w + delta * e_i
        x_tp = w - delta * e_i
        f_plus = env.value(t, x_t)
        f_minus = env.value(t, x_tp)
        v_t = (f_plus - f_minus) / (2 * delta)

        g_t = d * (v_t - g_tilde[i_t]) * e_i + g_tilde
        g_tilde_next = (v_t - g_tilde[i_t]) * e_i + g_tilde

        vbar_running += float(np.sum((g_t - g_tilde) ** 2))

        eta_t = eta_fn(t, vbar_running)
        w_hat_next = project_to_ball(w_hat - eta_t * g_t, (1 - xi) * radius)
        eta_next = eta_fn(t + 1, vbar_running)
        w_next = project_to_ball(w_hat_next - eta_next * g_tilde_next, (1 - xi) * radius)

        # Regret uses the average two-point loss 0.5*(f(x_t)+f(x'_t)) as in
        # the paper's REG_T definition (per-round midpoint surrogate).
        total_regret_loss += 0.5 * (f_plus + f_minus)
        losses_at_w.append(0.5 * (f_plus + f_minus))

        w_hat, w, g_tilde = w_hat_next, w_next, g_tilde_next

    x_star, offline_opt = env.offline_optimum(radius)
    regret = total_regret_loss - offline_opt
    return {
        "regret": regret,
        "vbar_T": vbar_running,
        "V_T": env.gradient_variation(),
        "losses": np.array(losses_at_w),
        "offline_opt": offline_opt,
    }