Replace Chapter 3 folder

Author: srikanthbaride

ch3_multi_armed_bandits/README.md

@@ -0,0 +1,20 @@
+# Chapter 3 — Multi-Armed Bandits
+
+Implements ε-Greedy, UCB1, and Thompson Sampling on Bernoulli bandits.
+
+## Run experiments
+```bash
+python -m ch3_multi_armed_bandits.experiments --K 10 --T 5000 --trials 50 --eps 0.1 --c 1.0
+```
+
+## Worked Examples
+See `examples/` for scripts reproducing the numerical examples:
+- Example 3.1: `ex1_regret_basic.py`
+- Example 3.2: `ex2_epsilon_update.py`
+- Example 3.3: `ex3_ucb_score.py`
+- Example 3.4: `ex4_thompson_update.py`
+
+## Tests
+```bash
+pytest -q ch3_multi_armed_bandits/tests
+```

ch3_multi_armed_bandits/__init__.py

@@ -1 +1 @@
-
+__all__ = ["bandits", "epsilon_greedy", "ucb", "thompson", "experiments"]

ch3_multi_armed_bandits/bandits.py

@@ -1,20 +1,32 @@
 from __future__ import annotations
-import numpy as np
 from dataclasses import dataclass
+from typing import Iterable, Optional
+import numpy as np
 
 @dataclass
 class BernoulliBandit:
-    probs: np.ndarray
+    p: Iterable[float]
+    seed: Optional[int] = None
 
     def __post_init__(self):
-        self.probs = np.array(self.probs, dtype=float)
-        assert self.probs.ndim == 1 and (0 <= self.probs).all() and (self.probs <= 1).all()
-        self.K = self.probs.shape[0]
-        self.opt_idx = int(np.argmax(self.probs))
-        self.opt_mean = float(self.probs[self.opt_idx])
+        self.p = np.asarray(list(self.p), dtype=float)
+        if np.any(self.p < 0) or np.any(self.p > 1):
+            raise ValueError("All probabilities must be in [0,1].")
+        self.K = int(self.p.size)
+        self._rng = np.random.default_rng(self.seed)
+
+    def step(self, arm: int) -> int:
+        if not (0 <= arm < self.K):
+            raise IndexError("Arm index out of range.")
+        return int(self._rng.random() < self.p[arm])
+
+    def reset(self, seed: Optional[int] = None):
+        self._rng = np.random.default_rng(seed)
 
-    def pull(self, arm: int, rng: np.random.Generator) -> float:
-        return float(rng.random() < self.probs[arm])
+def regret_from_choices(true_means: np.ndarray, choices: np.ndarray, rewards: np.ndarray) -> np.ndarray:
+    mu_star = float(np.max(true_means))
+    t = np.arange(1, rewards.size + 1, dtype=float)
+    return mu_star * t - np.cumsum(rewards)
 
-    def pseudo_regret(self, arm: int) -> float:
-        return self.opt_mean - float(self.probs[arm])
+def ensure_rng(seed: Optional[int]) -> np.random.Generator:
+    return np.random.default_rng(seed)

ch3_multi_armed_bandits/epsilon_greedy.py

@@ -1,20 +1,26 @@
 from __future__ import annotations
+from typing import Optional, Dict, Any
 import numpy as np
+from .bandits import BernoulliBandit, regret_from_choices, ensure_rng
 
-class EpsilonGreedy:
-    def __init__(self, K: int, epsilon: float = 0.1, rng: np.random.Generator | None = None):
-        self.K = K
-        self.epsilon = float(epsilon)
-        self.rng = rng or np.random.default_rng()
-        self.counts = np.zeros(K, dtype=int)
-        self.values = np.zeros(K, dtype=float)
-
-    def select_arm(self) -> int:
-        if self.rng.random() < self.epsilon:
-            return int(self.rng.integers(self.K))
-        return int(np.argmax(self.values))
-
-    def update(self, arm: int, reward: float):
-        self.counts[arm] += 1
-        n = self.counts[arm]
-        self.values[arm] += (reward - self.values[arm]) / n
+def run(true_means, epsilon: float, steps: int, seed: Optional[int] = None) -> Dict[str, Any]:
+    if not (0 <= float(epsilon) <= 1):
+        raise ValueError("epsilon must be in [0,1].")
+    K = len(true_means)
+    env = BernoulliBandit(true_means, seed=seed)
+    rng = ensure_rng(seed)
+    Q, N = np.zeros(K), np.zeros(K, dtype=int)
+    choices, rewards = np.zeros(steps, int), np.zeros(steps, float)
+    for t in range(steps):
+        if rng.random() < epsilon:
+            a = rng.integers(0, K)
+        else:
+            a = int(np.argmax(Q))
+        r = env.step(a)
+        N[a] += 1
+        Q[a] += (r - Q[a]) / N[a]
+        choices[t], rewards[t] = a, r
+    return {
+        "rewards": rewards, "choices": choices, "Q": Q, "N": N,
+        "cum_regret": regret_from_choices(np.asarray(true_means, float), choices, rewards),
+    }

