scripts/batch.py

"""Seed-batched implementations of Algorithm 6 and of the PABLO loop, so that hundreds of
independent replicates can be run in one vectorised pass.  Mathematically identical to the
reference implementations in pablo.py (checked in scripts/claim2_static.py)."""

from __future__ import annotations

import numpy as np


class BatchAlg6:
    """Algorithm 6 replicated over S independent seeds.  State shape (S, |S_eta|, d)."""

    EXP_CLIP = 300.0

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

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

    def update(self, g):  # g: (S, d)
        k, alpha, gamma = self.k, self.alpha, self.gamma
        etas = self.etas[None, :, None]
        nw = np.linalg.norm(self.W, axis=2, keepdims=True)
        theta = (
            (k / etas) * np.log(nw / alpha + 1.0) * (self.W / np.where(nw > 0, nw, 1.0))
        )
        theta = theta - g[:, None, :]
        nt = np.linalg.norm(theta, axis=2, keepdims=True)
        gg = np.sum(g * g, axis=1)[:, None, None]
        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.minimum(x, self.EXP_CLIP)
        mag = alpha * np.maximum(np.expm1(x), 0.0)
        self.W = mag * theta / np.where(nt > 0, nt, 1.0)


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

    def play(self):
        return self.w

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


def run_pablo_batch(losses, olo, d, eps, rng, S, comparators=None, track_w=False):
    """PABLO (Algorithm 1, isotropic H_t of Eq. (4)) over S independent seeds sharing one
    oblivious loss sequence.

    Returns per-seed arrays:
      cum_play   sum_t <l_t, wtilde_t>            (so R_T(u) = cum_play - <sum_t l_t, u> for
                                                   any static comparator u)
      dyn_play   sum_t <l_t, wtilde_t - u_t>      (dynamic regret, if comparators given)
      olo_reg    sum_t <ltilde_t, w_t - u_t>      (OLO regret of the subroutine on estimates)
      sum_lt_sq  sum_t ||ltilde_t||^2
      sum_w_sq   sum_t ||w_t||^2
      max_w      max_t ||w_t||
    """
    T = len(losses)
    cum_play = np.zeros(S)
    dyn_play = np.zeros(S)
    olo_reg = np.zeros(S)
    sum_lt_sq = np.zeros(S)
    sum_w_sq = np.zeros(S)
    max_w = np.zeros(S)
    sum_dir = np.zeros(S)   # sum_t <z_t, l_t> with z_t = wtilde_t/||wtilde_t||  (Eq. (7))
    sqd = np.sqrt(d)
    for t in range(T):
        w = olo.play()  # (S, d)
        ell = losses[t]  # (d,)
        nw = np.linalg.norm(w, axis=1)
        m = np.maximum(nw, eps)  # (S,)
        idx = rng.integers(d, size=S)
        sgn = rng.integers(2, size=S) * 2.0 - 1.0
        Smat = np.zeros((S, d))
        Smat[np.arange(S), idx] = sgn
        wt = w + (sqd * m)[:, None] * Smat
        obs = wt @ ell  # (S,) bandit feedback
        lt = (sqd / m)[:, None] * Smat * obs[:, None]
        cum_play += obs
        sum_dir += obs / np.maximum(np.linalg.norm(wt, axis=1), 1e-300)
        u = comparators[t] if comparators is not None else None
        if u is not None:
            dyn_play += obs - float(np.dot(ell, u))
            olo_reg += np.sum(lt * (w - u[None, :]), axis=1)
        else:
            olo_reg += np.sum(lt * w, axis=1)
        sum_lt_sq += np.sum(lt * lt, axis=1)
        sum_w_sq += nw**2
        max_w = np.maximum(max_w, nw)
        olo.update(lt)
    return dict(
        cum_play=cum_play,
        dyn_play=dyn_play,
        olo_reg=olo_reg,
        sum_lt_sq=sum_lt_sq,
        sum_w_sq=sum_w_sq,
        max_w=max_w,
        sum_dir=sum_dir,
        sum_losses=losses.sum(axis=0),
        V_T=float(np.sum(losses * losses)),
    )


def fit_exponent(xs, ys):
    """Least-squares slope of log y vs log x, with a standard error."""
    lx, ly = np.log(np.asarray(xs, float)), np.log(np.asarray(ys, float))
    A = np.vstack([lx, np.ones_like(lx)]).T
    coef, res, *_ = np.linalg.lstsq(A, ly, rcond=None)
    pred = A @ coef
    n = len(lx)
    if n > 2:
        s2 = float(np.sum((ly - pred) ** 2) / (n - 2))
        se = float(np.sqrt(s2 * np.linalg.inv(A.T @ A)[0, 0]))
    else:
        se = float("nan")
    return float(coef[0]), float(coef[1]), se


class BatchAlg5:
    """A single Algorithm-5 instance (one fixed step size eta) replicated over S seeds.
    Used as the *oracle-tuned* / *mistuned* baseline against Algorithm 6, which has to
    aggregate over the whole eta-grid because it does not know P_T."""

    EXP_CLIP = 300.0

    def __init__(self, S, d, T, L, epsilon, eta, k=4.0):
        self.eta, self.k = eta, k
        self.alpha = epsilon / T
        self.gamma = L / T
        self.w = np.zeros((S, d))
        self.clipped = 0

    def play(self):
        return self.w

    def update(self, g):
        k, eta, alpha, gamma = self.k, self.eta, self.alpha, self.gamma
        nw = np.linalg.norm(self.w, axis=1, keepdims=True)
        theta = (k / eta) * np.log(nw / alpha + 1.0) * (self.w / np.where(nw > 0, nw, 1.0))
        theta = theta - g
        nt = np.linalg.norm(theta, axis=1, keepdims=True)
        gg = np.sum(g * g, axis=1)[:, None]
        x = (eta / k) * (nt - 0.5 * eta * gg - gamma)
        if np.any(x > self.EXP_CLIP):
            self.clipped += int(np.sum(x > self.EXP_CLIP))
            x = np.minimum(x, self.EXP_CLIP)
        self.w = alpha * np.maximum(np.expm1(x), 0.0) * theta / np.where(nt > 0, nt, 1.0)