repro-bundle-v0/repro/network_env.py

"""
Synthetic "online shortest path with a long-term delay constraint" testbed,
built to mirror the structure of the paper's own experiments (Section 7 /
Appendix E: online shortest path over network measurements, comparing
against a fixed-horizon + doubling-trick baseline). We use a synthetic
layered graph and synthetic non-stationary/adversarial edge traffic instead
of the paper's RIPE-Atlas-derived dataset (no network dataset access is
required here), while keeping the same problem structure: a decision set
that is a probability simplex over candidate paths, linear round costs
(path latency-under-tolls) and linear round constraints (path delay budget).

Scale: N=16 nodes in 4 layers (width 4), 64 source-sink paths, K=48
candidate paths sub-sampled as the action set, T up to 40,000 rounds -- not
a toy size for a convex-optimization OCO experiment (this is exactly the
"vector of size K, T iterations" regime the paper's own Theorem statements
are about), and each full run costs well under a second of CPU time.
"""
from __future__ import annotations

import numpy as np


class PathNetwork:
    def __init__(self, layer_width: int = 4, n_layers: int = 3, n_paths: int = 48, seed: int = 0):
        rng = np.random.default_rng(seed)
        # layers: source(1) -> L1(w) -> L2(w) -> ... -> Ln(w) -> sink(1)
        self.layer_width = layer_width
        self.n_layers = n_layers
        layers = [1] + [layer_width] * n_layers + [1]
        # enumerate all source-sink paths as sequences of node indices per layer
        choices = [range(w) for w in layers[1:-1]]
        import itertools

        all_paths = list(itertools.product(*choices))
        rng.shuffle(all_paths)
        self.paths = all_paths[:n_paths]
        self.K = len(self.paths)
        self.layers = layers

        # base (mean) edge cost / delay per (layer, from_idx, to_idx)
        self.n_edge_layers = n_layers + 1
        self.base_cost = [rng.uniform(0.5, 2.0, size=(layers[i], layers[i + 1])) for i in range(self.n_edge_layers)]
        self.base_delay = [rng.uniform(0.5, 2.0, size=(layers[i], layers[i + 1])) for i in range(self.n_edge_layers)]
        # per-edge drift phase/frequency for smooth non-stationarity
        self.cost_phase = [rng.uniform(0, 2 * np.pi, size=b.shape) for b in self.base_cost]
        self.delay_phase = [rng.uniform(0, 2 * np.pi, size=b.shape) for b in self.base_delay]
        self.rng = rng

    def _edge_path_indices(self, path):
        # returns list of (layer, i, j) edges along this path (source and sink are singleton layers)
        seq = [0] + list(path) + [0]
        return list(enumerate(zip(seq[:-1], seq[1:])))

    def edge_values(self, t: int, noise_scale: float, spike_prob: float, spike_scale: float, rng: np.random.Generator):
        """Return (cost_edges, delay_edges): lists of matrices, one per edge-layer,
        for round t. Smooth drift + gaussian noise + rare adversarial spikes."""
        cost_edges = []
        delay_edges = []
        for li in range(self.n_edge_layers):
            drift_c = 0.5 * np.sin(0.01 * t + self.cost_phase[li])
            drift_d = 0.5 * np.sin(0.013 * t + self.delay_phase[li] + 1.0)
            noise_c = rng.normal(0, noise_scale, size=self.base_cost[li].shape)
            noise_d = rng.normal(0, noise_scale, size=self.base_delay[li].shape)
            spikes_c = (rng.uniform(size=self.base_cost[li].shape) < spike_prob) * spike_scale
            spikes_d = (rng.uniform(size=self.base_delay[li].shape) < spike_prob) * spike_scale
            c = np.clip(self.base_cost[li] + drift_c + noise_c + spikes_c, 0.05, None)
            d = np.clip(self.base_delay[li] + drift_d + noise_d + spikes_d, 0.05, None)
            cost_edges.append(c)
            delay_edges.append(d)
        return cost_edges, delay_edges

    def path_vectors(self, cost_edges, delay_edges):
        c_vec = np.zeros(self.K)
        d_vec = np.zeros(self.K)
        for k, path in enumerate(self.paths):
            edges = self._edge_path_indices(path)
            c_vec[k] = sum(cost_edges[li][i, j] for li, (i, j) in edges)
            d_vec[k] = sum(delay_edges[li][i, j] for li, (i, j) in edges)
        return c_vec, d_vec


