scripts/alg5.py

"""Algorithm 5 of the paper (Appendix E.2, 'Refined Dynamic Base Algorithm for OLO'),
implemented standalone so that the guarantee of Theorem E.6 can be checked numerically.

    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 ]_+

Theorem E.6 (untuned form):
    R_T(u_{1:T}) <= 2k[ Phi(||u_T-w1||, 1/alpha) + Phi-path(k/(eta*alpha*gamma)) ] / (2 eta)
                    + (eta/2) sum ||g_t||^2 ||u_t-w1||
                    + gamma sum ||u_t - w1||
                    + eta*alpha*sum ||g_t||^2
with Phi(x, lam) = x log(lam x + 1).
"""

from __future__ import annotations

import numpy as np


class Alg5:
    EXP_CLIP = 300.0

    def __init__(self, d, alpha, eta, gamma, k=4.0, w1=None):
        self.d = d
        self.alpha, self.eta, self.gamma, self.k = alpha, eta, gamma, k
        self.w1 = np.zeros(d) if w1 is None else np.asarray(w1, float)
        self.w = self.w1.copy()
        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
        z = self.w - self.w1
        nz = float(np.linalg.norm(z))
        grad_psi = (
            (k / eta) * np.log(nz / alpha + 1.0) * (z / nz) if nz > 0 else 0.0 * z
        )
        theta = grad_psi - g
        nt = float(np.linalg.norm(theta))
        x = (eta / k) * (nt - 0.5 * eta * float(np.dot(g, g)) - gamma)
        if x > self.EXP_CLIP:
            self.clipped += 1
            x = self.EXP_CLIP
        mag = alpha * max(np.expm1(x), 0.0)
        self.w = self.w1 + (mag * theta / nt if nt > 0 else 0.0 * theta)


def phi(x, lam):
    return x * np.log(lam * x + 1.0)


def thm_E6_bound(u_seq, g_seq, alpha, eta, gamma, k=4.0, w1=None):
    u = np.asarray(u_seq, float)
    g = np.asarray(g_seq, float)
    w1 = np.zeros(u.shape[1]) if w1 is None else w1
    nu = np.linalg.norm(u - w1[None, :], axis=1)
    gn2 = np.sum(g * g, axis=1)
    path = np.linalg.norm(u[1:] - u[:-1], axis=1)
    PhiT = phi(nu[-1], 1.0 / alpha)
    PT_phi = float(np.sum(phi(path, k / (eta * alpha * gamma))))
    return (
        2 * k * (PhiT + PT_phi) / (2 * eta)
        + 0.5 * eta * float(np.sum(gn2 * nu))
        + gamma * float(np.sum(nu))
        + eta * alpha * float(np.sum(gn2))
    )


def run_alg5(g_seq, u_seq, alpha, eta, gamma, k=4.0):
    d = g_seq.shape[1]
    a = Alg5(d, alpha, eta, gamma, k)
    reg = 0.0
    for t in range(len(g_seq)):
        w = a.play()
        reg += float(np.dot(g_seq[t], w - u_seq[t]))
        a.update(g_seq[t])
    return reg, a.clipped