To implement Thompson Sampling for binary rewards with Beta(1,1) priors, we keep per-arm Beta parameters (alpha, beta). Each selection samples theta_i ~ Beta(alpha_i, beta_i) and picks argmax. Update increments alpha by 1 on reward=1, beta by 1 on reward=0. This is O(k) per round (vectorized sampling with numpy is efficient for k up to ~100).python
import numpy as np
class ThompsonSampling:
def __init__(self, k, seed=None):
"""
k: number of arms
Priors: Beta(1,1) for each arm
"""
self.k = k
self.rng = np.random.default_rng(seed)
self.alpha = np.ones(k, dtype=np.int64) # successes + 1
self.beta = np.ones(k, dtype=np.int64) # failures + 1
self.counts = np.zeros(k, dtype=np.int64) # number of pulls
def select_arm(self):
"""Sample theta for each arm and return index of chosen arm."""
samples = self.rng.beta(self.alpha, self.beta)
return int(np.argmax(samples))
def update(self, chosen_arm, reward):
"""
Update Beta posterior for chosen arm.
reward: 1 or 0 (binary)
"""
if reward not in (0, 1):
raise ValueError("reward must be 0 or 1")
self.counts[chosen_arm] += 1
if reward == 1:
self.alpha[chosen_arm] += 1
else:
self.beta[chosen_arm] += 1
def simulate(self, n_rounds, true_rates):
"""
Run a simulation for n_rounds with given true_rates (list/array length k).
Returns history dict with pulls, rewards, cumulative_reward.
"""
true_rates = np.asarray(true_rates)
if true_rates.shape[0] != self.k:
raise ValueError("true_rates length must equal k")
rewards = np.empty(n_rounds, dtype=np.int64)
chosen = np.empty(n_rounds, dtype=np.int64)
for t in range(n_rounds):
a = self.select_arm()
chosen[t] = a
reward = int(self.rng.random() < true_rates[a])
rewards[t] = reward
self.update(a, reward)
return {
"chosen": chosen,
"rewards": rewards,
"cumulative_reward": int(rewards.sum()),
"counts": self.counts.copy(),
"alpha": self.alpha.copy(),
"beta": self.beta.copy()
}
# Example usage
if __name__ == "__main__":
k = 5
true_rates = [0.1, 0.15, 0.2, 0.05, 0.18]
ts = ThompsonSampling(k, seed=42)
hist = ts.simulate(5000, true_rates)
est_rates = (hist["alpha"] - 1) / (hist["alpha"] + hist["beta"] - 2) # posterior means
print("Pull counts:", hist["counts"])
print("Estimated rates (posterior means):", np.round(est_rates, 3))
print("Cumulative reward:", hist["cumulative_reward"])
Key points:- Time: O(k) per select_arm (vectorized sampling keeps constant factors low).- Space: O(k).- Edge cases: validate rewards are binary; ensure true_rates length matches k.- Alternatives: For large k or streaming, use incremental sampling or approximation methods (e.g., sampling only top candidates).