DRLMAPF_A3C_RNN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Pathfinding via Reinforcement and Imitation Multi-Agent Learning (PRIMAL)\n",
    "\n",
    "While training is taking place, statistics on agent performance are available from Tensorboard. To launch it use:\n",
    "\n",
    "`tensorboard --logdir train_primal`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#this should be the thing, right?\n",
    "from __future__ import division\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "import random\n",
    "import tensorflow as tf\n",
    "import tensorflow.contrib.layers as layers\n",
    "import matplotlib.pyplot as plt\n",
    "from od_mstar3 import cpp_mstar\n",
    "from od_mstar3.col_set_addition import OutOfTimeError,NoSolutionError\n",
    "import threading\n",
    "import time\n",
    "import scipy.signal as signal\n",
    "import os\n",
    "import GroupLock\n",
    "import multiprocessing\n",
    "%matplotlib inline\n",
    "import mapf_gym as mapf_gym\n",
    "import pickle\n",
    "import imageio\n",
    "from ACNet import ACNet\n",
    "\n",
    "from tensorflow.python.client import device_lib\n",
    "dev_list = device_lib.list_local_devices()\n",
    "print(dev_list)\n",
    "assert len(dev_list) > 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_gif(images, fname, duration=2, true_image=False,salience=False,salIMGS=None):\n",
    "    imageio.mimwrite(fname,images,subrectangles=True)\n",
    "    print(\"wrote gif\")\n",
    "\n",
    "# Copies one set of variables to another.\n",
    "# Used to set worker network parameters to those of global network.\n",
    "def update_target_graph(from_scope,to_scope):\n",
    "    from_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, from_scope)\n",
    "    to_vars = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, to_scope)\n",
    "\n",
    "    op_holder = []\n",
    "    for from_var,to_var in zip(from_vars,to_vars):\n",
    "        op_holder.append(to_var.assign(from_var))\n",
    "    return op_holder\n",
    "\n",
    "def discount(x, gamma):\n",
    "    return signal.lfilter([1], [1, -gamma], x[::-1], axis=0)[::-1]\n",
    "\n",
    "def good_discount(x, gamma):\n",
    "    return discount(x,gamma)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Worker Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Worker:\n",
    "    def __init__(self, game, metaAgentID, workerID, a_size, groupLock):\n",
    "        self.workerID = workerID\n",
    "        self.env = game\n",
    "        self.metaAgentID = metaAgentID\n",
    "        self.name = \"worker_\"+str(workerID)\n",
    "        self.agentID = ((workerID-1) % num_workers) + 1 \n",
    "        self.groupLock = groupLock\n",
    "\n",
    "        self.nextGIF = episode_count # For GIFs output\n",
    "        #Create the local copy of the network and the tensorflow op to copy global parameters to local network\n",
    "        self.local_AC = ACNet(self.name,a_size,trainer,True,GRID_SIZE,GLOBAL_NET_SCOPE)\n",
    "        self.pull_global = update_target_graph(GLOBAL_NET_SCOPE, self.name)\n",
    "\n",
    "    def synchronize(self):\n",
    "        #handy thing for keeping track of which to release and acquire\n",
    "        if(not hasattr(self,\"lock_bool\")):\n",
    "            self.lock_bool=False\n",
    "        self.groupLock.release(int(self.lock_bool),self.name)\n",
    "        self.groupLock.acquire(int(not self.lock_bool),self.name)\n",
    "        self.lock_bool=not self.lock_bool\n",
    "        \n",
    "    def train(self, rollout, sess, gamma, bootstrap_value, rnn_state0, imitation=False):\n",
    "        global episode_count\n",
    "        if imitation:\n",
    "            rollout=np.array(rollout)\n",
    "            #we calculate the loss differently for imitation\n",
    "            #if imitation=True the rollout is assumed to have different dimensions:\n",
    "            #[o[0],o[1],optimal_actions]\n",
    "            feed_dict={global_step:episode_count,\n",
    "                       self.local_AC.inputs:np.stack(rollout[:,0]),\n",
    "                       self.local_AC.goal_pos:np.stack(rollout[:,1]),\n",
    "                       self.local_AC.optimal_actions:np.stack(rollout[:,2]),\n",
    "                       self.local_AC.state_in[0]:rnn_state0[0],\n",
    "                       self.local_AC.state_in[1]:rnn_state0[1]\n",
    "                      }\n",
    "            _,i_l,_=sess.run([self.local_AC.policy,self.local_AC.imitation_loss,\n",
    "                              self.local_AC.apply_imitation_grads],\n",
    "                             feed_dict=feed_dict)\n",
    "            return i_l\n",
    "        rollout = np.array(rollout)\n",
    "        observations = rollout[:,0]\n",
    "        goals=rollout[:,-2]\n",
    "        actions = rollout[:,1]\n",
    "        rewards = rollout[:,2]\n",
    "        values = rollout[:,5]\n",
    "        valids = rollout[:,6]\n",
    "        blockings = rollout[:,10]\n",
    "        on_goals=rollout[:,8]\n",
    "        train_value = rollout[:,-1]\n",
    "\n",
    "        # Here we take the rewards and values from the rollout, and use them to \n",
    "        # generate the advantage and discounted returns. (With bootstrapping)\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 = good_discount(advantages,gamma)\n",
    "\n",
    "        num_samples = min(EPISODE_SAMPLES,len(advantages))\n",
    "        sampleInd = np.sort(np.random.choice(advantages.shape[0], size=(num_samples,), replace=False))\n",
    "\n",
    "        # Update the global network using gradients from loss\n",
    "        # Generate network statistics to periodically save\n",
    "        feed_dict = {\n",
    "            global_step:episode_count,\n",
    "            self.local_AC.target_v:np.stack(discounted_rewards),\n",
    "            self.local_AC.inputs:np.stack(observations),\n",
    "            self.local_AC.goal_pos:np.stack(goals),\n",
    "            self.local_AC.actions:actions,\n",
    "            self.local_AC.train_valid:np.stack(valids),\n",
    "            self.local_AC.advantages:advantages,\n",
    "            self.local_AC.train_value:train_value,\n",
    "            self.local_AC.target_blockings:blockings,\n",
    "            self.local_AC.target_on_goals:on_goals,\n",
    "            self.local_AC.state_in[0]:rnn_state0[0],\n",
    "            self.local_AC.state_in[1]:rnn_state0[1]\n",
    "        }\n",
    "        \n",
    "        v_l,p_l,valid_l,e_l,g_n,v_n,b_l,og_l,_ = sess.run([self.local_AC.value_loss,\n",
    "            self.local_AC.policy_loss,\n",
    "            self.local_AC.valid_loss,\n",
    "            self.local_AC.entropy,\n",
    "            self.local_AC.grad_norms,\n",
    "            self.local_AC.var_norms,\n",
    "            self.local_AC.blocking_loss,\n",
    "            self.local_AC.on_goal_loss,\n",
    "            self.local_AC.apply_grads],\n",
    "            feed_dict=feed_dict)\n",
    "        return v_l/len(rollout), p_l/len(rollout), valid_l/len(rollout), e_l/len(rollout), b_l/len(rollout), og_l/len(rollout), g_n, v_n\n",
    "\n",
    "    def shouldRun(self, coord, episode_count):\n",
    "        if TRAINING:\n",
    "            return (not coord.should_stop())\n",
    "        else:\n",
    "            return (episode_count < NUM_EXPS)\n",
    "\n",
    "    def parse_path(self,path):\n",
    "        '''needed function to take the path generated from M* and create the \n",
    "        observations and actions for the agent\n",
    "        path: the exact path ouput by M*, assuming the correct number of agents\n",
    "        returns: the list of rollouts for the \"episode\": \n",
    "                list of length num_agents with each sublist a list of tuples \n",
    "                (observation[0],observation[1],optimal_action,reward)'''\n",
    "        result=[[] for i in range(num_workers)]\n",
    "        for t in range(len(path[:-1])):\n",
    "            observations=[]\n",
    "            move_queue=list(range(num_workers))\n",
    "            for agent in range(1,num_workers+1):\n",
    "                observations.append(self.env._observe(agent))\n",
    "            steps=0\n",
    "            while len(move_queue)>0:\n",
    "                steps+=1\n",
    "                i=move_queue.pop(0)\n",
    "                o=observations[i]\n",
    "                pos=path[t][i]\n",
    "                newPos=path[t+1][i]#guaranteed to be in bounds by loop guard\n",
    "                direction=(newPos[0]-pos[0],newPos[1]-pos[1])\n",
    "                a=self.env.world.getAction(direction)\n",
    "                state, reward, done, nextActions, on_goal, blocking, valid_action=self.env._step((i+1,a))\n",
    "                if steps>num_workers**2:\n",
    "                    #if we have a very confusing situation where lots of agents move\n",
    "                    #in a circle (difficult to parse and also (mostly) impossible to learn)\n",
    "                    return None\n",
    "                if not valid_action:\n",
    "                    #the tie must be broken here\n",
    "                    move_queue.append(i)\n",
    "                    continue\n",
    "                result[i].append([o[0],o[1],a])\n",
    "        return result\n",
    "        \n",
    "    def work(self,max_episode_length,gamma,sess,coord,saver):\n",
    "        global episode_count, swarm_reward, episode_rewards, episode_lengths, episode_mean_values, episode_invalid_ops,episode_wrong_blocking #, episode_invalid_goals\n",
    "        total_steps, i_buf = 0, 0\n",
    "        episode_buffers, s1Values = [ [] for _ in range(NUM_BUFFERS) ], [ [] for _ in range(NUM_BUFFERS) ]\n",
    "\n",
    "        with sess.as_default(), sess.graph.as_default():\n",
    "            while self.shouldRun(coord, episode_count):\n",
    "                sess.run(self.pull_global)\n",
    "\n",
    "                episode_buffer, episode_values = [], []\n",
    "                episode_reward = episode_step_count = episode_inv_count = 0\n",
    "                d = False\n",
    "\n",
    "                # Initial state from the environment\n",
    "                if self.agentID==1:\n",
    "                    self.env._reset(self.agentID)\n",
    "                self.synchronize() # synchronize starting time of the threads\n",
    "                validActions          = self.env._listNextValidActions(self.agentID)\n",
    "                s                     = self.env._observe(self.agentID)\n",
    "                blocking              = False\n",
    "                p=self.env.world.getPos(self.agentID)\n",
    "                on_goal               = self.env.world.goals[p[0],p[1]]==self.agentID\n",
    "                s                     = self.env._observe(self.agentID)\n",
    "                rnn_state             = self.local_AC.state_init\n",
    "                rnn_state0            = rnn_state\n",
    "                RewardNb = 0 \n",
    "                wrong_blocking  = 0\n",
    "                wrong_on_goal=0\n",
    "\n",
    "                if self.agentID==1:\n",
    "                    global demon_probs\n",
    "                    demon_probs[self.metaAgentID]=np.random.rand()\n",
    "                self.synchronize() # synchronize starting time of the threads\n",
    "\n",
    "                # reset swarm_reward (for tensorboard)\n",
    "                swarm_reward[self.metaAgentID] = 0\n",
    "                if episode_count<PRIMING_LENGTH or demon_probs[self.metaAgentID]<DEMONSTRATION_PROB:\n",
    "                    #for the first PRIMING_LENGTH episodes, or with a certain probability\n",
    "                    #don't train on the episode and instead observe a demonstration from M*\n",
    "                    if self.workerID==1 and episode_count%100==0:\n",
    "                        saver.save(sess, model_path+'/model-'+str(int(episode_count))+'.cptk')\n",
    "                    global rollouts\n",
    "                    rollouts[self.metaAgentID]=None\n",
    "                    if(self.agentID==1):\n",
    "                        world=self.env.getObstacleMap()\n",
    "                        start_positions=tuple(self.env.getPositions())\n",
    "                        goals=tuple(self.env.getGoals())\n",
    "                        try:\n",
    "                            mstar_path=cpp_mstar.find_path(world,start_positions,goals,2,5)\n",
    "                            rollouts[self.metaAgentID]=self.parse_path(mstar_path)\n",
    "                        except OutOfTimeError:\n",
    "                            #M* timed out \n",
    "                            print(\"timeout\",episode_count)\n",
    "                        except NoSolutionError:\n",
    "                            print(\"nosol????\",episode_count,start_positions)\n",
    "                    self.synchronize()\n",
    "                    if rollouts[self.metaAgentID] is not None:\n",
    "                        i_l=self.train(rollouts[self.metaAgentID][self.agentID-1], sess, gamma, None, rnn_state0, imitation=True)\n",
    "                        episode_count+=1./num_workers\n",
    "                        if self.agentID==1:\n",
    "                            summary = tf.Summary()\n",
    "                            summary.value.add(tag='Losses/Imitation loss', simple_value=i_l)\n",
    "                            global_summary.add_summary(summary, int(episode_count))\n",
    "                            global_summary.flush()\n",
    "                        continue\n",
    "                    continue\n",
    "                saveGIF = False\n",
    "                if OUTPUT_GIFS and self.workerID == 1 and ((not TRAINING) or (episode_count >= self.nextGIF)):\n",
    "                    saveGIF = True\n",
    "                    self.nextGIF =episode_count + 64\n",
    "                    GIF_episode = int(episode_count)\n",
    "                    episode_frames = [ self.env._render(mode='rgb_array',screen_height=900,screen_width=900) ]\n",
    "                    \n",
    "                while (not self.env.finished): # Give me something!\n",
    "                    #Take an action using probabilities from policy network output.\n",
    "                    a_dist,v,rnn_state,pred_blocking,pred_on_goal = sess.run([self.local_AC.policy,\n",
    "                                                   self.local_AC.value,\n",
    "                                                   self.local_AC.state_out,\n",
    "                                                   self.local_AC.blocking,\n",
    "                                                    self.local_AC.on_goal], \n",
    "                                         feed_dict={self.local_AC.inputs:[s[0]],\n",
    "                                                    self.local_AC.goal_pos:[s[1]],\n",
    "                                                    self.local_AC.state_in[0]:rnn_state[0],\n",
    "                                                    self.local_AC.state_in[1]:rnn_state[1]})\n",
    "\n",
    "                    if(not (np.argmax(a_dist.flatten()) in validActions)):\n",
    "                        episode_inv_count += 1\n",
    "                    train_valid = np.zeros(a_size)\n",
    "                    train_valid[validActions] = 1\n",
    "\n",
    "                    valid_dist = np.array([a_dist[0,validActions]])\n",
    "                    valid_dist /= np.sum(valid_dist)\n",
    "\n",
    "                    if TRAINING:\n",
    "                        if (pred_blocking.flatten()[0] < 0.5) == blocking:\n",
    "                            wrong_blocking += 1\n",
    "                        if (pred_on_goal.flatten()[0] < 0.5) == on_goal:\n",
    "                            wrong_on_goal += 1\n",
    "                        a           = validActions[ np.random.choice(range(valid_dist.shape[1]),p=valid_dist.ravel()) ]\n",
    "                        train_val   = 1.\n",
    "                    else:\n",
    "                        a         = np.argmax(a_dist.flatten())\n",
    "                        if a not in validActions or not GREEDY:\n",
    "                            a     = validActions[ np.random.choice(range(valid_dist.shape[1]),p=valid_dist.ravel()) ]\n",
    "                        train_val = 1.\n",
    "\n",
    "                    _, r, _, _, on_goal,blocking,_ = self.env._step((self.agentID, a),episode=episode_count)\n",
    "\n",
    "                    self.synchronize() # synchronize threads\n",
    "\n",
    "                    # Get common observation for all agents after all individual actions have been performed\n",
    "                    s1           = self.env._observe(self.agentID)\n",
    "                    validActions = self.env._listNextValidActions(self.agentID, a,episode=episode_count)\n",
    "                    d            = self.env.finished\n",
    "\n",
    "                    if saveGIF:\n",
    "                        episode_frames.append(self.env._render(mode='rgb_array',screen_width=900,screen_height=900))\n",
    "\n",
    "                    episode_buffer.append([s[0],a,r,s1,d,v[0,0],train_valid,pred_on_goal,int(on_goal),pred_blocking,int(blocking),s[1],train_val])\n",
    "                    episode_values.append(v[0,0])\n",
    "                    episode_reward += r\n",
    "                    s = s1\n",
    "                    total_steps += 1\n",
    "                    episode_step_count += 1\n",
    "\n",
    "                    if r>0:\n",
    "                        RewardNb += 1\n",
    "                    if d == True:\n",
    "                        print('\\n{} Goodbye World. We did it!'.format(episode_step_count), end='\\n')\n",
    "\n",
    "                    # If the episode hasn't ended, but the experience buffer is full, then we\n",
    "                    # make an update step using that experience rollout.\n",
    "                    if TRAINING and (len(episode_buffer) % EXPERIENCE_BUFFER_SIZE == 0 or d):\n",
    "                        # Since we don't know what the true final return is, we \"bootstrap\" from our current value estimation.\n",
    "                        if len(episode_buffer) >= EXPERIENCE_BUFFER_SIZE:\n",
    "                            episode_buffers[i_buf] = episode_buffer[-EXPERIENCE_BUFFER_SIZE:]\n",
    "                        else:\n",
    "                            episode_buffers[i_buf] = episode_buffer[:]\n",
    "\n",
    "                        if d:\n",
    "                            s1Values[i_buf] = 0\n",
    "                        else:\n",
    "                            s1Values[i_buf] = sess.run(self.local_AC.value, \n",
    "                                 feed_dict={self.local_AC.inputs:np.array([s[0]])\n",
    "                                            ,self.local_AC.goal_pos:[s[1]]\n",
    "                                            ,self.local_AC.state_in[0]:rnn_state[0]\n",
    "                                            ,self.local_AC.state_in[1]:rnn_state[1]})[0,0]\n",
    "\n",
    "                        if (episode_count-EPISODE_START) < NUM_BUFFERS:\n",
    "                            i_rand = np.random.randint(i_buf+1)\n",
    "                        else:\n",
    "                            i_rand = np.random.randint(NUM_BUFFERS)\n",
    "                            tmp = np.array(episode_buffers[i_rand])\n",
    "                            while tmp.shape[0] == 0:\n",
    "                                i_rand = np.random.randint(NUM_BUFFERS)\n",
    "                                tmp = np.array(episode_buffers[i_rand])\n",
    "                        v_l,p_l,valid_l,e_l,b_l,og_l,g_n,v_n = self.train(episode_buffers[i_rand],sess,gamma,s1Values[i_rand],rnn_state0)\n",
    "\n",
    "                        i_buf = (i_buf + 1) % NUM_BUFFERS\n",
    "                        rnn_state0             = rnn_state\n",
    "                        episode_buffers[i_buf] = []\n",
    "\n",
    "                    self.synchronize() # synchronize threads\n",
    "                    # sess.run(self.pull_global)\n",
    "                    if episode_step_count >= max_episode_length or d:\n",
    "                        break\n",
    "\n",
    "                episode_lengths[self.metaAgentID].append(episode_step_count)\n",
    "                episode_mean_values[self.metaAgentID].append(np.nanmean(episode_values))\n",
    "                episode_invalid_ops[self.metaAgentID].append(episode_inv_count)\n",
    "                episode_wrong_blocking[self.metaAgentID].append(wrong_blocking)\n",
    "\n",
    "                # Periodically save gifs of episodes, model parameters, and summary statistics.\n",
    "                if episode_count % EXPERIENCE_BUFFER_SIZE == 0 and printQ:\n",
    "                    print('                                                                                   ', end='\\r')\n",
    "                    print('{} Episode terminated ({},{})'.format(episode_count, self.agentID, RewardNb), end='\\r')\n",
    "\n",
    "                swarm_reward[self.metaAgentID] += episode_reward\n",
    "\n",
    "                self.synchronize() # synchronize threads\n",
    "\n",
    "                episode_rewards[self.metaAgentID].append(swarm_reward[self.metaAgentID])\n",
    "\n",
    "                if not TRAINING:\n",
    "                    mutex.acquire()\n",
    "                    if episode_count < NUM_EXPS:\n",
    "                        plan_durations[episode_count] = episode_step_count\n",
    "                    if self.workerID == 1:\n",
    "                        episode_count += 1\n",
    "                        print('({}) Thread {}: {} steps, {:.2f} reward ({} invalids).'.format(episode_count, self.workerID, episode_step_count, episode_reward, episode_inv_count))\n",
    "                    GIF_episode = int(episode_count)\n",
    "                    mutex.release()\n",
    "                else:\n",
    "                    episode_count+=1./num_workers\n",
    "\n",
    "                    if episode_count % SUMMARY_WINDOW == 0:\n",
    "                        if episode_count % 100 == 0:\n",
    "                            print ('Saving Model', end='\\n')\n",
    "                            saver.save(sess, model_path+'/model-'+str(int(episode_count))+'.cptk')\n",
    "                            print ('Saved Model', end='\\n')\n",
    "                        SL = SUMMARY_WINDOW * num_workers\n",
    "                        mean_reward = np.nanmean(episode_rewards[self.metaAgentID][-SL:])\n",
    "                        mean_length = np.nanmean(episode_lengths[self.metaAgentID][-SL:])\n",
    "                        mean_value = np.nanmean(episode_mean_values[self.metaAgentID][-SL:])\n",
    "                        mean_invalid = np.nanmean(episode_invalid_ops[self.metaAgentID][-SL:])\n",
    "                        mean_wrong_blocking = np.nanmean(episode_wrong_blocking[self.metaAgentID][-SL:])\n",
    "                        current_learning_rate = sess.run(lr,feed_dict={global_step:episode_count})\n",
    "\n",
    "                        summary = tf.Summary()\n",
    "                        summary.value.add(tag='Perf/Learning Rate',simple_value=current_learning_rate)\n",
    "                        summary.value.add(tag='Perf/Reward', simple_value=mean_reward)\n",
    "                        summary.value.add(tag='Perf/Length', simple_value=mean_length)\n",
    "                        summary.value.add(tag='Perf/Valid Rate', simple_value=(mean_length-mean_invalid)/mean_length)\n",
    "                        summary.value.add(tag='Perf/Blocking Prediction Accuracy', simple_value=(mean_length-mean_wrong_blocking)/mean_length)\n",
    "\n",
    "                        summary.value.add(tag='Losses/Value Loss', simple_value=v_l)\n",
    "                        summary.value.add(tag='Losses/Policy Loss', simple_value=p_l)\n",
    "                        summary.value.add(tag='Losses/Blocking Loss', simple_value=b_l)\n",
    "                        summary.value.add(tag='Losses/On Goal Loss', simple_value=og_l)\n",
    "                        summary.value.add(tag='Losses/Valid Loss', simple_value=valid_l)\n",
    "                        summary.value.add(tag='Losses/Grad Norm', simple_value=g_n)\n",
    "                        summary.value.add(tag='Losses/Var Norm', simple_value=v_n)\n",
    "                        global_summary.add_summary(summary, int(episode_count))\n",
    "\n",
    "                        global_summary.flush()\n",
    "\n",
    "                        if printQ:\n",
    "                            print('{} Tensorboard updated ({})'.format(episode_count, self.workerID), end='\\r')\n",
    "\n",
    "                if saveGIF:\n",
    "                    # Dump episode frames for external gif generation (otherwise, makes the jupyter kernel crash)\n",
    "                    time_per_step = 0.1\n",
    "                    images = np.array(episode_frames)\n",
    "                    if TRAINING:\n",
    "                        make_gif(images, '{}/episode_{:d}_{:d}_{:.1f}.gif'.format(gifs_path,GIF_episode,episode_step_count,swarm_reward[self.metaAgentID]))\n",
    "                    else:\n",
    "                        make_gif(images, '{}/episode_{:d}_{:d}.gif'.format(gifs_path,GIF_episode,episode_step_count), duration=len(images)*time_per_step,true_image=True,salience=False)\n",
    "                if SAVE_EPISODE_BUFFER:\n",
    "                    with open('gifs3D/episode_{}.dat'.format(GIF_episode), 'wb') as file:\n",
    "                        pickle.dump(episode_buffer, file)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Learning parameters\n",
    "max_episode_length     = 256\n",
    "episode_count          = 0\n",
    "EPISODE_START          = episode_count\n",
    "gamma                  = .95 # discount rate for advantage estimation and reward discounting\n",
    "#moved network parameters to ACNet.py\n",
    "EXPERIENCE_BUFFER_SIZE = 128\n",
    "GRID_SIZE              = 10 #the size of the FOV grid to apply to each agent\n",
    "ENVIRONMENT_SIZE       = (10,70)#the total size of the environment (length of one side)\n",
    "OBSTACLE_DENSITY       = (0,.5) #range of densities\n",
    "DIAG_MVMT              = False # Diagonal movements allowed?\n",
    "a_size                 = 5 + int(DIAG_MVMT)*4\n",
    "SUMMARY_WINDOW         = 10\n",
    "NUM_META_AGENTS        = 3\n",
    "NUM_THREADS            = 8 #int(multiprocessing.cpu_count() / (2 * NUM_META_AGENTS))\n",
    "NUM_BUFFERS            = 1 # NO EXPERIENCE REPLAY int(NUM_THREADS / 2)\n",
    "EPISODE_SAMPLES        = EXPERIENCE_BUFFER_SIZE # 64\n",
    "LR_Q                   = 2.e-5 #8.e-5 / NUM_THREADS # default: 1e-5\n",
    "ADAPT_LR               = True\n",
    "ADAPT_COEFF            = 5.e-5 #the coefficient A in LR_Q/sqrt(A*steps+1) for calculating LR\n",
    "load_model             = False\n",
    "RESET_TRAINER          = False\n",
    "model_path             = 'model_primal'\n",
    "gifs_path              = 'gifs_primal'\n",
    "train_path             = 'train_primal'\n",
    "GLOBAL_NET_SCOPE       = 'global'\n",
    "\n",
    "#Imitation options\n",
    "PRIMING_LENGTH         = 0    # number of episodes at the beginning to train only on demonstrations\n",
    "DEMONSTRATION_PROB     = 0.5  # probability of training on a demonstration per episode\n",
    "\n",
    "# Simulation options\n",
    "FULL_HELP              = False\n",
    "OUTPUT_GIFS            = False\n",
    "SAVE_EPISODE_BUFFER    = False\n",
    "\n",
    "# Testing\n",
    "TRAINING               = True\n",
    "GREEDY                 = False\n",
    "NUM_EXPS               = 100\n",
    "MODEL_NUMBER           = 313000\n",
    "\n",
    "# Shared arrays for tensorboard\n",
    "episode_rewards        = [ [] for _ in range(NUM_META_AGENTS) ]\n",
    "episode_lengths        = [ [] for _ in range(NUM_META_AGENTS) ]\n",
    "episode_mean_values    = [ [] for _ in range(NUM_META_AGENTS) ]\n",
    "episode_invalid_ops    = [ [] for _ in range(NUM_META_AGENTS) ]\n",
    "episode_wrong_blocking = [ [] for _ in range(NUM_META_AGENTS) ]\n",
    "rollouts               = [ None for _ in range(NUM_META_AGENTS)]\n",
    "demon_probs=[np.random.rand() for _ in range(NUM_META_AGENTS)]\n",
    "# episode_steps_on_goal  = [ [] for _ in range(NUM_META_AGENTS) ]\n",
    "printQ                 = False # (for headless)\n",
    "swarm_reward           = [0]*NUM_META_AGENTS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "tf.reset_default_graph()\n",
    "print(\"Hello World\")\n",
    "if not os.path.exists(model_path):\n",
    "    os.makedirs(model_path)\n",
    "config = tf.ConfigProto(allow_soft_placement = True)\n",
    "config.gpu_options.allow_growth=True\n",
    "\n",
    "if not TRAINING:\n",
    "    plan_durations = np.array([0 for _ in range(NUM_EXPS)])\n",
    "    mutex = threading.Lock()\n",
    "    gifs_path += '_tests'\n",
    "    if SAVE_EPISODE_BUFFER and not os.path.exists('gifs3D'):\n",
    "        os.makedirs('gifs3D')\n",
    "\n",
    "#Create a directory to save episode playback gifs to\n",
    "if not os.path.exists(gifs_path):\n",
    "    os.makedirs(gifs_path)\n",
    "\n",
    "with tf.device(\"/gpu:0\"):\n",
    "    master_network = ACNet(GLOBAL_NET_SCOPE,a_size,None,False,GRID_SIZE,GLOBAL_NET_SCOPE) # Generate global network\n",
    "\n",
    "    global_step = tf.placeholder(tf.float32)\n",
    "    if ADAPT_LR:\n",
    "        #computes LR_Q/sqrt(ADAPT_COEFF*steps+1)\n",
    "        #we need the +1 so that lr at step 0 is defined\n",
    "        lr=tf.divide(tf.constant(LR_Q),tf.sqrt(tf.add(1.,tf.multiply(tf.constant(ADAPT_COEFF),global_step))))\n",
    "    else:\n",
    "        lr=tf.constant(LR_Q)\n",
    "    trainer = tf.contrib.opt.NadamOptimizer(learning_rate=lr, use_locking=True)\n",
    "\n",
    "    if TRAINING:\n",
    "        num_workers = NUM_THREADS # Set workers # = # of available CPU threads\n",
    "    else:\n",
    "        num_workers = NUM_THREADS\n",
    "        NUM_META_AGENTS = 1\n",
    "    \n",
    "    gameEnvs, workers, groupLocks = [], [], []\n",
    "    n=1#counter of total number of agents (for naming)\n",
    "    for ma in range(NUM_META_AGENTS):\n",
    "        num_agents=NUM_THREADS\n",
    "        gameEnv = mapf_gym.MAPFEnv(num_agents=num_agents, DIAGONAL_MOVEMENT=DIAG_MVMT, SIZE=ENVIRONMENT_SIZE, \n",
    "                                   observation_size=GRID_SIZE,PROB=OBSTACLE_DENSITY, FULL_HELP=FULL_HELP)\n",
    "        gameEnvs.append(gameEnv)\n",
    "\n",
    "        # Create groupLock\n",
    "        workerNames = [\"worker_\"+str(i) for i in range(n,n+num_workers)]\n",
    "        groupLock = GroupLock.GroupLock([workerNames,workerNames])\n",
    "        groupLocks.append(groupLock)\n",
    "\n",
    "        # Create worker classes\n",
    "        workersTmp = []\n",
    "        for i in range(ma*num_workers+1,(ma+1)*num_workers+1):\n",
    "            workersTmp.append(Worker(gameEnv,ma,n,a_size,groupLock))\n",
    "            n+=1\n",
    "        workers.append(workersTmp)\n",
    "\n",
    "    global_summary = tf.summary.FileWriter(train_path)\n",
    "    saver = tf.train.Saver(max_to_keep=2)\n",
    "\n",
    "    with tf.Session(config=config) as sess:\n",
    "        sess.run(tf.global_variables_initializer())\n",
    "        coord = tf.train.Coordinator()\n",
    "        if load_model == True:\n",
    "            print ('Loading Model...')\n",
    "            if not TRAINING:\n",
    "                with open(model_path+'/checkpoint', 'w') as file:\n",
    "                    file.write('model_checkpoint_path: \"model-{}.cptk\"'.format(MODEL_NUMBER))\n",
    "                    file.close()\n",
    "            ckpt = tf.train.get_checkpoint_state(model_path)\n",
    "            p=ckpt.model_checkpoint_path\n",
    "            p=p[p.find('-')+1:]\n",
    "            p=p[:p.find('.')]\n",
    "            episode_count=int(p)\n",
    "            saver.restore(sess,ckpt.model_checkpoint_path)\n",
    "            print(\"episode_count set to \",episode_count)\n",
    "            if RESET_TRAINER:\n",
    "                trainer = tf.contrib.opt.NadamOptimizer(learning_rate=lr, use_locking=True)\n",
    "\n",
    "        # This is where the asynchronous magic happens.\n",
    "        # Start the \"work\" process for each worker in a separate thread.\n",
    "        worker_threads = []\n",
    "        for ma in range(NUM_META_AGENTS):\n",
    "            for worker in workers[ma]:\n",
    "                groupLocks[ma].acquire(0,worker.name) # synchronize starting time of the threads\n",
    "                worker_work = lambda: worker.work(max_episode_length,gamma,sess,coord,saver)\n",
    "                print(\"Starting worker \" + str(worker.workerID))\n",
    "                t = threading.Thread(target=(worker_work))\n",
    "                t.start()\n",
    "                worker_threads.append(t)\n",
    "        coord.join(worker_threads)\n",
    "\n",
    "if not TRAINING:\n",
    "    print([np.mean(plan_durations), np.sqrt(np.var(plan_durations)), np.mean(np.asarray(plan_durations < max_episode_length, dtype=float))])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}