Module 3. Bandits, exploration, and one-step decisions
Read the explanations and code in the existing notebooks, or download the files. Code does not run on this site.
Download course files (.zip) 382.1 KBThe archive preserves the source folder structure, both languages, and supporting files.
Module 3. Bandits, exploration, and evaluation from logged feedback
Main module notebook
Download .ipynb module-03-bandits-ope.en.ipynb · 18.0 KB
Show notebook text
Module 3. Bandits, exploration, and evaluation from logged feedback
This notebook joins four levels introduced in the lecture: a stationary stochastic bandit, non-stationarity, contextual routing, and evaluation of a new policy from data collected by an older one. Every experiment is reproducible and runs on a CPU.
Profiles. smoke is intended for a fast learning pass. full increases the number of repetitions and log size so curves and estimates become more stable. The reported numbers are outputs of the course code under the declared protocol, not an external algorithm benchmark.
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)
1. Stationary bandit: compare a distribution of runs, not one lucky trace
A single seed can tell a misleading story. Each strategy below is therefore run repeatedly on the same five-arm Bernoulli environment. We measure pseudo-regret: each round adds the difference between the best arm's mean reward and the selected arm's mean reward.
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("Round")
plt.ylabel("Mean cumulative pseudo-regret")
plt.title("Stationary Bernoulli bandit: mean and 95% interval")
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}")
The plot does not establish universal superiority of one method. It shows how these strategies behave under fixed arm probabilities, horizon, tie-breaking rule, and seeds. A different problem may change the ordering.
2. Non-stationarity: accumulated history can become a liability
Halfway through the next experiment, the best arm changes. An ordinary sample average keeps the entire history with equal total influence. A constant step size forgets old rewards and adapts faster, at the cost of persistent estimation noise.
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("Round")
plt.ylabel("Median cumulative pseudo-regret")
plt.title("Non-stationary bandit: adaptation after the best arm changes")
plt.legend()
plt.tight_layout()
plt.show()
3. A contextual log and the target policy's true value
Actions now depend on request features. In this synthetic model we know every route's true success probability, but the log stores only the selected action's reward and its logging propensity. This lets us compare offline estimators with known truth without presenting a synthetic example as a production result.
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()
IPS uses only observed rewards and probability ratios. The direct method relies on a reward model. DR begins with model prediction and adds an importance-weighted residual. A correct formula still cannot rescue missing support: when the logging policy almost never selected an action favored by the target, weights concentrate on a few rows.
4. Support-overlap stress test
We increase target-policy sharpness. It increasingly concentrates on a narrow action set while the log came from a softer policy. We track estimator error together with maximum weight and effective sample size.
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("Target-policy sharpness")
plt.ylabel("Value-estimation error")
plt.title("OPE as support overlap deteriorates")
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("Target-policy sharpness")
plt.ylabel("Effective sample size")
plt.title("Importance-weight concentration")
plt.tight_layout()
plt.show()
print("last point:", rows[-1])
5. LLM-like routing: quality, cost, and latency
In a deployed system, reward is often not “true quality” but an explicitly defined utility. The three actions below stand for a cheap model, a stronger model, and a model with an additional tool. Utility subtracts cost and latency. Changing those coefficients changes the optimal policy; that is part of the problem definition, not incidental engineering.
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}")
6. What the notebook establishes
The result is not one “winning algorithm” but a set of testable engineering contracts: the exploration strategy must match environmental stationarity; the log must store the selected action's probability; IPS/DR require support overlap; ESS and the weight tail must accompany a single value estimate; and routing utility must explicitly include quality, cost, and latency. These are the contracts that transfer to RL for LLM systems.