image_output_1.png
import matplotlib.pyplot as plt

import numpy as np # linear algebra

import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

class Envierment:
    def __init__(self, K, horizon):
        self.K, self.horizon = (K, horizon)
        self.q_values = list()
        for k in range(self.K):
            self.q_values.append(np.random.randn())
    def give_reward(self, k):
        if k >= K:
            print('Invalid action!')
            raise ValueError()
        else:
            return np.random.randn() + self.q_values[k]
class Agent:
    def __init__(self, K, horizon):
        self.K, self.horizon = (K, horizon)
        self.action, self.reward = (0, 0)
        self.result = np.zeros((horizon, 1))
        self.result_df = pd.DataFrame()
    def take_action(self, k=None):
        if k == None:
            self.action = np.random.choice(self.K)
        else:
            self.action = k
        return self.action
    def record_reward(self, trial, n, reward):
        self.reward = reward
        if self.result.shape[1] < n + 1:
            self.result = np.c_[self.result, np.zeros(self.horizon)]
        self.result[n][trial] = self.reward
    def turn_result_into_df(self):
        self.result = np.cumsum(self.result, axis=0)
        for n in range(1, self.horizon):
            self.result[n] /= n
        self.result_df = pd.DataFrame(self.result, columns=['trial_' + str(x) for x in range(self.result.shape[1])])
K, horizon = (3, 100)
envioronment = Envierment(K, horizon)
agent = Agent(envioronment.K, envioronment.horizon)
trials = 10
Result = np.zeros((horizon, trials))
for trial in range(trials):
    for n in range(horizon):
        action = agent.take_action()
        reward = envioronment.give_reward(action)
        agent.record_reward(trial, n, reward)
agent.turn_result_into_df()
fig, ax = plt.subplots(figsize=(20, 10))
for trial in range(trials):
    ax.plot(agent.result_df['trial_' + str(trial)])
ax.legend(agent.result_df.columns, loc='upper right')
ax.set(title='Average reward for each trial')
print(f'K: {envioronment.K}')
print(f'action-value: {envioronment.q_values}')