Модуль 3. Бандиты, исследование и одношаговые решения

from __future__ import annotations

from dataclasses import replace
from pathlib import Path
import sys
import math
import numpy as np
import matplotlib.pyplot as plt


def find_course_root() -> Path:
    candidates = [Path.cwd(), *Path.cwd().parents]
    for candidate in candidates:
        if (candidate / "modules").is_dir() and (candidate / "COURSE_MODULE_REVISION_PROTOCOL.md").is_file():
            return candidate
    raise FileNotFoundError("Course root was not found. Run the notebook from the repository tree.")

ROOT = find_course_root()
np.set_printoptions(precision=4, suppress=True)
print(f"Course root: {ROOT}")
M3_DIR = ROOT / "modules" / "module-03"
if str(M3_DIR) not in sys.path:
    sys.path.insert(0, str(M3_DIR))

from module_3_reference import (
    simulate_many_bandits,
    summarize_many,
    ips_value,
    snips_value,
    doubly_robust_value,
    categorical_entropy,
)

PROFILE = "smoke"  # change to "full" for the larger protocol
assert PROFILE in {"smoke", "full"}
print("Profile:", PROFILE)

probabilities = np.array([0.20, 0.35, 0.42, 0.55, 0.60])
horizon = 1_200 if PROFILE == "smoke" else 3_000
n_runs = 120 if PROFILE == "smoke" else 500
algorithms = {
    "greedy": ("greedy", 0.0),
    "epsilon-greedy (0.1)": ("epsilon_greedy", 0.1),
    "UCB1": ("ucb1", 0.0),
    "Thompson": ("thompson", 0.0),
}
curves = {}
for label, (algorithm, epsilon) in algorithms.items():
    result = simulate_many_bandits(probabilities, horizon, n_runs, 20260804, algorithm, epsilon=epsilon)
    curves[label] = summarize_many(result)

plt.figure(figsize=(9, 5))
for label, summary in curves.items():
    plt.plot(summary.mean_regret, label=label)
    plt.fill_between(np.arange(horizon), summary.lower_95, summary.upper_95, alpha=0.12)
plt.xlabel("Раунд")
plt.ylabel("Среднее накопленное псевдосожаление")
plt.title("Стационарный Bernoulli-бандит: среднее и 95% интервал")
plt.legend()
plt.tight_layout()
plt.show()

for label, summary in curves.items():
    print(f"{label:24s} final mean pseudo-regret = {summary.mean_regret[-1]:8.2f}; "
          f"optimal-arm rate = {summary.mean_optimal_rate[-1]:.3f}")

def stable_sigmoid(x: np.ndarray) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    out = np.empty_like(x)
    positive = x >= 0
    out[positive] = 1.0 / (1.0 + np.exp(-x[positive]))
    z = np.exp(x[~positive])
    out[~positive] = z / (1.0 + z)
    return out


def row_softmax(logits: np.ndarray) -> np.ndarray:
    logits = np.asarray(logits, dtype=float)
    shifted = logits - np.max(logits, axis=1, keepdims=True)
    e = np.exp(shifted)
    return e / e.sum(axis=1, keepdims=True)


def contextual_probabilities(x: np.ndarray) -> np.ndarray:
    """True expected reward q(x, a) for three routing actions."""
    x = np.asarray(x, dtype=float)
    z = np.column_stack([
        0.35 + 0.70 * x[:, 0] - 0.25 * x[:, 1],
        0.10 - 0.35 * x[:, 0] + 0.80 * x[:, 1] + 0.25 * x[:, 2],
       -0.15 + 0.30 * x[:, 0] + 0.10 * x[:, 1] + 0.95 * x[:, 2],
    ])
    return stable_sigmoid(z)


def logging_policy(x: np.ndarray, floor: float = 0.08) -> np.ndarray:
    logits = np.column_stack([
        0.25 + 0.35 * x[:, 0],
        0.05 + 0.30 * x[:, 1],
       -0.10 + 0.25 * x[:, 2],
    ])
    raw = row_softmax(logits)
    k = raw.shape[1]
    return (1.0 - k * floor) * raw + floor


