Largely based on:

- https://github.com/SalahMouslih/Multi-Arms-Bandit-algorithms-simulation/blob/main/Multi-Armed%20Bandit%20algorithms.ipynb
- https://stackoverflow.com/questions/43445103/inline-animations-in-jupyter

In [1]:

    !pip install matplotlib pandas seaborn

    Requirement already satisfied: matplotlib in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (3.8.3)
    Collecting pandas
      Downloading pandas-2.2.1-cp312-cp312-win_amd64.whl.metadata (19 kB)
    Collecting seaborn
      Downloading seaborn-0.13.2-py3-none-any.whl.metadata (5.4 kB)
    Requirement already satisfied: contourpy>=1.0.1 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from matplotlib) (1.2.0)
    Requirement already satisfied: cycler>=0.10 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from matplotlib) (0.12.1)
    Requirement already satisfied: fonttools>=4.22.0 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from matplotlib) (4.49.0)
    Requirement already satisfied: kiwisolver>=1.3.1 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from matplotlib) (1.4.5)
    Requirement already satisfied: numpy<2,>=1.21 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from matplotlib) (1.26.4)
    Requirement already satisfied: packaging>=20.0 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from matplotlib) (23.2)
    Requirement already satisfied: pillow>=8 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from matplotlib) (10.2.0)
    Requirement already satisfied: pyparsing>=2.3.1 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from matplotlib) (3.1.2)
    Requirement already satisfied: python-dateutil>=2.7 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from matplotlib) (2.8.2)
    Collecting pytz>=2020.1 (from pandas)
      Downloading pytz-2024.1-py2.py3-none-any.whl.metadata (22 kB)
    Collecting tzdata>=2022.7 (from pandas)
      Downloading tzdata-2024.1-py2.py3-none-any.whl.metadata (1.4 kB)
    Requirement already satisfied: six>=1.5 in c:\users\vilia\data\projects\playground\.venv\lib\site-packages (from python-dateutil>=2.7->matplotlib) (1.16.0)
    Downloading pandas-2.2.1-cp312-cp312-win_amd64.whl (11.5 MB)
       ---------------------------------------- 0.0/11.5 MB ? eta -:--:--
       ---------------------------------------- 0.1/11.5 MB 3.5 MB/s eta 0:00:04
       - -------------------------------------- 0.5/11.5 MB 6.0 MB/s eta 0:00:02
       --- ------------------------------------ 1.0/11.5 MB 8.2 MB/s eta 0:00:02
       ------ --------------------------------- 1.8/11.5 MB 10.1 MB/s eta 0:00:01
       --------- ------------------------------ 2.8/11.5 MB 12.8 MB/s eta 0:00:01
       -------------- ------------------------- 4.2/11.5 MB 16.7 MB/s eta 0:00:01
       ------------------- -------------------- 5.5/11.5 MB 17.6 MB/s eta 0:00:01
       ---------------------------- ----------- 8.1/11.5 MB 22.4 MB/s eta 0:00:01
       ---------------------------------------  11.5/11.5 MB 40.9 MB/s eta 0:00:01
       ---------------------------------------- 11.5/11.5 MB 36.4 MB/s eta 0:00:00
    Downloading seaborn-0.13.2-py3-none-any.whl (294 kB)
       ---------------------------------------- 0.0/294.9 kB ? eta -:--:--
       --------------------------------------- 294.9/294.9 kB 19.0 MB/s eta 0:00:00
    Downloading pytz-2024.1-py2.py3-none-any.whl (505 kB)
       ---------------------------------------- 0.0/505.5 kB ? eta -:--:--
       ---------------------------------------- 505.5/505.5 kB ? eta 0:00:00
    Downloading tzdata-2024.1-py2.py3-none-any.whl (345 kB)
       ---------------------------------------- 0.0/345.4 kB ? eta -:--:--
       --------------------------------------- 345.4/345.4 kB 22.3 MB/s eta 0:00:00
    Installing collected packages: pytz, tzdata, pandas, seaborn
    Successfully installed pandas-2.2.1 pytz-2024.1 seaborn-0.13.2 tzdata-2024.1

In [2]:

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

In [14]:

    class Arm:
        def __init__(self, mean=0, seed=None) -> None:
            self.mean = mean
            self.rnd = np.random.default_rng(seed=seed)

        def draw(self) -> float:
            return self.rnd.normal(self.mean)

