diff --git a/k-bandits.py b/k-bandits.py
index 96a428e..1bfeba4 100644
--- a/k-bandits.py
+++ b/k-bandits.py
@@ -13,7 +13,8 @@ def bandit(a: int, bias: np.array):
 def test(_k_arms: int, _episode_len: int, _bandit: Callable, _param: tuple[float, float]) -> np.array:
 	_epsilon = _param[0]		# Anti-greediness (ability to explore)
 	_rho = _param[1]		# reduce epsilon with age
-	_r_mean = []
+	_r_list = []
+	_a_list = []
 
 	# Init Q and N
 	_qu = np.zeros(_k_arms)
@@ -24,26 +25,26 @@ 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]
+		_a = a_expl[j]
 		# choose action
 		if z[j] > _epsilon:
 			# Choose best action
-			a = np.argmax(_qu)
+			_a = np.argmax(_qu)
 
 		# get reward from bandit
-		r = _bandit(a)
-		_nu[a] = _nu[a] + 1
-		_qu[a] = _qu[a] + (r - _qu[a]) / _nu[a]
+		_r = _bandit(_a)
+		_nu[_a] = _nu[_a] + 1
+		_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)
 
 		# Statistics
-		_r_mean.append(r)
+		_a_list.append(_a)
+		_r_list.append(_r)
 
-	return np.array(_r_mean)
+	return np.array(_r_list), np.array(_a_list)
 
 
 if __name__ == '__main__':
@@ -64,7 +65,8 @@ if __name__ == '__main__':
 			# -> expected reward q*(a)
 			ql_star = np.random.uniform(-highest_reward, +highest_reward, k_arms)
 
-			r_mean += test(_k_arms=k_arms, _episode_len=episode_len, _bandit=lambda a: bandit(a, ql_star), _param=param)
+			r_list, a_list = test(_k_arms=k_arms, _episode_len=episode_len, _bandit=lambda a: bandit(a, ql_star), _param=param)
+			r_mean += r_list
 
 		pl.plot(r_mean/num_realisations)
 		legend.append(f"Param {param}")