bandit.py

class Bandit:
    
    def __init__(self,arms : int):
        self.arms = []
        for i in range(arms):
            self.arms.append(np.random.rand())
        self.action_count = arms
        
    def step(self, action: int):
        reward = 0
        roll = np.random.rand()
        if roll < self.arms[action]:
            reward = 1
        else:
            reward = -1
        return reward