refactored

k-bandits.py

@@ -1,5 +1,6 @@
 import numpy as np
 import matplotlib.pyplot as pl
+from pyparsing import alphas
 from tqdm import tqdm
 
 # For alternate implementation, see:
@@ -18,19 +19,13 @@ REWARD_VARIANCE = 1.0
 def simple_max(Q, N, t, _tie_break):
 	am = np.argmax(Q + _tie_break[t])
 	return am
-#	fm = Q == Q.max()
-#	ffm = np.flatnonzero(fm)
-#	return np.random.choice(ffm)  # breaking ties randomly
-
-def test(_k_arms: int, _episode_len: int, _param: tuple[float, float], ql_star, _tie_break) -> np.array:
-	_epsilon = _param[0]		# Anti-greediness (ability to explore)
-	_rho = _param[1]			# reduce epsilon with age
 
+def test(_k_arms: int, _episode_len: int, ql_star, _tie_break, _epsilon:float=0, _rho:float=0, _q_ic:float=5, _alpha:float=0) -> np.array:
 	rewards = np.zeros(_episode_len)
 	actions = np.zeros(_episode_len)
 
 	# Init Q and N
-	_qu = np.zeros(_k_arms)
+	_qu = np.zeros(_k_arms) + _q_ic
 	_nu = np.zeros(_k_arms)
 
 	# Calc z in advance
@@ -54,10 +49,13 @@ def test(_k_arms: int, _episode_len: int, _param: tuple[float, float], ql_star,
 		_reward = _reward_z[j] + ql_star[_a]
 
 		# calc
-		_nu[_a] = _nu[_a] + 1
-		_qu[_a] = _qu[_a] + (_reward - _qu[_a]) / _nu[_a]
+		if _alpha > 0:
+			_qu[_a] = _qu[_a] + _alpha * (_reward - _qu[_a])
+		else:
+			_nu[_a] = _nu[_a] + 1
+			_qu[_a] = _qu[_a] + (_reward - _qu[_a]) / _nu[_a]
 
-		# Reduce tendency to explore with number of steps (or with age for humans)
+# Reduce tendency to explore with number of steps (or with age for humans)
 		_epsilon = _epsilon * (1 - _rho)
 
 		# Statistics
@@ -68,6 +66,39 @@ def test(_k_arms: int, _episode_len: int, _param: tuple[float, float], ql_star,
 
 	return rewards, actions
 
+class Param:
+	def __init__(self, epsilon:float=0, rho:float=0, alpha:float =0, q_ic:float=0):
+
+		# Anti-greediness (ability to explore)
+		self.epsilon = epsilon
+
+		# Reduce epsilon with age
+		self.rho = rho
+
+		# Constant step size
+		self.alpha = alpha
+
+		# Initial condition Q
+		self.q_ic = q_ic
+
+	def __repr__(self):
+		return f"epsilon={self.epsilon}, rho={self.rho}, alpha={self.alpha}, q_ic={self.q_ic}"
+
+def batch(k_arms, num_episode, num_problems, _param: Param):
+	# Init stats
+	r_mean = np.zeros(episode_len)
+	a_mean = np.zeros(episode_len)
+	for k in tqdm(range(0, num_problems)):
+		# init bandits with different biases for shifting reward probability
+		# -> expected reward q*(a)
+		tie_break = 0.001 * np.random.normal(size=(num_episode, k_arms))
+		r, a = test(_k_arms=k_arms, _episode_len=num_episode, ql_star=q_star[k],
+					_tie_break=tie_break, _epsilon=_param.epsilon,
+					_rho=_param.rho, _q_ic=param.q_ic, _alpha=param.alpha)
+		r_mean += r
+		a_mean += a
+
+	return r_mean, a_mean
 
 if __name__ == '__main__':
 	q_star = np.random.normal(0, 1, (num_realisations, k_arms))
@@ -83,27 +114,18 @@ if __name__ == '__main__':
 	pl.violinplot(arms, positions=range(1, k_arms+1), showmedians=True)
 
 	pl.figure(figsize=(12, 8))
-	params = [(0.0, 0.0), (0.01, 0.0), (0.1, 0.0)]
+	params = [Param(epsilon=0.0), Param(epsilon=0.01), Param(epsilon=0.1)]
+#	params = [Param(epsilon=0.1), Param(alpha=0.2, q_ic=5)]
 	legend = []
 	for param in params:
-		# Init stats
-		r_mean = np.zeros(episode_len)
-		a_mean = np.zeros(episode_len)
-		for k in tqdm(range(0, num_realisations)):
-			# init bandits with different biases for shifting reward probability
-			# -> expected reward q*(a)
-			tie_break = 0.05 * np.random.normal(size=(episode_len, k_arms))
-			r, a = test(_k_arms=k_arms, _episode_len=episode_len, _param=param, ql_star=q_star[k], _tie_break=tie_break)
-			r_mean += r
-			a_mean += a
-
-		legend.append(f"Param {param}")
+		res_r, res_a = batch(k_arms, episode_len, num_realisations, param)
+		legend.append(param)
 
 		pl.subplot(2, 1, 1)
-		pl.plot(r_mean/num_realisations)
+		pl.plot(res_r/num_realisations)
 
 		pl.subplot(2, 1, 2)
-		pl.plot(100*a_mean/num_realisations)
+		pl.plot(100*res_a/num_realisations)
 
 	pl.subplot(2, 1, 1)
 	pl.title(f"E(R) over {num_realisations} realizations, num. bandit arms: K={k_arms}")