# k-armed bandit problem in Python

I implemented k-armed bandit problem in C#, MATLAB and Python. C# and Matlab code run fairly fast (With same settings of T = 2000 and nRun = 1000 the elapsed time is about 6sec). However, the Python version is ten times slower and takes about 60 seconds to complete. I am a beginner in Python but I guess it should be faster than this. I appreciate any tips to speed up this code.

import numpy as np
import matplotlib.pyplot as plt
import time

t1 = time.time()

plt.close('all')

nActions = 10
# np.random.seed(0)
sigma = np.ones(nActions)
mu = np.array([-0.4677,-0.1249,1.4790,-0.8608,0.7847,0.3086,-0.2339,-1.0570,-0.2841,-0.0867])

def Reward(action:int, mu:np.ndarray, sigma:np.ndarray):
return mu[action] + sigma[action]*np.random.normal()

def EpsGreedyPolicy(Q,eps):
p = np.random.rand()
if p<eps:
nAction = np.size(Q)
return np.random.randint(nAction)
else:
return GreedyPolicy(Q)

def GreedyPolicy(Q):
A = np.nonzero(Q == np.max(Q)) # A is a tuple
n = np.size(A)
if n == 1:
return A[0]
else:
j = np.random.randint(n)
return A[0][j] # Note A here is a tuple not array

# Reinforcement Learning

BestAction = np.argmax(mu)

T = 2000
nRun = 1000;

arrA = np.zeros((T,nRun))
arrR = np.zeros((T,nRun))

Q0 = 0 # Initial values

for j in range(nRun):
TotalReward = np.zeros((nActions))
Counter = np.zeros((nActions))

for t in range(T):

# Calculate action values
Q = TotalReward/Counter
Q[np.isnan(Q)] = Q0

# Apply Policy
a = EpsGreedyPolicy(Q,0.01)

# Commit Action
r = Reward(a, mu, sigma)

# Update relevant stats
Counter[a] += 1
TotalReward[a] += r

# Save results
arrA[t][j] = a
arrR[t][j] = r

#print('End of Run ', j)

print("end of run ",j)
# end of main loop

# Plot Results
BestActionSelected = (arrA == BestAction)
BestActionSelectedMean = np.average(BestActionSelected, axis=1)

plt.figure(2,figsize=(10,10))

plt.subplot(2,1,1)
plot1 = plt.plot(BestActionSelectedMean)
plt.title('best action selection')

plt.subplot(2,1,2)
plot2 = plt.plot(np.average(arrR, axis=1))
plt.title('average reward')

plt.show()

t2 = time.time()

print(t2-t1)


Python is notoriously slow with native loops; i.e.

for j in range(nRun):
...
for t in range(T):
...


is probably your bottleneck. I haven't tried to understand necessarily what you are coding but can you vectorise some calculations?

Something like

sigma = np.ones(nActions)
mu = np.array([-0.4677,-0.1249,1.4790,-0.8608,0.7847,0.3086,-0.2339,-1.0570,-0.2841,-0.0867])
def Reward(action:int, mu:np.ndarray, sigma:np.ndarray):
return mu[action] + sigma[action]*np.random.normal()


is better replaced by:

temp_rvs = np.random.rand(nRun * T, 10)
reward_arr = mu + sigma * temp_rvs


then you can just index an existing variable instead of generating random variables each time through a function in your loop, i.e. you can:

reward[(j+1)*T, a]


But this takes up $10 * nRun * T * 8$ bytes of RAM to store the random variables.

I answered a similar question to this here: How to simulate this Gamma expansion in a Python way perhaps it might give you some ideas.

• Thank you. I changed my code to generate all random numbers outside the main loop and access them whenever needed. This trick reduced my code run-time to under 30 seconds. Do you think further improvement in speed is possible? C# and Matlab version are still five time faster. Am I missing something here? – Omid Aug 19 '18 at 5:11
• Python probably won't match the speeds in C. You can investigate cython: en.wikipedia.org/wiki/Cython, I am surprised it is much slower than Matlab, unless Matlab functions are built-in in which case they probably use faster execution code, for example NumPy in python uses C-extensions to dramatically increase speed. – Attack68 Aug 19 '18 at 6:35