refactored

k-bandits.py

@@ -20,12 +20,33 @@ def simple_max(Q, N, t, _tie_break):
 	am = np.argmax(Q + _tie_break[t])
 	return am
 
-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:
+class Param:
+	def __init__(self, epsilon:float=0, rho:float=0, alpha:float =0, q_ic:float=0, _argmax_func=simple_max):
+
+		# 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
+
+		# argmax function
+		self.argmax_func=_argmax_func
+
+	def __repr__(self):
+		return f"epsilon={self.epsilon}, rho={self.rho}, alpha={self.alpha}, q_ic={self.q_ic}\nargmax={self.argmax_func.__repr__()}"
+
+def test(_k_arms: int, _episode_len: int, ql_star, _tie_break, _param: Param) -> np.array:
 	rewards = np.zeros(_episode_len)
 	actions = np.zeros(_episode_len)
 
 	# Init Q and N
-	_qu = np.zeros(_k_arms) + _q_ic
+	_qu = np.zeros(_k_arms) + _param.q_ic
 	_nu = np.zeros(_k_arms)
 
 	# Calc z in advance
@@ -40,23 +61,23 @@ def test(_k_arms: int, _episode_len: int, ql_star, _tie_break, _epsilon:float=0,
 	best_action = np.argmax(ql_star)
 	for j in range(0, _episode_len):
 		# choose action
-		if z[j] < _epsilon:
+		if z[j] < _param.epsilon:
 			_a = _a_expl[j]
 		else:
-			_a = simple_max(_qu, _nu, j, _tie_break)
+			_a = _param.argmax_func(_qu, _nu, j, _tie_break)
 
 		# get reward from bandit
 		_reward = _reward_z[j] + ql_star[_a]
 
 		# calc
-		if _alpha > 0:
-			_qu[_a] = _qu[_a] + _alpha * (_reward - _qu[_a])
+		if _param.alpha > 0:
+			_qu[_a] = _qu[_a] + _param.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)
-		_epsilon = _epsilon * (1 - _rho)
+		# Reduce tendency to explore with number of steps (or with age for humans)
+		_epsilon = _param.epsilon * (1 - _param.rho)
 
 		# Statistics
 		rewards[j] += _reward
@@ -66,24 +87,6 @@ def test(_k_arms: int, _episode_len: int, ql_star, _tie_break, _epsilon:float=0,
 
 	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)
@@ -93,8 +96,7 @@ def batch(k_arms, num_episode, num_problems, _param: Param):
 		# -> 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)
+					_tie_break=tie_break, _param=_param)
 		r_mean += r
 		a_mean += a
 
@@ -114,8 +116,8 @@ if __name__ == '__main__':
 	pl.violinplot(arms, positions=range(1, k_arms+1), showmedians=True)
 
 	pl.figure(figsize=(12, 8))
-	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)]
+#	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:
 		res_r, res_a = batch(k_arms, episode_len, num_realisations, param)