import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

def bandit(mode, K=5, N=1000):
    q_true, Q, N_a = np.random.randn(K), np.zeros(K), np.zeros(K)
    for t in range(1, N + 1):
        bonus = np.sqrt(2 * np.log(t) / np.maximum(N_a, 1e-9))
        a = np.argmax(Q + bonus) if mode == "ucb" else (np.random.randint(K) if np.random.rand() < 0.1 else np.argmax(Q))
        r = q_true[a] + np.random.randn()
        N_a[a] += 1
        Q[a] += (r - Q[a]) / N_a[a]
        yield r

pd.DataFrame({"UCB": bandit("ucb"), "Eps-Greedy": bandit("eps")}).expanding().mean().plot(title="Bandit Exploration")
plt.show()