Thompson Sampling Mock Bandit in Python

import random

class ThompsonSamplingBandit:
    def __init__(self, num_arms, alpha=1.0, beta=1.0):
        self.num_arms = num_arms
        self.alpha = [alpha] * num_arms
        self.beta = [beta] * num_arms

    def select_arm(self):
        samples = [random.betavariate(a, b) for a, b in zip(self.alpha, self.beta)]
        return max(range(self.num_arms), key=lambda i: samples[i])

    def update(self, arm, reward):
        if reward:
            self.alpha[arm] += 1
        else:
            self.beta[arm] += 1

if __name__ == "__main__":
    random.seed(42)
    true_probabilities = [0.2, 0.5, 0.8]
    bandit = ThompsonSamplingBandit(len(true_probabilities))
    total_reward = 0
    for _ in range(1000):
        arm = bandit.select_arm()
        reward = 1 if random.random() < true_probabilities[arm] else 0
        bandit.update(arm, reward)
        total_reward += reward
    print(f"Arm selections per arm: {[bandit.alpha[i] + bandit.beta[i] - 2 for i in range(len(true_probabilities))]}")
    print(f"Total reward: {total_reward}")
    print(f"Estimated probabilities: {[round(bandit.alpha[i] / (bandit.alpha[i] + bandit.beta[i]), 3) for i in range(len(true_probabilities))]}")