repro-bundle-v0/src/fab.py

"""Algorithm 1 (FAB) and its single-level special case (Push-Pull / AB), exactly
as specified in the paper's Eqs. (6)-(8) and Section 3.3.

FAB (bilevel, Eqs 6-8): each agent i keeps decision vars (x_i, y_i, z_i) and
tracking vars (t_{x,i}, t_{y,i}, t_{z,i}). Per iteration k:
  1. Decision update:  v_i^{k+1} = sum_j a_ij^k v_j^k - eta_v^k t_{v,i}^k   (v in {x,y,z})
  2. Local gradients:  d_{x,i}^{k+1} = grad_x L_i,  d_{y,i}^{k+1} = grad_y L_i,
                       d_{z,i}^{k+1} = -grad_z L_i     (L_i = f_i + lambda*(g_i(x,y)-g_i(x,z)))
  3. Tracking update:  t_{v,i}^{k+1} = sum_j b_ij^k t_{v,j}^k + d_{v,i}^{k+1} - d_{v,i}^k

Push-Pull (single-level, Sec 3.3): same scheme with only the x-block, on
f_i(x) directly (no y/z, no lambda).
"""
from __future__ import annotations

import numpy as np


def run_fab(problem, graphs, K: int, eta_x: float, eta_y: float, eta_z: float,
            lam: float, seed: int = 0, noise_std: float = 0.0):
    """Runs FAB for K iterations. Returns per-iteration x_bar, and running-min
    trackers for the true hypergradient norm^2 and the x/y/z consensus errors,
    as in Theorem 3.4 / Proposition 3.5 (all are `min_{0<=k<K-1}` statistics).

    `noise_std` adds i.i.d. Gaussian noise to each agent's local gradient
    (stochastic-gradient oracle) -- without it, this deterministic smooth
    instance converges geometrically fast and never exercises the theorem's
    sublinear worst-case rate; persistent noise is what keeps the problem
    "hard" across the whole horizon K, the same role played by the paper's
    own noise level omega (RL task) / mini-batch sampling (hyper-cleaning).
    """
    n, d = problem.n, problem.d
    rng = np.random.default_rng(seed)
    x = rng.normal(scale=0.1, size=(n, d))
    y = rng.normal(scale=0.1, size=(n, d))
    z = y.copy()

    def local_grads(x, y, z):
        dx = np.empty((n, d))
        dy = np.empty((n, d))
        dz = np.empty((n, d))
        for i in range(n):
            dx[i] = problem.grad_x_f_i(i, x[i], y[i]) + lam * (
                problem.grad_x_g_i(i, x[i], y[i]) - problem.grad_x_g_i(i, x[i], z[i]))
            dy[i] = problem.grad_y_f_i(i, x[i], y[i]) + lam * problem.grad_y_g_i(i, x[i], y[i])
            dz[i] = lam * problem.grad_y_g_i(i, x[i], z[i])  # d_z = -grad_z[L_i] = -grad_z[-lambda g_i(x,z)]
        if noise_std > 0:
            dx = dx + rng.normal(scale=noise_std, size=dx.shape)
            dy = dy + rng.normal(scale=noise_std, size=dy.shape)
            dz = dz + rng.normal(scale=noise_std, size=dz.shape)
        return dx, dy, dz

    d0x, d0y, d0z = local_grads(x, y, z)
    tx, ty, tz = d0x.copy(), d0y.copy(), d0z.copy()

    grad_norm_hist = np.empty(K)
    cons_x_hist = np.empty(K)
    cons_y_hist = np.empty(K)
    cons_z_hist = np.empty(K)

    for k in range(K):
        A, B = graphs.step()
        x_new = A @ x - eta_x * tx
        y_new = A @ y - eta_y * ty
        z_new = A @ z - eta_z * tz

        d1x, d1y, d1z = local_grads(x_new, y_new, z_new)
        tx = B @ tx + d1x - d0x
        ty = B @ ty + d1y - d0y
        tz = B @ tz + d1z - d0z
        d0x, d0y, d0z = d1x, d1y, d1z
        x, y, z = x_new, y_new, z_new

        x_bar = x.mean(axis=0)
        grad_norm_hist[k] = float(np.sum(problem.true_hypergradient(x_bar) ** 2))
        cons_x_hist[k] = float(np.mean(np.sum((x - x_bar) ** 2, axis=1)))
        cons_y_hist[k] = float(np.mean(np.sum((y - y.mean(axis=0)) ** 2, axis=1)))
        cons_z_hist[k] = float(np.mean(np.sum((z - z.mean(axis=0)) ** 2, axis=1)))

    return {
        "grad_norm_running_min": np.minimum.accumulate(grad_norm_hist),
        "cons_x_running_min": np.minimum.accumulate(cons_x_hist),
        "cons_y_running_min": np.minimum.accumulate(cons_y_hist),
        "cons_z_running_min": np.minimum.accumulate(cons_z_hist),
        "grad_norm_hist": grad_norm_hist,
    }


def run_pushpull(problem, graphs, K: int, eta_x: float, seed: int = 0, noise_std: float = 0.0):
    """Single-level Push-Pull (Theorem 3.6 / Eq. 9): min_x (1/n) sum f_i(x)."""
    n, d = problem.n, problem.d
    rng = np.random.default_rng(seed)
    x = rng.normal(scale=0.1, size=(n, d))

    def grads(x):
        g = np.array([problem.grad_i(i, x[i]) for i in range(n)])
        if noise_std > 0:
            g = g + rng.normal(scale=noise_std, size=g.shape)
        return g

    d0 = grads(x)
    tx = d0.copy()

    grad_norm_hist = np.empty(K)
    cons_hist = np.empty(K)

    for k in range(K):
        A, B = graphs.step()
        x_new = A @ x - eta_x * tx
        d1 = grads(x_new)
        tx = B @ tx + d1 - d0
        d0 = d1
        x = x_new

        x_bar = x.mean(axis=0)
        grad_norm_hist[k] = float(np.sum(problem.true_gradient(x_bar) ** 2))
        cons_hist[k] = float(np.mean(np.sum((x - x_bar) ** 2, axis=1)))

    return {
        "grad_norm_running_min": np.minimum.accumulate(grad_norm_hist),
        "cons_running_min": np.minimum.accumulate(cons_hist),
    }