| 16 | |
| 17 | |
| 18 | class Bandit: |
| 19 | def __init__(self, p): |
| 20 | # p: the win rate |
| 21 | self.p = p |
| 22 | self.p_estimate = 0. |
| 23 | self.N = 0. # num samples collected so far |
| 24 | |
| 25 | def pull(self): |
| 26 | # draw a 1 with probability p |
| 27 | return np.random.random() < self.p |
| 28 | |
| 29 | def update(self, x): |
| 30 | self.N += 1. |
| 31 | self.p_estimate = ((self.N - 1)*self.p_estimate + x) / self.N |
| 32 | |
| 33 | |
| 34 | def ucb(mean, n, nj): |