def make_round_sequence(
    net: PathNetwork,
    T: int,
    G: float,
    delay_budget_quantile: float = 0.55,
    noise_scale: float = 0.15,
    spike_prob: float = 0.01,
    spike_scale: float = 3.0,
    seed: int = 1,
):
    """Precompute T rounds of (c_t, d_t) path-cost/path-delay vectors, each
    rescaled to have L2 norm exactly G (so the Lipschitz constant of every
    round's linear cost/constraint function over the simplex is exactly G,
    letting us check the paper's constants exactly rather than up to unknown
    slack). Returns c_seq (T,K), d_seq (T,K), and the delay budget Delta.
    """
    rng = np.random.default_rng(seed)
    c_raw = np.zeros((T, net.K))
    d_raw = np.zeros((T, net.K))
    for t in range(1, T + 1):
        cost_edges, delay_edges = net.edge_values(t, noise_scale, spike_prob, spike_scale, rng)
        c_vec, d_vec = net.path_vectors(cost_edges, delay_edges)
        c_raw[t - 1] = c_vec
        d_raw[t - 1] = d_vec

    # delay budget: a quantile of the raw per-path delays, so the constraint
    # is neither vacuous nor infeasible for every path
    Delta = float(np.quantile(d_raw, delay_budget_quantile))

    # Guarantee path 0 is a "reliable low-delay route": its raw delay is kept
    # well under Delta in every single round. This guarantees the globally
    # feasible set {x : g_t(x)<=0 for all t} is non-empty (contains at least
    # the vertex e_0), so the static-regret LP comparator below is always
    # solvable -- without this, T rounds of independent random constraints
    # can jointly rule out every fixed action.
    safe_level = 0.15 * Delta
    d_raw[:, 0] = safe_level + rng.uniform(0, 0.02 * max(Delta, 1e-6), size=T)

    d_centered = d_raw - Delta

    def rescale(mat):
        norms = np.linalg.norm(mat, axis=1, keepdims=True)
        norms = np.maximum(norms, 1e-9)
        return mat / norms * G

    c_seq = rescale(c_raw)
    d_seq = rescale(d_centered)
    return c_seq, d_seq


def iid_adversarial_sequence(K: int, T: int, G: float, seed: int, violation_rate: float = 0.5):
    """A harder, structure-free instance: i.i.d. random unit-norm cost and
    constraint vectors each round (no smooth drift to exploit). This is the
    classic hard case for online linear optimization -- even the best FIXED
    action in hindsight is itself a noisy quantity, which is what forces
    Theta(sqrt(t)) regret for *any* algorithm and lets us see the claimed
    growth rate directly, rather than the near-zero regret a trackable,
    slowly-drifting environment produces."""
    rng = np.random.default_rng(seed)
    c_raw = rng.standard_normal((T, K))
    d_raw = rng.standard_normal((T, K))
    Delta = float(np.quantile(d_raw, violation_rate))
    d_raw[:, 0] = np.quantile(d_raw, 0.05) * np.ones(T)  # guaranteed-safe column
    d_centered = d_raw - Delta

    def rescale(mat):
        norms = np.linalg.norm(mat, axis=1, keepdims=True)
        return mat / np.maximum(norms, 1e-9) * G

    return rescale(c_raw), rescale(d_centered)


def block_adversarial_sequence(K: int, T: int, block_size: int, G: float, seed: int, violation_rate: float = 0.5):
    """Piecewise-constant adversarial sequence: cost/constraint vectors are
    redrawn i.i.d. every `block_size` rounds and held fixed within a block.
    This gives a direct, controllable knob on the comparator's path length
    P_T for the dynamic-regret experiment: the per-round best-response
    minimizer only moves between blocks, so smaller `block_size` (more,
    smaller blocks) yields larger P_T, while larger `block_size` yields
    smaller P_T -- without needing thousands of independent LP solves
    (only one LP per unique block is needed, since the optimum is constant
    within a block).
    """
    rng = np.random.default_rng(seed)
    n_blocks = int(np.ceil(T / block_size))
    c_blocks = rng.standard_normal((n_blocks, K))
    d_blocks = rng.standard_normal((n_blocks, K))
    Delta = float(np.quantile(d_blocks, violation_rate))
    d_blocks[:, 0] = np.quantile(d_blocks, 0.05)  # guaranteed-safe column, all blocks
    d_blocks = d_blocks - Delta

    def rescale(mat):
        norms = np.linalg.norm(mat, axis=1, keepdims=True)
        return mat / np.maximum(norms, 1e-9) * G

    c_blocks = rescale(c_blocks)
    d_blocks = rescale(d_blocks)

    block_id = np.repeat(np.arange(n_blocks), block_size)[:T]
    c_seq = c_blocks[block_id]
    d_seq = d_blocks[block_id]
    return c_seq, d_seq, block_id, c_blocks, d_blocks


def cost_constraint_fns(c_seq: np.ndarray, d_seq: np.ndarray):
    """Build cost_fn(t,x)->(f,grad) and constraint_fn(t,x)->(g,grad) closures
    for linear round functions f_t(x)=<c_t,x>, g_t(x)=<d_t,x> (already budget-centered)."""

    def cost_fn(t, x):
        c = c_seq[t - 1]
        return float(np.dot(c, x)), c

    def constraint_fn(t, x):
        d = d_seq[t - 1]
        return float(np.dot(d, x)), d

    return cost_fn, constraint_fn