"""Chapter 2 lab: bandit agents. Fill in every `# TODO` and delete the `raise NotImplementedError` lines. Work through the tasks in order; each one has its own tests: python -m pytest -q -k task1 # just Task 1 python -m pytest -q # everything Every agent has the same two methods, used by envs.run_bandit: agent.select_action() -> int which arm to pull now agent.update(action, reward) learn from the reward you just got Notation (same as the chapter): K arms, Q[a] is our estimate of the true mean q_*(a), N[a] is how many times arm a has been pulled, t is the total number of pulls so far. numpy reminders: `rng.random()` is a uniform number in [0, 1); `rng.integers(k)` is a uniform integer in {0, ..., k-1}; `rng.choice(array)` picks one element uniformly; `rng.beta(alpha, beta)` samples Beta variables (elementwise if you pass arrays); `np.flatnonzero(mask)` gives the indices where a boolean mask is True. """ from __future__ import annotations import numpy as np # ---------------------------------------------------------------- Task 1 def argmax_tiebreak(values, rng: np.random.Generator) -> int: """Index of the largest entry of `values`, breaking ties uniformly at random. Why not just np.argmax? np.argmax always returns the FIRST maximal index. At the start every Q is 0, so a greedy agent using np.argmax would always try arm 0 first, and in general would be biased toward low-numbered arms. (Exercise 2.2 in the chapter shows how much tie-breaking can matter.) Args: values: 1-D array-like of numbers, shape (K,). May contain np.inf. rng: a numpy Generator, used only to break ties. Returns: an int a such that values[a] == max(values); if several indices tie, each of them is returned with equal probability. """ # TODO: find all indices equal to the maximum, return one of them at random. raise NotImplementedError("Task 1: argmax_tiebreak") # ---------------------------------------------------------------- Task 2 class EpsilonGreedy: """ε-greedy action selection with incremental value estimates. With probability epsilon: explore, i.e. pick an arm uniformly at random (any of the K arms, including the greedy one). Otherwise: exploit, i.e. pick argmax_a Q[a] (ties broken at random with argmax_tiebreak). The same class gives three agents from the chapter: EpsilonGreedy(k, 0.0, rng) greedy EpsilonGreedy(k, 0.1, rng) ε-greedy EpsilonGreedy(k, 0.0, rng, initial_value=1.0, step_size=0.1) optimistic greedy Attributes: Q: float array, shape (K,), starts at initial_value N: int array, shape (K,), starts at 0 """ def __init__(self, k: int, epsilon: float, rng: np.random.Generator, initial_value: float = 0.0, step_size: float | None = None): self.k = k self.epsilon = epsilon self.rng = rng self.step_size = step_size # None means "sample average" (step 1/N) self.Q = np.full(k, float(initial_value)) self.N = np.zeros(k, dtype=int) def select_action(self) -> int: """Explore with probability epsilon, otherwise act greedily on Q. Use exactly one call to self.rng.random() for the coin flip, then self.rng.integers(self.k) to explore or argmax_tiebreak(self.Q, self.rng) to exploit. """ # TODO raise NotImplementedError("Task 2: EpsilonGreedy.select_action") def update(self, action: int, reward: float) -> None: """Incremental update of the pulled arm's estimate (the chapter's "estimate <- estimate + step x error" rule): N[a] <- N[a] + 1 Q[a] <- Q[a] + step * (reward - Q[a]) where step = 1 / N[a] (after incrementing) if step_size is None, which makes Q[a] exactly the sample average of a's rewards, or step = step_size (a constant α) otherwise. Only arm `action` changes. """ # TODO raise NotImplementedError("Task 2: EpsilonGreedy.update") # ---------------------------------------------------------------- Task 3 def ucb_scores(Q, N, t: int, c: float) -> np.ndarray: """UCB scores for every arm: score[a] = Q[a] + c * sqrt( ln(t) / N[a] ) if N[a] > 0 score[a] = +inf if N[a] == 0 (Untried arms get +inf so they are always tried first.) Args: Q: estimates, shape (K,) N: pull counts, shape (K,) t: total number of pulls so far (t >= 1) c: exploration strength; c = sqrt(2) is the textbook UCB1 for rewards in [0, 1] Returns: float array of shape (K,). Careful: don't divide by N where N == 0 (numpy would warn and produce inf/nan). """ # TODO raise NotImplementedError("Task 3: ucb_scores") class UCB1: """Upper-confidence-bound action selection. select_action: argmax of ucb_scores(Q, N, t, c), ties broken at random. (When t == 0 nothing has been tried and every score is inf; pass max(t, 1) so the log is defined.) update: t <- t + 1, N[a] <- N[a] + 1, Q[a] <- sample average (the same incremental rule as Task 2 with step 1/N[a]). """ def __init__(self, k: int, rng: np.random.Generator, c: float = np.sqrt(2.0)): self.k = k self.c = c self.rng = rng self.Q = np.zeros(k) self.N = np.zeros(k, dtype=int) self.t = 0 def select_action(self) -> int: # TODO raise NotImplementedError("Task 3: UCB1.select_action") def update(self, action: int, reward: float) -> None: # TODO raise NotImplementedError("Task 3: UCB1.update") # ---------------------------------------------------------------- Task 4 class ThompsonBernoulli: """Thompson sampling for 0/1 rewards with Beta posteriors. Each arm a has a Beta(alpha[a], beta[a]) belief about its win probability. The prior is Beta(alpha0, beta0); Beta(1, 1) is uniform on [0, 1]. select_action: draw one sample theta[a] ~ Beta(alpha[a], beta[a]) for every arm (one call: self.rng.beta(self.alpha, self.beta)), then pull argmax_a theta[a] (use argmax_tiebreak). update: a win (reward 1) adds 1 to alpha[action]; a loss (reward 0) adds 1 to beta[action]. (Writing alpha += r and beta += 1 - r does both at once.) """ def __init__(self, k: int, rng: np.random.Generator, alpha0: float = 1.0, beta0: float = 1.0): self.k = k self.rng = rng self.alpha = np.full(k, float(alpha0)) self.beta = np.full(k, float(beta0)) def select_action(self) -> int: # TODO raise NotImplementedError("Task 4: ThompsonBernoulli.select_action") def update(self, action: int, reward: float) -> None: # TODO raise NotImplementedError("Task 4: ThompsonBernoulli.update")