ch3_multi_armed_bandits/examples/ex1_regret_basic.py

@@ -0,0 +1,7 @@
+import numpy as np
+def compute():
+    mu = np.array([0.5,0.6,0.7]); T=100; chosen=0
+    mu_star = mu.max()
+    regret = T*mu_star - T*mu[chosen]
+    return {"optimal":np.argmax(mu)+1,"mu*":mu_star,"chosen_mu":mu[chosen],"regret":regret}
+if __name__=="__main__": print(compute())

ch3_multi_armed_bandits/examples/ex2_epsilon_update.py

@@ -0,0 +1,5 @@
+def compute():
+    N,Q,R=4,0.5,1; N_new=N+1
+    Q_new=Q+(R-Q)/N_new
+    return {"N":N,"N_new":N_new,"Q_new":Q_new}
+if __name__=="__main__": print(compute())

ch3_multi_armed_bandits/examples/ex3_ucb_score.py

@@ -0,0 +1,6 @@
+import numpy as np, math
+def compute():
+    Q=[0.6,0.7,0.4]; N=[5,10,2]; t=20;c=1.0
+    ucb=[Q[a]+c*math.sqrt(math.log(t)/N[a]) for a in range(3)]
+    return {"ucb":ucb,"selected":np.argmax(ucb)+1}
+if __name__=="__main__": print(compute())

ch3_multi_armed_bandits/examples/ex4_thompson_update.py

@@ -0,0 +1,4 @@
+def compute():
+    alpha=[1+3,1+1]; beta=[1+2,1+4]
+    return {"alpha":alpha,"beta":beta}
+if __name__=="__main__": print(compute())

ch3_multi_armed_bandits/examples/run_all_examples.py

@@ -0,0 +1,5 @@
+import ex1_regret_basic,ex2_epsilon_update,ex3_ucb_score,ex4_thompson_update
+print(ex1_regret_basic.compute())
+print(ex2_epsilon_update.compute())
+print(ex3_ucb_score.compute())
+print(ex4_thompson_update.compute())

ch3_multi_armed_bandits/experiments.py

@@ -1,70 +1,44 @@
-from __future__ import annotations
-import os
-import numpy as np
-import matplotlib.pyplot as plt
-from .bandits import BernoulliBandit
-from .epsilon_greedy import EpsilonGreedy
-from .ucb import UCB1
-from .thompson import ThompsonSamplingBernoulli
+import argparse, os, numpy as np, matplotlib.pyplot as plt
+from .epsilon_greedy import run as run_eps
+from .ucb import run as run_ucb
+from .thompson import run as run_ts
 
-def run_algorithm(env, algo, T: int, seed: int) -> dict:
-    rng = np.random.default_rng(seed)
-    rewards = np.zeros(T, dtype=float)
-    regret = np.zeros(T, dtype=float)
-    for t in range(T):
-        a = algo.select_arm()
-        r = env.pull(a, rng)
-        algo.update(a, r)
-        rewards[t] = r
-        regret[t] = env.pseudo_regret(a)
-    return {
-        "rewards": rewards,
-        "cum_rewards": np.cumsum(rewards),
-        "regret": regret,
-        "cum_regret": np.cumsum(regret),
-    }
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--K", type=int, default=10)
+    p.add_argument("--T", type=int, default=5000)
+    p.add_argument("--trials", type=int, default=50)
+    p.add_argument("--eps", type=float, default=0.1)
+    p.add_argument("--c", type=float, default=1.0)
+    p.add_argument("--seed", type=int, default=123)
+    p.add_argument("--outdir", type=str, default="ch3_multi_armed_bandits/plots")
+    return p.parse_args()
+
+def make_true_means(K, rng): return rng.uniform(0.1, 0.9, size=K)
 
