multi-armed bandit

'''
based on https://github.com/awjuliani/DeepRL-Agents

'''

# coding: utf-8

import tensorflow as tf
import tensorflow.contrib.slim as slim
import numpy as np
tf.reset_default_graph()
#List out our bandit arms. 
#Currently arm 4 (index #3) is set to most often provide a positive reward.
bandit_arms = np.array([0.2,0,-0.2,-2])
num_arms = len(bandit_arms)
def pullBandit(bandit,batch_size=1):
    #Get a random number.
    result = np.random.randn(batch_size)
    return 2*(result > bandit)-1
    
def single_train():
    #These two lines established the feed-forward part of the network. 
    weights = tf.Variable(tf.ones([num_arms]))
    output = tf.nn.softmax(weights)
    
    #The next six lines establish the training proceedure. We feed the reward and chosen action into the network
    #to compute the loss, and use it to update the network.
    reward_holder = tf.placeholder(shape=[1],dtype=tf.float32)
    action_holder = tf.placeholder(shape=[1],dtype=tf.int32)
    
    responsible_output = tf.slice(output,action_holder,[1])
    loss = -(tf.log(responsible_output)*reward_holder)
    optimizer = tf.train.AdamOptimizer(learning_rate=1e-3)
    
    
    update = optimizer.minimize(loss)
    
    total_episodes = 10000 #Set total number of episodes to train agent on.
    total_reward = np.zeros(num_arms) #Set scoreboard for bandit arms to 0.
    
    init = tf.global_variables_initializer()
    
    # Launch the tensorflow graph
    with tf.Session() as sess:
        sess.run(init)
        i = 0
        while i < total_episodes:
            
            #Choose action according to Boltzmann distribution.
            actions = sess.run(output)
            a = np.random.choice(actions,p=actions)
            action = np.argmax(actions == a)
    
            reward = pullBandit(bandit_arms[action]) #Get our reward from picking one of the bandit arms.
            
            #Update the network.
            _,resp,ww = sess.run([update,responsible_output,weights], feed_dict={reward_holder:reward,action_holder:[action]})
            
            #Update our running tally of scores.
            total_reward[action] += reward
            if i % 50 == 0:
                print("Running reward for the " + str(num_arms) + " arms of the bandit: " + str(total_reward))
            i+=1
    
        print('output: ',sess.run(output))
    print("\nThe agent thinks arm " + str(np.argmax(ww)+1) + " is the most promising....")
    if np.argmax(ww) == np.argmax(-np.array(bandit_arms)):
        print("...and it was right!")
    else:
        print("...and it was wrong!")


def batch_train():
    batch_size=5
    #These two lines established the feed-forward part of the network. 
    weights = tf.Variable(tf.ones([num_arms]))
    output = tf.nn.softmax(weights)
    
    #The next six lines establish the training proceedure. We feed the reward and chosen action into the network
    #to compute the loss, and use it to update the network.
    reward_holder = tf.placeholder(shape=[None],dtype=tf.float32)
    action_holder = tf.placeholder(shape=[None],dtype=tf.int32)
    
    responsible_output = tf.gather(output,action_holder)
    loss = tf.reduce_mean(-(tf.log(responsible_output)*reward_holder))
    optimizer = tf.train.AdamOptimizer(learning_rate=1e-3)
    
    
    update = optimizer.minimize(loss)
    
    total_episodes = 500 #Set total number of episodes to train agent on.
    total_reward = np.zeros(num_arms) #Set scoreboard for bandit arms to 0.
    
    init = tf.global_variables_initializer()
    
    # Launch the tensorflow graph
    with tf.Session() as sess:
        sess.run(init)
        i = 0
        while i < total_episodes:
            
            #Choose action according to Boltzmann distribution.
            actions = sess.run(output)
            action = np.random.choice(np.arange(num_arms),size=batch_size,p=actions)
    
            reward = pullBandit(bandit_arms[action],batch_size) #Get our reward from picking one of the bandit arms.
            
            #Update the network.
            _,resp,ww = sess.run([update,responsible_output,weights], feed_dict={reward_holder:reward,action_holder:action})
            
            #Update our running tally of scores.
            total_reward[action] += reward
            if i % 50 == 0:
                print("Running reward for the " + str(num_arms) + " arms of the bandit: " + str(total_reward))
            i+=1
    
        print('output: ',sess.run(output))
    print("\nThe agent thinks arm " + str(np.argmax(ww)+1) + " is the most promising....")
    if np.argmax(ww) == np.argmax(-np.array(bandit_arms)):
        print("...and it was right!")
    else:
        print("...and it was wrong!")



if __name__ == '__main__':
    single_train()
    #batch_train()