def target_policy(x: np.ndarray, sharpness: float = 2.0) -> np.ndarray:
    q = contextual_probabilities(x)
    # The target prefers reward but includes small action-dependent deployment costs.
    costs = np.array([0.03, 0.08, 0.17])
    return row_softmax(sharpness * (q - costs))


def sample_log(n: int, seed: int, floor: float = 0.08) -> dict[str, np.ndarray]:
    rng = np.random.default_rng(seed)
    x = rng.normal(size=(n, 3))
    mu = logging_policy(x, floor=floor)
    q = contextual_probabilities(x)
    u = rng.random(n)
    cdf = np.cumsum(mu, axis=1)
    actions = (u[:, None] > cdf).sum(axis=1)
    rewards = (rng.random(n) < q[np.arange(n), actions]).astype(float)
    return {"x": x, "mu": mu, "q": q, "actions": actions, "rewards": rewards}


def features(x: np.ndarray) -> np.ndarray:
    x = np.asarray(x, dtype=float)
    return np.column_stack([
        np.ones(len(x)),
        x,
        x ** 2,
        x[:, 0] * x[:, 1],
        x[:, 0] * x[:, 2],
        x[:, 1] * x[:, 2],
    ])


def cross_fitted_reward_model(x: np.ndarray, actions: np.ndarray, rewards: np.ndarray,
                              folds: int = 5, ridge: float = 2.0, seed: int = 0) -> np.ndarray:
    """Per-action ridge reward model with grouped cross-fitting."""
    rng = np.random.default_rng(seed)
    n = len(x)
    order = rng.permutation(n)
    fold_id = np.empty(n, dtype=int)
    fold_id[order] = np.arange(n) % folds
    phi = features(x)
    pred = np.empty((n, 3), dtype=float)
    eye = np.eye(phi.shape[1])
    eye[0, 0] = 0.0  # do not regularize the intercept
    for fold in range(folds):
        train = fold_id != fold
        valid = ~train
        for action in range(3):
            rows = train & (actions == action)
            if rows.sum() < phi.shape[1]:
                raise RuntimeError("Insufficient action coverage for reward-model fitting")
            lhs = phi[rows].T @ phi[rows] + ridge * eye
            rhs = phi[rows].T @ rewards[rows]
            coef = np.linalg.solve(lhs, rhs)
            pred[valid, action] = phi[valid] @ coef
    return np.clip(pred, 0.0, 1.0)


def ope_estimates(log: dict[str, np.ndarray], sharpness: float = 2.0,
                  clip: float | None = None, seed: int = 0) -> dict[str, float]:
    x, mu, actions, rewards = log["x"], log["mu"], log["actions"], log["rewards"]
    pi = target_policy(x, sharpness=sharpness)
    p_logged = mu[np.arange(len(x)), actions]
    p_target = pi[np.arange(len(x)), actions]
    weights = p_target / p_logged
    if clip is not None:
        weights_used = np.minimum(weights, clip)
        p_target_used = weights_used * p_logged
    else:
        weights_used = weights
        p_target_used = p_target
    rhat = cross_fitted_reward_model(x, actions, rewards, seed=seed)
    dm_rows = np.sum(pi * rhat, axis=1)
    ips = float(np.mean(weights_used * rewards))
    snips = float(np.sum(weights_used * rewards) / np.sum(weights_used))
    dr = float(np.mean(dm_rows + weights_used * (rewards - rhat[np.arange(len(x)), actions])))
    ess = float(weights_used.sum() ** 2 / np.sum(weights_used ** 2))
    truth = float(np.mean(np.sum(pi * log["q"], axis=1)))
    return {
        "truth": truth,
        "IPS": ips,
        "SNIPS": snips,
        "DM": float(np.mean(dm_rows)),
        "DR": dr,
        "ESS": ess,
        "max_weight": float(np.max(weights_used)),
        "mean_weight": float(np.mean(weights_used)),
    }


