Solving Bernoulli multi-armed bandit with Thompson sampling

beta_bernoulli_bandit.py

def make_bandits(params):
    def pull(arm, size=None):
        while True:
            # Bernoulli distributed rewards
            reward = np.random.binomial(n=1, p=params[arm], size=size)
            yield reward
    
    return pull, len(params)


def bayesian_strategy(pull, num_bandits):
    num_rewards = np.zeros(num_bandits)
    num_trials = np.zeros(num_bandits)
    
    while True:
        # Sample from the bandits' priors, and choose largest
        choice = np.argmax(np.random.beta(a=2+num_rewards,
                                          b=2+num_trials-num_rewards))
        
        # Sample the chosen bandit
        reward = next(pull(choice))
        
        # Update
        num_rewards[choice] += reward
        num_trials[choice] += 1
        
        yield choice, reward, num_rewards, num_trials


if __name__ == '__main__':
    pull, num_bandits = make_bandits([0.2, 0.5, 0.7])
    play = bayesian_strategy(pull, num_bandits)
    
    for _ in range(100):
        choice, reward, num_rewards, num_trials = next(play)