-def average_over_runs(env, algo_ctor, T: int, n_runs: int, base_seed: int = 0) -> dict:
-    cum_regrets = []
-    for run in range(n_runs):
-        algo = algo_ctor()
-        result = run_algorithm(env, algo, T, seed=base_seed + run)
-        cum_regrets.append(result["cum_regret"])
-    cum_regrets = np.array(cum_regrets)
-    mean = cum_regrets.mean(axis=0)
-    se = cum_regrets.std(axis=0, ddof=1) / np.sqrt(n_runs)
-    return {"mean": mean, "se": se}
+def run_all(true_means, T, trials, eps, c, seed):
+    rng = np.random.default_rng(seed)
+    avg_regret = {"eps": np.zeros(T), "ucb": np.zeros(T), "ts": np.zeros(T)}
+    for _ in range(trials):
+        s = int(rng.integers(0, 2**31-1))
+        avg_regret["eps"] += run_eps(true_means, eps, T, s)["cum_regret"]
+        avg_regret["ucb"] += run_ucb(true_means, c, T, s)["cum_regret"]
+        avg_regret["ts"]  += run_ts(true_means, T, s)["cum_regret"]
+    for k in avg_regret: avg_regret[k] /= trials
+    return avg_regret
 
-def plot_regret(curves: dict, title: str, fname: str | None):
-    fig, ax = plt.subplots()
-    for label, stats in curves.items():
-        ax.plot(stats["mean"], label=label)
-    ax.set_xlabel("Time")
-    ax.set_ylabel("Average cumulative pseudo-regret")
-    ax.set_title(title)
-    ax.legend()
-    if fname:
-        out_dir = os.path.dirname(fname)
-        if out_dir and not os.path.exists(out_dir):
-            os.makedirs(out_dir, exist_ok=True)
-        fig.savefig(fname, bbox_inches="tight")
-    else:
-        plt.show()
+def plot(xs, series, ylabel, title, outpath):
+    plt.figure()
+    for label,y in series: plt.plot(xs,y,label=label)
+    plt.xlabel("Time"); plt.ylabel(ylabel); plt.title(title); plt.legend()
+    os.makedirs(os.path.dirname(outpath), exist_ok=True)
+    plt.savefig(outpath, dpi=300); plt.close()
 
 def main():
-    probs = np.array([0.2, 0.25, 0.3, 0.35, 0.5])
-    env = BernoulliBandit(probs=probs)
-    T = 2000
-    n_runs = 200
-    curves = {}
-    curves["ε-greedy(0.10)"] = average_over_runs(env, lambda: EpsilonGreedy(env.K, 0.10), T, n_runs, 123)
-    curves["ε-greedy(0.01)"] = average_over_runs(env, lambda: EpsilonGreedy(env.K, 0.01), T, n_runs, 223)
-    curves["UCB1(c=0.5)"] = average_over_runs(env, lambda: UCB1(env.K, c=0.5), T, n_runs, 323)
-    curves["Thompson (Beta-Bernoulli)"] = average_over_runs(env, lambda: ThompsonSamplingBernoulli(env.K), T, n_runs, 423)
-    here = os.path.dirname(__file__)
-    out_path = os.path.join(here, "plots", "regret_bernoulli.png")
-    plot_regret(curves, "Multi-Armed Bandits: Average Cumulative Pseudo-Regret", out_path)
-    print(f"Saved plot to {out_path}")
-
-if __name__ == "__main__":
-    main()
+    a = parse_args()
+    true_means = make_true_means(a.K, np.random.default_rng(a.seed))
+    xs = np.arange(1, a.T+1)
+    reg = run_all(true_means,a.T,a.trials,a.eps,a.c,a.seed)
+    plot(xs,[("ε-Greedy",reg["eps"]),("UCB1",reg["ucb"]),("Thompson",reg["ts"])],
+         "Cumulative Regret","Regret vs Time",os.path.join(a.outdir,"regret.png"))
+if __name__=="__main__": main()