Fast&Slow.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZxMFW3MlFU93",
    "outputId": "6d476016-e649-474a-9682-f487b48e50f9"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "zsh:1: no matches found: stable_baselines3[extra]\n",
      "Collecting sb3-contrib\n",
      "  Downloading sb3_contrib-1.7.0-py3-none-any.whl (79 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m79.1/79.1 kB\u001b[0m \u001b[31m2.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting stable-baselines3>=1.7.0\n",
      "  Downloading stable_baselines3-1.7.0-py3-none-any.whl (171 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m171.8/171.8 kB\u001b[0m \u001b[31m5.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: pandas in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from stable-baselines3>=1.7.0->sb3-contrib) (1.4.4)\n",
      "Requirement already satisfied: cloudpickle in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from stable-baselines3>=1.7.0->sb3-contrib) (2.0.0)\n",
      "Collecting torch>=1.11\n",
      "  Downloading torch-1.13.1-cp39-none-macosx_10_9_x86_64.whl (135.3 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m135.3/135.3 MB\u001b[0m \u001b[31m17.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hCollecting importlib-metadata~=4.13\n",
      "  Downloading importlib_metadata-4.13.0-py3-none-any.whl (23 kB)\n",
      "Collecting gym==0.21\n",
      "  Downloading gym-0.21.0.tar.gz (1.5 MB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.5/1.5 MB\u001b[0m \u001b[31m24.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25h  Preparing metadata (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25hRequirement already satisfied: numpy in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from stable-baselines3>=1.7.0->sb3-contrib) (1.21.5)\n",
      "Requirement already satisfied: matplotlib in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from stable-baselines3>=1.7.0->sb3-contrib) (3.5.2)\n",
      "Requirement already satisfied: zipp>=0.5 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from importlib-metadata~=4.13->stable-baselines3>=1.7.0->sb3-contrib) (3.8.0)\n",
      "Requirement already satisfied: typing-extensions in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from torch>=1.11->stable-baselines3>=1.7.0->sb3-contrib) (4.3.0)\n",
      "Requirement already satisfied: cycler>=0.10 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from matplotlib->stable-baselines3>=1.7.0->sb3-contrib) (0.11.0)\n",
      "Requirement already satisfied: packaging>=20.0 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from matplotlib->stable-baselines3>=1.7.0->sb3-contrib) (21.3)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from matplotlib->stable-baselines3>=1.7.0->sb3-contrib) (1.4.2)\n",
      "Requirement already satisfied: pyparsing>=2.2.1 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from matplotlib->stable-baselines3>=1.7.0->sb3-contrib) (3.0.9)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from matplotlib->stable-baselines3>=1.7.0->sb3-contrib) (4.25.0)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from matplotlib->stable-baselines3>=1.7.0->sb3-contrib) (2.8.2)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from matplotlib->stable-baselines3>=1.7.0->sb3-contrib) (9.2.0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from pandas->stable-baselines3>=1.7.0->sb3-contrib) (2022.1)\n",
      "Requirement already satisfied: six>=1.5 in /Users/tanchongmin/opt/anaconda3/lib/python3.9/site-packages (from python-dateutil>=2.7->matplotlib->stable-baselines3>=1.7.0->sb3-contrib) (1.16.0)\n",
      "Building wheels for collected packages: gym\n",
      "  Building wheel for gym (setup.py) ... \u001b[?25ldone\n",
      "\u001b[?25h  Created wheel for gym: filename=gym-0.21.0-py3-none-any.whl size=1616799 sha256=0cad8d8e8e02200db4a5a631dc5e226046bcb661d926251c5ea9fb63e4255ea9\n",
      "  Stored in directory: /Users/tanchongmin/Library/Caches/pip/wheels/b3/50/6c/0a82c1358b4da2dbd9c1bb17e0f89467db32812ab236dbf6d5\n",
      "Successfully built gym\n",
      "Installing collected packages: torch, importlib-metadata, gym, stable-baselines3, sb3-contrib\n",
      "  Attempting uninstall: importlib-metadata\n",
      "    Found existing installation: importlib-metadata 4.11.3\n",
      "    Uninstalling importlib-metadata-4.11.3:\n",
      "      Successfully uninstalled importlib-metadata-4.11.3\n",
      "Successfully installed gym-0.21.0 importlib-metadata-4.13.0 sb3-contrib-1.7.0 stable-baselines3-1.7.0 torch-1.13.1\n"
     ]
    }
   ],
   "source": [
    "!pip install stable_baselines3[extra]\n",
    "!pip install sb3-contrib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "lNhZqpya_csr",
    "outputId": "6f087b34-8308-4977-88ab-60c32e91899e"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from collections import defaultdict\n",
    "from queue import PriorityQueue\n",
    "import threading\n",
    "from threading import Thread\n",
    "\n",
    "import stable_baselines3 as sb3\n",
    "from stable_baselines3 import PPO, A2C, DQN, HerReplayBuffer\n",
    "from sb3_contrib import ARS, TRPO, QRDQN\n",
    "import torch\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "from gym import spaces\n",
    "\n",
    "print(tf.__version__)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "iwgO42H0wtii"
   },
   "source": [
    "# Define some sample next states and actions to try to learn by supervised learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jRo0RG_a_xR_"
   },
   "outputs": [],
   "source": [
    "def getnextaction(state):\n",
    "  cur_x, cur_y, goal_x, goal_y = state\n",
    "  if cur_x > goal_x:\n",
    "    return 0\n",
    "  elif cur_x < goal_x:\n",
    "    return 1\n",
    "  elif cur_y > goal_y:\n",
    "    return 2\n",
    "  elif cur_y < goal_y:\n",
    "    return 3\n",
    "  else:\n",
    "    return 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WRKZnh3NYNW2"
   },
   "outputs": [],
   "source": [
    "# this is used for the alternate formulation\n",
    "def getnextaction2(state):\n",
    "  cur_x, cur_y, goal_x, goal_y = state\n",
    "  if cur_y > goal_y:\n",
    "    return 0\n",
    "  elif cur_y < goal_y:\n",
    "    return 1\n",
    "  elif cur_x > goal_x:\n",
    "    return 2\n",
    "  elif cur_x < goal_x:\n",
    "    return 3\n",
    "  else:\n",
    "    return 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "z3voJT7bQ7ZD"
   },
   "outputs": [],
   "source": [
    "def getnextstate(state):\n",
    "  cur_x, cur_y, goal_x, goal_y = state\n",
    "  if cur_x > goal_x:\n",
    "    return (cur_x-1, cur_y)\n",
    "  elif cur_x < goal_x:\n",
    "    return (cur_x+1, cur_y)\n",
    "  elif cur_y > goal_y:\n",
    "    return (cur_x, cur_y - 1)\n",
    "  elif cur_y < goal_y:\n",
    "    return (cur_x, cur_y+1)\n",
    "  else:\n",
    "    return (cur_x, cur_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "luYbhWRnYjVV"
   },
   "outputs": [],
   "source": [
    "# this is used for the alternate formulation\n",
    "def getnextstate2(state):\n",
    "  cur_x, cur_y, goal_x, goal_y = state\n",
    "  if cur_y > goal_y:\n",
    "    return (cur_x, cur_y - 1)\n",
    "  elif cur_y < goal_y:\n",
    "    return (cur_x, cur_y+1)\n",
    "  elif cur_x > goal_x:\n",
    "    return (cur_x-1, cur_y)\n",
    "  elif cur_x < goal_x:\n",
    "    return (cur_x+1, cur_y)\n",
    "  else:\n",
    "    return (cur_x, cur_y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SHbFgbR-AMiT"
   },
   "outputs": [],
   "source": [
    "def GetSamples(grid_size = 10, seed = 0):\n",
    "  np.random.seed(seed)\n",
    "  states = np.random.randint(0, grid_size, (1000, 4))\n",
    "  actions = np.array([getnextaction(state) for state in states])\n",
    "  next_states = np.array([getnextstate(state) for state in states])\n",
    "\n",
    "  actions2 = np.array([getnextaction2(state) for state in states])\n",
    "  next_states2 = np.array([getnextstate2(state) for state in states])\n",
    "\n",
    "  return states, actions, next_states, actions2, next_states2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oNxVMm4hXN_v"
   },
   "source": [
    "# Experiment 1:\n",
    "# Input: Current State, Goal State\n",
    "# Output: Action\n",
    "\n",
    "This shows that learning an action given the current state and conditioned on the goal state is easy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "VPN7ipwMXUeD",
    "outputId": "a8aedd1a-f54e-4df6-e8d4-fc17779de555"
   },
   "outputs": [],
   "source": [
    "grid_size = 10\n",
    "inputs = tf.keras.layers.Input(shape = (4,))\n",
    "# x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), strides=(1,1), padding='same', activation = tf.nn.relu)(inputs)\n",
    "# x = tf.keras.layers.MaxPooling2D(pool_size = (2,2))(x)\n",
    "# x = tf.keras.layers.Flatten()(x)\n",
    "x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(inputs)\n",
    "x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(x)\n",
    "output1 = tf.keras.layers.Dense(5, activation = tf.nn.softmax)(x)\n",
    " \n",
    "model = tf.keras.Model(inputs = inputs, outputs = output1)\n",
    "\n",
    "model.compile(optimizer = 'adam',\n",
    "              loss = 'sparse_categorical_crossentropy', metrics = ['accuracy'])\n",
    "\n",
    "states, actions, next_states, actions2, next_states2 = GetSamples(grid_size = 10)\n",
    "\n",
    "datasize = len(states)\n",
    "split = datasize//5 * 4\n",
    "# reshape the input to add a channel dimension at the end for CNN processing\n",
    "history = model.fit(states[:split], actions[:split], epochs=50)\n",
    "\n",
    "## swap x and y priority for the next state\n",
    "datasize = len(states)\n",
    "split = datasize//5 * 4\n",
    "# reshape the input to add a channel dimension at the end for CNN processing\n",
    "history2 = model.fit(states[:split], actions2[:split], epochs=50)\n",
    "\n",
    "action10 = history.history['accuracy']+history2.history['accuracy']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Qp18xjRPdTUz",
    "outputId": "fa4f7374-0756-4f39-cd53-bd5e65f1719d"
   },
   "outputs": [],
   "source": [
    "grid_size = 20\n",
    "inputs = tf.keras.layers.Input(shape = (4,))\n",
    "# x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), strides=(1,1), padding='same', activation = tf.nn.relu)(inputs)\n",
    "# x = tf.keras.layers.MaxPooling2D(pool_size = (2,2))(x)\n",
    "# x = tf.keras.layers.Flatten()(x)\n",
    "x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(inputs)\n",
    "x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(x)\n",
    "output1 = tf.keras.layers.Dense(5, activation = tf.nn.softmax)(x)\n",
    " \n",
    "model = tf.keras.Model(inputs = inputs, outputs = output1)\n",
    "\n",
    "model.compile(optimizer = 'adam',\n",
    "              loss = 'sparse_categorical_crossentropy', metrics = ['accuracy'])\n",
    "\n",
    "states, actions, next_states, actions2, next_states2 = GetSamples(grid_size = 20)\n",
    "\n",
    "datasize = len(states)\n",
    "split = datasize//5 * 4\n",
    "history = model.fit(states[:split], actions[:split], epochs=50)\n",
    "\n",
    "## swap x and y priority for the next state\n",
    "datasize = len(states)\n",
    "split = datasize//5 * 4\n",
    "history2 = model.fit(states[:split], actions2[:split], epochs=50)\n",
    "\n",
    "action20 = history.history['accuracy']+history2.history['accuracy']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 296
    },
    "id": "1-R55JAbt-VR",
    "outputId": "b2dabd3d-92b1-412a-df23-8accfc834787"
   },
   "outputs": [],
   "source": [
    "plt.plot(action10, 'b', label = 'Next Action Prediction 10x10')\n",
    "plt.plot(action20, 'r', label = 'Next Action Prediction 20x20')\n",
    "plt.xlabel('Number of Epochs')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.xticks([i*10 for i in range(11)])\n",
    "plt.grid()\n",
    "plt.legend(loc = 'lower right')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nhdFm7rK_tBr"
   },
   "source": [
    "# Experiment 2\n",
    "# Input: Current State, Goal State\n",
    "# Output: Next State\n",
    "\n",
    "This shows that learning the next state conditioned on current state and goal state is difficult."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DIojL9Ed_lyR",
    "outputId": "e09c5366-3858-4536-b5c0-b4a11559a5fb"
   },
   "outputs": [],
   "source": [
    "grid_size = 10\n",
    "inputs = tf.keras.layers.Input(shape = (4,))\n",
    "# x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), strides=(1,1), padding='same', activation = tf.nn.relu)(inputs)\n",
    "# x = tf.keras.layers.MaxPooling2D(pool_size = (2,2))(x)\n",
    "# x = tf.keras.layers.Flatten()(x)\n",
    "x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(inputs)\n",
    "x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(x)\n",
    "output1 = tf.keras.layers.Dense(grid_size, activation = tf.nn.softmax, name = 'state_x')(x)\n",
    "output2 = tf.keras.layers.Dense(grid_size, activation = tf.nn.softmax, name = 'state_y')(x)\n",
    " \n",
    "model2 = tf.keras.Model(inputs = inputs, outputs = [output1, output2])\n",
    "\n",
    "model2.compile(optimizer = 'adam',\n",
    "              loss = 'sparse_categorical_crossentropy', metrics = ['accuracy'])\n",
    "\n",
    "states, actions, next_states, actions2, next_states2 = GetSamples(grid_size = 10)\n",
    "\n",
    "datasize = len(states)\n",
    "split = datasize//5 * 4\n",
    "history = model2.fit(states[:split], [next_states[:split][:,0], next_states[:split][:,1]], epochs=200)\n",
    "\n",
    "## swap x and y priority for next state\n",
    "datasize = len(states)\n",
    "split = datasize//5 * 4\n",
    "history2 = model2.fit(states[:split], [next_states2[:split][:,0], next_states2[:split][:,1]], epochs=200)\n",
    "\n",
    "acc_x = history.history['state_x_accuracy']+history2.history['state_x_accuracy']\n",
    "acc_y = history.history['state_y_accuracy']+history2.history['state_y_accuracy']\n",
    "nextstate10 = (np.array(acc_x)+np.array(acc_y))/2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "yBKTgEqjwHRm",
    "outputId": "b68e2528-b582-443e-fa55-070fc1b55082"
   },
   "outputs": [],
   "source": [
    "grid_size = 20\n",
    "inputs = tf.keras.layers.Input(shape = (4,))\n",
    "# x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), strides=(1,1), padding='same', activation = tf.nn.relu)(inputs)\n",
    "# x = tf.keras.layers.MaxPooling2D(pool_size = (2,2))(x)\n",
    "# x = tf.keras.layers.Flatten()(x)\n",
    "x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(inputs)\n",
    "x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(x)\n",
    "output1 = tf.keras.layers.Dense(grid_size, activation = tf.nn.softmax, name = 'state_x')(x)\n",
    "output2 = tf.keras.layers.Dense(grid_size, activation = tf.nn.softmax, name = 'state_y')(x)\n",
    " \n",
    "model2 = tf.keras.Model(inputs = inputs, outputs = [output1, output2])\n",
    "\n",
    "model2.compile(optimizer = 'adam',\n",
    "              loss = 'sparse_categorical_crossentropy', metrics = ['accuracy'])\n",
    "\n",
    "states, actions, next_states, actions2, next_states2 = GetSamples(grid_size = 20)\n",
    "\n",
    "datasize = len(states)\n",
    "split = datasize//5 * 4\n",
    "history = model2.fit(states[:split], [next_states[:split][:,0], next_states[:split][:,1]], epochs=200)\n",
    "\n",
    "## swap x and y priority for next state\n",
    "datasize = len(states)\n",
    "split = datasize//5 * 4\n",
    "history2 = model2.fit(states[:split], [next_states2[:split][:,0], next_states2[:split][:,1]], epochs=200)\n",
    "\n",
    "acc_x = history.history['state_x_accuracy']+history2.history['state_x_accuracy']\n",
    "acc_y = history.history['state_y_accuracy']+history2.history['state_y_accuracy']\n",
    "nextstate20 = (np.array(acc_x)+np.array(acc_y))/2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 296
    },
    "id": "soX0Y8AueBUu",
    "outputId": "e006de69-8b47-40d8-b48e-7206fff28fb2"
   },
   "outputs": [],
   "source": [
    "plt.plot(nextstate10, 'b', label = 'Next State Prediction 10x10')\n",
    "plt.plot(nextstate20, 'r', label = 'Next State Prediction 20x20')\n",
    "plt.xticks([i*50 for i in range(9)])\n",
    "plt.grid()\n",
    "plt.xlabel('Number of Epochs')\n",
    "plt.ylabel('Accuracy')\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "D5j9QyU8heZk"
   },
   "source": [
    "# Create Grid Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "LKNoqxse2swm"
   },
   "outputs": [],
   "source": [
    "num_actions = 4\n",
    "\n",
    "class GridEnv(gym.Env):\n",
    "    \"\"\"Custom Environment that follows gym interface.\"\"\"\n",
    "\n",
    "    metadata = {\"render.modes\": [\"human\"]}\n",
    "\n",
    "    def __init__(self, size=5, start_state = None, goal_state = None, forbidden_squares = [], seed = None):\n",
    "        super().__init__()\n",
    "\n",
    "        # initialize action and observation space\n",
    "        self.action_space = spaces.Discrete(4)\n",
    "        self.observation_space = spaces.Box(low=0, high=size-1, shape=(4,), dtype=np.uint8)\n",
    "\n",
    "        self.env_size = size\n",
    "        self.forbidden_squares = forbidden_squares\n",
    "\n",
    "        self.steps = 0\n",
    "\n",
    "        if seed is not None:\n",
    "          np.random.seed(seed)\n",
    "\n",
    "        self.start_state = (np.random.randint(0, self.env_size), np.random.randint(0, self.env_size))\n",
    "        self.goal_state = (np.random.randint(0, self.env_size), np.random.randint(0, self.env_size))\n",
    "\n",
    "        # make sure we don't spawn in the forbidden squares\n",
    "        while self.start_state in self.forbidden_squares:\n",
    "          self.start_state = (np.random.randint(0, self.env_size), np.random.randint(0, self.env_size))\n",
    "        # make sure we don't spawn in forbidden squares, and that the goal state is not the start state\n",
    "        while self.goal_state in self.forbidden_squares or self.goal_state == self.start_state:\n",
    "          self.goal_state = (np.random.randint(0, self.env_size), np.random.randint(0, self.env_size))\n",
    "        if start_state is not None:\n",
    "          self.start_state = start_state\n",
    "        if goal_state is not None:\n",
    "          self.goal_state = goal_state\n",
    "        self.cur_state = self.start_state\n",
    "        self.min_steps = self.get_min()\n",
    "\n",
    "    def get_min(self):\n",
    "        ''' Gets the minimum number of moves from start state to goal state using breadth first search '''\n",
    "        visited = [self.start_state]\n",
    "        cur_nodes = [self.start_state]\n",
    "        count = 0\n",
    "        while self.goal_state not in cur_nodes and count <= self.env_size * self.env_size:\n",
    "          future_nodes = []\n",
    "          count += 1 # increment 1 to the visited counts\n",
    "          while len(cur_nodes) > 0:\n",
    "            node = cur_nodes.pop()\n",
    "            for future_node in [(node[0]-1, node[1]), (node[0]+1, node[1]), (node[0], node[1]-1), (node[0], node[1]+1)]:\n",
    "              if future_node not in visited and 0 <= future_node[0] < self.env_size and 0 <= future_node[1] < self.env_size:\n",
    "                visited.append(future_node)\n",
    "                future_nodes.append(future_node)\n",
    "          cur_nodes = future_nodes.copy()\n",
    "        return count\n",
    "\n",
    "    def step(self, action):\n",
    "        # only step if episode is not terminated\n",
    "        ''' Takes an action and also returns 1 if the goal state is reached '''\n",
    "        initial_state = self.cur_state\n",
    "        # left\n",
    "        if action == 0 and self.cur_state[0] > 0:\n",
    "          self.cur_state = (self.cur_state[0]-1, self.cur_state[1])\n",
    "        # up\n",
    "        elif action == 1 and self.cur_state[1] > 0:\n",
    "          self.cur_state = (self.cur_state[0], self.cur_state[1]-1)\n",
    "        # right\n",
    "        elif action == 2 and self.cur_state[0]+1 < self.env_size:\n",
    "          self.cur_state = (self.cur_state[0]+1, self.cur_state[1])\n",
    "        # down\n",
    "        elif action == 3 and self.cur_state[1]+1 < self.env_size:\n",
    "          self.cur_state = (self.cur_state[0], self.cur_state[1]+1)\n",
    "        # otherwise don't move\n",
    "        # check if we hit the forbidden states. If so, reposition back to the beginning\n",
    "        if self.cur_state in self.forbidden_squares:\n",
    "          self.cur_state = initial_state\n",
    "      \n",
    "        # increment number of steps by 1\n",
    "        observation = self.cur_state + self.goal_state\n",
    "        self.steps += 1\n",
    "        done = self.steps >= self.env_size * self.env_size # only stop once we reached end of the training steps\n",
    "\n",
    "        reward = 0\n",
    "        # terminate if we have reached the goal\n",
    "        if self.cur_state == self.goal_state:\n",
    "          reward = 1\n",
    "          done = True\n",
    "\n",
    "        return observation, reward, done, {}\n",
    "\n",
    "    def reset(self):\n",
    "        self.cur_state = self.start_state # reset start state\n",
    "        observation = self.cur_state + self.goal_state\n",
    "        self.steps = 0 # reset step count\n",
    "        return observation  # reward, done, info can't be included\n",
    "\n",
    "    def render(self, mode=\"human\"):\n",
    "        print(self.cur_state+self.goal_state)\n",
    "\n",
    "    def close(self):\n",
    "        ''' Don't have to close anything as we are not rendering anything '''\n",
    "        return"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2IPHtfFfZhct"
   },
   "source": [
    "## Helper Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "fAYWKVDAZjOf"
   },
   "outputs": [],
   "source": [
    "def minion():\n",
    "    ''' This references existing memory and tries to find a trajectory towards the goal state '''\n",
    "    # this stores the list of states the minion experiences\n",
    "    minion_state_history = [cur_state]\n",
    "    minion_action_history = []\n",
    "    for depth in range(lookahead_depth):\n",
    "      # look for matches\n",
    "      matches = []\n",
    "\n",
    "      for action in range(num_actions):\n",
    "        matches.extend([(state, action) for state in memory[minion_state_history[-1]][action]])\n",
    "\n",
    "      # no match, then discontinue this minion's search\n",
    "      if len(matches) == 0:\n",
    "        break\n",
    "\n",
    "      next_minion_state, next_minion_action = matches[np.random.randint(len(matches))]\n",
    "      # update minion's value\n",
    "      minion_state_history.append(next_minion_state)\n",
    "      if depth==0:\n",
    "        best_action = next_minion_action\n",
    "      minion_action_history.append(next_minion_action)\n",
    "      \n",
    "      # check if goal state reached\n",
    "      if minion_state_history[-1] == goal_state:\n",
    "        sem.acquire()\n",
    "        best_minions.put((depth, best_action, minion_state_history[:-1], minion_action_history, next_minion_action))\n",
    "        sem.release()\n",
    "        break\n",
    "\n",
    "def GetModel(seed = 0):\n",
    "    ''' This takes in a random seed and builds the desired TensorFlow model '''\n",
    "\n",
    "    # set random seeds\n",
    "    tf.random.set_seed(seed)\n",
    "    np.random.seed(seed)\n",
    "\n",
    "    inputs = tf.keras.layers.Input(shape = (4,))\n",
    "    # x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3,3), strides=(1,1), padding='same', activation = tf.nn.relu)(inputs)\n",
    "    # x = tf.keras.layers.MaxPooling2D(pool_size = (2,2))(x)\n",
    "    # x = tf.keras.layers.Flatten()(x)\n",
    "    x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(inputs)\n",
    "    x = tf.keras.layers.Dense(128, activation = tf.nn.relu)(x)\n",
    "    output1 = tf.keras.layers.Dense(num_actions, activation = tf.nn.softmax)(x)\n",
    "    \n",
    "    model = tf.keras.Model(inputs = inputs, outputs = output1)\n",
    "\n",
    "    model.compile(optimizer = 'adam',\n",
    "                  loss = 'sparse_categorical_crossentropy', metrics = ['accuracy'])\n",
    "    \n",
    "    return model\n",
    "\n",
    "def GetGrid(episode, grid_size = 10):\n",
    "  ''' This takes in an episode number and returns the desired forbidden squares'''\n",
    "  # Define what is traversible\n",
    "  if episode < 50:\n",
    "    forbidden_squares = [(x, grid_size//2) for x in range(grid_size) if x!=grid_size//2]\n",
    "  else:\n",
    "    forbidden_squares = [(grid_size//2, x) for x in range(grid_size) if x!=grid_size//2]\n",
    "\n",
    "  return forbidden_squares"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kEMAgs80rjnV"
   },
   "source": [
    "# Experiment 1: Solving the grid with fixed start and end points\n",
    "\n",
    "- We define a 10x10 grid and the agent is supposed to go from (0, 0) to (grid_size-1, grid_size - 1)\n",
    "- No obstacles"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FhKU2r5XEZI9"
   },
   "source": [
    "## Model 1: PPO, TRPO, A2C\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "xf5EApSiEXhJ",
    "outputId": "673f267e-998a-4387-b333-d84832d695c3"
   },
   "outputs": [],
   "source": [
    "for ALGO in [PPO, TRPO, A2C]:\n",
    "    print (ALGO.__name__)\n",
    "    grid_size = 10\n",
    "    # Custom MLP policy of two layers of size 128 each with Relu activation function\n",
    "    policy_kwargs = dict(activation_fn=torch.nn.ReLU, net_arch=[128, 128])\n",
    "    # Create the environment\n",
    "    env = GridEnv(size = grid_size, start_state = (0, 0), goal_state = (grid_size-1, grid_size-1), forbidden_squares = [])\n",
    "    # Create the agent\n",
    "    # set random seed\n",
    "    sb3.common.utils.set_random_seed(0)\n",
    "    if ALGO == A2C:\n",
    "      model = ALGO(\"MlpPolicy\", env, policy_kwargs=policy_kwargs, learning_rate = 0.0001)\n",
    "    else:\n",
    "      model = ALGO(\"MlpPolicy\", env, policy_kwargs=policy_kwargs)\n",
    "    model.save(\"my_env\")\n",
    "\n",
    "    best_steps = []\n",
    "    actual_steps = []\n",
    "    for i in range(100):\n",
    "      # learn for each episode\n",
    "      env = GridEnv(size = grid_size, start_state = (0, 0), goal_state = (grid_size-1, grid_size-1), forbidden_squares = [])\n",
    "      model = ALGO.load(\"my_env\", env=env)\n",
    "      model.learn(total_timesteps=100)\n",
    "      model.save(\"my_env\")\n",
    "      obs = env.reset()\n",
    "      for steps in range(100):\n",
    "        action, _states = model.predict(obs)\n",
    "        obs, rewards, done, info = env.step(action)\n",
    "        if done:\n",
    "          if rewards == 1 and steps+1 < grid_size*grid_size:\n",
    "            print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "\n",
    "          else:\n",
    "            print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "          best_steps.append(env.min_steps)\n",
    "          actual_steps.append(steps+1)\n",
    "          break\n",
    "\n",
    "    # save the best steps\n",
    "    if ALGO == PPO:\n",
    "      ppo_steps = actual_steps\n",
    "    elif ALGO == TRPO:\n",
    "      trpo_steps = actual_steps\n",
    "    else:\n",
    "      a2c_steps = actual_steps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Lue-S4r2H5GC"
   },
   "source": [
    "## Model 2: Q-Learning Approach to solving the grid with fixed start and end points\n",
    "$Q(s,a) = Q(s,a) + \\alpha(r + \\gamma\\max_a Q(s',a) - Q(s,a))$\n",
    "\n",
    "The maximum next state value only applies if it is not a terminal state.\n",
    "\n",
    "Epsilon-greedy action selection, with greedy action being $argmax_a Q(s,a)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 381
    },
    "id": "9SMFn29AH7K2",
    "outputId": "0002db81-e4b6-4097-fe8d-47d6bd3672e7"
   },
   "outputs": [],
   "source": [
    "grid_size = 10\n",
    "alpha = 1\n",
    "gamma = 0.99\n",
    "eps = 1\n",
    "num_actions = 4\n",
    "num_episodes = 100\n",
    "\n",
    "best_steps = []\n",
    "actual_steps = []\n",
    "q_values = defaultdict(lambda: 0)\n",
    "np.random.seed(0)\n",
    "for i in range(num_episodes):\n",
    "  env = GridEnv(grid_size, start_state = (0, 0), goal_state = (grid_size-1, grid_size-1), forbidden_squares = [])\n",
    "  for steps in range(grid_size * grid_size):\n",
    "    cur_state = env.cur_state\n",
    "    action = 0\n",
    "    # take random action if less than desired episode count\n",
    "    if i<75:\n",
    "      action = np.random.randint(num_actions)\n",
    "    # otherwise choose the maximum action\n",
    "    else:\n",
    "      action_values = [q_values[(cur_state, a)] for a in range(num_actions)]\n",
    "      action = np.argmax(action_values)\n",
    "\n",
    "    obs, rewards, done, _ = env.step(action)\n",
    "    next_state = (obs[0], obs[1])\n",
    "\n",
    "    # update the Q value tables\n",
    "    future_action_values = [q_values[(next_state, a)] for a in range(num_actions)]\n",
    "    old_state_value = q_values[(cur_state, action)]\n",
    "    if rewards == 1:\n",
    "      q_values[(cur_state, action)] = old_state_value + alpha*(rewards - old_state_value)\n",
    "    else:\n",
    "      q_values[(cur_state, action)] = old_state_value + alpha*(rewards + gamma*np.max(future_action_values) - old_state_value)\n",
    "\n",
    "    if done:\n",
    "      if rewards == 1 and steps+1 < grid_size*grid_size:\n",
    "        print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "\n",
    "      else:\n",
    "        print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "      best_steps.append(env.min_steps)\n",
    "      actual_steps.append(steps+1)\n",
    "      break\n",
    "\n",
    "qlearning_steps = actual_steps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FSturl9OrQeE"
   },
   "source": [
    "## Model 3: Memory-based approach "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "x5w5BBv1wG_9"
   },
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "use_fast = True\n",
    "use_slow = True\n",
    "\n",
    "grid_size = 10\n",
    "exploration_const = 1\n",
    "parallel_threads = 100\n",
    "lookahead_depth = 20\n",
    "\n",
    "# create model\n",
    "model = GetModel()\n",
    "\n",
    "# reset memory\n",
    "memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "# store only first action\n",
    "best_steps = []\n",
    "actual_steps = []\n",
    "\n",
    "for i in range(100):\n",
    "  env = GridEnv(grid_size, start_state = (0, 0), goal_state = (grid_size-1, grid_size-1), forbidden_squares = [])\n",
    "  statehistory = []\n",
    "  actionhistory = []\n",
    "  episode_memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "  for steps in range(grid_size*grid_size):\n",
    "    cur_state = env.cur_state\n",
    "    goal_state = env.goal_state\n",
    "\n",
    "    # initialize training array\n",
    "    start_states = []\n",
    "    start_action = []\n",
    "    # STEP 1: Fast neural network to get general cue for direction\n",
    "    if use_fast:\n",
    "      action_prob = model.predict(np.array(cur_state+goal_state).reshape(-1,4), verbose = 0)\n",
    "      action_probs = action_prob[0]\n",
    "    else:\n",
    "      # set to 0s if not using the fast neural network\n",
    "      action_probs = np.array([0 for _ in range(num_actions)], dtype = 'float')\n",
    "\n",
    "    # STEP 2: Do explore-exploit to choose action\n",
    "\n",
    "    # calculate the exploration parameter, take action_probs and add in count values from instances in memory\n",
    "    action_probs -= [exploration_const*np.sqrt(len(episode_memory[env.cur_state][i])) for i in range(num_actions)]\n",
    "    \n",
    "    # choose the best action via explore-exploit tradeoff\n",
    "    action = np.argmax(action_probs)\n",
    "\n",
    "    # STEP 3: Forward Replay: find out if there is a path in memory to the goal (akin to Model-Based Planning)\n",
    "    if use_slow:\n",
    "      best_minions = PriorityQueue()\n",
    "      best_action = 0\n",
    "\n",
    "      threads = []\n",
    "      sem = threading.Semaphore()\n",
    "      for _ in range(parallel_threads):\n",
    "        t = Thread(target=minion())\n",
    "        threads.append(t)\n",
    "        t.start()\n",
    "\n",
    "      for t in threads:\n",
    "        t.join()\n",
    "      \n",
    "      # Do memory action if found in memory. Also, do a replay to learn the transition of the start state to the end state\n",
    "      if not best_minions.empty():\n",
    "        depth, action, minion_state_history, minion_action_history, final_action = best_minions.get()\n",
    "        # print(f'Found best action {action}. Depth {depth}. Current state {env.cur_state}.')\n",
    "\n",
    "        if depth > 0:\n",
    "          # update fast neural network with each state on future state trajectory as start state, and the final goal state as the goal state, with each action on future action trajectory\n",
    "          start_states.extend([x+goal_state for x in minion_state_history])\n",
    "          start_action.extend([int(x) for x in minion_action_history])\n",
    "\n",
    "    # perform action\n",
    "    obs, rewards, done, _= env.step(action)\n",
    "    next_state = (obs[0], obs[1])\n",
    "\n",
    "    # STEP 4: store transition in memory, delete irrelevant states if environment has changed\n",
    "    for mem in [memory, episode_memory]:\n",
    "      if len(mem[cur_state][action]) > 0 and mem[cur_state][action][-1] != next_state:\n",
    "        mem[cur_state][action] = [next_state]\n",
    "      else:\n",
    "        mem[cur_state][action].append(next_state)\n",
    "\n",
    "    # append current state and action to history\n",
    "    statehistory.append(cur_state)\n",
    "    actionhistory.append(action)\n",
    "\n",
    "    # terminate if completed\n",
    "    if done:\n",
    "      if rewards == 1 and steps+1 < grid_size*grid_size:\n",
    "        print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "\n",
    "      else:\n",
    "        print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "      best_steps.append(env.min_steps)\n",
    "      actual_steps.append(steps+1)\n",
    "\n",
    "      break\n",
    "\n",
    "    # STEP 5: Backward Replay: update fast neural network with next state as goal state, and states before as start state, with current action\n",
    "    if use_fast:\n",
    "      start_states.extend([x+next_state for x in statehistory])\n",
    "      start_action.extend([int(x) for x in actionhistory])\n",
    "      model.fit(start_states, start_action, epochs=1, verbose = False)\n",
    "\n",
    "memory_steps = actual_steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 279
    },
    "id": "3LCpGsr0NNEq",
    "outputId": "79518cbc-7769-4e50-af3b-294f53b14000"
   },
   "outputs": [],
   "source": [
    "# Across all states\n",
    "plt.plot(memory_steps, 'b-', label = 'Fast & Slow')\n",
    "plt.plot(ppo_steps, 'r-', label = 'PPO')\n",
    "plt.plot(trpo_steps, 'y-', label = 'TRPO')\n",
    "plt.plot(a2c_steps, 'g-', label = 'A2C')\n",
    "plt.plot(qlearning_steps, 'c-', label = 'Q-Learning')\n",
    "# plt.plot(best_steps, 'm--', label = 'Best')\n",
    "plt.xlabel('Episode Number')\n",
    "plt.ylabel('Number of Steps')\n",
    "plt.legend(loc = 'upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 279
    },
    "id": "DJbeNNxTpQNp",
    "outputId": "1750c15d-7594-4a84-92c6-90d0426b03a3"
   },
   "outputs": [],
   "source": [
    "# Only first 10 states\n",
    "plt.plot(memory_steps[:10], 'b-', label = 'Fast & Slow')\n",
    "plt.plot(ppo_steps[:10], 'r-', label = 'PPO')\n",
    "plt.plot(trpo_steps[:10], 'y-', label = 'TRPO')\n",
    "plt.plot(a2c_steps[:10], 'g-', label = 'A2C')\n",
    "plt.plot(qlearning_steps[:10], 'c-', label = 'Q-Learning')\n",
    "# plt.plot(best_steps[:10], 'm--', label = 'Best')\n",
    "plt.xlabel('Episode Number')\n",
    "plt.ylabel('Number of Steps')\n",
    "plt.legend(loc = 'upper right')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Co3sCyF3o2Nv",
    "outputId": "bf4147be-824a-460b-afcc-709eb85c3c16"
   },
   "outputs": [],
   "source": [
    "# This shows us the solve rate and steps above minimum for all methods\n",
    "for steps in [memory_steps, ppo_steps, trpo_steps, a2c_steps, qlearning_steps]:\n",
    "  actual_steps = np.array(steps)\n",
    "  print(f'Number of episodes solved overall is {sum(actual_steps<grid_size*grid_size)}')\n",
    "  print(f'Number of steps above minimum overall is {sum(actual_steps)-sum(best_steps)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LaVZr4tnRDWO"
   },
   "source": [
    "# Experiment 2: Solving the grid with random start and end points\n",
    "- Start and end goal differs each episode\n",
    "- Obstacles will change every 50 episodes\n",
    "- Agent has no pre-built model of the environment and has to lookahead without any external help\n",
    "- Agent can rely on stored memory from start of the episodes, and memory can also carry over between episodes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VcRSVBR-h8_H"
   },
   "source": [
    "## Model 1: PPO, TRPO, A2C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "hSl2vxy-h_aC",
    "outputId": "3a4303b6-c0e7-4c30-f0a3-e5264d6201d2"
   },
   "outputs": [],
   "source": [
    "for ALGO in [PPO, TRPO, A2C]:\n",
    "    print (ALGO.__name__)\n",
    "    grid_size = 40\n",
    "    # Custom MLP policy of two layers of size 128 each with Relu activation function\n",
    "    policy_kwargs = dict(activation_fn=torch.nn.ReLU, net_arch=[128, 128])\n",
    "    # Create the environment\n",
    "    i = 0\n",
    "    env = GridEnv(grid_size, start_state = None, goal_state = None, forbidden_squares = GetGrid(i, grid_size), seed = i)\n",
    "    # Create the agent\n",
    "    # set random seed\n",
    "    sb3.common.utils.set_random_seed(0)\n",
    "    if ALGO == A2C:\n",
    "      model = ALGO(\"MlpPolicy\", env, policy_kwargs=policy_kwargs, learning_rate = 0.0001)\n",
    "    else:\n",
    "      model = ALGO(\"MlpPolicy\", env, policy_kwargs=policy_kwargs)\n",
    "    model.save(\"my_env\")\n",
    "\n",
    "    best_steps = []\n",
    "    actual_steps = []\n",
    "    for i in range(100):\n",
    "      # learn for each episode\n",
    "      env = GridEnv(grid_size, start_state = None, goal_state = None, forbidden_squares = GetGrid(i, grid_size), seed = i)\n",
    "      model = ALGO.load(\"my_env\", env=env)\n",
    "      model.learn(total_timesteps=100)\n",
    "      model.save(\"my_env\")\n",
    "      obs = env.reset()\n",
    "      for steps in range(grid_size*grid_size):\n",
    "        action, _states = model.predict(obs)\n",
    "        obs, rewards, done, info = env.step(action)\n",
    "        if done:\n",
    "          if rewards == 1 and steps+1 < grid_size*grid_size:\n",
    "            print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "\n",
    "          else:\n",
    "            print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "          best_steps.append(env.min_steps)\n",
    "          actual_steps.append(steps+1)\n",
    "\n",
    "          break\n",
    "\n",
    "    plt.plot(best_steps, 'b-', label = 'Minimum Steps per Episode')\n",
    "    plt.plot(actual_steps, 'g-', label = 'Actual Steps per Episode')\n",
    "    plt.xlabel('Episode number')\n",
    "    plt.ylabel('Number of Steps')\n",
    "    plt.xticks([i*10 for i in range(11)])\n",
    "    # plt.legend(loc = 'upper right')\n",
    "    plt.show()\n",
    "\n",
    "    actual_steps = np.array(actual_steps)\n",
    "    print(f'Number of steps above minimum for first 50 epochs is {sum(actual_steps[:50])-sum(best_steps[:50])}')\n",
    "    print(f'Number of steps above minimum for last 50 epochs is {sum(actual_steps[50:])-sum(best_steps[50:])}')\n",
    "    print(f'Number of steps above minimum overall is {sum(actual_steps)-sum(best_steps)}')\n",
    "    print(f'Number of episodes solved for first 50 epochs is {sum(actual_steps[:50]<grid_size*grid_size)}')\n",
    "    print(f'Number of episodes solved for last 50 epochs is {sum(actual_steps[50:]<grid_size*grid_size)}')\n",
    "    print(f'Number of episodes solved overall is {sum(actual_steps<grid_size*grid_size)}')\n",
    "\n",
    "    # save the best steps\n",
    "    if ALGO == PPO:\n",
    "      ppo_best_steps = best_steps\n",
    "    elif ALGO == TRPO:\n",
    "      trpo_steps = best_steps\n",
    "    else:\n",
    "      a2c_steps = best_steps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SWokzVQjiFYV"
   },
   "source": [
    "## Model 2: Fast and Slow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 381
    },
    "id": "I6eqmieJrmDI",
    "outputId": "9f48087c-22c6-4522-d125-edf071a57e02"
   },
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "use_fast = True\n",
    "use_slow = True\n",
    "\n",
    "grid_size = 10\n",
    "exploration_const = 1\n",
    "parallel_threads = 100\n",
    "lookahead_depth = 20\n",
    "\n",
    "# create model\n",
    "model = GetModel()\n",
    "\n",
    "# reset memory\n",
    "memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "# store only first action\n",
    "best_steps = []\n",
    "actual_steps = []\n",
    "\n",
    "for i in range(100):\n",
    "  env = GridEnv(grid_size, start_state = None, goal_state = None, forbidden_squares = GetGrid(i, grid_size), seed = i)\n",
    "  statehistory = []\n",
    "  actionhistory = []\n",
    "  episode_memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "  for steps in range(grid_size*grid_size):\n",
    "    cur_state = env.cur_state\n",
    "    goal_state = env.goal_state\n",
    "\n",
    "    # initialize training array\n",
    "    start_states = []\n",
    "    start_action = []\n",
    "    # STEP 1: Fast neural network to get general cue for direction\n",
    "    if use_fast:\n",
    "      action_prob = model.predict(np.array(cur_state+goal_state).reshape(-1,4), verbose = 0)\n",
    "      action_probs = action_prob[0]\n",
    "    else:\n",
    "      # set to 0s if not using the fast neural network\n",
    "      action_probs = np.array([0 for _ in range(num_actions)], dtype = 'float')\n",
    "\n",
    "    # STEP 2: Do explore-exploit to choose action\n",
    "\n",
    "    # calculate the exploration parameter, take action_probs and add in count values from instances in memory\n",
    "    action_probs -= [exploration_const*np.sqrt(len(episode_memory[env.cur_state][i])) for i in range(num_actions)]\n",
    "    \n",
    "    # choose the best action via explore-exploit tradeoff\n",
    "    action = np.argmax(action_probs)\n",
    "\n",
    "    # STEP 3: Forward Replay: find out if there is a path in memory to the goal (akin to Model-Based Planning)\n",
    "    if use_slow:\n",
    "      best_minions = PriorityQueue()\n",
    "      best_action = 0\n",
    "\n",
    "      threads = []\n",
    "      sem = threading.Semaphore()\n",
    "      for _ in range(parallel_threads):\n",
    "        t = Thread(target=minion())\n",
    "        threads.append(t)\n",
    "        t.start()\n",
    "\n",
    "      for t in threads:\n",
    "        t.join()\n",
    "      \n",
    "      # Do memory action if found in memory. Also, do a replay to learn the transition of the start state to the end state\n",
    "      if not best_minions.empty():\n",
    "        depth, action, minion_state_history, minion_action_history, final_action = best_minions.get()\n",
    "        # print(f'Found best action {action}. Depth {depth}. Current state {env.cur_state}.')\n",
    "\n",
    "        if depth > 0:\n",
    "          # update fast neural network with each state on future state trajectory as start state, and the final goal state as the goal state, with each action on future action trajectory\n",
    "          start_states.extend([x+goal_state for x in minion_state_history])\n",
    "          start_action.extend([int(x) for x in minion_action_history])\n",
    "\n",
    "    # perform action\n",
    "    obs, rewards, done, _= env.step(action)\n",
    "    next_state = (obs[0], obs[1])\n",
    "\n",
    "    # STEP 4: store transition in memory, delete irrelevant states if environment has changed\n",
    "    for mem in [memory, episode_memory]:\n",
    "      if len(mem[cur_state][action]) > 0 and mem[cur_state][action][-1] != next_state:\n",
    "        mem[cur_state][action] = [next_state]\n",
    "      else:\n",
    "        mem[cur_state][action].append(next_state)\n",
    "\n",
    "    # append current state and action to history\n",
    "    statehistory.append(cur_state)\n",
    "    actionhistory.append(action)\n",
    "\n",
    "    # terminate if completed\n",
    "    if done:\n",
    "      if rewards == 1 and steps+1 < grid_size*grid_size:\n",
    "        print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "\n",
    "      else:\n",
    "        print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "      best_steps.append(env.min_steps)\n",
    "      actual_steps.append(steps+1)\n",
    "\n",
    "      break\n",
    "\n",
    "    # STEP 5: Backward Replay: update fast neural network with next state as goal state, and states before as start state, with current action\n",
    "    if use_fast:\n",
    "      start_states.extend([x+next_state for x in statehistory])\n",
    "      start_action.extend([int(x) for x in actionhistory])\n",
    "      model.fit(start_states, start_action, epochs=1, verbose = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XnnrP_Gh3y1D"
   },
   "outputs": [],
   "source": [
    "plt.plot(best_steps, 'b-', label = 'Minimum Steps per Episode')\n",
    "plt.plot(actual_steps, 'g-', label = 'Actual Steps per Episode')\n",
    "plt.xlabel('Episode number')\n",
    "plt.ylabel('Number of Steps')\n",
    "plt.xticks([i*10 for i in range(11)])\n",
    "# plt.legend(loc = 'upper right')\n",
    "plt.show()\n",
    "\n",
    "actual_steps = np.array(actual_steps)\n",
    "print(f'Number of steps above minimum for first 50 epochs is {sum(actual_steps[:50])-sum(best_steps[:50])}')\n",
    "print(f'Number of steps above minimum for last 50 epochs is {sum(actual_steps[50:])-sum(best_steps[50:])}')\n",
    "print(f'Number of steps above minimum overall is {sum(actual_steps)-sum(best_steps)}')\n",
    "print(f'Number of episodes solved for first 50 epochs is {sum(actual_steps[:50]<grid_size*grid_size)}')\n",
    "print(f'Number of episodes solved for last 50 epochs is {sum(actual_steps[50:]<grid_size*grid_size)}')\n",
    "print(f'Number of episodes solved overall is {sum(actual_steps<grid_size*grid_size)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Gel08FwGsTly"
   },
   "source": [
    "## Model 3: Without slow memory mechanism"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Nu4JbRTgtcmZ",
    "outputId": "c139a5cc-fe00-4c6c-9823-cc5f6279f8b7"
   },
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "use_fast = True\n",
    "use_slow = False\n",
    "\n",
    "grid_size = 10\n",
    "exploration_const = 1\n",
    "parallel_threads = 100\n",
    "lookahead_depth = 20\n",
    "\n",
    "# create model\n",
    "model = GetModel()\n",
    "\n",
    "# reset memory\n",
    "memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "# store only first action\n",
    "best_steps = []\n",
    "actual_steps = []\n",
    "\n",
    "for i in range(100):\n",
    "  env = GridEnv(grid_size, start_state = None, goal_state = None, forbidden_squares = GetGrid(i, grid_size), seed = i)\n",
    "  statehistory = []\n",
    "  actionhistory = []\n",
    "  episode_memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "  for steps in range(grid_size*grid_size):\n",
    "    cur_state = env.cur_state\n",
    "    goal_state = env.goal_state\n",
    "\n",
    "    # initialize training array\n",
    "    start_states = []\n",
    "    start_action = []\n",
    "    # STEP 1: Fast neural network to get general cue for direction\n",
    "    if use_fast:\n",
    "      action_prob = model.predict(np.array(cur_state+goal_state).reshape(-1,4), verbose = 0)\n",
    "      action_probs = action_prob[0]\n",
    "    else:\n",
    "      # set to 0s if not using the fast neural network\n",
    "      action_probs = np.array([0 for _ in range(num_actions)], dtype = 'float')\n",
    "\n",
    "    # STEP 2: Do explore-exploit to choose action\n",
    "\n",
    "    # calculate the exploration parameter, take action_probs and add in count values from instances in memory\n",
    "    action_probs -= [exploration_const*np.sqrt(len(episode_memory[env.cur_state][i])) for i in range(num_actions)]\n",
    "    \n",
    "    # choose the best action via explore-exploit tradeoff\n",
    "    action = np.argmax(action_probs)\n",
    "\n",
    "    # STEP 3: Forward Replay: find out if there is a path in memory to the goal (akin to Model-Based Planning)\n",
    "    if use_slow:\n",
    "      best_minions = PriorityQueue()\n",
    "      best_action = 0\n",
    "\n",
    "      threads = []\n",
    "      sem = threading.Semaphore()\n",
    "      for _ in range(parallel_threads):\n",
    "        t = Thread(target=minion())\n",
    "        threads.append(t)\n",
    "        t.start()\n",
    "\n",
    "      for t in threads:\n",
    "        t.join()\n",
    "      \n",
    "      # Do memory action if found in memory. Also, do a replay to learn the transition of the start state to the end state\n",
    "      if not best_minions.empty():\n",
    "        depth, action, minion_state_history, minion_action_history, final_action = best_minions.get()\n",
    "        # print(f'Found best action {action}. Depth {depth}. Current state {env.cur_state}.')\n",
    "\n",
    "        if depth > 0:\n",
    "          # update fast neural network with each state on future state trajectory as start state, and the final goal state as the goal state, with each action on future action trajectory\n",
    "          start_states.extend([x+goal_state for x in minion_state_history])\n",
    "          start_action.extend([int(x) for x in minion_action_history])\n",
    "\n",
    "    # perform action\n",
    "    obs, rewards, done, _= env.step(action)\n",
    "    next_state = (obs[0], obs[1])\n",
    "\n",
    "    # STEP 4: store transition in memory, delete irrelevant states if environment has changed\n",
    "    for mem in [memory, episode_memory]:\n",
    "      if len(mem[cur_state][action]) > 0 and mem[cur_state][action][-1] != next_state:\n",
    "        mem[cur_state][action] = [next_state]\n",
    "      else:\n",
    "        mem[cur_state][action].append(next_state)\n",
    "\n",
    "    # append current state and action to history\n",
    "    statehistory.append(cur_state)\n",
    "    actionhistory.append(action)\n",
    "\n",
    "    # terminate if completed\n",
    "    if done:\n",
    "      if rewards == 1 and steps+1 < grid_size*grid_size:\n",
    "        print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "\n",
    "      else:\n",
    "        print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "      best_steps.append(env.min_steps)\n",
    "      actual_steps.append(steps+1)\n",
    "\n",
    "      break\n",
    "\n",
    "    # STEP 5: Backward Replay: update fast neural network with next state as goal state, and states before as start state, with current action\n",
    "    if use_fast:\n",
    "      start_states.extend([x+next_state for x in statehistory])\n",
    "      start_action.extend([int(x) for x in actionhistory])\n",
    "      model.fit(start_states, start_action, epochs=1, verbose = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 383
    },
    "id": "o0eZI_F933Mv",
    "outputId": "d1e4ae03-0ed8-4940-e587-5bd1a6fd7d8d"
   },
   "outputs": [],
   "source": [
    "plt.plot(best_steps, 'b-', label = 'Minimum Steps per Episode')\n",
    "plt.plot(actual_steps, 'g-', label = 'Actual Steps per Episode')\n",
    "plt.xlabel('Episode number')\n",
    "plt.ylabel('Number of Steps')\n",
    "plt.xticks([i*10 for i in range(11)])\n",
    "# plt.legend(loc = 'upper right')\n",
    "plt.show()\n",
    "\n",
    "actual_steps = np.array(actual_steps)\n",
    "print(f'Number of steps above minimum for first 50 epochs is {sum(actual_steps[:50])-sum(best_steps[:50])}')\n",
    "print(f'Number of steps above minimum for last 50 epochs is {sum(actual_steps[50:])-sum(best_steps[50:])}')\n",
    "print(f'Number of steps above minimum overall is {sum(actual_steps)-sum(best_steps)}')\n",
    "print(f'Number of episodes solved for first 50 epochs is {sum(actual_steps[:50]<grid_size*grid_size)}')\n",
    "print(f'Number of episodes solved for last 50 epochs is {sum(actual_steps[50:]<grid_size*grid_size)}')\n",
    "print(f'Number of episodes solved overall is {sum(actual_steps<grid_size*grid_size)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dccTifE6tsM1"
   },
   "source": [
    "## Model 4: Without fast goal-directed mechanism"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RX2mIbqCtjbk",
    "outputId": "bb51810a-f970-4737-9d30-117e24a6386d"
   },
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "use_fast = False\n",
    "use_slow = True\n",
    "\n",
    "grid_size = 10\n",
    "exploration_const = 1\n",
    "parallel_threads = 100\n",
    "lookahead_depth = 20\n",
    "\n",
    "# create model\n",
    "model = GetModel()\n",
    "\n",
    "# reset memory\n",
    "memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "# store only first action\n",
    "best_steps = []\n",
    "actual_steps = []\n",
    "\n",
    "for i in range(100):\n",
    "  env = GridEnv(grid_size, start_state = None, goal_state = None, forbidden_squares = GetGrid(i, grid_size), seed = i)\n",
    "  statehistory = []\n",
    "  actionhistory = []\n",
    "  episode_memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "  for steps in range(grid_size*grid_size):\n",
    "    cur_state = env.cur_state\n",
    "    goal_state = env.goal_state\n",
    "\n",
    "    # initialize training array\n",
    "    start_states = []\n",
    "    start_action = []\n",
    "    # STEP 1: Fast neural network to get general cue for direction\n",
    "    if use_fast:\n",
    "      action_prob = model.predict(np.array(cur_state+goal_state).reshape(-1,4), verbose = 0)\n",
    "      action_probs = action_prob[0]\n",
    "    else:\n",
    "      # set to 0s if not using the fast neural network\n",
    "      action_probs = np.array([0 for _ in range(num_actions)], dtype = 'float')\n",
    "\n",
    "    # STEP 2: Do explore-exploit to choose action\n",
    "\n",
    "    # calculate the exploration parameter, take action_probs and add in count values from instances in memory\n",
    "    action_probs -= [exploration_const*np.sqrt(len(episode_memory[env.cur_state][i])) for i in range(num_actions)]\n",
    "    \n",
    "    # choose the best action via explore-exploit tradeoff\n",
    "    action = np.argmax(action_probs)\n",
    "\n",
    "    # STEP 3: Forward Replay: find out if there is a path in memory to the goal (akin to Model-Based Planning)\n",
    "    if use_slow:\n",
    "      best_minions = PriorityQueue()\n",
    "      best_action = 0\n",
    "\n",
    "      threads = []\n",
    "      sem = threading.Semaphore()\n",
    "      for _ in range(parallel_threads):\n",
    "        t = Thread(target=minion())\n",
    "        threads.append(t)\n",
    "        t.start()\n",
    "\n",
    "      for t in threads:\n",
    "        t.join()\n",
    "      \n",
    "      # Do memory action if found in memory. Also, do a replay to learn the transition of the start state to the end state\n",
    "      if not best_minions.empty():\n",
    "        depth, action, minion_state_history, minion_action_history, final_action = best_minions.get()\n",
    "        # print(f'Found best action {action}. Depth {depth}. Current state {env.cur_state}.')\n",
    "\n",
    "        if depth > 0:\n",
    "          # update fast neural network with each state on future state trajectory as start state, and the final goal state as the goal state, with each action on future action trajectory\n",
    "          start_states.extend([x+goal_state for x in minion_state_history])\n",
    "          start_action.extend([int(x) for x in minion_action_history])\n",
    "\n",
    "    # perform action\n",
    "    obs, rewards, done, _= env.step(action)\n",
    "    next_state = (obs[0], obs[1])\n",
    "\n",
    "    # STEP 4: store transition in memory, delete irrelevant states if environment has changed\n",
    "    for mem in [memory, episode_memory]:\n",
    "      if len(mem[cur_state][action]) > 0 and mem[cur_state][action][-1] != next_state:\n",
    "        mem[cur_state][action] = [next_state]\n",
    "      else:\n",
    "        mem[cur_state][action].append(next_state)\n",
    "\n",
    "    # append current state and action to history\n",
    "    statehistory.append(cur_state)\n",
    "    actionhistory.append(action)\n",
    "\n",
    "    # terminate if completed\n",
    "    if done:\n",
    "      if rewards == 1 and steps+1 < grid_size*grid_size:\n",
    "        print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "\n",
    "      else:\n",
    "        print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "      best_steps.append(env.min_steps)\n",
    "      actual_steps.append(steps+1)\n",
    "\n",
    "      break\n",
    "\n",
    "    # STEP 5: Backward Replay: update fast neural network with next state as goal state, and states before as start state, with current action\n",
    "    if use_fast:\n",
    "      start_states.extend([x+next_state for x in statehistory])\n",
    "      start_action.extend([int(x) for x in actionhistory])\n",
    "      model.fit(start_states, start_action, epochs=1, verbose = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 383
    },
    "id": "Ms-m5fVm35sH",
    "outputId": "6200ba64-7b49-4c78-9aaa-ebc457df2df1"
   },
   "outputs": [],
   "source": [
    "plt.plot(best_steps, 'b-', label = 'Minimum Steps per Episode')\n",
    "plt.plot(actual_steps, 'g-', label = 'Actual Steps per Episode')\n",
    "plt.xlabel('Episode number')\n",
    "plt.ylabel('Number of Steps')\n",
    "plt.xticks([i*10 for i in range(11)])\n",
    "# plt.legend(loc = 'upper right')\n",
    "plt.show()\n",
    "\n",
    "actual_steps = np.array(actual_steps)\n",
    "print(f'Number of steps above minimum for first 50 epochs is {sum(actual_steps[:50])-sum(best_steps[:50])}')\n",
    "print(f'Number of steps above minimum for last 50 epochs is {sum(actual_steps[50:])-sum(best_steps[50:])}')\n",
    "print(f'Number of steps above minimum overall is {sum(actual_steps)-sum(best_steps)}')\n",
    "print(f'Number of episodes solved for first 50 epochs is {sum(actual_steps[:50]<grid_size*grid_size)}')\n",
    "print(f'Number of episodes solved for last 50 epochs is {sum(actual_steps[50:]<grid_size*grid_size)}')\n",
    "print(f'Number of episodes solved overall is {sum(actual_steps<grid_size*grid_size)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "PXanUK4rvno4"
   },
   "source": [
    "## Model 5: Without slow memory and fast goal-directed mechanism\n",
    "Just explore-exploit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "t_FzdkHHuOqJ",
    "outputId": "d203f392-9f5a-437a-fc69-315b5dd3316e"
   },
   "outputs": [],
   "source": [
    "# hyperparameters\n",
    "use_fast = False\n",
    "use_slow = False\n",
    "\n",
    "grid_size = 10\n",
    "exploration_const = 1\n",
    "parallel_threads = 100\n",
    "lookahead_depth = 20\n",
    "\n",
    "# create model\n",
    "model = GetModel()\n",
    "\n",
    "# reset memory\n",
    "memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "# store only first action\n",
    "best_steps = []\n",
    "actual_steps = []\n",
    "\n",
    "for i in range(100):\n",
    "  env = GridEnv(grid_size, start_state = None, goal_state = None, forbidden_squares = GetGrid(i, grid_size), seed = i)\n",
    "  statehistory = []\n",
    "  actionhistory = []\n",
    "  episode_memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "  for steps in range(grid_size*grid_size):\n",
    "    cur_state = env.cur_state\n",
    "    goal_state = env.goal_state\n",
    "\n",
    "    # initialize training array\n",
    "    start_states = []\n",
    "    start_action = []\n",
    "    # STEP 1: Fast neural network to get general cue for direction\n",
    "    if use_fast:\n",
    "      action_prob = model.predict(np.array(cur_state+goal_state).reshape(-1,4), verbose = 0)\n",
    "      action_probs = action_prob[0]\n",
    "    else:\n",
    "      # set to 0s if not using the fast neural network\n",
    "      action_probs = np.array([0 for _ in range(num_actions)], dtype = 'float')\n",
    "\n",
    "    # STEP 2: Do explore-exploit to choose action\n",
    "\n",
    "    # calculate the exploration parameter, take action_probs and add in count values from instances in memory\n",
    "    action_probs -= [exploration_const*np.sqrt(len(episode_memory[env.cur_state][i])) for i in range(num_actions)]\n",
    "    \n",
    "    # choose the best action via explore-exploit tradeoff\n",
    "    action = np.argmax(action_probs)\n",
    "\n",
    "    # STEP 3: Forward Replay: find out if there is a path in memory to the goal (akin to Model-Based Planning)\n",
    "    if use_slow:\n",
    "      best_minions = PriorityQueue()\n",
    "      best_action = 0\n",
    "\n",
    "      threads = []\n",
    "      sem = threading.Semaphore()\n",
    "      for _ in range(parallel_threads):\n",
    "        t = Thread(target=minion())\n",
    "        threads.append(t)\n",
    "        t.start()\n",
    "\n",
    "      for t in threads:\n",
    "        t.join()\n",
    "      \n",
    "      # Do memory action if found in memory. Also, do a replay to learn the transition of the start state to the end state\n",
    "      if not best_minions.empty():\n",
    "        depth, action, minion_state_history, minion_action_history, final_action = best_minions.get()\n",
    "        # print(f'Found best action {action}. Depth {depth}. Current state {env.cur_state}.')\n",
    "\n",
    "        if depth > 0:\n",
    "          # update fast neural network with each state on future state trajectory as start state, and the final goal state as the goal state, with each action on future action trajectory\n",
    "          start_states.extend([x+goal_state for x in minion_state_history])\n",
    "          start_action.extend([int(x) for x in minion_action_history])\n",
    "\n",
    "    # perform action\n",
    "    obs, rewards, done, _= env.step(action)\n",
    "    next_state = (obs[0], obs[1])\n",
    "\n",
    "    # STEP 4: store transition in memory, delete irrelevant states if environment has changed\n",
    "    for mem in [memory, episode_memory]:\n",
    "      if len(mem[cur_state][action]) > 0 and mem[cur_state][action][-1] != next_state:\n",
    "        mem[cur_state][action] = [next_state]\n",
    "      else:\n",
    "        mem[cur_state][action].append(next_state)\n",
    "\n",
    "    # append current state and action to history\n",
    "    statehistory.append(cur_state)\n",
    "    actionhistory.append(action)\n",
    "\n",
    "    # terminate if completed\n",
    "    if done:\n",
    "      if rewards == 1:\n",
    "        print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "        best_steps.append(env.min_steps)\n",
    "        actual_steps.append(steps+1)\n",
    "\n",
    "      else:\n",
    "        print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "        best_steps.append(env.min_steps)\n",
    "        actual_steps.append(steps+1)\n",
    "\n",
    "      break\n",
    "\n",
    "    # STEP 5: Backward Replay: update fast neural network with next state as goal state, and states before as start state, with current action\n",
    "    if use_fast:\n",
    "      start_states.extend([x+next_state for x in statehistory])\n",
    "      start_action.extend([int(x) for x in actionhistory])\n",
    "      model.fit(start_states, start_action, epochs=1, verbose = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 383
    },
    "id": "tENS4Ncnv01Q",
    "outputId": "a454023f-fed5-4fe7-a828-e145fc543270"
   },
   "outputs": [],
   "source": [
    "plt.plot(best_steps, 'b-', label = 'Minimum Steps per Episode')\n",
    "plt.plot(actual_steps, 'g-', label = 'Actual Steps per Episode')\n",
    "plt.xlabel('Episode number')\n",
    "plt.ylabel('Number of Steps')\n",
    "plt.xticks([i*10 for i in range(11)])\n",
    "# plt.legend(loc = 'upper right')\n",
    "plt.show()\n",
    "\n",
    "actual_steps = np.array(actual_steps)\n",
    "print(f'Number of steps above minimum for first 50 epochs is {sum(actual_steps[:50])-sum(best_steps[:50])}')\n",
    "print(f'Number of steps above minimum for last 50 epochs is {sum(actual_steps[50:])-sum(best_steps[50:])}')\n",
    "print(f'Number of steps above minimum overall is {sum(actual_steps)-sum(best_steps)}')\n",
    "print(f'Number of episodes solved for first 50 epochs is {sum(actual_steps[:50]<grid_size*grid_size)}')\n",
    "print(f'Number of episodes solved for last 50 epochs is {sum(actual_steps[50:]<grid_size*grid_size)}')\n",
    "print(f'Number of episodes solved overall is {sum(actual_steps<grid_size*grid_size)}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "jmpii43rV7T1"
   },
   "source": [
    "# Ablation Studies\n",
    "- Vary lookahead depth\n",
    "- Vary number of parallel branch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "PPWl-W6SfhXL",
    "outputId": "7636bc5e-27ee-419d-ab4b-07cc4622f0a4"
   },
   "outputs": [],
   "source": [
    "for lookahead_depth in [5, 10, 50]:\n",
    "\n",
    "    # hyperparameters\n",
    "    use_fast = True\n",
    "    use_slow = True\n",
    "\n",
    "    grid_size = 10\n",
    "    exploration_const = 1\n",
    "    parallel_threads = 100\n",
    "    # lookahead_depth = 20\n",
    "\n",
    "    # create model\n",
    "    model = GetModel()\n",
    "\n",
    "    # reset memory\n",
    "    memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "    # store only first action\n",
    "    best_steps = []\n",
    "    actual_steps = []\n",
    "\n",
    "    for i in range(100):\n",
    "      env = GridEnv(grid_size, start_state = None, goal_state = None, forbidden_squares = GetGrid(i, grid_size), seed = i)\n",
    "      statehistory = []\n",
    "      actionhistory = []\n",
    "      episode_memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "      for steps in range(grid_size*grid_size):\n",
    "        cur_state = env.cur_state\n",
    "        goal_state = env.goal_state\n",
    "\n",
    "        # initialize training array\n",
    "        start_states = []\n",
    "        start_action = []\n",
    "        # STEP 1: Fast neural network to get general cue for direction\n",
    "        if use_fast:\n",
    "          action_prob = model.predict(np.array(cur_state+goal_state).reshape(-1,4), verbose = 0)\n",
    "          action_probs = action_prob[0]\n",
    "        else:\n",
    "          # set to 0s if not using the fast neural network\n",
    "          action_probs = np.array([0 for _ in range(num_actions)], dtype = 'float')\n",
    "\n",
    "        # STEP 2: Do explore-exploit to choose action\n",
    "\n",
    "        # calculate the exploration parameter, take action_probs and add in count values from instances in memory\n",
    "        action_probs -= [exploration_const*np.sqrt(len(episode_memory[env.cur_state][i])) for i in range(num_actions)]\n",
    "        \n",
    "        # choose the best action via explore-exploit tradeoff\n",
    "        action = np.argmax(action_probs)\n",
    "\n",
    "        # STEP 3: Forward Replay: find out if there is a path in memory to the goal (akin to Model-Based Planning)\n",
    "        if use_slow:\n",
    "          best_minions = PriorityQueue()\n",
    "          best_action = 0\n",
    "\n",
    "          threads = []\n",
    "          sem = threading.Semaphore()\n",
    "          for _ in range(parallel_threads):\n",
    "            t = Thread(target=minion())\n",
    "            threads.append(t)\n",
    "            t.start()\n",
    "\n",
    "          for t in threads:\n",
    "            t.join()\n",
    "          \n",
    "          # Do memory action if found in memory. Also, do a replay to learn the transition of the start state to the end state\n",
    "          if not best_minions.empty():\n",
    "            depth, action, minion_state_history, minion_action_history, final_action = best_minions.get()\n",
    "            # print(f'Found best action {action}. Depth {depth}. Current state {env.cur_state}.')\n",
    "\n",
    "            if depth > 0:\n",
    "              # update fast neural network with each state on future state trajectory as start state, and the final goal state as the goal state, with each action on future action trajectory\n",
    "              start_states.extend([x+goal_state for x in minion_state_history])\n",
    "              start_action.extend([int(x) for x in minion_action_history])\n",
    "\n",
    "        # perform action\n",
    "        obs, rewards, done, _= env.step(action)\n",
    "        next_state = (obs[0], obs[1])\n",
    "\n",
    "        # STEP 4: store transition in memory, delete irrelevant states if environment has changed\n",
    "        for mem in [memory, episode_memory]:\n",
    "          if len(mem[cur_state][action]) > 0 and mem[cur_state][action][-1] != next_state:\n",
    "            mem[cur_state][action] = [next_state]\n",
    "          else:\n",
    "            mem[cur_state][action].append(next_state)\n",
    "\n",
    "        # append current state and action to history\n",
    "        statehistory.append(cur_state)\n",
    "        actionhistory.append(action)\n",
    "\n",
    "        # terminate if completed\n",
    "        if done:\n",
    "          if rewards == 1:\n",
    "            print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "            best_steps.append(env.min_steps)\n",
    "            actual_steps.append(steps+1)\n",
    "\n",
    "          else:\n",
    "            print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "            best_steps.append(env.min_steps)\n",
    "            actual_steps.append(steps+1)\n",
    "\n",
    "          break\n",
    "\n",
    "        # STEP 5: Backward Replay: update fast neural network with next state as goal state, and states before as start state, with current action\n",
    "        if use_fast:\n",
    "          start_states.extend([x+next_state for x in statehistory])\n",
    "          start_action.extend([int(x) for x in actionhistory])\n",
    "          model.fit(start_states, start_action, epochs=1, verbose = False)\n",
    "\n",
    "    plt.plot(best_steps, 'b-', label = 'Minimum Steps per Episode')\n",
    "    plt.plot(actual_steps, 'g-', label = 'Actual Steps per Episode')\n",
    "    plt.xlabel('Episode number')\n",
    "    plt.ylabel('Number of Steps')\n",
    "    plt.xticks([i*10 for i in range(11)])\n",
    "    # plt.legend(loc = 'upper right')\n",
    "    plt.show()\n",
    "\n",
    "    actual_steps = np.array(actual_steps)\n",
    "    print(f'Number of steps above minimum for first 50 epochs is {sum(actual_steps[:50])-sum(best_steps[:50])}')\n",
    "    print(f'Number of steps above minimum for last 50 epochs is {sum(actual_steps[50:])-sum(best_steps[50:])}')\n",
    "    print(f'Number of steps above minimum overall is {sum(actual_steps)-sum(best_steps)}')\n",
    "    print(f'Number of episodes solved for first 50 epochs is {sum(actual_steps[:50]<grid_size*grid_size)}')\n",
    "    print(f'Number of episodes solved for last 50 epochs is {sum(actual_steps[50:]<grid_size*grid_size)}')\n",
    "    print(f'Number of episodes solved overall is {sum(actual_steps<grid_size*grid_size)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "31ia4OcsWJF1",
    "outputId": "f86fef90-3d61-4393-bdb8-0e4f7589b553"
   },
   "outputs": [],
   "source": [
    "for parallel_threads in [10, 50, 200]:\n",
    "\n",
    "    # hyperparameters\n",
    "    use_fast = True\n",
    "    use_slow = True\n",
    "\n",
    "    grid_size = 10\n",
    "    exploration_const = 1\n",
    "    # parallel_threads = 100\n",
    "    lookahead_depth = 20\n",
    "\n",
    "    # create model\n",
    "    model = GetModel()\n",
    "\n",
    "    # reset memory\n",
    "    memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "    # store only first action\n",
    "    best_steps = []\n",
    "    actual_steps = []\n",
    "\n",
    "    for i in range(100):\n",
    "      env = GridEnv(grid_size, start_state = None, goal_state = None, forbidden_squares = GetGrid(i, grid_size), seed = i)\n",
    "      statehistory = []\n",
    "      actionhistory = []\n",
    "      episode_memory = defaultdict(lambda: defaultdict(lambda: []))\n",
    "\n",
    "      for steps in range(grid_size*grid_size):\n",
    "        cur_state = env.cur_state\n",
    "        goal_state = env.goal_state\n",
    "\n",
    "        # initialize training array\n",
    "        start_states = []\n",
    "        start_action = []\n",
    "        # STEP 1: Fast neural network to get general cue for direction\n",
    "        if use_fast:\n",
    "          action_prob = model.predict(np.array(cur_state+goal_state).reshape(-1,4), verbose = 0)\n",
    "          action_probs = action_prob[0]\n",
    "        else:\n",
    "          # set to 0s if not using the fast neural network\n",
    "          action_probs = np.array([0 for _ in range(num_actions)], dtype = 'float')\n",
    "\n",
    "        # STEP 2: Do explore-exploit to choose action\n",
    "\n",
    "        # calculate the exploration parameter, take action_probs and add in count values from instances in memory\n",
    "        action_probs -= [exploration_const*np.sqrt(len(episode_memory[env.cur_state][i])) for i in range(num_actions)]\n",
    "        \n",
    "        # choose the best action via explore-exploit tradeoff\n",
    "        action = np.argmax(action_probs)\n",
    "\n",
    "        # STEP 3: Forward Replay: find out if there is a path in memory to the goal (akin to Model-Based Planning)\n",
    "        if use_slow:\n",
    "          best_minions = PriorityQueue()\n",
    "          best_action = 0\n",
    "\n",
    "          threads = []\n",
    "          sem = threading.Semaphore()\n",
    "          for _ in range(parallel_threads):\n",
    "            t = Thread(target=minion())\n",
    "            threads.append(t)\n",
    "            t.start()\n",
    "\n",
    "          for t in threads:\n",
    "            t.join()\n",
    "          \n",
    "          # Do memory action if found in memory. Also, do a replay to learn the transition of the start state to the end state\n",
    "          if not best_minions.empty():\n",
    "            depth, action, minion_state_history, minion_action_history, final_action = best_minions.get()\n",
    "            # print(f'Found best action {action}. Depth {depth}. Current state {env.cur_state}.')\n",
    "\n",
    "            if depth > 0:\n",
    "              # update fast neural network with each state on future state trajectory as start state, and the final goal state as the goal state, with each action on future action trajectory\n",
    "              start_states.extend([x+goal_state for x in minion_state_history])\n",
    "              start_action.extend([int(x) for x in minion_action_history])\n",
    "\n",
    "        # perform action\n",
    "        obs, rewards, done, _= env.step(action)\n",
    "        next_state = (obs[0], obs[1])\n",
    "\n",
    "        # STEP 4: store transition in memory, delete irrelevant states if environment has changed\n",
    "        for mem in [memory, episode_memory]:\n",
    "          if len(mem[cur_state][action]) > 0 and mem[cur_state][action][-1] != next_state:\n",
    "            mem[cur_state][action] = [next_state]\n",
    "          else:\n",
    "            mem[cur_state][action].append(next_state)\n",
    "\n",
    "        # append current state and action to history\n",
    "        statehistory.append(cur_state)\n",
    "        actionhistory.append(action)\n",
    "\n",
    "        # terminate if completed\n",
    "        if done:\n",
    "          if rewards == 1:\n",
    "            print(f'Run {i+1} completed in {steps+1} steps. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "            best_steps.append(env.min_steps)\n",
    "            actual_steps.append(steps+1)\n",
    "\n",
    "          else:\n",
    "            print(f'Run {i+1} not completed. Minimum steps: {env.min_steps}, Start state: {env.start_state}, Goal state: {env.goal_state}')\n",
    "            best_steps.append(env.min_steps)\n",
    "            actual_steps.append(steps+1)\n",
    "\n",
    "          break\n",
    "\n",
    "        # STEP 5: Backward Replay: update fast neural network with next state as goal state, and states before as start state, with current action\n",
    "        if use_fast:\n",
    "          start_states.extend([x+next_state for x in statehistory])\n",
    "          start_action.extend([int(x) for x in actionhistory])\n",
    "          model.fit(start_states, start_action, epochs=1, verbose = False)\n",
    "\n",
    "    plt.plot(best_steps, 'b-', label = 'Minimum Steps per Episode')\n",
    "    plt.plot(actual_steps, 'g-', label = 'Actual Steps per Episode')\n",
    "    plt.xlabel('Episode number')\n",
    "    plt.ylabel('Number of Steps')\n",
    "    plt.xticks([i*10 for i in range(11)])\n",
    "    # plt.legend(loc = 'upper right')\n",
    "    plt.show()\n",
    "\n",
    "    actual_steps = np.array(actual_steps)\n",
    "    print(f'Number of steps above minimum for first 50 epochs is {sum(actual_steps[:50])-sum(best_steps[:50])}')\n",
    "    print(f'Number of steps above minimum for last 50 epochs is {sum(actual_steps[50:])-sum(best_steps[50:])}')\n",
    "    print(f'Number of steps above minimum overall is {sum(actual_steps)-sum(best_steps)}')\n",
    "    print(f'Number of episodes solved for first 50 epochs is {sum(actual_steps[:50]<grid_size*grid_size)}')\n",
    "    print(f'Number of episodes solved for last 50 epochs is {sum(actual_steps[50:]<grid_size*grid_size)}')\n",
    "    print(f'Number of episodes solved overall is {sum(actual_steps<grid_size*grid_size)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0tvzGm3NXuQs"
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}