Add Chapter 3 multi-armed bandits: env, algorithms, examples, tests

Author: srikanthbaride

ch3_multi_armed_bandits/__init__.py

@@ -0,0 +1,2 @@
+from .bandit_env import MultiArmedBanditBernoulli
+from .algorithms import EpsilonGreedy, UCB1, ThompsonSamplingBeta, simulate

ch3_multi_armed_bandits/algorithms.py

@@ -1,2 +1,87 @@
-def epsilon_greedy(*args, **kwargs):
-    return 0
+from __future__ import annotations
+import numpy as np
+from dataclasses import dataclass
+from typing import Optional, Dict, Any
+
+@dataclass
+class EpsilonGreedy:
+    n_arms: int
+    epsilon: float = 0.1
+    init: float = 0.0
+    def __post_init__(self):
+        self.counts = np.zeros(self.n_arms, dtype=int)
+        self.values = np.full(self.n_arms, float(self.init), dtype=float)
+        self.rng = np.random.default_rng()
+    def select_arm(self) -> int:
+        if self.rng.random() < self.epsilon:
+            return int(self.rng.integers(self.n_arms))
+        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
+
+@dataclass
+class UCB1:
+    n_arms: int
+    c: float = 2.0
+    def __post_init__(self):
+        self.counts = np.zeros(self.n_arms, dtype=int)
+        self.values = np.zeros(self.n_arms, dtype=float)
+        self.total = 0
+        self.rng = np.random.default_rng()
+    def select_arm(self) -> int:
+        # pull each arm once
+        for a in range(self.n_arms):
+            if self.counts[a] == 0:
+                return a
+        ucb = self.values + self.c * np.sqrt(np.log(self.total) / self.counts)
+        return int(np.argmax(ucb))
+    def update(self, arm: int, reward: float):
+        self.total += 1
+        self.counts[arm] += 1
+        n = self.counts[arm]
+        self.values[arm] += (reward - self.values[arm]) / n
+
+@dataclass
+class ThompsonSamplingBeta:
+    n_arms: int
+    a0: float = 1.0
+    b0: float = 1.0
+    def __post_init__(self):
+        self.a = np.full(self.n_arms, self.a0, dtype=float)
+        self.b = np.full(self.n_arms, self.b0, dtype=float)
+        self.rng = np.random.default_rng()
+    def select_arm(self) -> int:
+        samples = self.rng.beta(self.a, self.b)
+        return int(np.argmax(samples))
+    def update(self, arm: int, reward: float):
+        self.a[arm] += reward
+        self.b[arm] += 1 - reward
+
+def simulate(env, agent, steps: int, seed: Optional[int] = None) -> Dict[str, Any]:
+    """Run interaction loop and return history stats."""
+    if seed is not None:
+        try:
+            agent.rng = np.random.default_rng(seed)
+        except AttributeError:
+            pass
+        if hasattr(env, 'rng'):
+            env.rng = np.random.default_rng(seed+1 if seed is not None else None)
+    rewards = np.zeros(steps, dtype=float)
+    pulls = np.zeros(env.k, dtype=int)
+    choices = np.zeros(steps, dtype=int)
+    for t in range(steps):
+        a = agent.select_arm()
+        r = env.pull(a)
+        agent.update(a, r)
+        rewards[t] = r
+        pulls[a] += 1
+        choices[t] = a
+    return {
+        "avg_reward": float(rewards.mean()),
+        "cum_reward": float(rewards.sum()),
+        "pulls": pulls,
+        "choices": choices,
+        "rewards": rewards,
+    }

ch3_multi_armed_bandits/bandit_env.py

