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&lt;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.