Counter_Factual_Multi_Agent_Policy_Gradient_Implementation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Counter Factual Multi-Agent Policy Gradient Implementation.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GA5-Us75vTI4",
        "colab_type": "text"
      },
      "source": [
        "# COMA Implementation\n",
        "\n",
        "This is my implementation of the Counter Factual Multi-Agent Policy Gradients system - COMA. This system was completed entirely by myself with my own research."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bBcOf9ENzRSt",
        "colab_type": "text"
      },
      "source": [
        "The Combat environment is used in this implementation - the wiki is given below and linked [here](https://github.com/koulanurag/ma-gym/wiki/). However, as is shown in the README - I found that the wiki is somewhat incorrect about the reward system."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eAePzQjrcriF",
        "colab_type": "text"
      },
      "source": [
        "### Environment Description\n",
        "![Combat](https://github.com/koulanurag/ma-gym/raw/master/static/gif/Combat-v0.gif)\n",
        "\n",
        "\n",
        "Combat environment simulates a battle involving two opposing teams in a `20×20` grid. Each team consists of `m = 10` agents and their initial positions are sampled uniformly in a `5×5` square around the team center, which is picked uniformly in the grid.\n",
        "\n",
        "#### Action Space\n",
        "At each time step, an agent can perform one of the following actions:\n",
        "*   move one cell in one of four directions;\n",
        "*   attack another agent by specifying its ID `j` (there are `m` attack actions, each corresponding to one enemy agent);\n",
        "*   do nothing.\n",
        "\n",
        "#### Transition Dynamics\n",
        "If agent A attacks agent B, then B’s health point will be reduced by 1, but only if B is inside the firing range of A (its surrounding `3×3` area). Agents need one-time step of cooling down after an attack, during which they cannot attack. All agents start with three health points and die when their health reaches 0. A team will win if all agents in the other team die. The simulation ends when one team wins, or neither of the teams win within 40 time steps (a draw).\n",
        "\n",
        "#### Observation Space\n",
        "When the input to a model, each agent is represented by a set of one-hot binary vectors `{i, t, l, h, c}` encoding its unique ID, team ID, location, health points, and cooldown. A model controlling an agent also sees other agents in its visual range (`3×3` surrounding area).\n",
        "\n",
        "#### Reward Settings\n",
        "The model gets a reward of -1 if the team loses or draws at the end of the game. In addition, it also get a reward of −0.1 times the total health points of the enemy team, which encourages it to attack enemy bots.\n",
        "\n",
        "#### Enemy Settings\n",
        "Your model controls one team during training, and the other team consists of bots that follow a hardcoded policy. The bot policy is to attack the nearest enemy agent if it is within its firing range. If not, it approaches the nearest visible enemy agent within the visual range. An agent is visible to all bots if it is inside the visual range of any individual bot. This shared vision gives an advantage to the bot team."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2bRNAhv1QMC1",
        "colab_type": "text"
      },
      "source": [
        "### Download Requirements and Set the Environment\n",
        "The following command will download the required scripts and set up the environment. \n",
        "\n",
        "Execute the following block until you see a WARNING and a button asking to **RESTART RUNTIME**. Click the button and wait for the runtime to restart, **then proceed to the next block of code.**\n",
        "\n",
        "Note: The below installs are written by the TAs from UCL's Multi-Agent Artificial Intelligence module, however, the actual implementation is written by me - this is indicated below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "N0MTkFoNkLi2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!rm -rf /content/ma-gym  \n",
        "!git clone https://github.com/koulanurag/ma-gym.git\n",
        "%cd /content/ma-gym \n",
        "!pip install -q -e ."
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58wskOfR5xv-",
        "colab_type": "text"
      },
      "source": [
        "**Run the following block of code to continue environment set-up AFTER you have restarted the RUNTIME from the previous block.**"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IZx9b4dS35AI",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "!apt-get install -y xvfb python-opengl x11-utils > /dev/null 2>&1\n",
        "!pip install pyvirtualdisplay > /dev/null 2>&1\n",
        "!apt-get install x11-utils\n",
        "!apt-get update > /dev/null 2>&1\n",
        "!apt-get install cmake > /dev/null 2>&1\n",
        "!pip install --upgrade setuptools 2>&1\n",
        "!pip install ez_setup > /dev/null 2>&1\n",
        "!pip install -U gym[atari] > /dev/null 2>&1"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3sYf6AK2d6kK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "%matplotlib inline\n",
        "import gym\n",
        "import ma_gym\n",
        "from ma_gym.wrappers import Monitor\n",
        "from ma_gym.envs.combat.combat import Combat\n",
        "import matplotlib.pyplot as plt\n",
        "import glob\n",
        "import io\n",
        "import base64\n",
        "from IPython.display import HTML\n",
        "from IPython import display as ipythondisplay\n",
        "\n",
        "from pyvirtualdisplay import Display\n",
        "display = Display(visible=0, size=(1400, 900))\n",
        "display.start()\n",
        "\n",
        "\"\"\"\n",
        "Utility functions to enable video recording of gym environment and displaying it\n",
        "To enable video, just do \"env = wrap_env(env)\"\"\n",
        "\"\"\"\n",
        "\n",
        "def show_video():\n",
        "  mp4list = glob.glob('video/*.mp4')\n",
        "  if len(mp4list) > 0:\n",
        "    mp4 = mp4list[0]\n",
        "    video = io.open(mp4, 'r+b').read()\n",
        "    encoded = base64.b64encode(video)\n",
        "    ipythondisplay.display(HTML(data='''<video alt=\"test\" autoplay \n",
        "                loop controls style=\"height: 400px;\">\n",
        "                <source src=\"data:video/mp4;base64,{0}\" type=\"video/mp4\" />\n",
        "             </video>'''.format(encoded.decode('ascii'))))\n",
        "  else: \n",
        "    print(\"Could not find video\")\n",
        "    \n",
        "\n",
        "def wrap_env(env):\n",
        "  env = Monitor(env, './video', force=True)\n",
        "  return env\n",
        "\n",
        "\n",
        "from functools import reduce\n",
        "from operator import add"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jyhtZjqwJtCu",
        "colab_type": "text"
      },
      "source": [
        "# COMA Implementation\n",
        "All code written from here is exclusively by myself and is the entire COMA implementation.\n",
        "\n",
        "COMA - [Counterfactual Multi-Agent Policy Gradients (COMA)](https://www.aaai.org/ocs/index.php/AAAI/AAAI18/paper/view/17193/16614)\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fxRIenrMvTJ1",
        "colab_type": "text"
      },
      "source": [
        "Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Krleu7MZvShr",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "#The below is to stick the TensorFlow version, since TF version2\n",
        "#does not have some required funtions (eg .get_session())\n",
        "%tensorflow_version 1.x\n",
        "import tensorflow.compat.v1 as tf\n",
        "tf.disable_v2_behavior()\n",
        "\n",
        "\n",
        "import keras\n",
        "from keras.models import Sequential\n",
        "from keras.layers import Dense\n",
        "from keras.models import Sequential\n",
        "from keras.layers import Dense, Dropout, Conv2D, MaxPooling2D, Activation, Flatten\n",
        "from keras.callbacks import TensorBoard\n",
        "from keras.optimizers import Adam\n",
        "from keras import optimizers\n",
        "from keras import backend\n",
        "from collections import deque\n",
        "from keras.layers import Activation, Dense\n",
        "from collections import deque\n",
        "from copy import deepcopy\n",
        "import random\n",
        "import numpy as np\n",
        "import sys\n",
        "import time\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-aMlruU0aWHc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Since this method uses the \"global state\" and has a lot of variables,\n",
        "#it is best to create a new class for the buffer.\n",
        "\n",
        "class Replay_Buffer():\n",
        "  def __init__(self, buffer_size):\n",
        "    self.max_length = buffer_size\n",
        "    # self.data = deque(maxlen = buffer_size)\n",
        "\n",
        "    \"\"\"\n",
        "    Originally, a deque was used as the buffer container.\n",
        "    However,l the rest of the coe has been created to not use random sample\n",
        "    (as this would be exceedingly difficult to process due to the\n",
        "    new COMA critic setup) and deques have O(n) complexity\n",
        "    when indexing middle values - a regular buffer size\n",
        "    is 100,000 values and thus this hugely slows the system down.\n",
        "    Therefore, a dictionary (which has O(1) complexity for indexing \n",
        "    middle values) has been selected, and a custom addition to the \n",
        "    add_buffer function has been included to implement a max length\n",
        "    attribute.\n",
        "\n",
        "    Key: index\n",
        "    Value: List of Data from that step\n",
        "    \n",
        "     \"\"\"\n",
        "    self.data = {}\n",
        "\n",
        "    self.tracker = 0 #this tracks additions to the dictionary\n",
        "\n",
        "  def Add_to_Buffer(self, prev_obs, prev_act, obs, curr_act, rewards, done):\n",
        "    \"\"\" Since the overall state, as well as the observation,\n",
        "    of this system is so large, it is best to organise a reasonable\n",
        "    addition to the buffer in this separate function\n",
        "    \n",
        "    The contents of the buffer should be as follows:\n",
        "    - Previous State\n",
        "    - Previous Actions\n",
        "    - Current State\n",
        "    - Current actions\n",
        "    - Group reward (for that step)\n",
        "    - Done \n",
        "\n",
        "    Timestep not included\n",
        "    \"\"\"\n",
        "\n",
        "    addition_list = [] #the list which will be added to the buffer\n",
        "\n",
        "    #add the previous state\n",
        "    addition_list.append(prev_obs)\n",
        "\n",
        "    #add the previous actions taken\n",
        "    addition_list.append(prev_act)\n",
        "\n",
        "    #add the current state\n",
        "    addition_list.append(obs)\n",
        "\n",
        "    #add the current actions\n",
        "    addition_list.append(curr_act)\n",
        "\n",
        "    #add the whole system reward\n",
        "    addition_list.append(sum(rewards)/10)\n",
        "    #Taking the average reward, instead of the summative, which is far higher\n",
        "\n",
        "    #add done\n",
        "    addition_list.append(done)\n",
        "\n",
        "    #append to the deque object\n",
        "    # self.data.append(addition_list)\n",
        "\n",
        "\n",
        "\n",
        "    #now, if the tracker is potentially going to go over the maximum\n",
        "    #size for the buffer, it should be reset to zero\n",
        "\n",
        "    if self.tracker >= self.max_length: \n",
        "      self.tracker = 0\n",
        "\n",
        "    #add as a new key-value pair to the dictionary\n",
        "    self.data[self.tracker] = addition_list\n",
        "\n",
        "    self.tracker += 1\n",
        "\n",
        "\n",
        "  def Return_all_current_actions(self, index):\n",
        "    #access and return the current actions of the specified index\n",
        "    return self.data[index][3]\n",
        "\n",
        "  def Return_all_current_actions_excl_agent(self,index,agent):\n",
        "    return [x for i, x in enumerate(self.data[index][3]) if i != agent-1] #-1 for indexing\n",
        "\n",
        "  def Return_current_state(self,index):\n",
        "    return self.data[index][2]\n",
        "\n",
        "  def Return_current_agent_state(self,index,agent):\n",
        "    return self.data[index][2][agent-1]\n",
        "\n",
        "  def Return_all_previous_actions(self,index):\n",
        "    return self.data[index][1]\n",
        "  \n",
        "  def Return_reward(self,index):\n",
        "    return self.data[index][4]\n",
        "\n",
        "  def Return_done(self,index):\n",
        "    return self.data[index][5]\n",
        "\n",
        "  def Return_previous_agent_action(self,index,agent):\n",
        "    return self.data[index][1][agent-1]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vS5uQ0NLun9B",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Coma_System:\n",
        "\n",
        "  def __init__(self):\n",
        "    self.env = wrap_env(Combat(grid_shape=(20, 20), n_agents=10, n_opponents=10))\n",
        "    self.sess = tf.Session() #this instantiates a TensorFlow session for us to use\n",
        "    keras.backend.set_session(self.sess)\n",
        "    print(\"action meanings\",self.env.get_action_meanings())\n",
        "    print(\"self.env.n_agents\",self.env.n_agents)\n",
        "    #each agent needs its own action-observation history, which conditions its own policy.\n",
        "    self.input_shape = self.env.n_agents*150\n",
        "    self.trace_decay = 0.8 #lambda in the paper, for TD error - the credit assignment variable.\n",
        "    self.discount_factor = 0.99 #gamma\n",
        "    self.n_of_actions = 15\n",
        "    \n",
        "    self.max_buffer_size = 100000\n",
        "    \n",
        "    self.buffer = Replay_Buffer(self.max_buffer_size)\n",
        "    self.minibatch_size = 32\n",
        "\n",
        "    #instantiating networks\n",
        "    self.actor_learning_rate = 1e-6\n",
        "    self.critic_learning_rate = 1e-5\n",
        "    #critic should have higher LR than actor - it is teaching\n",
        "    self.critic_model = self._Create_Critic(1670,15) #input shape and action space\n",
        "    self.actors = self._Instantiate_All_Actors(10,152,15) #dictionary of actors\n",
        "    #1670 and 152 above are the networks' input shapes, derived heuristically.\n",
        "\n",
        "    self.target_network = keras.models.clone_model(self.critic_model)\n",
        "    self.target_network_update_steps = 100\n",
        "    self.actor_critic_update_frequency = 32\n",
        "\n",
        "    self.training_episodes = 2000 \n",
        "    self.exploration_episodes = 200\n",
        "    self.epsilon_min = 0.1\n",
        "    self.epsilon_max = 0.99\n",
        "    self.epsilon_projection = -0.2 #what the linear interpolation \"aims\" for.\n",
        "    self.testing_episodes = 100\n",
        "\n",
        "    \n",
        "\n",
        "    self.episodes_trained = []\n",
        "    self.critic_loss = []\n",
        "    \n",
        "    self.sess.run(tf.global_variables_initializer())\n",
        "\n",
        "  \n",
        "  def _Create_Critic(self, input_dimensions,n_of_actions):\n",
        "\n",
        "    #INPUT to this system defined in COMA paper\n",
        "\n",
        "    #OUTPUT to this system must be the (15) Q values for all possible \n",
        "    #actions which the agent we're considering could take.\n",
        "\n",
        "    \"\"\" From Paper: use a centralised critic that conditions on the true global states,\n",
        "    if it is available,or the joint action-observation histories τ otherwise. \"\"\"\n",
        "\n",
        "    keras.backend.set_session(self.sess)\n",
        "    model = Sequential() #neuron numbers gained from linked GitHub (oxwhirl)\n",
        "    model.add(Dense(128, input_dim=input_dimensions, activation=\"relu\"))\n",
        "    model.add(Dense(128,activation = \"relu\"))\n",
        "    model.add(Dense(n_of_actions,activation = \"linear\"))\n",
        "    adam_opt = keras.optimizers.Adam(learning_rate = self.critic_learning_rate,clipnorm=1.)\n",
        "    model.compile(loss='mse',optimizer=adam_opt)\n",
        "    return model\n",
        "\n",
        "  def _Build_Critic_Input(self,index,Id = None): #agent id is none by default (during training)\n",
        "    \"\"\"This function builds the ijnput whicht eh critic should be trained on.\n",
        "    This function should be used when analysing individual components\n",
        "    of the replay buffer, ie its input should be a replay buffer index set.\n",
        "\n",
        "\n",
        "    This input shoiuld include:\n",
        "    \n",
        "\n",
        "    - selected agent ID(during training, we can randomise this, becomes useful in baseline)\n",
        "    - ALL actions of the agents EXCLUDING the selected agent (14 total)\n",
        "    - The global state \n",
        "    - The selected agent's observation\n",
        "    - All agents' actions from the previous timestep\n",
        "      \"\"\"\n",
        "\n",
        "    \"\"\" \n",
        "    I need a current and a \"future\" critic input.\n",
        "    We need the future one in order to utilise TD(1) learning.\n",
        "    TD(lambda) is implemented in the COMA paper, but TD(1) is implemented\n",
        "    here to create a workflow.\n",
        "    \"\"\"\n",
        "\n",
        "    critic_input = []\n",
        "\n",
        "    if Id == None:\n",
        "      Id = random.randint(1, self.env.n_agents)\n",
        "    critic_input.append(Id)\n",
        "    #attach all current actions excluding selected agent\n",
        "    critic_input += self.buffer.Return_all_current_actions_excl_agent(index,Id)\n",
        "\n",
        "    #The global state is currently assumed to be just a list of ALL observations\n",
        "    #attach the global state\n",
        "    critic_input += self._Unravel_Observation(self.buffer.Return_current_state(index))\n",
        "\n",
        "    #attach current agent observation\n",
        "    critic_input += self.buffer.Return_current_agent_state(index,Id)\n",
        "\n",
        "    #attach all previous observations\n",
        "    critic_input += self.buffer.Return_all_previous_actions(index)\n",
        "\n",
        "    #######Define future critic input\n",
        "    \"\"\" TODO: Check if there IS a future input - the system could just be done.\n",
        "    This could be a downfall of the system.\n",
        "\n",
        "    Q: if there is no future input, what should be done? This is what prevented\n",
        "    this feature being implemented\"\"\"\n",
        "    future_critic_input = []\n",
        "\n",
        "    #attach id\n",
        "    future_critic_input.append(Id)\n",
        "\n",
        "    #attach all current actions excl selected agent\n",
        "    future_critic_input += self.buffer.Return_all_current_actions_excl_agent(index+1,Id)\n",
        "\n",
        "    #attach global state\n",
        "    future_critic_input += self._Unravel_Observation(self.buffer.Return_current_state(index+1))\n",
        "\n",
        "    #attach current agent observation\n",
        "    future_critic_input += self.buffer.Return_current_agent_state(index+1,Id)\n",
        "\n",
        "    #attach all previous observations\n",
        "    future_critic_input += self.buffer.Return_all_previous_actions(index+1)\n",
        "\n",
        "    \"\"\"Now, in order to find the value-based targets we need the future rewards and future done\n",
        "    although these are not inputs to the critic network and just should be returned separately\n",
        "    (see return statement)\"\"\"\n",
        "\n",
        "    future_reward = self.buffer.Return_reward(index+1)\n",
        "    future_done = self.buffer.Return_done(index+1)\n",
        "\n",
        "\n",
        "\n",
        "    return critic_input, future_critic_input, future_reward, future_done\n",
        "  \n",
        "\n",
        "  def _Train_Critic(self):\n",
        "    \"\"\" Note that right now, I am actually implementing TD(1) (next-state\n",
        "    TD) instead of TD(lambda), as is implemented in the paper.\n",
        "    This system is simpler to implement and will allow me to discretise a \n",
        "    workflow more efficiently than being stuck on the TD(lambda) implementation\n",
        "    for a long time, as it is not one of the key contributions to this paper\n",
        "    and thus is not the focus point of this implementation - change this if \n",
        "    I have time.\n",
        "    \n",
        "    \n",
        "    The current difference between this training and the training which was\n",
        "    done in my DQN algorithm is that I am bootstrapping the predicted Q values\n",
        "    (from a target network too!) rather than the bootstrapping each value, V.\n",
        "    \"\"\"\n",
        "\n",
        "\n",
        "    \"\"\"\n",
        "    Steps:\n",
        "    - Predict future Q values from future state (AFTER the first action)\n",
        "    - Find the target for each of these Q values, using the global reward.\n",
        "      Abive target = reward+self.discount_rate*predicted_Q value\n",
        "      for each Q value.\n",
        "      Q values are predicted by target network.\n",
        "    \n",
        "    \n",
        "     \"\"\"\n",
        "\n",
        "\n",
        "    #In order to train the critic, we could just randomyl select an agent that we are\n",
        "    #\"considering\" at that time.This will then allow us to find targets with reference to this agent\n",
        "\n",
        "    #first thing is to randomly sample from the deque and build the critic input\n",
        "\n",
        "    for i in range(self.minibatch_size): \n",
        "      index = random.randint(0, len(self.buffer.data)-3)\n",
        "\n",
        "      critic_input_row, future_critic_input_row, future_reward_row, future_done_row = self._Build_Critic_Input(index)\n",
        "\n",
        "      if i == 0:\n",
        "        critic_input = critic_input_row\n",
        "        future_critic_input = future_critic_input_row\n",
        "        future_reward = future_reward_row\n",
        "        future_done = future_done_row\n",
        "      else:\n",
        "        critic_input = np.vstack((critic_input,critic_input_row))\n",
        "        future_critic_input = np.vstack((future_critic_input,future_critic_input_row))\n",
        "        future_reward = np.vstack((future_reward,future_reward_row))\n",
        "        future_done = np.vstack((future_done,future_done_row))\n",
        "\n",
        "    #Now we must use the target network to get the predicted future Q values for the specified agent (see TD(0) eqn): \n",
        "    #https://medium.com/@violante.andre/simple-reinforcement-learning-temporal-difference-learning-e883ea0d65b0\n",
        "\n",
        "    Predicted_Q_Values = self.target_network.predict(future_critic_input) #shape batch_sizex15 (action numbers)\n",
        "\n",
        "    #Now we must define the targets for the above predicted Q values\n",
        "\n",
        "    for i in range(self.minibatch_size):\n",
        "      #since we are using Q as the targets here, instead of V\n",
        "      #(as the paper uses Q and says that it gives better results)\n",
        "      Q_row = Predicted_Q_Values[i,:] #one state's predicted Q values\n",
        "      done_row = future_done[i,:]  #corresponding to above state, same as reward too\n",
        "\n",
        "      if all(done_row):\n",
        "        \n",
        "        target_row = future_reward[i] #should be 15-long - add the [0] as the there was a bug with the sizes.\n",
        "        #since we are using Q, we need the output of this to be 15 long\n",
        "        target_row = list(target_row) * self.n_of_actions #distribute average reward across action space\n",
        "      else:\n",
        "        target_row = future_reward[i] + self.discount_factor * Q_row # Using Q, not V\n",
        "\n",
        "      #now create the targets matrix\n",
        "      if i == 0:\n",
        "        targets = target_row\n",
        "      else:\n",
        "        targets = np.vstack((targets,target_row))\n",
        "\n",
        "    class LossHistory(keras.callbacks.Callback): #This class is from the Keras website and allows for the loss to be recorded\n",
        "      def on_train_begin(self, logs={}):\n",
        "          self.losses = []\n",
        "\n",
        "      def on_batch_end(self, batch, logs={}):\n",
        "          self.losses.append(logs.get('loss'))\n",
        "\n",
        "    history = LossHistory()\n",
        "    \n",
        "    #now that we have the targets and the input, let's fit the model.\n",
        "    history = self.critic_model.fit(critic_input,targets, verbose = 0, batch_size = self.minibatch_size, epochs = 200)\n",
        "    mean_loss = np.mean(history.history['loss'])\n",
        "    self.critic_loss.append(mean_loss)\n",
        "\n",
        "\n",
        "\n",
        "  \"\"\"We make the below custom loss function in order to use Keras, Keras.backend and tensorFlow\n",
        "    to find our loss of gradients (needed in policy network optimisation and A2C), without finding them\n",
        "    manually, which can lead to error. \"\"\"\n",
        "\n",
        "  def _my_log_loss(self,y_true,y_pred):\n",
        "    # This allows us to input actions and advantages into the Keras API.\n",
        "    action_taken = y_true[:,1:]\n",
        "    advantages = y_true[:,0]\n",
        "\n",
        "    loss = keras.backend.log(keras.backend.sum(y_pred * action_taken,axis=1,\n",
        "                                               keepdims=True) + 1e-5) * keras.backend.stop_gradient(advantages)\n",
        "    #It is highly important that the action_taken is the ONE-HOT version:\n",
        "    # eg. [0,0,0,1,0,0,0,0], where 1 shows the index of the action taken\n",
        "    \n",
        "    return -loss #we want gradient ascent, not descent\n",
        "\n",
        "  \n",
        "\n",
        "  def _Create_Actor(self,input_dimensions,n_of_actions):\n",
        "    #Each agent should have its own actor\n",
        "\n",
        "    keras.backend.set_session(self.sess)\n",
        "    model = Sequential()\n",
        "    model.add(Dense(128, input_dim=input_dimensions, activation=\"relu\"))\n",
        "    model.add(Dense(128,activation = \"relu\"))\n",
        "    model.add(Dense(n_of_actions,activation = \"softmax\")) #Must be softmax since it outputs probabilities!\n",
        "    adam_opt = keras.optimizers.Adam(learning_rate=self.actor_learning_rate)  #default lr is 0.001\n",
        "\n",
        "    model.compile(loss= self._my_log_loss,optimizer=adam_opt)\n",
        "    return model\n",
        "\n",
        "  def _Instantiate_All_Actors(self,n_of_agents,input_dimensions,n_of_actions):\n",
        "    \"\"\" \n",
        "    COMA requires N actor networks to work for N agents,\n",
        "    therefore these will be stored in a dictionary\n",
        "    key: agent ID\n",
        "    value: actor network\n",
        "    \"\"\"\n",
        "    actors = {}\n",
        "\n",
        "    for i in range(n_of_agents): #loop's i goes from 0 to 9\n",
        "      model = self._Create_Actor(input_dimensions,n_of_actions)\n",
        "      actors[i+1] = model #i+1 so we get keys 1 - 10, coinciding with agent IDs\n",
        "    return actors\n",
        "  \n",
        "  def _Build_Actor_Input(self,index,agent):\n",
        "    \"\"\" \n",
        "    The actor input consists of the following:\n",
        "\n",
        "    - Current Agent Observation\n",
        "    - Current Agent ID\n",
        "    - Current Agent's Previous action\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    actor_input = []\n",
        "\n",
        "    #add the current agent observation\n",
        "    actor_input += self.buffer.Return_current_agent_state(index,agent)\n",
        "\n",
        "    #add the current Agent ID\n",
        "    actor_input.append(agent)\n",
        "\n",
        "    #add the current agent's previous action\n",
        "    actor_input.append(self.buffer.Return_previous_agent_action(index,agent))\n",
        "\n",
        "    return actor_input\n",
        "\n",
        "  def _Build_Actor_Input_Without_Buffer(self,agent_observation,agent,previous_agent_action):\n",
        "    \"\"\" \n",
        "    The actor input consists of the following:\n",
        "\n",
        "    - Current Agent Observation\n",
        "    - Current Agent ID\n",
        "    - Current Agent's Previous action\n",
        "\n",
        "    \"\"\"\n",
        "\n",
        "    actor_input = []\n",
        "\n",
        "    #add the current agent observation\n",
        "    actor_input += agent_observation\n",
        "\n",
        "    #add the current Agent ID\n",
        "    actor_input.append(agent)\n",
        "\n",
        "    #add the current agent's previous action\n",
        "    actor_input.append(previous_agent_action)\n",
        "\n",
        "    return actor_input\n",
        "\n",
        "\n",
        "\n",
        "  def _Define_Advantage(self,index,agent): \n",
        "    \"\"\"when we train the actor, we should train it via minibatch gradient ascent, to keep the data i.i.d\n",
        "    Q: would this create a problem in the temporal nature of the actor optimisation?\n",
        "    (ie it may be optimising from data of its PREVIOUS policy)\n",
        "    The Q values here are based off of the whole true space \"\"\"\n",
        "\n",
        "    #First part is to return all of the Q values from the critic network\n",
        "    critic_input_row, _, _, _ = self._Build_Critic_Input(index,agent) #only need the current critic input!\n",
        "    \n",
        "    critic_input_row = np.array([critic_input_row])# this must be defined as a batch, even though it is just of size 1\n",
        "\n",
        "    Q_values = self.critic_model.predict(critic_input_row) #these are whole Q values\n",
        "\n",
        "    #Now we must define the action which was actually taken after assessing these Q values\n",
        "    actual_agent_action_taken = self.buffer.Return_all_current_actions(index)[agent-1]\n",
        "\n",
        "    #Now define the actual whole Q value for the system when this action was picked\n",
        "    Q_s_u = Q_values[0][actual_agent_action_taken] #this is what the actual actor is PICKING to do!\n",
        "\n",
        "    #Now we must use the current actor model to define the probability that each action will be chosen:\n",
        "    current_actor = self.actors[agent]\n",
        "\n",
        "    current_actor_input = self._Build_Actor_Input(index,agent)\n",
        "    current_actor_input = np.array([current_actor_input])\n",
        "\n",
        "    probabilities = current_actor.predict(current_actor_input) \n",
        "    #Now this output should be 15-long list of probabilitites that each action will be picked\n",
        "    \n",
        "    taken_action_probability = probabilities[0][actual_agent_action_taken]\n",
        "    #the probability that the ACTUAL action taken was to be taken.\n",
        "    \n",
        "    \n",
        "    #Now we must loop through all of the available actions:\n",
        "\n",
        "    baseline_summation = 0\n",
        "    possible_actions = np.arange(0,15)\n",
        "    Advantage = []\n",
        "    \n",
        "\n",
        "\n",
        "    #This is where much of the innovation of the paper is - creating the counter-factual baseline.\n",
        "    # Advantage = Q - Baseline\n",
        "    #Q values come from the critic!\n",
        "\n",
        "\n",
        "    for i in range(len(Q_values[0])):\n",
        "      current_Q_value = Q_values[0][i] #q value for the current action\n",
        "\n",
        "      current_action_probability = probabilities[0][i]\n",
        "\n",
        "      baseline_summation += current_action_probability * current_Q_value\n",
        "\n",
        "    Advantage_component = Q_s_u - baseline_summation \n",
        "\n",
        "    return Advantage_component, actual_agent_action_taken, Q_values[0]\n",
        "    \n",
        "\n",
        "\n",
        "  def _Train_Actors(self):\n",
        "    \"\"\" \n",
        "    This code is used to train all of the actors during the training stage.\n",
        "    During the execution stage, this is not needed. This just updates the policies.\n",
        "    \"\"\"\n",
        "    for agent in self.actors: #loop through all agents - their models are indexed in the advantage function\n",
        "      gradient_sum = 0\n",
        "\n",
        "      #Now specify the current agent's model\n",
        "      current_agent_model = self.actors[agent]\n",
        "\n",
        "      for i in range(self.minibatch_size): #we are training through minibatch gradient ASCENT\n",
        "        \n",
        "        #Firstly, a random index from the buffer must be selected\n",
        "        index = random.randint(0, len(self.buffer.data)-3) #selecting the final value can lead to errors\n",
        "\n",
        "        #now find the input to the actor system\n",
        "        actor_input_row = self._Build_Actor_Input(index,agent)\n",
        "        advantage_row, agent_action, agent_Q_values_row = self._Define_Advantage(index,agent)\n",
        "        one_hot_agent_actions_row = [0] * self.n_of_actions\n",
        "        one_hot_agent_actions_row[agent_action] = 1\n",
        "        if i == 0:\n",
        "          actor_input = actor_input_row\n",
        "          advantage = advantage_row\n",
        "          one_hot_agent_actions = one_hot_agent_actions_row\n",
        "          agent_Q_values = agent_Q_values_row\n",
        "        else:\n",
        "          actor_input = np.vstack((actor_input,actor_input_row))\n",
        "          advantage = np.vstack((advantage,advantage_row))\n",
        "          one_hot_agent_actions = np.vstack((one_hot_agent_actions,one_hot_agent_actions_row))\n",
        "          agent_Q_values = np.vstack((agent_Q_values,agent_Q_values_row))\n",
        "          #per agent we should have a scalar advantage\n",
        "\n",
        "      actor_input = np.matrix(actor_input).reshape((self.minibatch_size,len(actor_input[0])))\n",
        "\n",
        "\n",
        "      # This allows us to input the actions and advantages directly into the Keras API.\n",
        "\n",
        "      advantages_and_actions = np.concatenate([advantage,one_hot_agent_actions],axis = 1) \n",
        "      #axis = 1 means that they will be side by side - first value is advantage!\n",
        "\n",
        "      current_agent_model.fit(actor_input, advantages_and_actions , verbose = 0, batch_size = self.minibatch_size, epochs = 200)\n",
        "\n",
        "\n",
        "  def Train_Policy(self):\n",
        "\n",
        "    env = self.env\n",
        "\n",
        "    actions = [4] * self.env.n_agents \n",
        "    #this is an initial (NOOP) action, needed since the action acts recurseively (to maintain buffer properties)\n",
        "\n",
        "    overall_steps = 0 #counts the steps across all episodes and iterations\n",
        "    self.total_rewards = []\n",
        "    start_time = time.time()\n",
        "    beginning = time.time()\n",
        "    for episode in range(self.training_episodes):\n",
        "      elapsed_time = time.time() - start_time\n",
        "      total_time = time.time() - beginning\n",
        "      start_time = time.time()\n",
        "      # print(\"elapsed time for this episode\",elapsed_time)\n",
        "      # print(\"total elapsed time\",total_time)\n",
        "      if episode % 100 == 0:\n",
        "        print(\"episode \",episode, \"out of\", self.training_episodes)\n",
        "      # print(\"buffer length\",len(self.buffer.data))\n",
        "      \n",
        "      ep_reward = 0\n",
        "      observations = env.reset()\n",
        "      done = [False for _ in range(env.n_agents)] #instantiation for the loop\n",
        "\n",
        "\n",
        "      #ANNEALING EPSILON\n",
        "      if episode < self.exploration_episodes:\n",
        "          self.epsilon = 1.1\n",
        "      else: #don't anneal epsilon during random exploration \n",
        "          max_ep = self.training_episodes-self.exploration_episodes\n",
        "          ep = episode-self.exploration_episodes\n",
        "          epsilon = ((max_ep-ep)/(max_ep))*(self.epsilon_max-self.epsilon_projection)+self.epsilon_projection\n",
        "          self.epsilon = np.max([self.epsilon_min,epsilon])\n",
        "      \n",
        "      # print(\"epsilon\",self.epsilon)\n",
        "      timestep_in_episode = 0\n",
        "      actions = [4] * self.env.n_agents\n",
        "      #this is an initial (NOOP) action, needed since the action acts recurseively (to maintain buffer properties)\n",
        "\n",
        "      while not all(done):\n",
        "        timestep_in_episode += 1\n",
        "        overall_steps += 1\n",
        "\n",
        "        previous_observations = observations #for buffer\n",
        "        previous_actions = actions #for buffer\n",
        "  \n",
        "        if random.random() < self.epsilon:\n",
        "          actions = env.action_space.sample() #random action\n",
        "\n",
        "        else:\n",
        "          actions = []\n",
        "          for agent, current_actor in self.actors.items():\n",
        "\n",
        "            agent_observation = observations[agent-1]\n",
        "            previous_agent_action = previous_actions[agent - 1] \n",
        "            #-1 becuse agents in my dict are 1 - 10, not index values (0 - 9)\n",
        "\n",
        "            actor_input = self._Build_Actor_Input_Without_Buffer(agent_observation,agent,previous_agent_action)\n",
        "            \n",
        "            actor_input = np.array([actor_input])\n",
        "            probabilities = current_actor.predict(actor_input)\n",
        "            probabilities = probabilities[0]\n",
        "            action = np.random.choice(self.n_of_actions,1,p = probabilities) #pick an action based on the outputted probabilities\n",
        "            actions.append(action)\n",
        "\n",
        "\n",
        "        if not isinstance(actions[0],int): \n",
        "          #to fix a bug where action numbers are put into arrays, caused by custom loss function?\n",
        "          actions = [x[0] for x in actions]\n",
        "\n",
        "        observations, rewards, done, info = env.step(actions)\n",
        "\n",
        "        \"\"\"\n",
        "        The description given at the top of this section, about the combat environment is taken\n",
        "        directly from the ma gym wiki. HOWEVER, from going through the ma gym source code, I am \n",
        "        pretty certain the description is incorrect. The source code states that the reward \n",
        "        is ACTUALLY += 1 if our team hits an enemy and -= 1 if the enemy team hits us.\n",
        "        This does NOT encourage the agents to fight, but ijn fact encourages our team to run\n",
        "        away, as any random exchange which could potentially encourage fighting is likely won\n",
        "        by the hardcoded team.\n",
        "        \n",
        "        The below code attempts to fix this issue as a supplementary reward, however,\n",
        "        without a way to get rid of the in-built reward, this supplementary method actually detriments\n",
        "        performance and hence has been omitted, but kept here for method reference\n",
        "        \"\"\"\n",
        "\n",
        "\n",
        "        self.buffer.Add_to_Buffer(previous_observations, previous_actions, observations, actions, rewards, done)\n",
        "\n",
        "        ep_reward += sum(rewards)\n",
        "\n",
        "        #exploration episodes to properly populate the buffer, not train the networks\n",
        "        if overall_steps % self.actor_critic_update_frequency == 0 and episode > self.exploration_episodes: \n",
        "          #update the actor and the critic routinely, not constantly after each step\n",
        "          self._Train_Critic()\n",
        "          self._Train_Actors()\n",
        "          self.episodes_trained.append(episode) #To plot the critic loss\n",
        "\n",
        "        if overall_steps % self.target_network_update_steps == 0 and episode > self.exploration_episodes:\n",
        "          #update the target network\n",
        "          self.target_network = keras.models.clone_model(self.critic_model)\n",
        "      # print(\"ep reward: \",ep_reward)\n",
        "      # print(\"enemy healths: \",env.opp_health.values())\n",
        "      self.total_rewards.append(ep_reward) #to plot the reward\n",
        "\n",
        "    env.render()\n",
        "    env.close()\n",
        "    # To improve the training efficiency, render() is not necessary during the training.\n",
        "    # We provide the render and video code here just want to demonstrate how to debugging and analysis.\n",
        "    # show_video()\n",
        "    #show_video() CAN crash the system on Google Collab due to usage restrictions\n",
        "\n",
        "\n",
        "\n",
        "  def Test_Policy(self): \n",
        "    #majoritively the same as training, except no random actions or training - just using actor networks.\n",
        "    env = self.env\n",
        "\n",
        "    actions = [4] * self.env.n_agents\n",
        "    overall_steps = 0 #counts the steps across all episodes and iterations\n",
        "    self.test_rewards = []\n",
        "    start_time = time.time()\n",
        "    beginning = time.time()\n",
        "    for episode in range(self.testing_episodes):\n",
        "      elapsed_time = time.time() - start_time\n",
        "      total_time = time.time() - beginning\n",
        "      start_time = time.time()\n",
        "      # print(\"elapsed time for this episode\",elapsed_time)\n",
        "      # print(\"total elapsed time\",total_time)\n",
        "      if episode % 50 == 0:\n",
        "        print(\"testing episode \",episode, \"out of\", self.testing_episodes)\n",
        "      # print(\"buffer length\",len(self.buffer.data))\n",
        "      \n",
        "      ep_reward = 0\n",
        "      observations = env.reset()\n",
        "      done = [False for _ in range(env.n_agents)] #instantiation for the loop\n",
        "\n",
        "      timestep_in_episode = 0\n",
        "      actions = [4] * self.env.n_agents \n",
        "      while not all(done):\n",
        "        timestep_in_episode += 1\n",
        "        overall_steps += 1\n",
        "\n",
        "        previous_observations = observations #for buffer\n",
        "        previous_actions = actions #for buffer \n",
        "        \n",
        "        actions = []\n",
        "        for agent, current_actor in self.actors.items():\n",
        "          agent_observation = observations[agent-1]\n",
        "          previous_agent_action = previous_actions[agent - 1] \n",
        "          #-1 becuse agents in my dict are 1 - 10, not index values (0 - 9)\n",
        "\n",
        "          actor_input = self._Build_Actor_Input_Without_Buffer(agent_observation,agent,previous_agent_action)\n",
        "          \n",
        "          actor_input = np.array([actor_input])\n",
        "          probabilities = current_actor.predict(actor_input)\n",
        "          probabilities = probabilities[0]\n",
        "          action = np.random.choice(self.n_of_actions,1,p = probabilities) \n",
        "          #pick an action based on the outputted probabilities\n",
        "\n",
        "          actions.append(action)\n",
        "\n",
        "        if not isinstance(actions[0],int): \n",
        "          #to fix a bug where action numbers are put into arrays, caused by custom loss function?\n",
        "          actions = [x[0] for x in actions]\n",
        "\n",
        "        observations, rewards, done, info = env.step(actions)\n",
        "\n",
        "        ep_reward += sum(rewards)\n",
        "        \n",
        "      # print(\"ep reward: \",ep_reward)\n",
        "      \n",
        "      self.test_rewards.append(ep_reward)\n",
        "    env.render()\n",
        "    env.close()\n",
        "    # To improve the training efficiency, render() is not necessary during the training.\n",
        "    # We provide the render and video code here just want to demonstrate how to debugging and analysis.\n",
        "    # show_video()\n",
        "    #show_video() CAN crash the system on Google Collab due to usage restrictions\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "    \n",
        "  def _Unravel_Observation(self,obs): #this function works!\n",
        "    #The critic runs on the entire state of our team, and thus the observation must be unravelled into\n",
        "    #one long string, to input into a single NN\n",
        "\n",
        "    #Since we do not have a global state available,this must be achieved by concatenating observations\n",
        "    #Source: COMA paper (see report below).\n",
        "\n",
        "    result_obs = [0 for _ in range(len(obs[0])*len(obs))] #zeros list of correct length\n",
        "\n",
        "    #now loop through all values methodically\n",
        "    step = 0\n",
        "    for i in range(len(obs)):\n",
        "      for j in range(len(obs[0])): #obs[0] has same length as obs[1] etc etc\n",
        "        result_obs[step] = obs[i][j]\n",
        "        step += 1\n",
        "    return result_obs\n",
        "\n",
        "\n",
        "COMA = Coma_System()\n",
        "COMA.Train_Policy() #execute\n",
        "\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0nMHoGAbuoVb",
        "colab_type": "text"
      },
      "source": [
        "#### Plot the Learning Curve"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "d9Sr8mHKurcE",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "fig, ax = plt.subplots()\n",
        "\n",
        "number_of_episodes = len(COMA.total_rewards)\n",
        "episode_matrix = np.arange(0,number_of_episodes)\n",
        "\n",
        "ax.plot(episode_matrix, COMA.total_rewards)\n",
        "ax.set_title(\"Learning Curve\")\n",
        "ax.set_ylabel(\"Episodic Reward (Cumulative for All Agents)\")\n",
        "ax.set_xlabel(\"Episodes\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "74_nj0VsiXR3",
        "colab_type": "text"
      },
      "source": [
        "Plot the Critic Loss"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Kny_pcngie9K",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "fig, ax = plt.subplots()\n",
        "print(\"critic_loss\",COMA.critic_loss)\n",
        "print(\"episodes trained\",COMA.episodes_trained)\n",
        "\n",
        "ax.plot(COMA.episodes_trained,COMA.critic_loss)\n",
        "\n",
        "ax.set_title(\"Critic Loss\")\n",
        "ax.set_ylabel(\"Loss\")\n",
        "ax.set_xlabel(\"Episodes Trained\")\n",
        "\n",
        "\"\"\" \n",
        "My report (below) mentions the sporadic critic loss, this can be seen on this graph.\n",
        "\"\"\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mcOjASlXihdU",
        "colab_type": "text"
      },
      "source": [
        "Testing Actor/Agent/Policy System:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r5LprplvilQf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "\"\"\" \n",
        "This tests the trained policy / actor networks.\n",
        "It is placed after the training and train-plots because\n",
        "strenuous computational work is more likely to crash Google Collab\n",
        "than plotting etc, which after training is heavily disadvantageous.\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "COMA.Test_Policy()\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BJzDqlaNiq5Q",
        "colab_type": "text"
      },
      "source": [
        "Plot the Testing Curve"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gruDsqvUitkK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "fig, ax = plt.subplots()\n",
        "\n",
        "number_of_episodes = len(COMA.test_rewards)\n",
        "episode_matrix = np.arange(0,number_of_episodes)\n",
        "\n",
        "ax.plot(episode_matrix, COMA.test_rewards)\n",
        "ax.set_title(\"Testing Curve\")\n",
        "ax.set_ylabel(\"Episodic Reward (Cumulative for All Agents)\")\n",
        "ax.set_xlabel(\"Episodes\")\n",
        "\n",
        "\"\"\"\n",
        "System rewsrds can vary due to randomness (spawn, instantiations etc),\n",
        "hence the neeed for this curve, as the reward is not always necessarily the \n",
        "same, even with the same policy/actor networks.\n",
        "\"\"\""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T1yh_c3bj4io",
        "colab_type": "text"
      },
      "source": [
        "Show Video"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kw2y-Kw1j3hM",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\"\"\"\n",
        "This shows the behaviour of the trained network.\n",
        "It is outside of the main training and testing loop\n",
        "since it often causes Google Collab crashes due\n",
        "to usage limits. \n",
        "\"\"\"\n",
        "\n",
        "show_video()"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}