def nonstationary_epsilon_greedy(horizon: int, seed: int, alpha: float | None) -> np.ndarray:
    """Regret when the best arm changes halfway through the run.

    alpha=None uses the full sample average; a finite alpha uses constant-step
    recency weighting and therefore forgets old rewards.
    """
    rng = np.random.default_rng(seed)
    means_before = np.array([0.65, 0.45, 0.25])
    means_after = np.array([0.25, 0.45, 0.70])
    estimates = np.zeros(3)
    counts = np.zeros(3, dtype=int)
    regret = np.zeros(horizon)
    running = 0.0
    for t in range(horizon):
        means = means_before if t < horizon // 2 else means_after
        if rng.random() < 0.10:
            arm = int(rng.integers(3))
        else:
            best = np.flatnonzero(np.isclose(estimates, estimates.max()))
            arm = int(rng.choice(best))
        reward = float(rng.random() < means[arm])
        counts[arm] += 1
        step = (1.0 / counts[arm]) if alpha is None else alpha
        estimates[arm] += step * (reward - estimates[arm])
        running += float(means.max() - means[arm])
        regret[t] = running
    return regret

ns_horizon = 1_500
seeds = np.arange(80 if PROFILE == "smoke" else 300) + 9000
sample_average = np.stack([nonstationary_epsilon_greedy(ns_horizon, int(s), None) for s in seeds])
recency_weighted = np.stack([nonstationary_epsilon_greedy(ns_horizon, int(s), 0.10) for s in seeds])

plt.figure(figsize=(9, 5))
plt.plot(np.median(sample_average, axis=0), label="sample average")
plt.plot(np.median(recency_weighted, axis=0), label="constant step size 0.10")
plt.axvline(ns_horizon // 2, linestyle="--", label="best arm changes")
plt.xlabel("Раунд")
plt.ylabel("Медианное накопленное псевдосожаление")
plt.title("Нестационарный бандит: адаптация после смены лучшей руки")
plt.legend()
plt.tight_layout()
plt.show()

n = 6_000 if PROFILE == "smoke" else 25_000
log = sample_log(n, seed=20260804, floor=0.08)
base = ope_estimates(log, sharpness=2.0, seed=44)
for key, value in base.items():
    print(f"{key:12s}: {value:.6f}")
assert base["ESS"] > 0
assert np.isfinite(list(base.values())).all()

sharpness_grid = np.linspace(0.5, 8.0, 16)
rows = [ope_estimates(log, sharpness=float(s), seed=55) for s in sharpness_grid]

plt.figure(figsize=(9, 5))
for name in ["IPS", "SNIPS", "DM", "DR"]:
    error = [row[name] - row["truth"] for row in rows]
    plt.plot(sharpness_grid, error, marker="o", label=name)
plt.axhline(0.0, linewidth=1)
plt.xlabel("Резкость целевой политики")
plt.ylabel("Ошибка оценки ценности")
plt.title("OPE при ухудшении перекрытия носителей")
plt.legend()
plt.tight_layout()
plt.show()

plt.figure(figsize=(9, 5))
plt.plot(sharpness_grid, [row["ESS"] for row in rows], marker="o", label="ESS")
plt.xlabel("Резкость целевой политики")
plt.ylabel("Эффективный размер выборки")
plt.title("Концентрация importance weights")
plt.tight_layout()
plt.show()

print("last point:", rows[-1])

rng = np.random.default_rng(777)
x_router = rng.normal(size=(12_000, 3))
quality = contextual_probabilities(x_router)
cost = np.array([0.02, 0.11, 0.28])
latency = np.array([0.10, 0.35, 0.80])

for cost_weight, latency_weight in [(0.2, 0.1), (0.8, 0.4), (1.8, 0.8)]:
    utility = quality - cost_weight * cost - latency_weight * latency
    chosen = np.argmax(utility, axis=1)
    shares = np.bincount(chosen, minlength=3) / len(chosen)
    print(f"cost_weight={cost_weight:3.1f}, latency_weight={latency_weight:3.1f}, route shares={shares}")