In [34]:

    arms = [ Arm(mean=m) for m in [0.4,0.6,0.0,0.5]]

In [160]:

    def sim_UCB1(arms, num_rounds):
        df = pd.DataFrame(columns = ['Chosen','Reward'] + [f'm_{i}' for i in range(len(arms))] + [f'b_{i}' for i in range(len(arms))])

        mean = np.array([0.0 for _ in range(len(arms))])
        N = np.array([0 for _ in range(len(arms))])
        beta = np.array([np.nan for _ in range(len(arms))])

        for t in range(len(arms)):
            mean[t] = arms[t].draw()
            N[t] += 1
            df.loc[len(df)] = [t,mean[t]] + mean.tolist() + beta.tolist()
        #
        for t in range(len(arms),num_rounds):
            if np.isnan(beta[0]): beta = np.sqrt(2*np.log(t)/N)
            pi_t = np.argmax(mean+beta)

            r = arms[pi_t].draw()
            mean[pi_t] = ((mean[pi_t]*N[pi_t])+r)/(N[pi_t]+1)
            N[pi_t] += 1

            beta = np.sqrt(2*np.log(t)/N)
            df.loc[len(df)] = [pi_t,r] + mean.tolist() + beta.tolist()

        return df

In [163]:

    df = sim_UCB1(arms, 500)
    df['Chosen'] = df['Chosen'].astype('int64')
    print(df)

         Chosen    Reward       m_0       m_1       m_2       m_3       b_0  \
    0         0 -0.240790 -0.240790  0.000000  0.000000  0.000000       NaN
    1         1  1.186736 -0.240790  1.186736  0.000000  0.000000       NaN
    2         2  0.594749 -0.240790  1.186736  0.594749  0.000000       NaN
    3         3  1.418498 -0.240790  1.186736  0.594749  1.418498       NaN
    4         3  0.682473 -0.240790  1.186736  0.594749  1.050486  1.665109
    ..      ...       ...       ...       ...       ...       ...       ...
    495       1  0.522401  0.254973  0.658312 -0.075245  0.428030  0.543558
    496       1 -0.383651  0.254973  0.655352 -0.075245  0.428030  0.543646
    497       1 -0.659148  0.254973  0.651628 -0.075245  0.428030  0.543734
    498       1  1.028480  0.254973  0.652692 -0.075245  0.428030  0.543822
    499       1  0.599777  0.254973  0.652543 -0.075245  0.428030  0.543910

              b_1       b_2       b_3
    0         NaN       NaN       NaN
    1         NaN       NaN       NaN
    2         NaN       NaN       NaN
    3         NaN       NaN       NaN
    4    1.665109  1.665109  1.177410
    ..        ...       ...       ...
    495  0.188026  0.909546  0.375517
    496  0.187789  0.909694  0.375578
    497  0.187553  0.909842  0.375639
    498  0.187318  0.909989  0.375699
    499  0.187084  0.910136  0.375760

    [500 rows x 10 columns]

In [96]:

    import matplotlib.animation
    import matplotlib.pyplot as plt

In [164]:

    %matplotlib inline
    plt.rcParams["animation.html"] = "jshtml"
    plt.rcParams['figure.dpi'] = 150
    plt.ioff()
    fig, ax = plt.subplots()

    x= np.linspace(0,10,100)
    def animate(t):
        plt.cla()
        plt.ylim((-2,3))
        plt.xlim((-0.2,3.2))
        plt.title(f"Picked A{df['Chosen'][t]} and received r={df['Reward'][t]}")
        x = np.array(range(len(arms)))
        tics = [f'A{i}' for i in range(len(arms))]
        m = np.array(df.loc[t][2:len(arms)+2])
        ci = np.array(df.loc[t][len(arms)+2:])
        plt.xticks(x, [f'A{i}' for i in range(len(arms))])
        plt.plot([x, x],[m+ci,m-ci], color='black')
        plt.plot([x-0.05, x+0.05],[m+ci,m+ci], color='black')
        plt.plot([x-0.05, x+0.05],[m-ci,m-ci], color='black')
        plt.plot(x,m,'o')

    matplotlib.animation.FuncAnimation(fig, animate, frames=len(df))

Out[164]:
Once Loop Reflect
In [ ]: