arxiv.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Two-step-task as described in \"Learning to Reinforcement Learn\", Wang. et al., 2016\n",
    "\n",
    "This iPython notebook includes an implementation of the two-step task described here: [arxiv pdf](https://arxiv.org/pdf/1611.05763.pdf).\n",
    "\n",
    "The episodes are 20 steps (10 trials) long. \n",
    "\n",
    "### First training to get an idea of the number of episodes (not given by the paper)\n",
    "\n",
    "We first did a training with an open-ended number of episode to see after how many episodes it seemed to converge (not mentionned in their arxiv pre-print):\n",
    "\n",
    "![reward curves for 137k episode](results/arxiv/arxiv_137k/train/reward_curve.png)\n",
    "\n",
    "**Result**: __The convergence seems to be after ~40k episodes__\n",
    "\n",
    "We then launched another training of 40k episodes.\n",
    "\n",
    "![reward curves for 40k](results/arxiv/arxiv_40k/train/reward_curve.png)\n",
    "\n",
    "Here is how the plot of the probability of repeating the same action (in function of the transition and previous reward) evolved during those 40 episodes of training:\n",
    "\n",
    "![gif of evolution of training](results/arxiv/arxiv_40k/train/training_40k_gif.gif)\n",
    "\n",
    "And finally, a side by side comparaison with the results from the arxiv pre-print (Figure 5)c)):\n",
    "\n",
    "![side by side comparison](results/arxiv/arxiv_40k/test/side_by_side.png)\n",
    "\n",
    "We can see that there is a clear difference in our results. Both of the plots have the same shape, but on the results from the paper their common-rewarded and uncommon-not-rewarded stay probabilities are clearly closer to zero, whereas our probabilities are closer to 0.5.\n",
    "\n",
    "The hyper-parameters were the one mentionned in \"Learning to Reinforcement Learn\", Table 1, Exp. 5.\n",
    "\n",
    "**Reasons for not having the same results**: in the paper they did not give the exact number of episodes in training and test, so we had to guess. Also, for the learning rate and the discount factor they put \"tuned\", so we also had to guess those ones, using for instance the parameters from their biorxiv pre-print.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "outputs": [],
   "source": [
    "# Saving path name\n",
    "from datetime import datetime\n",
    "\n",
    "dir_name = \"train_\" + datetime.now().strftime(\"%m%d-%H%M%S\")\n",
    "model_path = dir_name+'/model_meta_context'\n",
    "frame_path = dir_name+'/frames'\n",
    "plot_path = dir_name+'/plots'\n",
    "load_model_path = \"results/arxiv/arxiv_40k/train/model_meta_context/model-40000\"\n",
    "\n",
    "# encoding of the higher stages\n",
    "S_1 = 0\n",
    "S_2 = 1\n",
    "S_3 = 2\n",
    "nb_states = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class two_step_task():\n",
    "    def __init__(self):\n",
    "        # start in S_1\n",
    "        self.state = S_1\n",
    "        \n",
    "        # defines what is the stage with the highest expected reward. Initially random\n",
    "        self.highest_reward_second_stage = np.random.choice([S_2,S_3])\n",
    "        \n",
    "        self.num_actions = 2\n",
    "        self.reset()\n",
    "        \n",
    "        # initialization of plotting variables\n",
    "        common_prob = 0.75\n",
    "        self.transitions = np.array([\n",
    "            [common_prob, 1-common_prob],\n",
    "            [1-common_prob, common_prob]\n",
    "        ])\n",
    "        self.transition_count = np.zeros((2,2,2))\n",
    "        \n",
    "        self.last_action = None\n",
    "        self.last_state = None\n",
    "    \n",
    "    def get_state(self):\n",
    "        one_hot_array = np.zeros(nb_states)\n",
    "        one_hot_array[self.state] = 1\n",
    "        return one_hot_array\n",
    "\n",
    "    def possible_switch(self):\n",
    "        if (np.random.uniform() < 0.5):\n",
    "            # switches which of S_2 or S_3 has expected reward of 0.9\n",
    "            self.highest_reward_second_stage = S_2 if (self.highest_reward_second_stage == S_3) else S_3\n",
    "            \n",
    "    def get_rprobs(self):\n",
    "        \"\"\"\n",
    "        probability of reward of states S_2 and S_3, in the form [[p, 1-p], [1-p, p]]\n",
    "        \"\"\"\n",
    "        if (self.highest_reward_second_stage == S_2):\n",
    "            r_prob = 0.9\n",
    "        else:\n",
    "            r_prob = 0.1\n",
    "        \n",
    "        rewards = np.array([\n",
    "            [r_prob, 1-r_prob],\n",
    "            [1-r_prob, r_prob]\n",
    "        ])\n",
    "        return rewards\n",
    "\n",
    "            \n",
    "    def isCommon(self,action,state):\n",
    "        if self.transitions[action][state] >= 1/2:\n",
    "            return True\n",
    "        return False\n",
    "        \n",
    "    def updateStateProb(self,action):\n",
    "        if self.last_is_rewarded: #R\n",
    "            if self.last_is_common: #C\n",
    "                if self.last_action == action: #Rep\n",
    "                    self.transition_count[0,0,0] += 1\n",
    "                else: #URep\n",
    "                    self.transition_count[0,0,1] += 1\n",
    "            else: #UC\n",
    "                if self.last_action == action: #Rep\n",
    "                    self.transition_count[0,1,0] += 1\n",
    "                else: #URep\n",
    "                    self.transition_count[0,1,1] += 1\n",
    "        else: #UR\n",
    "            if self.last_is_common:\n",
    "                if self.last_action == action:\n",
    "                    self.transition_count[1,0,0] += 1\n",
    "                else:\n",
    "                    self.transition_count[1,0,1] += 1\n",
    "            else:\n",
    "                if self.last_action == action:\n",
    "                    self.transition_count[1,1,0] += 1\n",
    "                else:\n",
    "                    self.transition_count[1,1,1] += 1\n",
    "                    \n",
    "        \n",
    "    def stayProb(self):\n",
    "        print(self.transition_count)\n",
    "        row_sums = self.transition_count.sum(axis=-1)\n",
    "        stay_prob = self.transition_count / row_sums[:,:,np.newaxis] \n",
    "       \n",
    "        return stay_prob\n",
    "\n",
    "    def reset(self):\n",
    "        self.timestep = 0\n",
    "        \n",
    "        # for the two-step task plots\n",
    "        self.last_is_common = None\n",
    "        self.last_is_rewarded = None\n",
    "        self.last_action = None\n",
    "        self.last_state = None\n",
    "        \n",
    "        # come back to S_1 at the end of an episode\n",
    "        self.state = S_1\n",
    "        \n",
    "        return self.get_state()\n",
    "        \n",
    "    def step(self,action):\n",
    "        self.timestep += 1\n",
    "        self.last_state = self.state\n",
    "        \n",
    "        # get next stage\n",
    "        if (self.state == S_1):\n",
    "            # get reward\n",
    "            reward = 0\n",
    "            # update stage\n",
    "            self.state = S_2 if (np.random.uniform() < self.transitions[action][0]) else S_3\n",
    "            # keep track of stay probability after first action\n",
    "            if (self.last_action != None):    \n",
    "                self.updateStateProb(action)\n",
    "            self.last_action = action\n",
    "            # book-keeping for plotting\n",
    "            self.last_is_common = self.isCommon(action,self.state-1)\n",
    "            \n",
    "        else:# case S_2 or S_3\n",
    "            # get probability of reward in stage\n",
    "            r_prob = 0.9 if (self.highest_reward_second_stage == self.state) else 0.1\n",
    "            # get reward\n",
    "            reward = 1 if np.random.uniform() < r_prob else 0\n",
    "            # update stage\n",
    "            self.state = S_1\n",
    "            # book-keeping for plotting\n",
    "            self.last_is_rewarded = reward\n",
    "\n",
    "        # new state after the decision\n",
    "        new_state = self.get_state()\n",
    "        if self.timestep >= 20: \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": {},
   "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,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.nn.rnn_cell.LSTMCell(48,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, 48])\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.05 * 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": {},
   "outputs": [],
   "source": [
    "\n",
    "class Worker():\n",
    "    def __init__(self,game,name,a_size,trainer,model_path,global_episodes,make_gif=False):\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(dir_name)\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",
    "        self.make_gif = make_gif\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",
    "        \n",
    "        episode_count = sess.run(self.global_episodes)\n",
    "        # set count to zero when loading a model\n",
    "        if not train:\n",
    "            episode_count = 0\n",
    "            \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() and episode_count <= num_episodes:\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",
    "                if (self.env.state == S_1):\n",
    "                    self.env.possible_switch()\n",
    "                \n",
    "                while d == False:\n",
    "                    #possible switch of S_2 <-> S_3 with probability 2.5% at the beginning of a trial (every two steps)\n",
    "                    #if (self.env.state == S_1):\n",
    "                    #    self.env.possible_switch()\n",
    "                    \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.step(a)                        \n",
    "                    episode_buffer.append([s,a,r,t,d,v[0,0]])\n",
    "                    episode_values.append(v[0,0])\n",
    "                    \n",
    "                    if episode_count % 1000 == 0 and self.name == 'worker_0':\n",
    "                        if self.make_gif and self.env.last_state == S_2 or self.env.last_state == S_3:\n",
    "                            episode_frames.append(make_frame(frame_path,self.env.transitions,\n",
    "                                                                    self.env.get_rprobs(), \n",
    "                                                                    t, action=self.env.last_action, \n",
    "                                                                    final_state=self.env.last_state, \n",
    "                                                                    reward=r))\n",
    "                \n",
    "                            \n",
    "                    \n",
    "                    \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 % 50 == 0 and episode_count != 0:\n",
    "                    if episode_count % 1000 == 0 and self.name == 'worker_0':\n",
    "                        if train == True:\n",
    "                            # save model\n",
    "                            os.makedirs(model_path+'/model-'+str(episode_count))\n",
    "                            saver.save(sess,model_path+'/model-'+str(episode_count)+\n",
    "                                       '/model-'+str(episode_count)+'.cptk')\n",
    "                            print (\"Saved Model\")\n",
    "                        \n",
    "                            if episode_count % 1000 == 0:\n",
    "                                # generate plot\n",
    "                                self.plot(episode_count)\n",
    "                                print (\"Saved Plot\")\n",
    "\n",
    "                        if self.make_gif:\n",
    "                            # generate gif\n",
    "                            make_gif(episode_frames,frame_path+\"/test_\"+str(episode_count)+'.gif')    \n",
    "                            print (\"Saved Gif\")\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\n",
    "                if (episode_count % 10 == 0):\n",
    "                    print(\"episode_count is: \", episode_count)\n",
    "        \n",
    "        if not train:\n",
    "            self.plot(episode_count-1)\n",
    "\n",
    "    def plot(self, episode_count):\n",
    "        fig, ax = plt.subplots()\n",
    "        x = np.arange(2)\n",
    "        ax.set_ylim([0, 1.2])\n",
    "        ax.set_ylabel('Stay Probability')\n",
    "        \n",
    "        stay_probs = self.env.stayProb()\n",
    "        \n",
    "        common = [stay_probs[0,0,0],stay_probs[1,0,0]]\n",
    "        uncommon = [stay_probs[0,1,0],stay_probs[1,1,0]]\n",
    "        \n",
    "        ax.set_xticks([1.3,3.3])\n",
    "        ax.set_xticklabels(['Last trial rewarded', 'Last trial not rewarded'])\n",
    "        \n",
    "        c = plt.bar([1,3],  common, color='b', width=0.5)\n",
    "        uc = plt.bar([1.8,3.8], uncommon, color='r', width=0.5)\n",
    "        ax.legend( (c[0], uc[0]), ('common', 'uncommon') )\n",
    "        plt.savefig(plot_path +\"/\"+ str(episode_count) + \".png\")\n",
    "        self.env.transition_count = np.zeros((2,2,2))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma = .9 # discount rate for advantage estimation and reward discounting\n",
    "a_size = 2 \n",
    "load_model = True\n",
    "train = False\n",
    "num_episodes = 3000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "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(frame_path):\n",
    "    os.makedirs(frame_path)\n",
    "    \n",
    "if not os.path.exists(plot_path):\n",
    "    os.makedirs(plot_path)\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.RMSPropOptimizer(learning_rate=7e-4)\n",
    "    master_network = AC_Network(a_size,'global',None) # Generate global network\n",
    "    num_workers = 1\n",
    "    workers = []\n",
    "    # Create worker classes\n",
    "    for i in range(num_workers):\n",
    "        workers.append(Worker(two_step_task(),i,a_size,trainer,model_path,global_episodes, make_gif=True))\n",
    "    saver = tf.train.Saver(max_to_keep=5)\n",
    "\n",
    "with tf.Session() as sess:\n",
    "    tf.set_random_seed(1)\n",
    "    coord = tf.train.Coordinator()\n",
    "    if load_model == True:\n",
    "        print ('Loading Model...')\n",
    "        ckpt = tf.train.get_checkpoint_state(load_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 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}