A3C-Meta-Context.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Meta Reinforcement Learning with A3C - Rainbow Bandit\n",
    "\n",
    "This iPython notebook includes an implementation of the [A3C algorithm capable of Meta-RL](https://arxiv.org/pdf/1611.05763.pdf).\n",
    "\n",
    "For more information see the accompanying [Medium post](https://medium.com/p/b15b592a2ddf)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import threading\n",
    "import multiprocessing\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import tensorflow as tf\n",
    "import tensorflow.contrib.slim as slim\n",
    "import scipy.signal\n",
    "from PIL import Image\n",
    "from PIL import ImageDraw \n",
    "from PIL import ImageFont\n",
    "%matplotlib inline\n",
    "from helper import *\n",
    "\n",
    "from random import choice\n",
    "from time import sleep\n",
    "from time import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class contextual_bandit():\n",
    "    def __init__(self):\n",
    "        self.num_actions = 2\n",
    "        self.reset()\n",
    "        \n",
    "    def get_state(self):\n",
    "        self.internal_state = np.random.permutation(self.choices)\n",
    "        self.state = np.concatenate(np.reshape(self.internal_state,[2,1,1,3]),axis=1)\n",
    "        return self.state\n",
    "                \n",
    "    def reset(self):\n",
    "        self.timestep = 0\n",
    "        color = [np.random.uniform(),np.random.uniform(),np.random.uniform()]\n",
    "        a = [np.reshape(np.array(color),[1,1,3]),np.reshape(1-np.array(color),[1,1,3])]\n",
    "        self.true = a[0]\n",
    "        self.choices = a\n",
    "        return self.get_state()\n",
    "        \n",
    "    def pullArm(self,action):\n",
    "        self.timestep += 1\n",
    "        if (self.internal_state[action] == self.true).all() == True: \n",
    "            reward = 1.0\n",
    "        else:\n",
    "            reward = 0.0\n",
    "        new_state = self.get_state()\n",
    "        if self.timestep > 99: \n",
    "            done = True\n",
    "        else: \n",
    "            done = False\n",
    "        return new_state,reward,done,self.timestep"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Actor-Critic Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class AC_Network():\n",
    "    def __init__(self,a_size,scope,trainer):\n",
    "        with tf.variable_scope(scope):\n",
    "            #Input and visual encoding layers\n",
    "            self.state = tf.placeholder(shape=[None,1,2,3],dtype=tf.float32)\n",
    "            self.prev_rewards = tf.placeholder(shape=[None,1],dtype=tf.float32)\n",
    "            self.prev_actions = tf.placeholder(shape=[None],dtype=tf.int32)\n",
    "            self.timestep = tf.placeholder(shape=[None,1],dtype=tf.float32)\n",
    "            self.prev_actions_onehot = tf.one_hot(self.prev_actions,a_size,dtype=tf.float32)\n",
    "\n",
    "            hidden = tf.concat([slim.flatten(self.state),self.prev_rewards,self.prev_actions_onehot,self.timestep],1)\n",
    "            #Recurrent network for temporal dependencies\n",
    "            lstm_cell = tf.contrib.rnn.BasicLSTMCell(64,state_is_tuple=True)\n",
    "            c_init = np.zeros((1, lstm_cell.state_size.c), np.float32)\n",
    "            h_init = np.zeros((1, lstm_cell.state_size.h), np.float32)\n",
    "            self.state_init = [c_init, h_init]\n",
    "            c_in = tf.placeholder(tf.float32, [1, lstm_cell.state_size.c])\n",
    "            h_in = tf.placeholder(tf.float32, [1, lstm_cell.state_size.h])\n",
    "            self.state_in = (c_in, h_in)\n",
    "            rnn_in = tf.expand_dims(hidden, [0])\n",
    "            step_size = tf.shape(self.prev_rewards)[:1]\n",
    "            state_in = tf.contrib.rnn.LSTMStateTuple(c_in, h_in)\n",
    "            lstm_outputs, lstm_state = tf.nn.dynamic_rnn(\n",
    "                lstm_cell, rnn_in, initial_state=state_in, sequence_length=step_size,\n",
    "                time_major=False)\n",
    "            lstm_c, lstm_h = lstm_state\n",
    "            self.state_out = (lstm_c[:1, :], lstm_h[:1, :])\n",
    "            rnn_out = tf.reshape(lstm_outputs, [-1, 64])\n",
    "            \n",
    "            self.actions = tf.placeholder(shape=[None],dtype=tf.int32)\n",
    "            self.actions_onehot = tf.one_hot(self.actions,a_size,dtype=tf.float32)\n",
    "                        \n",
    "            #Output layers for policy and value estimations\n",
    "            self.policy = slim.fully_connected(rnn_out,a_size,\n",
    "                activation_fn=tf.nn.softmax,\n",
    "                weights_initializer=normalized_columns_initializer(0.01),\n",
    "                biases_initializer=None)\n",
    "            self.value = slim.fully_connected(rnn_out,1,\n",
    "                activation_fn=None,\n",
    "                weights_initializer=normalized_columns_initializer(1.0),\n",
    "                biases_initializer=None)\n",
    "            \n",
    "            #Only the worker network need ops for loss functions and gradient updating.\n",
    "            if scope != 'global':\n",
    "                self.target_v = tf.placeholder(shape=[None],dtype=tf.float32)\n",
    "                self.advantages = tf.placeholder(shape=[None],dtype=tf.float32)\n",
    "                \n",
    "                self.responsible_outputs = tf.reduce_sum(self.policy * self.actions_onehot, [1])\n",
    "\n",
    "                #Loss functions\n",
    "                self.value_loss = 0.5 * tf.reduce_sum(tf.square(self.target_v - tf.reshape(self.value,[-1])))\n",
    "                self.entropy = - tf.reduce_sum(self.policy * tf.log(self.policy + 1e-7))\n",
    "                self.policy_loss = -tf.reduce_sum(tf.log(self.responsible_outputs + 1e-7)*self.advantages)\n",
    "                self.loss = 0.5 * self.value_loss + self.policy_loss - self.entropy * 0.05\n",
    "\n",
    "                #Get gradients from local network using local losses\n",
    "                local_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, scope)\n",
    "                self.gradients = tf.gradients(self.loss,local_vars)\n",
    "                self.var_norms = tf.global_norm(local_vars)\n",
    "                grads,self.grad_norms = tf.clip_by_global_norm(self.gradients,999.0)\n",
    "                \n",
    "                #Apply local gradients to global network\n",
    "                global_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'global')\n",
    "                self.apply_grads = trainer.apply_gradients(zip(grads,global_vars))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Worker Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "class Worker():\n",
    "    def __init__(self,game,name,a_size,trainer,model_path,global_episodes):\n",
    "        self.name = \"worker_\" + str(name)\n",
    "        self.number = name        \n",
    "        self.model_path = model_path\n",
    "        self.trainer = trainer\n",
    "        self.global_episodes = global_episodes\n",
    "        self.increment = self.global_episodes.assign_add(1)\n",
    "        self.episode_rewards = []\n",
    "        self.episode_lengths = []\n",
    "        self.episode_mean_values = []\n",
    "        self.summary_writer = tf.summary.FileWriter(\"train_\"+str(self.number))\n",
    "\n",
    "        #Create the local copy of the network and the tensorflow op to copy global paramters to local network\n",
    "        self.local_AC = AC_Network(a_size,self.name,trainer)\n",
    "        self.update_local_ops = update_target_graph('global',self.name)        \n",
    "        self.env = game\n",
    "        \n",
    "    def train(self,rollout,sess,gamma,bootstrap_value):\n",
    "        rollout = np.array(rollout)\n",
    "        states = rollout[:,0]\n",
    "        actions = rollout[:,1]\n",
    "        rewards = rollout[:,2]\n",
    "        timesteps = rollout[:,3]\n",
    "        prev_rewards = [0] + rewards[:-1].tolist()\n",
    "        prev_actions = [0] + actions[:-1].tolist()\n",
    "        values = rollout[:,5]\n",
    "        \n",
    "        self.pr = prev_rewards\n",
    "        self.pa = prev_actions\n",
    "        # Here we take the rewards and values from the rollout, and use them to \n",
    "        # generate the advantage and discounted returns. \n",
    "        # The advantage function uses \"Generalized Advantage Estimation\"\n",
    "        self.rewards_plus = np.asarray(rewards.tolist() + [bootstrap_value])\n",
    "        discounted_rewards = discount(self.rewards_plus,gamma)[:-1]\n",
    "        self.value_plus = np.asarray(values.tolist() + [bootstrap_value])\n",
    "        advantages = rewards + gamma * self.value_plus[1:] - self.value_plus[:-1]\n",
    "        advantages = discount(advantages,gamma)\n",
    "\n",
    "        # Update the global network using gradients from loss\n",
    "        # Generate network statistics to periodically save\n",
    "        rnn_state = self.local_AC.state_init\n",
    "        feed_dict = {self.local_AC.target_v:discounted_rewards,\n",
    "            self.local_AC.state:np.stack(states,axis=0),\n",
    "            self.local_AC.prev_rewards:np.vstack(prev_rewards),\n",
    "            self.local_AC.prev_actions:prev_actions,\n",
    "            self.local_AC.actions:actions,\n",
    "            self.local_AC.timestep:np.vstack(timesteps),\n",
    "            self.local_AC.advantages:advantages,\n",
    "            self.local_AC.state_in[0]:rnn_state[0],\n",
    "            self.local_AC.state_in[1]:rnn_state[1]}\n",
    "        v_l,p_l,e_l,g_n,v_n,_ = sess.run([self.local_AC.value_loss,\n",
    "            self.local_AC.policy_loss,\n",
    "            self.local_AC.entropy,\n",
    "            self.local_AC.grad_norms,\n",
    "            self.local_AC.var_norms,\n",
    "            self.local_AC.apply_grads],\n",
    "            feed_dict=feed_dict)\n",
    "        return v_l / len(rollout),p_l / len(rollout),e_l / len(rollout), g_n,v_n\n",
    "        \n",
    "    def work(self,gamma,sess,coord,saver,train):\n",
    "        episode_count = sess.run(self.global_episodes)\n",
    "        total_steps = 0\n",
    "        print \"Starting worker \" + str(self.number)\n",
    "        with sess.as_default(), sess.graph.as_default():                 \n",
    "            while not coord.should_stop():\n",
    "                sess.run(self.update_local_ops)\n",
    "                episode_buffer = []\n",
    "                episode_values = []\n",
    "                episode_frames = []\n",
    "                episode_reward = 0\n",
    "                episode_step_count = 0\n",
    "                d = False\n",
    "                r = 0\n",
    "                a = 0\n",
    "                t = 0\n",
    "                s = self.env.reset()\n",
    "                rnn_state = self.local_AC.state_init\n",
    "                \n",
    "                while d == False:\n",
    "                    #Take an action using probabilities from policy network output.\n",
    "                    a_dist,v,rnn_state_new = sess.run([self.local_AC.policy,self.local_AC.value,self.local_AC.state_out], \n",
    "                        feed_dict={\n",
    "                        self.local_AC.state:[s],\n",
    "                        self.local_AC.prev_rewards:[[r]],\n",
    "                        self.local_AC.timestep:[[t]],\n",
    "                        self.local_AC.prev_actions:[a],\n",
    "                        self.local_AC.state_in[0]:rnn_state[0],\n",
    "                        self.local_AC.state_in[1]:rnn_state[1]})\n",
    "                    a = np.random.choice(a_dist[0],p=a_dist[0])\n",
    "                    a = np.argmax(a_dist == a)\n",
    "                    \n",
    "                    rnn_state = rnn_state_new\n",
    "                    s1,r,d,t = self.env.pullArm(a)                        \n",
    "                    episode_buffer.append([s,a,r,t,d,v[0,0]])\n",
    "                    episode_values.append(v[0,0])\n",
    "                    episode_frames.append(set_image_context(self.env.true,s,episode_reward,a,t))\n",
    "                    episode_reward += r\n",
    "                    total_steps += 1\n",
    "                    episode_step_count += 1\n",
    "                    s = s1\n",
    "                                            \n",
    "                self.episode_rewards.append(episode_reward)\n",
    "                self.episode_lengths.append(episode_step_count)\n",
    "                self.episode_mean_values.append(np.mean(episode_values))\n",
    "                \n",
    "                # Update the network using the experience buffer at the end of the episode.\n",
    "                if len(episode_buffer) != 0 and train == True:\n",
    "                    v_l,p_l,e_l,g_n,v_n = self.train(episode_buffer,sess,gamma,0.0)\n",
    "            \n",
    "                    \n",
    "                # Periodically save gifs of episodes, model parameters, and summary statistics.\n",
    "                if episode_count % 20 == 0 and episode_count != 0:\n",
    "                    if episode_count % 500 == 0 and self.name == 'worker_0' and train == True and len(self.episode_rewards) != 0:\n",
    "                        saver.save(sess,self.model_path+'/model-'+str(episode_count)+'.cptk')\n",
    "                        print \"Saved Model\"\n",
    "                        \n",
    "                    if episode_count % 40 == 0 and self.name == 'worker_0':\n",
    "                        self.images = np.array(episode_frames)\n",
    "                        make_gif(self.images,'./frames/image'+str(episode_count)+'.gif',\n",
    "                            duration=len(self.images)*0.1,true_image=True)\n",
    "\n",
    "\n",
    "                    mean_reward = np.mean(self.episode_rewards[-10:])\n",
    "                    mean_length = np.mean(self.episode_lengths[-10:])\n",
    "                    mean_value = np.mean(self.episode_mean_values[-10:])\n",
    "                    summary = tf.Summary()\n",
    "                    summary.value.add(tag='Perf/Reward', simple_value=float(mean_reward))\n",
    "                    summary.value.add(tag='Perf/Length', simple_value=float(mean_length))\n",
    "                    summary.value.add(tag='Perf/Value', simple_value=float(mean_value))\n",
    "                    if train == True:\n",
    "                        summary.value.add(tag='Losses/Value Loss', simple_value=float(v_l))\n",
    "                        summary.value.add(tag='Losses/Policy Loss', simple_value=float(p_l))\n",
    "                        summary.value.add(tag='Losses/Entropy', simple_value=float(e_l))\n",
    "                        summary.value.add(tag='Losses/Grad Norm', simple_value=float(g_n))\n",
    "                        summary.value.add(tag='Losses/Var Norm', simple_value=float(v_n))\n",
    "                    self.summary_writer.add_summary(summary, episode_count)\n",
    "\n",
    "                    self.summary_writer.flush()\n",
    "                if self.name == 'worker_0':\n",
    "                    sess.run(self.increment)\n",
    "                episode_count += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "gamma = .8 # discount rate for advantage estimation and reward discounting\n",
    "a_size = 2 \n",
    "load_model = True\n",
    "train = False\n",
    "model_path = './model_meta_context'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "\n",
    "if not os.path.exists(model_path):\n",
    "    os.makedirs(model_path)\n",
    "    \n",
    "if not os.path.exists('./frames'):\n",
    "    os.makedirs('./frames')\n",
    "    \n",
    "with tf.device(\"/cpu:0\"): \n",
    "    global_episodes = tf.Variable(0,dtype=tf.int32,name='global_episodes',trainable=False)\n",
    "    trainer = tf.train.AdamOptimizer(learning_rate=1e-3)\n",
    "    master_network = AC_Network(a_size,'global',None) # Generate global network\n",
    "    num_workers = multiprocessing.cpu_count() # Set workers ot number of available CPU threads\n",
    "    workers = []\n",
    "    # Create worker classes\n",
    "    for i in range(num_workers):\n",
    "        workers.append(Worker(contextual_bandit(),i,a_size,trainer,model_path,global_episodes))\n",
    "    saver = tf.train.Saver(max_to_keep=5)\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    coord = tf.train.Coordinator()\n",
    "    if load_model == True:\n",
    "        print 'Loading Model...'\n",
    "        ckpt = tf.train.get_checkpoint_state(model_path)\n",
    "        saver.restore(sess,ckpt.model_checkpoint_path)\n",
    "    else:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        \n",
    "    worker_threads = []\n",
    "    for worker in workers:\n",
    "        worker_work = lambda: worker.work(gamma,sess,coord,saver,train)\n",
    "        thread = threading.Thread(target=(worker_work))\n",
    "        thread.start()\n",
    "        worker_threads.append(thread)\n",
    "    coord.join(worker_threads)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}