From 7c7a1000b8c6cdc06993ecc5a7db36be91b29531 Mon Sep 17 00:00:00 2001
From: Jens Ahrensfeld <jens@jayfield.org>
Date: Wed, 2 Apr 2025 19:31:08 +0200
Subject: [PATCH] - refactored - calculate optimal actions

---
 k-bandits.py | 40 +++++++++++++++++++++-------------------
 1 file changed, 21 insertions(+), 19 deletions(-)

diff --git a/k-bandits.py b/k-bandits.py
index 1bfeba4..83ef1ab 100644
--- a/k-bandits.py
+++ b/k-bandits.py
@@ -5,17 +5,14 @@ from typing import Callable
 float_formatter = "{:.3f}".format
 np.set_printoptions(formatter={'float_kind': float_formatter})
 
-def bandit(a: int, bias: np.array):
-	zn = np.random.normal() + bias[a]
-	return zn
+REWARD_VARIANCE = 0.1
 
-
-def test(_k_arms: int, _episode_len: int, _bandit: Callable, _param: tuple[float, float]) -> np.array:
+def test(_k_arms: int, _episode_len: int, _bias: np.array, _param: tuple[float, float]) -> np.array:
 	_epsilon = _param[0]		# Anti-greediness (ability to explore)
-	_rho = _param[1]		# reduce epsilon with age
+	_rho = _param[1]			# reduce epsilon with age
 	_r_list = []
 	_a_list = []
-
+	_opt_a_list = []
 	# Init Q and N
 	_qu = np.zeros(_k_arms)
 	_nu = np.zeros(_k_arms)
@@ -26,29 +23,33 @@ def test(_k_arms: int, _episode_len: int, _bandit: Callable, _param: tuple[float
 	# 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):
-		_a = a_expl[j]
+
 		# choose action
-		if z[j] > _epsilon:
-			# Choose best action
-			_a = np.argmax(_qu)
+		_a = np.argmax(_qu) if z[j] > _epsilon else a_expl[j]
 
 		# get reward from bandit
-		_r = _bandit(_a)
+		_reward_list = REWARD_VARIANCE*np.random.normal(size=_k_arms) + _bias
+
+		# get reward from action
+		_reward = _reward_list[_a]
+
+		# calc
 		_nu[_a] = _nu[_a] + 1
-		_qu[_a] = _qu[_a] + (_r - _qu[_a]) / _nu[_a]
+		_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)
 
 		# Statistics
 		_a_list.append(_a)
-		_r_list.append(_r)
+		_r_list.append(_reward)
+		opt_a = 1 if _a == np.argmax(_reward_list) else 0
+		_opt_a_list.append(opt_a)
 
-	return np.array(_r_list), np.array(_a_list)
+	return np.array(_r_list), np.array(_a_list), np.array(_opt_a_list)
 
 
 if __name__ == '__main__':
-
 	# Init parameters
 	k_arms = 10
 	highest_reward = 1.5
@@ -58,15 +59,16 @@ if __name__ == '__main__':
 	params = [(0.0, 0.0), (0.01, 0.002), (0.1, 0.002)]
 	legend = []
 	for param in params:
-		# Init r_mean
+		# Init stats
 		r_mean = np.zeros(episode_len)
+		o_mean = np.zeros(episode_len)
 		for realization in 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 = test(_k_arms=k_arms, _episode_len=episode_len, _bandit=lambda a: bandit(a, ql_star), _param=param)
+			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
 
 		pl.plot(r_mean/num_realisations)
 		legend.append(f"Param {param}")