@@ -1,2 +1,24 @@
-class MultiArmedBandit:
-    pass
+from __future__ import annotations
+import numpy as np
+from typing import Sequence, Optional
+
+class MultiArmedBanditBernoulli:
+    """K-armed bandit with Bernoulli rewards.
+    probs: list/array of success probabilities for each arm (0..1).
+    reward is 0/1. RNG can be injected for reproducibility.
+    """
+    def __init__(self, probs: Sequence[float], rng: Optional[np.random.Generator] = None):
+        self.probs = np.asarray(probs, dtype=float)
+        assert np.all((0 <= self.probs) & (self.probs <= 1)), "probs must be in [0,1]"
+        self.k = len(self.probs)
+        self.rng = rng if rng is not None else np.random.default_rng()
+
+    def pull(self, arm: int) -> int:
+        p = self.probs[arm]
+        return int(self.rng.random() < p)
+
+    def best_arm(self) -> int:
+        return int(np.argmax(self.probs))
+
+    def optimal_mean(self) -> float:
+        return float(np.max(self.probs))

ch3_multi_armed_bandits/examples/demo_bandit.py

@@ -1 +1,19 @@
-print('Bandit demo placeholder')
+from ch3_multi_armed_bandits.bandit_env import MultiArmedBanditBernoulli
+from ch3_multi_armed_bandits.algorithms import EpsilonGreedy, UCB1, ThompsonSamplingBeta, simulate
+
+def main():
+    probs = [0.1, 0.2, 0.5, 0.4]
+    env = MultiArmedBanditBernoulli(probs)
+    best = env.best_arm()
+    steps = 5000
+
+    for agent in [
+        EpsilonGreedy(n_arms=len(probs), epsilon=0.1),
+        UCB1(n_arms=len(probs)),
+        ThompsonSamplingBeta(n_arms=len(probs)),
+    ]:
+        stats = simulate(env, agent, steps, seed=123)
+        print(f"{agent.__class__.__name__}: avg_reward={stats['avg_reward']:.3f}, best_arm_pulled={stats['pulls'][best]} times" )
+
+if __name__ == "__main__":
+    main()

ch3_multi_armed_bandits/requirements.txt

@@ -1,2 +1,3 @@
 numpy>=1.24
 pytest>=7.0
+pytest-cov>=4.1

ch3_multi_armed_bandits/tests/test_bandit.py

@@ -0,0 +1,47 @@
+import numpy as np
+from ch3_multi_armed_bandits.bandit_env import MultiArmedBanditBernoulli
+from ch3_multi_armed_bandits.algorithms import EpsilonGreedy, UCB1, ThompsonSamplingBeta, simulate
+
+PROBS = [0.1, 0.2, 0.5, 0.4]
+BEST = int(np.argmax(PROBS))
+STEPS = 4000
+
+def random_baseline_avg(probs, steps, seed=7):
+    rng = np.random.default_rng(seed)
+    k = len(probs)
+    rewards = []
+    for t in range(steps):
+        a = int(rng.integers(k))
+        r = int(rng.random() < probs[a])
+        rewards.append(r)
+    return float(np.mean(rewards))
+
+def run_and_check(agent, steps=STEPS):
+    env = MultiArmedBanditBernoulli(PROBS)
+    stats = simulate(env, agent, steps, seed=42)
+    return stats
+
+def test_algorithms_beat_random_baseline():
+    baseline = random_baseline_avg(PROBS, STEPS)
+    agents = [
+        EpsilonGreedy(n_arms=len(PROBS), epsilon=0.1),
+        UCB1(n_arms=len(PROBS)),
+        ThompsonSamplingBeta(n_arms=len(PROBS)),
+    ]
+    for agent in agents:
+        stats = run_and_check(agent)
+        # Should beat random baseline by a margin
+        assert stats['avg_reward'] >= baseline + 0.05, (agent.__class__.__name__, stats['avg_reward'], baseline)
+
+def test_learn_best_arm_frequently():
+    agents = [
+        EpsilonGreedy(n_arms=len(PROBS), epsilon=0.1),
+        UCB1(n_arms=len(PROBS)),
+        ThompsonSamplingBeta(n_arms=len(PROBS)),
+    ]
+    for agent in agents:
+        stats = run_and_check(agent)
+        pulls = stats['pulls']
+        assert pulls[BEST] == pulls.max()  # best arm most selected
+        # At least 50% of pulls go to best arm after learning
+        assert pulls[BEST] >= STEPS * 0.5