From c92fc9c13a9586f3033214b928f9299286d1f2ce Mon Sep 17 00:00:00 2001
From: Jens Ahrensfeld <jens@jayfield.org>
Date: Fri, 4 Apr 2025 14:21:51 +0200
Subject: [PATCH] - refactored - added violin plot for q_star - corrected
 calculation of results - increased speed

---
 k-bandits.py | 119 +++++++++++++++++++++++++++++++++++----------------
 1 file changed, 81 insertions(+), 38 deletions(-)

diff --git a/k-bandits.py b/k-bandits.py
index 83ef1ab..f9a260f 100644
--- a/k-bandits.py
+++ b/k-bandits.py
@@ -1,18 +1,34 @@
 import numpy as np
 import matplotlib.pyplot as pl
-from typing import Callable
+from tqdm import tqdm
+
+# For alternate implementation, see:
+# https://www.kaggle.com/code/parsasam/reinforcement-learning-notes-multi-armed-bandits
 
 float_formatter = "{:.3f}".format
 np.set_printoptions(formatter={'float_kind': float_formatter})
 
-REWARD_VARIANCE = 0.1
+# Init parameters
+k_arms = 10
+num_realisations = 2000
+episode_len = 1000
 
-def test(_k_arms: int, _episode_len: int, _bias: np.array, _param: tuple[float, float]) -> np.array:
+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
-	_r_list = []
-	_a_list = []
-	_opt_a_list = []
+
+	rewards = np.zeros(_episode_len)
+	actions = np.zeros(_episode_len)
+
 	# Init Q and N
 	_qu = np.zeros(_k_arms)
 	_nu = np.zeros(_k_arms)
@@ -20,18 +36,22 @@ def test(_k_arms: int, _episode_len: int, _bias: np.array, _param: tuple[float,
 	# Calc z in advance
 	z = np.random.uniform(size=_episode_len)
 
-	# Calc a_expl in advance
-	a_expl = np.random.randint(low=0, high=_k_arms, size=_episode_len)
-	for j in range(0, _episode_len):
+	# Calc a in advance
+	_a_expl = np.random.randint(_k_arms, size=_episode_len)
 
+	# Calc rewards in advance
+	_reward_z = REWARD_VARIANCE * np.random.normal(size=_episode_len)
+
+	best_action = np.argmax(ql_star)
+	for j in range(0, _episode_len):
 		# choose action
-		_a = np.argmax(_qu) if z[j] > _epsilon else a_expl[j]
+		if z[j] < _epsilon:
+			_a = _a_expl[j]
+		else:
+			_a = simple_max(_qu, _nu, j, _tie_break)
 
 		# get reward from bandit
-		_reward_list = REWARD_VARIANCE*np.random.normal(size=_k_arms) + _bias
-
-		# get reward from action
-		_reward = _reward_list[_a]
+		_reward = _reward_z[j] + ql_star[_a]
 
 		# calc
 		_nu[_a] = _nu[_a] + 1
@@ -41,45 +61,68 @@ def test(_k_arms: int, _episode_len: int, _bias: np.array, _param: tuple[float,
 		_epsilon = _epsilon * (1 - _rho)
 
 		# Statistics
-		_a_list.append(_a)
-		_r_list.append(_reward)
-		opt_a = 1 if _a == np.argmax(_reward_list) else 0
-		_opt_a_list.append(opt_a)
+		rewards[j] += _reward
 
-	return np.array(_r_list), np.array(_a_list), np.array(_opt_a_list)
+		if _a == best_action:
+			actions[j] += 1
+
+	return rewards, actions
 
 
 if __name__ == '__main__':
-	# Init parameters
-	k_arms = 10
-	highest_reward = 1.5
-	num_realisations = 2000
-	episode_len = 1000
+	q_star = np.random.normal(0, 1, (num_realisations, k_arms))
+	arms = np.zeros((num_realisations, k_arms))
+	for i in range(k_arms):
+		arms[:,i] = np.random.normal(q_star[0, i], 1, num_realisations)  # first problem as a sample
 
-	params = [(0.0, 0.0), (0.01, 0.002), (0.1, 0.002)]
+	pl.figure(figsize=(12, 8))
+	pl.ylabel('Rewards distribution')
+	pl.xlabel('Actions')
+	pl.xticks(range(1, k_arms+1))
+	pl.yticks(np.arange(-5, 5, 0.5))
+	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)]
 	legend = []
 	for param in params:
 		# Init stats
 		r_mean = np.zeros(episode_len)
-		o_mean = np.zeros(episode_len)
-		for realization in range(0, num_realisations):
+		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)
-			ql_star = np.random.uniform(-highest_reward, +highest_reward, k_arms)
-			r_list, a_list, o_list = test(_k_arms=k_arms, _episode_len=episode_len, _bias=ql_star, _param=param)
-			r_mean += r_list
-			o_mean += o_list
+			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
 
-		pl.plot(r_mean/num_realisations)
 		legend.append(f"Param {param}")
 
-	pl.grid()
-	pl.title(f"Norm. E(R), num. realizations: Z={num_realisations}, num. bandit arms: K={k_arms}")
-	pl.xlabel("Episode")
-	pl.ylabel("E(R)/Z")
+		pl.subplot(2, 1, 1)
+		pl.plot(r_mean/num_realisations)
+
+		pl.subplot(2, 1, 2)
+		pl.plot(100*a_mean/num_realisations)
+
+	pl.subplot(2, 1, 1)
+	pl.title(f"E(R) over {num_realisations} realizations, num. bandit arms: K={k_arms}")
 	pl.legend(legend)
+	pl.grid()
 	ax = pl.gca()
-	ax.set_xlim([-episode_len/10, episode_len])
-	ax.set_ylim([0, highest_reward])
+	ax.set_xlim([-episode_len / 10, episode_len])
+	ax.set_ylim([0, 1.6])
+	pl.ylabel("E(R)")
+
+	pl.subplot(2, 1, 2)
+	pl.legend(legend)
+	pl.grid()
+	ax = pl.gca()
+	ax.set_xlim([-episode_len / 10, episode_len])
+	ax.set_ylim([0, 100])
+	pl.ylabel("Percentage of optimal actions")
+	pl.ylabel("%")
+
+	pl.xlabel("Episode")
 
 	pl.show()