- refactored

- fixed evaluation of mean reward

k-bandits.py

@@ -1,5 +1,6 @@
 import numpy as np
 import matplotlib.pyplot as pl
+from typing import Callable
 
 float_formatter = "{:.3f}".format
 np.set_printoptions(formatter={'float_kind': float_formatter})
@@ -7,28 +8,16 @@ np.set_printoptions(formatter={'float_kind': float_formatter})
 K = 10
 
 
-def bandit(bias):
-	zn = np.random.uniform() + bias
+def bandit(a: int, bias: np.array):
+	zn = np.random.uniform() + bias[a]
 	return zn
 
 
-if __name__ == '__main__':
-
-	epsilon = 0.0		# Anti-greediness
-	rho = 0.01			# reduce epsilon with age
-
-	qu = np.zeros(K)
-	nu = np.zeros(K)
-	r_vec = []
-	# init bandits with different biases for shifting reward probability
-	# -> expected reward q*(a)
-	ql_star = [-0.4, -0.3, -0.2, -0.1, 0.0, +0.1, +0.2, +0.3, +0.4, +0.5]
-
-	N = 1000
-	n_ages = 1000
-	for age in range(0, n_ages):
+def test(episode_len: int, epsilon: float, qu: np.array, nu: np.array, bandit: Callable) -> np.array:
 	r_sum = 0
-		for j in range(0, N):
+	r_mean = []
+	n = 0
+	for j in range(0, episode_len):
 		# choose action
 		z = np.random.uniform()
 		if z <= epsilon:
@@ -39,18 +28,44 @@ if __name__ == '__main__':
 			a = np.argmax(qu)
 
 		# get reward from bandit
-			r = bandit(ql_star[a])
+		r = bandit(a)
 		nu[a] = nu[a] + 1
-			qu[a] = qu[a] + (r - qu[a])/nu[a]
-			r_sum += r
-		r_vec.append(r_sum/N)
+		qu[a] = qu[a] + (r - qu[a]) / nu[a]
+
 		# Reduce tendency to explore with number of steps (or with age for humans)
-		epsilon = epsilon*(1-rho)
+		epsilon = epsilon * (1 - rho)
+
+		# Statistics
+		r_sum += r
+		n += 1
+		r_mean.append(r_sum/n)
+
+	return np.array(r_mean)
+
+
+if __name__ == '__main__':
+
+	epsilon = 0.1		# Anti-greediness
+	rho = 0.01		# reduce epsilon with age
+
+	# init bandits with different biases for shifting reward probability
+	# -> expected reward q*(a)
+	ql_star = np.array([-0.4, -0.3, -0.2, -0.1, 0.0, +0.1, +0.2, +0.3, +0.4, +0.5])
+
+	num_realisations = 10
+	episode_len = 1000
+	r_mean = np.zeros(episode_len)
+	for age in range(0, num_realisations):
+		# Init Q and N
+		qu = np.zeros(K)
+		nu = np.zeros(K)
+
+		r_mean += test(episode_len, epsilon, qu, nu, bandit=lambda a: bandit(a, ql_star))
 
 	print(f"ql_star  = {ql_star}")
 	print(f"qu       = {qu}")
 	print(f"nu       = {nu}")
 
-	pl.plot(np.array(r_vec))
+	pl.plot(r_mean/num_realisations)
 	pl.grid()
 	pl.show()