GCSLDemo.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "GCSLDemo.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyMrw2yPptrb93hgMoY8B3jr",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/notdibya/gcsl/blob/master/GCSLDemo.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m0j9Anlrb8zV",
        "colab_type": "text"
      },
      "source": [
        "# Goal-Conditioned Supervised Learning (GCSL)\n",
        "\n",
        "<a href=\"https://ibb.co/gz3HWFD\"><img src=\"https://i.ibb.co/J5K0pFz/gcsl-figure-arxiv.png\" alt=\"gcsl-figure-arxiv\" border=\"0\" /></a>\n",
        "\n",
        "This notebook provides a brief introduction to goal-conditioned supervised learning (GCSL), as introduced in \n",
        "> Dibya Ghosh, Abhishek Gupta, Justin Fu, Ashwin Reddy, Coline Devin, Benjamin Eysenbach, Sergey Levine. *Learning To Reach Goals Without Reinforcement Learning.*\n",
        "\n",
        "In this notebook, we develop a simple instantiation of the GCSL algorithm for solving goal-reaching tasks on a simple 2d navigation task with discrete actions to explain the algorithm. For a codebase with more complete features, please see the Github repository linked below.\n",
        "\n",
        "Other useful links:\n",
        "- [ArXiv paper](https://arxiv.org/abs/1912.06088)\n",
        "- [Github](https://github.com/notdibya/gcsl/)\n",
        "- [Interactive Colab](https://colab.research.google.com/github/notdibya/gcsl/blob/master/GCSLDemo.ipynb) (follow along with the code)\n",
        "\n",
        "Let's get started!\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p3sBsJDaXpel",
        "colab_type": "text"
      },
      "source": [
        "## The Goal-Reaching Problem\n",
        "\n",
        "Goal-conditioned supervised learning (GCSL) is an algorithm for solving the problem of goal-reaching in an environment (formally, an MDP). \n",
        "\n",
        "At the beginning of each episode, a goal state $g$ is selected randomly in an environment and shown to the agent. The agent's goal is to reach this goal state by the end of the episode, which terminates after some fixed amount of time has passed.\n",
        "\n",
        "Each step, the agent must choose an action to take, given three pieces of information:\n",
        "\n",
        "1. The **state** the agent is currently at\n",
        "2. The **goal** the agent is trying to get to\n",
        "3. The amount of time the agent has left to reach the goal (the **horizon**)\n",
        "\n",
        "For example, here would be the pseudocode of an agent that takes actions randomly.\n",
        "\n",
        "```python\n",
        "class RandomAgent:\n",
        "  def get_action(self, state, goal, horizon):\n",
        "    return np.random.choice(n_actions) # Choose randomly\n",
        "```\n",
        "\n",
        "How can we train an agent $\\pi: (s, g, h) \\mapsto a$ to learn to reach goals in the environment? \n",
        "\n",
        "**RL**: The dominant paradigm is reinforcement learning (RL), where you define a *reward* function $r(s, g) = \\begin{cases} 1 & s=g \\\\ 0 & \\text{otherwise}\\end{cases}$, and optimize this agent using a value-based RL algorithm (DQN, DDPG, TD3, SAC). This can get quite complicated, and is often difficult to learn in practice.\n",
        "\n",
        "\n",
        "**Imitation Learning**: If instead, we had access (somehow) to an optimal policy $\\pi^*$, that given any state $s$, goal $g$ and time horizon $h$, could return the optimal action $a^*$, we could just use supervised learning to try to match the actions of the expert (MLE). If actions were discrete, this would correspond to a classification problem, and if actions were continuous, this would correspond to a regression problem."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-K0xJzWQdKOz",
        "colab_type": "text"
      },
      "source": [
        "## Goal-Conditioned Supervised Learning (GCSL)\n",
        "\n",
        "Of course, we don't have access to the optimal policy $\\pi^*$ (wouldn't life be simple, then?). What can we do in the absence of this optimal policy? GCSL leverages the insight that even if we don't have access to the optimal policy, we can use **the principle of hindsight** to turn data that the agent has collected into \"optimal\" data.\n",
        "\n",
        "\n",
        "**The Principle of Hindsight**: Suppose the agent has just executed the following episode. It was commanded to reach the goal $g^*$, and ends up seeing the following states and taking the following actions\n",
        "\n",
        "$$s_0, a_0, s_1, a_1, s_2, \\dots s_{50}, a_{50}$$\n",
        "\n",
        "The agent might have done a terrible job at reaching $g^*$. It could be that the final state $s_{50}$ and $g^*$ are on opposite sides of the room. However, we know that the agent was able to successfully reach $s_{50}$, and we can use the data collected in the trajectory to teach the agent how to better get to $s_{50}$. \n",
        "\n",
        "For that matter, the agent was also able to reach $s_{49}, s_{48}, \\dots$, and we can use this data to learn how to reach all of these states  more accurately.\n",
        "\n",
        "More precisely, consider any two timesteps $t_1$ and $t_2$ ($t_1 < t_2$). We know that after the agent took action $a_{t_1}$ in state $s_{t_1}$, it reached the state $s_{t_2}$ after $t_2-t_1$ timesteps. This means that we can treat $a_{t_1}$ as being **optimal** behavior at $s_{t_1}$ when trying to reach $s_{t_2}$ in $t_2 - t_1$ timesteps.\n",
        "\n",
        "**GCSL:** The algorithm performs a simple three-step procedure:\n",
        "\n",
        "1. Collects a trajectory with the current agent\n",
        "2. Picks many choices of $t_1,t_2$ to create a dataset $\\mathcal{D}$ of optimal data $\\mathbf{x} = (s=s_{t_1}, g=s_{t_2}, h=t_2-t_1)$ and $\\mathbf{y} = a_{t_1}$.\n",
        "3. Performs supervised learning (classification or regression depending on discrete actions) to have the policy produce output $\\mathbf{y}$ when provided input $\\mathbf{x}$.\n",
        "\n",
        "Formally, the policy is performing maximum-likelihood estimation (MLE) via supervised learning.\n",
        "\n",
        "$$\\arg\\max_\\theta \\mathbb{E}_{((s,g,h), a) \\sim \\mathcal{D}}[\\log \\pi_\\theta(a | s=s, g=g, h=h)]$$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E_6kifzpf9iq",
        "colab_type": "text"
      },
      "source": [
        "## GCSL in Code\n",
        "\n",
        "Now having provided a general overview of the algorithm, let's see a simple implementation of GCSL with a neural network agent trained in Pytorch, learning to reach goals in a pointmass navigation environment. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AaRJqidtb7_V",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import gym"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y3mSscqSiUBu",
        "colab_type": "text"
      },
      "source": [
        "## The Environment\n",
        "We use a navigation task in a room to illustrate GCSL. The agent, which gets to observe its $(x,y)$ coordinates, can take one of 5 actions: go up, right, down, left or stay still.\n",
        "\n",
        "\n",
        "The environment defined below is simply a normal Gym env, but with one new function:\n",
        "- `env.sample_goal()`: To define the goal reaching task, we need a notion of what goals the agent should try to reach. This function samples one such *desired goal*. In this case, a goal position is chosen randomly in the room.\n",
        "\n",
        "We also implement a helper function for visualizing data from this environment. The agent starts in the blue square, and attempts to reach the orange star."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iYMzlA9kd1C-",
        "colab_type": "code",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 265
        },
        "outputId": "e2f47170-c705-4647-d848-ac76078de48c"
      },
      "source": [
        "# @title Environment Definition\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import gym\n",
        "\n",
        "class PointmassEnv(gym.Env):\n",
        "  def __init__(self):\n",
        "    # 2-d coordinates\n",
        "    self.observation_space = gym.spaces.Box(-1, 1, shape=(2,), dtype=np.float32)\n",
        "    self.position = np.zeros(2) \n",
        "    # Up, Right, Down, Left, Stay\n",
        "    self.action_space = gym.spaces.Discrete(5)\n",
        "\n",
        "  def reset(self):\n",
        "    self.position = np.zeros(2)\n",
        "    return self.position\n",
        "\n",
        "  def action_to_direction(self, a):\n",
        "    actions = [\n",
        "               np.array([0, 1]), # Up\n",
        "               np.array([1, 0]), # Right\n",
        "               np.array([0, -1]), # Down\n",
        "               np.array([-1, 0]), # Left,\n",
        "               np.array([0, 0]), # Stay still\n",
        "    ]\n",
        "    return actions[a]\n",
        "  \n",
        "  def step(self, a):\n",
        "    direction = self.action_to_direction(a)\n",
        "    step = 0.05 * direction + np.random.randn() * 0.01 # Take a noisy step in direction\n",
        "    self.position = self.position + step\n",
        "    self.position = np.clip(self.position, -1, 1) # Clip to prevent object from escaping\n",
        "    return (self.position, # State\n",
        "          0,     # Reward (not necessary for GCSL)\n",
        "          False, # Done flag\n",
        "          dict()) # Additional info \n",
        "\n",
        "  def sample_goal(self):\n",
        "    return np.random.rand(2) * 2 - 1 # Sample uniformly from [-1, 1]\n",
        "\n",
        "\n",
        "def plot_trajectory(trajectory, ax=None):\n",
        "  if ax is None:\n",
        "    ax = plt.gca()\n",
        "  # Draw path\n",
        "  ax.plot(*trajectory['states'].T)\n",
        "  # Draw goal\n",
        "  ax.scatter(0, 0, s=200, marker='s')\n",
        "  ax.scatter(trajectory['desired_goal'][0], trajectory['desired_goal'][1], s=400, marker='*')\n",
        "  # Draw boundary\n",
        "  ax.vlines([-1, 1], -1, 1)\n",
        "  ax.hlines([-1, 1], -1, 1)\n",
        "\n",
        "  ax.set_xlim(-1.05, 1.05)\n",
        "  ax.set_ylim(-1.05, 1.05)\n",
        "  ax.axis('off')\n",
        "\n",
        "env = PointmassEnv()\n",
        "plot_trajectory(dict(states=env.reset()[None], desired_goal=env.sample_goal()))\n",
        "plt.axis('on');"
      ],
      "execution_count": 24,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYUAAAD4CAYAAAAD6PrjAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAXZklEQVR4nO3dfbBcdZ3n8ffHQILO1IzBZJkISGCNIhZbQe8w7lo1ovIQrVmCK2qwLIODkx1HZrfW0hKKKp1lxh2c+YN5kF1NYQR1FnAy5Zop4zLhIVq1inKtjTxZkGtch2RQroBUKRoI+e4ffXJsLvf59O2b4PtV1dXn/B7O/ebcTn/u6XO6O1WFJEkAz1vsAiRJhw9DQZLUMhQkSS1DQZLUMhQkSa2jFruA+VixYkWtXr16scuQpCPKt7/97R9X1crpxhyRobB69WpGR0cXuwxJOqIk+cFMYwby8lGSLUkeTnLPFP1J8jdJxpLcleRVfX0bk+xubhsHUY8kaX4GdU7hOmDdNP1vAtY0t03A/wBIcizwUeB3gDOBjyZZPqCaJElzNJBQqKqvAY9OM2Q98NnquQN4YZJVwHnAjqp6tKoeA3YwfbhIkhbQsM4pHA882Le+t2mbqr2Ts846q+smJOmIs3Pnzs7bOGIuSU2yKcloktHx8fEpx5111lns2rVriJVJ0uLbtWvXQP4gHtaRwj7gxL71E5q2fcBZE9p3TraBqtoMbAYYGRmZ9lP81q5dO5DElKQjxaBeIRnWkcI24N3NVUivAR6vqoeAm4FzkyxvTjCf27RJkhbBQI4UktxA7y/+FUn20rui6GiAqvoksB14MzAGPAG8p+l7NMmfAnc2m7qyqqY7YS1JWkADCYWqumiG/gLeP0XfFmDLIOqQJHVzxJxoliQtPENBktQyFCRJLUNBktQyFCRJLUNBktQyFCRJLUNBktQyFCRJLUNBktQyFCRJLUNBktQyFCRJLUNBktQyFCRJLUNBktQyFCRJrYGEQpJ1Se5PMpbkskn6r06yq7k9kOQnfX1P9/VtG0Q9kqT56fx1nEmWANcA5wB7gTuTbKuq+w6Nqar/0jf+j4Ez+jbx86pa27UOSVJ3gzhSOBMYq6o9VfUkcCOwfprxFwE3DODnSpIGbBChcDzwYN/63qbtWZKcBJwM3NbXfEyS0SR3JLlgqh+SZFMzbnR8fHwAZUuSJhr2ieYNwNaqerqv7aSqGgHeCfxVkn892cSq2lxVI1U1snLlymHUKkm/cgYRCvuAE/vWT2jaJrOBCS8dVdW+5n4PsJNnnm+QJA3RIELhTmBNkpOTLKX3xP+sq4iSnAosB77R17Y8ybJmeQXwWuC+iXMlScPR+eqjqjqQ5FLgZmAJsKWq7k1yJTBaVYcCYgNwY1VV3/RXAJ9KcpBeQF3Vf9WSJGm4OocCQFVtB7ZPaPvIhPU/mWTe14HTB1GDJKk739EsSWoZCpKklqEgSWoZCpKklqEgSWoZCpKklqEgSWoZCpKklqEgSWoZCpKklqEgSWoZCpKklqEgSWoZCpKklqEgSWoZCpKklqEgSWoNJBSSrEtyf5KxJJdN0n9xkvEku5rbe/v6NibZ3dw2DqIeSdL8dP46ziRLgGuAc4C9wJ1Jtk3yXcs3VdWlE+YeC3wUGAEK+HYz97GudUmS5m4QRwpnAmNVtaeqngRuBNbPcu55wI6qerQJgh3AugHUJEmah0GEwvHAg33re5u2id6a5K4kW5OcOMe5JNmUZDTJ6Pj4+ADKliRNNKwTzf8IrK6qf0PvaOD6uW6gqjZX1UhVjaxcuXLgBUqSBhMK+4AT+9ZPaNpaVfVIVe1vVq8FXj3buZKk4RlEKNwJrElycpKlwAZgW/+AJKv6Vs8Hvtss3wycm2R5kuXAuU2bJGkRdL76qKoOJLmU3pP5EmBLVd2b5EpgtKq2Af8pyfnAAeBR4OJm7qNJ/pResABcWVWPdq1JkjQ/nUMBoKq2A9sntH2kb/ly4PIp5m4BtgyiDklSN76jWZLUMhQkSS1DQZLUMhQkSS1DQZLUMhQkSS1DQZLUMhQkSS1DQZLUMhQkSS1DQZLUMhQkSS1DQZLUMhQkSS1DQZLUMhQkSS1DQZLUGkgoJFmX5P4kY0kum6T/A0nuS3JXkluTnNTX93SSXc1t28S5kqTh6fx1nEmWANcA5wB7gTuTbKuq+/qG/V9gpKqeSPI+4C+AdzR9P6+qtV3rkCR1N4gjhTOBsaraU1VPAjcC6/sHVNXtVfVEs3oHcMIAfq4kacAGEQrHAw/2re9t2qZyCfCVvvVjkowmuSPJBVNNSrKpGTc6Pj7erWJJ0qQ6v3w0F0neBYwAr+trPqmq9iU5Bbgtyd1V9b2Jc6tqM7AZYGRkpIZSsCT9ihnEkcI+4MS+9ROatmdIcjZwBXB+Ve0/1F5V+5r7PcBO4IwB1CRJmodBhMKdwJokJydZCmwAnnEVUZIzgE/RC4SH+9qXJ1nWLK8AXgv0n6CWJA1R55ePqupAkkuBm4ElwJaqujfJlcBoVW0D/hL4deDvkwD8c1WdD7wC+FSSg/QC6qoJVy1JkoZoIOcUqmo7sH1C20f6ls+eYt7XgdMHUYMkqTvf0SxJahkKkqSWoSBJahkKkqSWoSBJahkKkqSWoSBJahkKkqSWoSBJahkKkqSWoSBJahkKkqSWoSBJag31m9ek54JXfvR/87P9T8953q8tW8K9/3XdAlQkDY5HCtIczScQusyThslQkCS1DAVJUmsgoZBkXZL7k4wluWyS/mVJbmr6v5lkdV/f5U37/UnOG0Q9kqT56RwKSZYA1wBvAk4DLkpy2oRhlwCPVdVLgauBjzdzTwM2AK8E1gH/vdmeJGkRDOJI4UxgrKr2VNWTwI3A+glj1gPXN8tbgTcmSdN+Y1Xtr6rvA2PN9iRJi2AQoXA88GDf+t6mbdIxVXUAeBx40SznApBkU5LRJKPj4+MDKFuSNNERc6K5qjZX1UhVjaxcuXKxy5Gk56RBhMI+4MS+9ROatknHJDkK+E3gkVnOlSQNySBC4U5gTZKTkyyld+J424Qx24CNzfKFwG1VVU37hubqpJOBNcC3BlCTJGkeOn/MRVUdSHIpcDOwBNhSVfcmuRIYraptwKeBzyUZAx6lFxw0474A3AccAN5fVb7tU5IWyUA++6iqtgPbJ7R9pG/5F8Dbppj7MeBjg6hDktTNEXOiWZK08AwFaY5+bdn83l8533nSMPnR2dIc+fHXei7zSEGS1DIUJEktQ0GS1DIUJEktQ0GS1DIUJEktQ0GS1DIUJEktQ0GS1DIUJEktQ0GS1DIUJEktQ0GS1DIUJEmtTqGQ5NgkO5Lsbu6XTzJmbZJvJLk3yV1J3tHXd12S7yfZ1dzWdqlHktRN1yOFy4Bbq2oNcGuzPtETwLur6pXAOuCvkrywr/9DVbW2ue3qWI8kqYOuobAeuL5Zvh64YOKAqnqgqnY3y/8CPAys7PhzJUkLoGsoHFdVDzXLPwSOm25wkjOBpcD3+po/1rysdHWSZdPM3ZRkNMno+Ph4x7IlSZOZMRSS3JLknklu6/vHVVUBNc12VgGfA95TVQeb5suBU4HfBo4FPjzV/KraXFUjVTWycqUHGpK0EGb8juaqOnuqviQ/SrKqqh5qnvQfnmLcbwBfBq6oqjv6tn3oKGN/ks8AH5xT9ZKkger68tE2YGOzvBH40sQBSZYCXwQ+W1VbJ/Stau5D73zEPR3rkSR10DUUrgLOSbIbOLtZJ8lIkmubMW8Hfhe4eJJLT/8uyd3A3cAK4M861iNJ6mDGl4+mU1WPAG+cpH0UeG+z/Hng81PMf0OXny9JGizf0SxJahkKkqSWoSBJahkKkqSWoSBJahkKkqSWoSBJahkKkqSWoSBJahkKkqSWoSBJahkKkqSWoSBJahkKkqSWoSBJahkKkgbvF4/DJ87s3euIYihIGrwHboYf3w8P/NNiV6I56hQKSY5NsiPJ7uZ++RTjnu77Ks5tfe0nJ/lmkrEkNzXf5yzpSLfrfz7zXkeMrkcKlwG3VtUa4NZmfTI/r6q1ze38vvaPA1dX1UuBx4BLOtYjabHt/yn84Ou95R/8H3jyZ4tbj+akayisB65vlq8HLpjtxCQB3gBsnc98SYepsR2w5Oje8pKjYfeOxa1Hc9I1FI6rqoea5R8Cx00x7pgko0nuSHLoif9FwE+q6kCzvhc4fqoflGRTs43R8fHxjmVLWjDfuQGe/Glv+cmf9tZ1xDhqpgFJbgF+a5KuK/pXqqqS1BSbOamq9iU5Bbgtyd3AnC5LqKrNwGaAkZGRqX6OpIV28CA88ePJ+55+CvZ89Zlte3bC4/t+efQw0QtWwPO85uVwMWMoVNXZU/Ul+VGSVVX1UJJVwMNTbGNfc78nyU7gDOAfgBcmOao5WjgB2DePf4OkYbrvf8HW90CWTP5EnyXPXv/bVz173NNPQT0Nb7sOXvmWBSlVc9c1nrcBG5vljcCXJg5IsjzJsmZ5BfBa4L6qKuB24MLp5ks6zLzyLfB7V8NRS+HAfjjwi2fenppwYvmpnz17zIH9vfm/dzWc5qnEw0nXULgKOCfJbuDsZp0kI0mubca8AhhN8h16IXBVVd3X9H0Y+ECSMXrnGD7dsR5JCy2Bkd+HP9gJy0+Co54/t/lHPb837w929raTLESVmqcZXz6aTlU9ArxxkvZR4L3N8teB06eYvwc4s0sNkhbJvzoV/ugO2P4huGcrPPXzmecc/Xw4/W3wpr+Eo49Z+Bo1Z57dkTR/Rz8f1n8C3vrp3vK0Y18Ab90C5/+tgXAYMxQkdfeydc8+wTxRngcvO2849WjeDAVJ3R16B/NM/vkbC1uHOjMUJHV3103w1BO/XF+y9Jn30Ou/6wvDrUtzZihI6ubgQfjuNqiDvfWjXwCnvB7+8129+6Nf0Guvg733OBw8uHi1akaGgqRu9t4JB58G0jvZvO7P4Z039S47fedNcN5/a05Cpzdu3+hiV6xpGAqSurlna++TUJevhk1fhVdf/Mv3HiQw8p5e+/LVvXF3b51mY1pshoKkbvbeCa/e2HvPwsqXTz5m5ct7/a/eCHu/Ndz6NCed3rwmSWzaObtxRx8D//6vF7ISDYBHCpKklqEgSWoZCpKklqEgSWoZCpKklqEgSWoZCpKklqEgSWp1CoUkxybZkWR3c798kjGvT7Kr7/aLJBc0fdcl+X5f39ou9UiSuul6pHAZcGtVrQFubdafoapur6q1VbUWeAPwBPBPfUM+dKi/qnZ1rEeS1EHXUFgPXN8sXw9cMMP4C4GvVNUTM4yTJC2CrqFwXFU91Cz/EDhuhvEbgBsmtH0syV1Jrk6ybKqJSTYlGU0yOj4+3qFkSdJUZgyFJLckuWeS2/r+cVVVQE2znVXA6cDNfc2XA6cCvw0cC3x4qvlVtbmqRqpqZOXKlTOVLUmahxk/JbWqzp6qL8mPkqyqqoeaJ/2Hp9nU24EvVtVTfds+dJSxP8lngA/Osm5J0gLo+vLRNmBjs7wR+NI0Yy9iwktHTZCQJPTOR9zTsR5JUgddQ+Eq4Jwku4Gzm3WSjCS59tCgJKuBE4GvTpj/d0nuBu4GVgB/1rEeSVIHnb5kp6oeAd44Sfso8N6+9f8HHD/JuDd0+fmSpMHyHc2SpJahIElqGQqSpJahIElqGQqSpJahIElqGQqSpJahIElqGQqSpJahIElqGQqSpJahIElqGQqSpJahIElqGQqSpJahIElqdQqFJG9Lcm+Sg0lGphm3Lsn9ScaSXNbXfnKSbzbtNyVZ2qUeSVI3XY8U7gH+A/C1qQYkWQJcA7wJOA24KMlpTffHgaur6qXAY8AlHeuRJHXQKRSq6rtVdf8Mw84ExqpqT1U9CdwIrE8S4A3A1mbc9cAFXeqRJHXT6TuaZ+l44MG+9b3A7wAvAn5SVQf62p/1Pc7zsWvXLs4666xBbEqSjgi7du1i7dq1nbczYygkuQX4rUm6rqiqL3WuYJaSbAI2AbzkJS+ZctzOnTsNBEm/ctauXcvOnTs7b2fGUKiqszv+jH3AiX3rJzRtjwAvTHJUc7RwqH2qOjYDmwFGRkZquh84iB0jSb+KhnFJ6p3AmuZKo6XABmBbVRVwO3BhM24jMLQjD0nSs3W9JPUtSfYC/xb4cpKbm/YXJ9kO0BwFXArcDHwX+EJV3dts4sPAB5KM0TvH8Oku9UiSuknvD/Yjy8jISI2Oji52GZJ0REny7aqa8j1lcISGQpJx4AfTDFkB/HhI5cyH9c3f4VwbWF9X1tfNTPWdVFUrp9vAERkKM0kyOlMaLibrm7/DuTawvq6sr5tB1OdnH0mSWoaCJKn1XA2FzYtdwAysb/4O59rA+rqyvm461/ecPKcgSZqf5+qRgiRpHgwFSVLriA2Fw/kLfpIcm2RHkt3N/fJJxrw+ya6+2y+SXND0XZfk+3193T/6cI71NeOe7qthW1/7gn450iz339ok32geA3cleUdf34Lsv6keS339y5r9Mdbsn9V9fZc37fcnOW8Q9cyjvg8kua/ZX7cmOamvb9Lf9ZDruzjJeF8d7+3r29g8HnYn2bgItV3dV9cDSX7S1zeMfbclycNJ7pmiP0n+pqn/riSv6uub276rqiPyBrwCeDmwExiZYswS4HvAKcBS4DvAaU3fF4ANzfIngfcNsLa/AC5rli8DPj7D+GOBR4EXNOvXARcu4L6bVX3AT6doX7B9N9v6gJcBa5rlFwMPAS9cqP033WOpb8wfAZ9sljcANzXLpzXjlwEnN9tZsgj1vb7vMfa+Q/VN97secn0XA5+YZO6xwJ7mfnmzvHyYtU0Y/8fAlmHtu+Zn/C7wKuCeKfrfDHwFCPAa4Jvz3XdH7JFCHd5f8LO+2eZst30h8JWqemKANUxnrvW1hrDvYBb1VdUDVbW7Wf4X4GFg2ndqdjTpY2nCmP66twJvbPbXeuDGqtpfVd8HxprtDbW+qrq97zF2B71PJh6W2ey/qZwH7KiqR6vqMWAHsG4Ra7sIuGGAP39GVfU1en84TmU98NnquYPeJ1CvYh777ogNhVma7At+jmcBv+CncVxVPdQs/xA4bobxG3j2g+xjzWHg1UmWDbC2udR3TJLRJHccemmLhd93c6kPgCRn0vsL73t9zYPef1M9liYd0+yfx+ntr9nMHUZ9/S6h95flIZP9rhejvrc2v7etSQ595P5C779Zb795ye1k4La+5oXed7Mx1b9hzvtuGN+8Nm85TL7gZzLT1da/UlWVZMrrfps0P53ep8gecjm9J8Ol9K47/jBw5SLUd1JV7UtyCnBbkrvpPdF1NuD99zlgY1UdbJo777/nsiTvAkaA1/U1P+t3XVXfm3wLC+YfgRuqan+S/0jvqOsNQ65hJhuArVX1dF/b4bDvBuawDoU6TL7gZ661JflRklVV9VDzpPXwNJt6O/DFqnqqb9uH/kren+QzwAfnUtug6quqfc39niQ7gTOAf6DjvhtUfUl+A/gyvT8S7ujbduf9N4mpHkuTjdmb5CjgN+k91mYzdxj1keRsesH7uqraf6h9it/1IJ/YZqyvqh7pW72W3rmlQ3PPmjB35zBr67MBeH9/wxD23WxM9W+Y8757rr98tFhf8LOt2eZstv2s1yebJ8JDr99fAEx6xcFC1pdk+aGXXZKsAF4L3DeEfTfb+pYCX6T3OurWCX0Lsf8mfSxNU/eFwG3N/toGbEjv6qSTgTXAtwZQ05zqS3IG8Cng/Kp6uK990t/1ItS3qm/1fHrfvwK9o+hzmzqXA+fyzCPrBa+tqe9Ueidrv9HXNox9NxvbgHc3VyG9Bni8+eNo7vtuoc+aL9QNeAu918f2Az8Cbm7aXwxs7xv3ZuABesl9RV/7KfT+Y44Bfw8sG2BtLwJuBXYDtwDHNu0jwLV941bTS/LnTZh/G3A3vSezzwO/PuB9N2N9wL9ravhOc3/JMPbdHOp7F/AUsKvvtnYh999kjyV6L0ud3ywf0+yPsWb/nNI394pm3v3Amxbo/8RM9d3S/F85tL+2zfS7HnJ9fw7c29RxO3Bq39zfb/brGPCeYdfWrP8JcNWEecPadzfQu8LuKXrPe5cAfwj8YdMf4Jqm/rvpuyJzrvvOj7mQJLWe6y8fSZLmwFCQJLUMBUlSy1CQJLUMBUlSy1CQJLUMBUlS6/8Dg+03QFFfKO0AAAAASUVORK5CYII=\n",
            "text/plain": [
              "<Figure size 432x288 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d3kny1ATEvD8",
        "colab_type": "text"
      },
      "source": [
        "## The Agent\n",
        "\n",
        "Our agent must choose an action to take, given three pieces of information\n",
        "1. The **state** the agent is currently at\n",
        "2. The **goal** the agent is trying to get to\n",
        "3. The amount of time the agent has left to reach the goal (the **horizon**)\n",
        "\n",
        "Our agent's policy will be determined by a neural network. This agent takes in the state, goal, and horizon; concatenates them all together; and feeds this input into a neural network with two hidden layers of size 100 each. The neural network outputs logits for each action. The agent will sample according to the logits to output the action.\n",
        "\n",
        "Neural Network Schema (see code below for more details):\n",
        "\n",
        "```\n",
        "self.net = nn.Sequential(\n",
        "      nn.Linear(5, 100), # Input 2 (state) + 2 (goal) + 1 (horizon)\n",
        "      nn.ReLU(),\n",
        "      nn.Linear(100, 100),\n",
        "      nn.ReLU(),\n",
        "      nn.Linear(100, 5), # Output: logits for 5 actions\n",
        ")\n",
        "```\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ADPSIvYTggLq",
        "colab_type": "code",
        "colab": {},
        "cellView": "form"
      },
      "source": [
        "#@title Neural Network Agent Definition\n",
        "\n",
        "class NNAgent(nn.Module):\n",
        "  def __init__(self):\n",
        "    super().__init__()\n",
        "    self.net = nn.Sequential(\n",
        "      nn.Linear(5, 100), # Input 2 (state) + 2 (goal) + 1 (horizon)\n",
        "      nn.ReLU(),\n",
        "      nn.Linear(100, 100),\n",
        "      nn.ReLU(),\n",
        "      nn.Linear(100, 5), # Output: 5 actions\n",
        "      )\n",
        "  \n",
        "  def forward(self, state, goal, horizon):\n",
        "    horizon = horizon / 50 # Normalize between [0, 1]\n",
        "    x = torch.cat([state, goal, horizon], -1)\n",
        "    logits = self.net(x)\n",
        "    return logits\n",
        "  \n",
        "  def get_action(self, state, goal, horizon):\n",
        "    # Put into PyTorch Notation\n",
        "    state_torch, goal_torch, horizon_torch = to_torch(state, goal, horizon)\n",
        "    logits_torch = self.forward(state_torch, goal_torch, horizon_torch)[0]\n",
        "    probabilities_torch = torch.softmax(logits_torch, -1)\n",
        "    probabilities = probabilities_torch.detach().numpy()\n",
        "    \n",
        "    return np.random.choice(5, p=probabilities)\n",
        "  \n",
        "\n",
        "def to_torch(state, goal, horizon):\n",
        "  state_torch = torch.tensor(state, dtype=torch.float32)[None]\n",
        "  goal_torch = torch.tensor(goal, dtype=torch.float32)[None]\n",
        "  horizon_torch = torch.tensor(horizon, dtype=torch.float32)[None, None]\n",
        "  return state_torch, goal_torch, horizon_torch"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BuOS3Qu0JJiO",
        "colab_type": "text"
      },
      "source": [
        "## Acting in the Environment\n",
        "\n",
        "At the beginning of each episode, we choose a desired goal (according to `env.sample_goal()`). The agent interacts with the environment (in the typical RL fashion) trying to reach this goal by the last timestep of this episode (total 50 timesteps).\n",
        "\n",
        "The function `sample_trajectory(env, agent)` executes an episode according to the above logic and returns the states seen, the actions the agent took, and what the desired goal was.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j81noQG7JBTz",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title Agent-Environment Interaction\n",
        "\n",
        "def sample_trajectory(env, agent, T=50):\n",
        "  # Sample a target goal (fixed for episode)\n",
        "  desired_goal = env.sample_goal()\n",
        "\n",
        "  # Default control loop\n",
        "  state = env.reset()\n",
        "  states = []\n",
        "  actions = []\n",
        "  for i in range(T):\n",
        "    states.append(state)\n",
        "\n",
        "    action = agent.get_action(state=state,\n",
        "                              goal=desired_goal,\n",
        "                              horizon=np.array(T-i, dtype=float))\n",
        "    actions.append(action)\n",
        "\n",
        "    state, _, _, _ = env.step(action)\n",
        "\n",
        "  return {\n",
        "      'states': np.array(states),  \n",
        "      'actions': np.array(actions),\n",
        "      'desired_goal': desired_goal,\n",
        "  }\n",
        "\n"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jjv2HE_Bmsh5",
        "colab_type": "text"
      },
      "source": [
        "We also write some code to evaluate the performance of the agent. We will measure performance according to how far the agent was from the desired goal at the final state in the episode.\n",
        "\n",
        "- `evaluate_agent` collects 50 episodes, and prints the median distance to goal for the agent.\n",
        "- `visualize_agent` visualizes 5 different episodes, each with a different goal, for the agent.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sSe7Z8ebmsG5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def evaluate_agent(env, agent, n=50):\n",
        "  distances = []\n",
        "  for _ in range(n):\n",
        "    trajectory = sample_trajectory(env, agent)\n",
        "    distances.append(np.linalg.norm(trajectory['states'][-1] - trajectory['desired_goal']))\n",
        "  print('Median Distance to Goal:  %.3f'% np.median(distances))\n",
        "  print('Min Distance to Goal:  %.3f'% np.min(distances))\n",
        "  print('Max Distance to Goal:  %.3f'% np.max(distances))\n",
        "\n",
        "def visualize_agent(env, agent):\n",
        "  fig, axes = plt.subplots(1, 5, figsize=(15, 3))\n",
        "  for ax in axes:\n",
        "    plot_trajectory(sample_trajectory(env, agent), ax=ax)\n",
        "  "
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zUYOFv236K-y",
        "colab_type": "code",
        "outputId": "f39249f4-83b2-42db-be20-c6ed2920889f",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 245
        }
      },
      "source": [
        "env = PointmassEnv()\n",
        "agent = NNAgent()\n",
        "\n",
        "evaluate_agent(env, agent)\n",
        "visualize_agent(env, agent)"
      ],
      "execution_count": 18,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Median Distance to Goal:  0.780\n",
            "Min Distance to Goal:  0.195\n",
            "Max Distance to Goal:  1.566\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAA1MAAACxCAYAAAAh3OeIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAVrUlEQVR4nO3df5AkZX3H8ffuAQfccYMgXI4f0oBihCNS8YDDgKKFIVQLGj0IeBGtAEqpEH8moyZWW0mZrooGNVZS4USjQaJBUwauKX8kCoqooPKj5CCHgRY9DzwQB7hfwO3kj6eHHZb9Mdu7Mz2z835VbfV0zzPnt+Sp3fk8/fTzjDSbTSRJkiRJszNadQGSJEmSNIgMU5IkSZJUgmFKkiRJkkowTEmSJElSCYYpSZIkSSrBMCVJkiRJJexWdQGSNKOkdjZwBbAP8ChwAUnjy9UWJUmShp13piQNgvMJQQpgWXEuSZJUqRE37ZVUuaQ2Chwwxbu7A/cAe7Zd2wG8AHhyis9sIWmMzV+BkiRJzzarMDUyMnJ990rRQtBsNk+tuoZ29tnBsGLpyAFH7T96NMBYk8lCUHN0hEWtk7Emu4CRiY1GR8Ld9o0Pj23Y/HhzSyf/2/ZZDRr7rAZRP/Vb+6w60Wmf7XiaX9HxjitZj4bDcf30C8o+Ozg2P97ccu8jYxvHmoyNjjA6yc+i9vajIyyarN1Yk7F7Hxnb2GmQwj6rwWOf1SDqm35rn1WHOu6zs12A4rZ+GllQf+mXX5QT2GcHSVI7GrgWWAHsNYtPbgc2A2eSNDZ0+iH7rAaNfVaDqA/7rX1W05pNn3UBCkn9IwShlcAXgG0dfmobcCVwzGyClCRJ0ly5NLqk/pI0tgMXkdTWA1cBe0/TehtwHknjmp7UJkmS1MY7U5L61Xpg1wxtdgFZD2qRJEl6FsOUpH51SoftTu5qFZIkSVMwTEnqV2uBJW3nOyccIUwBXNuziiRJktoYpiT1n7CJ7xrGf0dtBb4JHF4ctxbXFwFnF+0lSZJ6yi8gkvrRakJQGiMsMvFO4CySRg6cBbyruD5WtDuxmjI1NJJajaR2J0mtVnUpkqT+YZiS1I/OBZYC9wGrSBqfJmk0AUgaTZLGOmBV8f7Sor3UTa8GjgbiqguRJPWP8mHKUTpJ3bMaWAesJGncNWmLcH1l0e6k3pWmIfWmCUdJkua0z1T7KN1V81OOJAFJ44QO2+0A3trdYjT0ktpS4GXF2ctJaktIGlun+4gkaTjMZZqfo3SSpGFwBvBE8XpncS5JUskwNdkonSRJC9P5wD7F62XFuSRJpaf5tUbpFjM+Svfl+SpKkqSeCUvrHzDFu7sDp0249iqS2iHAk1N8ZgtJY2y+ypMk9a+yYWqyUTrDlCRpEK0BvgTsYnw6X7unJjnfOEm7PQhL9Z8DXD2fBUqS+tP0YSqpLW+9XLyI3Ytrh+AonSRp4bga2Be4DNgLGJmh/dJJro0BOwh7oDm4KElDYqY7U5soRulOOHjR4uLaRhylkyQtFGEPs8tJajcC1wIrCKGqU9uBzcCZJI0NXahQktSnZlqAYiewJ7DX6AijoyOMEv7ATByVW1pcb/9pPU/1VhylkyT1uxCEVgJfALZ1+KltwJXAMQYpSX3JvWG7aqY7U8fjKJ0kaVgkje3ARSS19YQ9FPeepvU24DySxjU9qU2SynFv2C6a/s5U2yjdWJNOn3lylE6SNOjWExakmM4uIOtBLZI0F+4N20Uz7zOVNLaTNC7a+PDYhg4CVWuU7i0kjR3zUqEkSb13SoftTu5qFZI0F+4N23Udb9r74Nbmw0BzhmaO0kmSFoK1QPuXjp0TjhCmAK7tWUWSNHutvWFhfG9YzaOOw9T+e43s22FTR+kkSYMrbOK7hvG/kVuBbwKHF8etxfVFwNlFe0nqR5PtDat51PGmvcuXjhw4OsKitks7GV+xr7VsemuU7oZ5q1CSpN5aTQhKrb2j3glcQdJoktTOAi4EPk5Y7XYRcCLw/YpqlTTMwmDOAVO8uzvuDdt1nY2mJbXR/fYaaf8P5SidJGmhOpew5cd9wCqSxqeLvajCnlRJYx2wqnh/adFekqqwBniAsDfsfRN+ptsbdmLbTcW/8/qeVL2AdBp6Vo+M7wi/jTBKdxZJIwfOIuz4vo0witcapZMkaRCtBtYBK0kad03aIlxfWbQ7qXelSdIzXE3Y0/XpvWEn/Lg3bJd1Os3v3NERFu3cxQ7CKN34H5cwWreubef4IwijdE55kCQNnqRxQoftdhC+fEhSNcL38Mvbvoe7N2yPdXxnasu25uabN+26xVE6SZIkqY+07Q1LmC3WCfeGnQed3ZlKGids+PDI9R20c5ROkiRJ6rWksR24iKS2HriKsDDcVFp7w17Tk9oWMBeKkCRJkhaO9YS9X6fj3rDzxDAlSZIkLRyndNjOvWHngWFKkiRJWjjWAkvazndOOML43rCaI8OUJEmStBCEvV7XMP4d371hu8z/AyVJkqSFYTUhKI3h3rA90ek+U5IkSZL627mEjXnvJewd5d6wXeadKUmSJGlhWE3Y83Wle8P2hnemJEmSpIUgaZzQYTv3hp0n3pmSJEmSpBIMU5IkSZJUgmFKkiRJkkowTEmSJElSCYYpSZIkSSrBMCVJkiRJJRimJEmSJKkEw5QkSZIklWCYkiRJkqQSDFOSJEmSVIJhSpIkSZJKMExJkiRJUgmGKUmSJEkqwTAlSZIkSSXsVnUBkqoR1bMlwOPAy/M0/k7V9UiSgqiejQIvBb6fp/GuquuRNDXDlDSEij/UjxenN0T1rPXWSXka/6CaqiRJhSuANwMxcF21pUiajtP8pOH0tSmuvyOqZ4f2tBJJaklqNZLanSS1WtWlVCWqZ28jBCmAF1RYiqQOGKakIRPVs+cCr5ri7bXA/VE9O7aHJUldEdWzUx0cGDivBo4m3JEZVn8B3ADsAKJqS5E0E8OUNGTyNH4IWA28bppmb+pROVJXRPVsd+DbhC+lGhxvmnAcKlE92w04BLgR+D/g8GorkjQTn5mShlCexj8EiOrZh4E9gb8s3voEYST0jVE9ewr4neLn1jyN319BqVJZLyqOfhkdFEltKfCy4uzlJLUlJI2tVZZUgYOBRcDPgfsJwUpSHzNMSUMsT+MEIKpnVxKm+H0AOAN4DfBu4AFgX+AowDClQfL7rRdRPfsQ4YvpfcA9wK7ip5Gn8ZPVlKdJnAE8ASwGdhbnX660ot57XnH8ObCJtn4sqT8ZpiSRp/FPGQ9L1xXLpm/P07gZ1bNvAMuqq07qTFTP9gFeD9wNnNf21oen+MgPCVNe1R/OB/YpXi8rzoc1TO0C9gaWR/VsjzyNn6iwJknTMExJepY8jbe1nR5IGNWX+t1lwAUztLkduBw4E3hlVM9G8zQe63plgqQ2Chwwxbu7A6dNuPYqktohwFR3D7eQNBbMf7uonp0PfK44/e/i+BA+3y71NcOUpElF9WyEsCzvwcCPKy5HmlZUzz7BzEEK4G/yNP5KsdfaHxG+3D/Y1eLUsgb4EuGuy2R3Wp6a5HzjJO32IDxXdA5w9XwWWLGT2l6/D/gucHuexjsqqkdSBwxTkp5WrID2h4RnFc4AjijeurWyoqTO/AhoAj8DTiaM5m8GfgN8nfFpf+dE9ewY4ITi/FAMU71yNeEZzMuAvYCRGdovneTaGGHJ8Hex8KYAvgNoAO/O0/ijVRcjqTOGKUnA03eiNjE+DefXwKVAlqfxvZUVJnUgT+N/i+rZDuBKwkITtxRv7cczn586pzg+Cmwg9HP1QtJoApeT1G4ErgVWEEJVp7YTAvKZJI0NXaiwUnka7yqe+3us6lokdc4wJanlH3nm8wwHAhikNCjyNL46qmf/SxgEeGNx+WFg/7ZmK4CHXcWvQkljA0ltJfBJ4A2EhRZmsg34AnApSaOvp71F9ex5wAeBt+dpPHHq4kwWMfUzYpL6kGFKUsvbi+ODwF8B64BfVVeONHt5Gt8BXBjVs/cSpqneRng+ZxFwUJ7GD1RZnwpJYztwEUltPXAV0weqbcB5JI1relLb3F0HHAP8E2HBk9kYwwUnpIFimJpnUT07FthQ3K4/mrBi1EfzNN5VcWlSpzYRptK0XksDJ0/j3wI/ierZQYQg9fY8jTfP8DH13nrCghTT2QVkPahl1qJ6djrh7tp2YCthAOqY4u2flfgndxH6q6QBYZiaR1E9OxH4AfBe4GOE0anDCHP3v1VhadJs7E+YCgXjoUoaVIcVx59XWoWmckqH7U4GbuhmIbMV1bMXA1+b6v08jbcW7ZYBrwWOBCLC/lkAnyEssrGMMK16efHT6FrRkuadYWp+/UFxPLw43k34Q34ehikNjksJz08BvBX4QIW1SHPV2gTVvdL601pgSdv5TmBx2xHCFMC19FmY4pn7Yj0EPLf9zWJBlMVtl5o8827/6cDjhAUnHgDuIEyz/l43ipXUHYap+dUKUXlxvIvwy3JNVM8udqqf+twRwImEP+atL6CXFav8HQgcX/x8NU9jl0rXoPDOVL8Km/iuYfwZoa3At4FLCAM6ryAErUXA2SS1i/tsk97WSoQNwia7nyfMSGlZPLF9nsY7e1GYpN4xTM0gqmePAvvM8mN/H9WzDwF/XZzvS9jd3TClnijZbyGMkrZWyvo7wijq7hPaHM74NBVpXsyhz+4Efi9P48k2d4UwbWpbnsaPli5O3bKaEJRae0e9E7iCpNEkqZ0FXAh8HNizaHci8P2Kap1M6298DTi3+JnOjqieATyWp/GybhYmTWYOv2fts9MwTM2sTKdrfe7gtvPJdnuXuqVsv11KuLO6BHj/hPe+A3wR+Fz5sqQple2ziwlfsPef4v09CYsDqP+cS/idcy9h76i7nn4n7Em1rm1PqiOK9pWHqaieLQcuBvYo+U+U7evSXM3lO62mYJjqrve1vT6Z8GVU6mt5Gh8b1bPXAJcTpvd9Kk/jSyouS5rOdBu/7sn43Vb1l9WELRj+fMq9o5LGXcWeVJ8ATuphbc8S1bPTCHfKjpmpraThYZjqnfMxTGlw/CtheiqEUWOpnz3jOZqonn0QeBlwH+F37y+qKEozSBondNhuB2ExnMpE9ex1wFeqrEFSfzJM9c4jVRcgzUL7SL9hSv1u4qIEFwD7AS8h/J27u+cVaaG5Cfhn4GbCRrz34d91SRimeunhqguQOhHVs1U8cxUqR/XV7x6acP4c4PN5Gl8a1bN9CKvESaXlafwA8Lb2a8ViEpKGnGGq++4kzK/+TdWFSB26ZcK5S/mq361qvYjq2SLCFNVHAPI0fqyqoiRJC59hqvta00/umraV1D8mbj7ZT/u6SM+Sp/FvAIo90V5aXHYAS5KAqJ7tz7Pv4GuejM7cRHN0LGHK1B1RPTtspsZSH7hswrlhSn0tqmfLo3r2HuAOwkI/O4AfV1uVJPWNl1RdwEJmmOqNXwC/BT5bdSHSTPI0/gjwirZLPmStfrcJ+Cjh2aiLgRV5Gt9YbUmS1B/yNP5GnsYjVdexUDnNrzeuBTYCt1ZdiNShm4DrgKvzNP511cVIM1hUHO/L0/hfKq1EkjRUDFPddQ9wT57GF1ZdiDQbeRo/AcRV1yHN0k+rLkCSNFwMU931XOB/qi5CkhYyp69IkqpimOqu5+DUPkmSJPVYVM+WAX8KHAUcWXE5C5ZhamaPAfuU+NwTwEkYplSNsv3WPXlUFfusBo19Vv3uEuBvgW3A3nP4d+yz0zBMzSBP42UztYnq2TE8c67+kXka39u9qqTpddJvpX5in9Wgsc9qALwQ2AIsB25h+iXSX5un8X/1pKoFxqXR58fEB/XzKoqQJEmSCkcCG/I0buZpvIqw7+nK4r0c+BPCitMAP+p9eQuDYWp+XNB+kqexm5xKkiSpSkcCP2udFCv1HlqcvjlP4/8A7iY8mrK59+UtDE7zm6Oonj2H8GAfwOnAzRWWI0mSpCEX1bMlhOl9r47q2Z3A0cAG4Jqiye3F8TDgfm8ElGeYmrvT215/N0/j7ZVVIkmSJMEO4IuElaUfB0YIgepFQJ6n8W+LdocBT0X1bBVwe57GT1ZR7CAzTM3dGcXxZoOUJEmSqpan8S7gvPZrUT27BPgkcFvb5V8CJxIWqNgW1bPz8zT+Ss8KXQB8ZmruXlgc3ZxXkiRJ/epTwHuBf2hdyNN4DfA84FxgDHhlNaUNLu9Mzd2JxfH6KouQJEmSppKncRP42CTXfwF8KapnHwQO6XlhA847U/PnpqoLkCRJkkr6JXBw1UUMGu9MzZM8jR+vugZJkiSppAbjK1Q/Lapnzwf+GFgBHESYCviWPI2/2tvy+pNhau4uJCR5SZIkaVAtBx6c5PpHgLOBbcDexbWoRzX1PcPUHOVpfEXVNUiSJElztAK4c5LrhwLfAk4Djgd+SNtmwMPOZ6YkSZIkrQA2T3L9IOCXxQIWS4trj/Wsqj7nnSlJkiRpiEX1bDegBrw+qmdjwNcJj7E8QrF0elTPfgT8bvERw1TBO1OSJEnSEMvT+CngDcBPgLcAGXA7cH/RZA/Chr9vA54EflVBmX3JMCVJkiQNuTyN/x3YBOw5ydvHAwcQpvwdmKfxA72srZ85zU+SJEkaclE9eyPhrtRkbgJuz9P4+B6WNBC8MyVJkiQNsaievRT4TNuln7a9/jCwETisp0UNCMOUJEmSNNwOBEbazle2vf4esAW4p6cVDQjDlCRJkjTE8jT+KnAU8NlJ3r4UeD7uLTUpw5QkSZI05PI0vjdP4z8Djpvw1i7gELwzNSnDlCRJkqSWTW2v7wBOLV7f2vtS+p9hSpIkSVLLw8VxXZ7GLwb2Aw4GrquupP7l0uiSJEmSAMjTuBnVs0eAHcX5GG7SOyXvTEmSJElqt4PJN+/VBIYpSZIkSe12AourLmIQOM1PkiRJUrv3AA9WXcQgMExJkiRJelqexv9ZdQ2Dwml+kiRJklSCYUqSJEmSSjBMSZIkSVIJhilJkiRJKsEwJUmSJEklGKYkSZIkqQTDlCRJkiSVYJiSJEmSpBIMU5IkSZJUgmFKkiRJkkowTEmSJElSCYYpSZIkSSrBMCVJkiRJJRimJEmSJKkEw5QkSZIklWCYkiRJkqQSDFOSJEmSVIJhSpIkSZJKMExJkiRJUgmGKUmSJEkqwTAlSZIkSSUYpiRJkiSpBMOUJEmSJJVgmJIkSZKkEgxTkiRJklSCYUqSJEmSSjBMSZIkSVIJhilJkiRJKsEwJUmSJEklGKYkSZIkqQTDlCRJkiSVYJiSJEmSpBIMU5IkSZJUgmFKkiRJkkowTEmSJElSCYYpSZIkSSrBMCVJkiRJJRimJEmSJKkEw5QkSZIklWCYkiRJkqQSDFOSJEmSVMJus2x/3MjIyPXdKEQLwnHAbVUXMYF9VtOxz2rQ2Gc1iPqt39pnNZOO++xIs9ns+F+142kmzWbz1KpraGef1Uzssxo09lkNon7qt/ZZdaLTPjurMCVJkiRJCnxmSpIkSZJKMExJkiRJUgmGKUmSJEkqwTAlSZIkSSUYpiRJkiSpBMOUJEmSJJVgmJIkSZKkEgxTkiRJklSCYUqSJEmSSvh/bc0BQ/4vdDwAAAAASUVORK5CYII=\n",
            "text/plain": [
              "<Figure size 1080x216 with 5 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KFE4V2GGC5LF",
        "colab_type": "text"
      },
      "source": [
        "## Learning a Policy using GCSL (A Code Description)\n",
        "\n",
        "Now that we've established the basics of what an *agent* is, how it interacts with the environment, and how to evaluate the agent, let's actually train a policy!\n",
        "\n",
        "Suppose that the agent has already collected some trajectories in its buffer.\n",
        "```\n",
        "buffer = []\n",
        "for i in range(10):\n",
        "  buffer.append(sample_trajectory(env, agent))\n",
        "```\n",
        "\n",
        "When training the policy, we will sample a trajectory randomly from the dataset\n",
        "```\n",
        "trajectory = buffer[np.random.choice(len(buffer))]\n",
        "```\n",
        "Choose two time-steps randomly on this path $t_1 < t_2$ (remember, hindsight works for any trajectory and any such pair).\n",
        "```\n",
        "t1, t2 = np.random.randint(0, 50, size=2)\n",
        "t1, t2 = min([t1, t2]), max([t1, t2])\n",
        "```\n",
        "\n",
        "If we define\n",
        "- `s` to be the state at t1\n",
        "- `a` to be the action at t1\n",
        "- `g` to be the state at t2\n",
        "- `h` to be t2 - t1\n",
        "\n",
        "```\n",
        "s = trajectory['states'][t1]\n",
        "a = trajectory['actions'][t1]\n",
        "g = trajectory['states'][t2]\n",
        "h = t2 - t1\n",
        "```\n",
        "\n",
        "\n",
        "Then `a` is a good action to go from `s` to `g` in `h` timesteps.\n",
        "\n",
        "Because of this, we can treat it as expert data for `agent(state=s, goal=g, horizon=h)`, and try to match it using the classification loss.\n",
        "\n",
        "\n",
        "```\n",
        "  loss = nn.functional.cross_entropy(agent(s, g, h), a)\n",
        "```\n",
        "\n",
        "This leads to the following code-block that fully describes the GCSL algorithm:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zRhhZTINC20w",
        "colab_type": "code",
        "outputId": "df698b8d-a3aa-4814-a7df-4d589806bc6e",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 602
        }
      },
      "source": [
        "# Create agent and optimizer\n",
        "\n",
        "agent = NNAgent()\n",
        "learning_rate = 1e-4\n",
        "optimizer = torch.optim.Adam(agent.parameters(), lr=learning_rate)\n",
        "\n",
        "buffer = []\n",
        "n_episodes = 201\n",
        "n_steps_per_episode = 500\n",
        "\n",
        "for episode in range(n_episodes):\n",
        "  # Collect more data and put it in the replay buffer\n",
        "  new_trajectory = sample_trajectory(env, agent)\n",
        "  buffer.append(new_trajectory)\n",
        "\n",
        "  # GCSL optimization  \n",
        "  for step in range(n_steps_per_episode):\n",
        "\n",
        "    # Sample a trajectory and timesteps\n",
        "    trajectory = buffer[np.random.choice(len(buffer))]\n",
        "    t1, t2 = np.random.randint(0, 50, size=2)\n",
        "    t1, t2 = min([t1, t2]), max([t1, t2])\n",
        "\n",
        "    # Create optimal ((s, g, h), a) data\n",
        "    s = trajectory['states'][t1]\n",
        "    g = trajectory['states'][t2]\n",
        "    h = t2 - t1\n",
        "    a = trajectory['actions'][t1]\n",
        "\n",
        "    s, g, h = to_torch(s, g, h)\n",
        "    a = torch.tensor(a)[None]\n",
        "\n",
        "    # Optimize agent(s, g, h) to imitate action a \n",
        "    optimizer.zero_grad()\n",
        "    loss = nn.functional.cross_entropy(agent(s, g, h), a)\n",
        "    loss.backward()\n",
        "    optimizer.step()\n",
        "\n",
        "  # Print agent performance once in a while\n",
        "  if episode % 40 == 0:\n",
        "    print('##### Episode %d #####'%episode)\n",
        "    evaluate_agent(env, agent)\n",
        "\n",
        "visualize_agent(env, agent)"
      ],
      "execution_count": 20,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "##### Episode 0 #####\n",
            "Median Distance to Goal:  0.756\n",
            "Min Distance to Goal:  0.098\n",
            "Max Distance to Goal:  1.400\n",
            "##### Episode 40 #####\n",
            "Median Distance to Goal:  0.357\n",
            "Min Distance to Goal:  0.037\n",
            "Max Distance to Goal:  1.073\n",
            "##### Episode 80 #####\n",
            "Median Distance to Goal:  0.135\n",
            "Min Distance to Goal:  0.016\n",
            "Max Distance to Goal:  0.500\n",
            "##### Episode 120 #####\n",
            "Median Distance to Goal:  0.128\n",
            "Min Distance to Goal:  0.003\n",
            "Max Distance to Goal:  0.283\n",
            "##### Episode 160 #####\n",
            "Median Distance to Goal:  0.096\n",
            "Min Distance to Goal:  0.010\n",
            "Max Distance to Goal:  0.372\n",
            "##### Episode 200 #####\n",
            "Median Distance to Goal:  0.100\n",
            "Min Distance to Goal:  0.015\n",
            "Max Distance to Goal:  0.251\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAA1MAAACxCAYAAAAh3OeIAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjEsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+j8jraAAAcV0lEQVR4nO3de5QkZZ2n8SfrRt+rkZsoQgAqCK0IKjQKNAu4ioGMF1CUXZjRQVxmVVhmMdZRSUcZwnUQcQ8zCo5XQBA8KhBehlGBQVAUBQV6AIXgJiMNNEnfu6sq9483ysouqqqrk8prPZ9z8kTGLfvXp9+THd98I963VK1WkSRJkiRtnZ5WFyBJkiRJncgwJUmSJEl1MExJkiRJUh0MU5IkSZJUB8OUJEmSJNXBMCVJkiRJdTBMSZIkSVIdDFOSJEmSVAfDlCRJkiTVoW9rDi6VStc3qA51iWq1enira6hlm9WW2GbVaWyz6kTt1G5ts5qO6bbZafdMFQ3vlXXWo9nhle30BWWb1TTYZtVpbLPqRG3Tbm2zmqZpt9mt6pkCbm+nXxbUXtrli3Ic26wmZZtVp7HNqhO1Ybu1zWpKW9NmfWZKkiRJkupgmJIkSZKkOhimJEmSJKkOhilJkiRJqoNhSpIkSZLqYJiSJEmSpDoYpiRJkiSpDoYpSZIkSaqDYUqSJEmS6mCYkiRJkqQ6GKYkSZIkqQ6GKUmSJEmqg2FKkiRJkupgmJIkSZKkOhimJEmSJKkOhilJkiRJqoNhSpIkSZLqYJiSJEmSpDoYpiRJkiSpDoYpSZIkSaqDYUqSJEmS6mCYkiRJkqQ6GKYkSZIkqQ6GKUmSJEmqg2FKkiRJkupgmJIkSZKkOhimJEmSJKkOhilJkiRJqoNhSpIkSZLqYJiSJEmSpDoYpiRJkiSpDoYpSZIkSaqDYUqSJEmS6mCYkiRJkqQ6GKYkSZIkqQ59rS5AkiRJUo3y4CBwM/BaypVKq8vR5OyZkiRJktrLMcA+QNzqQjQ1w5QkSZLUXk4et1SbMkxJkiRJ7aI8uAA4rFhbRnlwfivL0dQMU5IkSVL7OBrYWLzfUKyrTRmmJEmSpPZxErCweL+oWFebcjQ/SZIkqVnKgz3ADpPs7QeOGrft9ZQHdwE2TXLOCsqVkZkqT1vHMCVJkiQ1z3HAFcAwY7fz1RqaYP3eCY4bAHqBdwBXzmSBmj7DlCRJktQ8VwKLgfOBuUBpC8cvmGDbCLAeOAO4akar01YxTEmSJEnNUq5UgYsoD94EXAPsTAhV07UOeAx4M+XK3Q2oUFvBASgkSZKkZgtBaAlwKbB2mmetBS4B9jVItQd7piRJkqRWKFfWAadQHrwWuAyYN8XRa4F3Ua5c3ZTaNC32TEmSJEmtdS1hQIqpDANZE2rRVrBnSpIkSWqtQ7d0wLrqAG/aeO5VDyTZE8AzwErg63kaP9Tw6jQpw5QkSZLUWidWq8z/yNB72a/0B17Xc9fGF/WsGAA2ANsAPF5dPO+B6s5vIdzuV2Js0IpPtaZkgWFKkiRJap0wie9xKxjs+fHwAXyTIwEG+hn6/XcHPvbQvj0PHgTM37X0eO9iVlUrLLi8SumLwC+A5a0sXT4zJUmSJLXSUqB3x1Jl5Bfb/M3a/913+XlQ/fgm+naMN/7DyHC1dAaw9nG2HXl+aWV1DhuPAF5VnHtb68oW2DMlSZIktcz3hg8+40ObPrCwWJ0HnAkMAX1QOmrPDZe+dBs2PryBgb0AotJj2+TVnV8FPAk82JqqNcqeKUmSJKlFfj6yzyETbK7t8Nh1AwNrehj+/qX951z1k4EzHyH0TN2Wp3G1OVVqMvZMSZIkSS1y7jmf3fmbSbYI2B3YA3gxsBfwamC/4rChEXrfcOKmv5tHGHxiNfCP4z8rSrITgbcTeqz2Bn4EXGDoahzDlCRJktRCeRo/A9xRvACIkmwH4PFi9UDCrX+LgH0J1/C3135GlGTzgUvGffQbgf/HluewUp0MU5IkSVL7WTRuvQ9YUbP+zijJDiH0aEWEHq1RXyIEs8vzNDZINZBhSpIkSWo/W5o/6q1ABXgAuA/4AfC3wPV5Gp/S4NpUMExJkiRJbSRKstcAJ0xxyC+B1+VpvGnceYcDL4mS7BPA3cB1eRo/1bBC5Wh+kiRJUptZDpwFHAYcBVw3bv+V44NU4XJgLfDR4v2FjSxS9kxJkiRJbSVP49XAZ0bXoyR7CLi3WH0Y+F2UZKXxo/TlaXwecF6UZHOAm4HtavdHSbYXcA2wE7AR2ARckqfxWY36u3Q7e6YkSZKkNpan8X2M9U69iPB81FenOH494XmquQBRks2Jkmwp8ENgEPgK8C1CL9YxDSt8FrBnSpIkSWp/byLMP3UAcAohEE1lHXBwlGS/AZYQrvvXAIfnafwrgCjJ/pkwL5XqZJiSJEmS2lyexkPAXcXrG9M45ZeE4PUnQk/Wr4Gf5Wn8WM0xTwDbRUnWk6fxyAyXPCsYpiRJkqQuk6fx2cDZWzhsPeGxnzmEW/60lXxmSpIkSZqdXgw8lqexQapOhilJkiRpdtqHMB+V6mSYkiRJkmaZKMlKGKaeM5+Z2oIoyZ4BFtZx6qo8jRfNdD3SdNhu1Wlss+o0tll1mina7AeiJPvAFKfaZqdgmNqyer4on8t50kyw3arT2GbVaWyzantRkm0HPB/YAdtsQximJEmSpC4TJdnhwE9bXUe385kpSZIkqfv8EvjXVhfR7eyZkiRJkrpMnsZrgDcAREk2B1jX2oq6kz1TkiRJUhfL03h9q2voVoYpSZIkSaqDt/m1SJRkg8D+wI15Go+0uh5JkiRJW8cw1UBRkh0B7Alcm6fxY1GS9QJHAGcDrysOO5DwgKAkSZKkDmKYaqxzgKUAUZJNdszTTatGkiRJ0owxTDXWk8CDwN3A0TXb1wNVYA6wMkqynYHnAQuAW/M0rja7UElqpCjJTgV2Az6Vp/HaVtcjSdJMMEw11n8CB+Rp/KYoyQbyNN44uiNKsm8DbwNWjDtnGXBjE2uUpIaKkuytwBeK1VOiJLsT+C1wD7Ac+PdW1SZJ0nNhmGqs5wHbRklWGg1SUZLtB/wVIUjVeppwseFFhaRucziwAfgL4KJi/SBgbrH/tpZUJUnSc2SYaqxjgW8AC6IkO5kQog4ANhX7rwTel6exz01J6mYPAtsAvwD2JfzQ9DDwAsKgPJ9pXWmSJNXPMLVlq4CFdZy3jvCr69cJFwqnArcDHzyp90fX/H3/1zLgFMqVyoxVKo2pt92uGn0TJdliwmiUewIvBfYqXi8FLsvT+LQZqFOzw4PFMsrT+HZgdbH+KPCNKMmuJdwWPVDHZ6/a8iFSQzzn71mpyWyzDWCY2oI8jRdNti9KsnnAtoQBJUZfvwX2Ae4tlj8HYsItLgfkaVyl/O4Ti30xcFlD/wKalaZqt+NFSfZJ4KPF6q+iJPsp8ApC70GtRwjPuNwHvDdKso/lafzkTNSrrpcXy90JPyptvjONVxJ6rqSOMdH3bHFbfzVKso8CnwQG8jTe9OyzpebbmmsDTZ9h6jkoRqTabFSqKMn2Kd7uVyzPIfRQbQPsCPwJOLnYdzKGKbVeqVj+C7AEeAlwPXAL8IfRV57GawCiJHsFcAeQRkn2B2AX4IXAMLASeArI8jR2IBWNyotlNHqx2cpipEaIkuxI4JooyX4P9BSb+xi7tV9SFzJMzbwzgf8ARieW+p9Af/F+N8qDa4DDivVllAfnU66saXKNUq0y8Jk8jad1y2mexr+NkuwW4K+LTSuBPxIuHrYFdiI8G/j6mS9VHeopQtj+LOGHpr9saTVSHaIk2xs4LE/jiyY55DWEH08fBfYvlkNNKk9Si/Rs+RBtjTyNP0uYQwrCg9U7Ae8hjNS3nDDf1OgQ6RvYfP4pqenyNB6abpCqcTSwN7AwT+Pn5Wm8JE/jfQgDCjxN6M3aalGS9UVJtihKst56zld7KnqiRv9NT2plLdJz8E3gi1GS7TrJ/l2Ap/M0PjpP4+cDu3qLn9T97JlqjL2K5W+LZwG+UrygzEmMPfy3iHBhcVWzC5SeiyJ8TRTAnk/onbprS59RBKYDgAOBVxavJYTJrImSbB1hoIJ1hNtp1417/+08jb/2XP8uaprjCN913291IVKdRm/rj4F/nmD/Cwm9UQDkaTzSjKIktZZhaqaVB3sWcvHLVzPvmbf33rjHqR+59WUb6e8bpqfnnL4v/+FFPRw17ozXUx7chcnvqV5BueIXsjrFvsXyWWEqSrIewgAERwL/ldBzu22x+0nCwAQXAo8BC2pec4vXvJrlq4A9AMNUh8jT+NtRkv0HIQhLnWhFsfxAlGRzCc+OrgJ2K16vZho/JEnqLoapmXfc/j2/f8+NI/tx1fCyW2t3XDNy8KbTeq7eMO74IcLIf+MNEG6LeQdhPiqpE4zOF3RwlGSvIoSrCNiVcAvM6PODjwDfBa4DbgIe2ZpBCaIkuxA4YYZqVvNsD/iMqDpClGQlwlQQ2xN++HlxsWtb4LwJTlmFPa/SrGOYmnlX3jiy3xUA21Hh7P6v088QZ276HzxeXdzP2MXkqAUTfMYIYZj1M/AWQHWIKMl2ItyqB/CpYvlH4H7CyIAPEUZ1uwG45zmO6PZH4HlRki120uvOECXZboSL0l+3uhZpmv4S+PK4bV/O0/i9UZLtSBhMZS5hHrUH/S6SZifD1AyL1l82Ouw5TzLIBzd9gM/1X8hiVrOqOnc6H7GOcJvTmylX7m5UnVID1M4TdDjwuzyNn2rQn5UXy5VRkj0K3A1cAVyRp/HqKMn6Cb8ivx94Jk/jjzWoDk3fsmJ5fSuLkLbCdsXyLYT/l1cCDwDkafw4oWdd0ixnmJpBUZJty+hAEzW+OvSGkWF6elazxTC1FrgU+CDlyvoGlCg10sOjb/I0vqHBf9aVhN7bvQm34RwEfAn4UpRkzzo4SrLzCL8yvwz4HmGQi6MplaDqlEdNsoxwMXpnqwuRpmn0y+EneRqvamklktqWYWqGFEOl/utE+26vvrgHICr951QfsRZ4F+XK1TNfndR4eRpXJwoyDfqzNgLfHl0vnm1YBvx0klNWMPZ9977RjT3zFv9sZM1K54FpjsOAf3eEM3WQ0dvyz4ySbDVhpL7LnXRaUi3D1Mw5l7Eh0Z/lhN6fbNqGof5bhl/Gwb3LJzpkmLGJfqVO9XHGhv5vmuLi5vooyfYEXkEY8W974GpgR+B4wgTac2rP6527qM8w1XhRku1DuO1youGkpXZ1brE8u2bbeuA7LahFUpsyTM2c9xFGJ/sW4Qv4H4GfEB5Q5fLhI/oBPj/8Nu7r+e/0l4Yn+oxDCA/nSx0pT+NPtvjPv58w4MVmm4FboyQ7i9A7chbwI+CCnrkL/Q5sgCjJ+oDXAW8nzC+1c7HrhCjJLimeN5Ha3RcIz10CfI4wiI53j0jajBcSMyRP4zXFHCoQHlbNgY0Ab+y59c6XlB7Zt680XHptz930l4Y3EB7WH11C+CX9RAxTUkMUvVc3ADdESXY4QM+cBX4HzrBiMubJ5s17DXBflGQvz9P4oSaWJdXj44yFqVvyNP5WK4uR1J56Wl1Al7kTeC8hRH0R2G2Q1S/4wsDndjmz/6rSh/q+w2t67llDGAFo92I5OudKL3A85UH/TaRmcfCJGZen8TBQJkzCXAZOqtl9JLCIMGCI1NbyNF5Rszrh7SSS5K+yM6j45fvLUZJ9hXCLyw53zHnf7oSgNDp31OnAv1CuVCkPHgv8NeH2gTnFcQcR5uSR1Dh9ANWRYdNUA+Rp/AngE6PrUZJ9vXj7m2K5S9OLkupzHnAm8KtWFyKpPRmmGqAIVTcBUObzhIl57yfMHTU2+kS5UgUupjx4E3ANsAdwAoYpqdF6ATBMNdvThJFLdwGIkmwR8HfAJ/I0XtvKwqRJLAaeIDwvJUnPYphqvKXAxcCHJp07qlxZTnlwCXABcHATa5NmqxCmMEs1yVKgUgyf/yjwomL7+wgDgqyKkuwOwvffasJkqb/P09jR/9RwUZL1ANVJhjx/NXCbw6FLmoxhqtHKlQOnedx64NTGFiOpsAqgZ2Beb6sLmQ3yNP5FzeqjjI3u92SxfDnhearDa88rRv5zslQ1TJRku1BMOB4l2b2EQWqWEtrkqDtaUJqkDmGYkjQbrQTY4a0fecUzv/pe3uJaZps/MtYDPzqB7zsmOfZtwNcaXpFms9Nq3r+0eI13YZNqkdSBDFOSZqOnR98sevVfRC2sYzZ6BNg9SrL7GAtTEC5YNxEG6Rl1dpRkxxMmXt4OSPM0vrhplWo2+Hvgx8DfAG8l3GZ6JHAbocfqB3ka39q68iS1O8OUpNnoceD26vDQfj431XQXEkLUHoQpIm4CjsjTeFOUZPMZC1M/BPYBXgCsIMxRtW/zy1U3y9N4PSFM/Xj8vijJVgLbNr0oSR3FMCVp9igPDgI353N4LeXK/tvHZzw8b+9Dd97ieZoxxWS9/2eSfWuiJBtdfX+exg8CREm2LfAUjqim5noKeHGUZMcA64CH8zS+t8U1SWozThAraTY5htDbEQOUevp6GBmya6q9fKdY1ganXYvlg02uRbPbA4SBKK4B/g24J0qyxa0tSVK7sWdK0mxycs3ystLA3N7qpo3DrSxIz/JOYG4xjHovEFGEX+yZUnOdAvxfYC7wYUI7dHRJSZsxTEmaHcqDC4DDirVllAfn9w2eM3do1Yp1rSxLY6Ik24kwPPp/iZLsYGAvYJti91rg9y0qTbNQnsYbgDsBoiQbBO7M09gfXyRtxjAlabY4GthIuDjfABw9sOPui1pbkgCiJDsV+CDhFkwIv/7fDFwHLC9ed+VpXGlNhRKvAK5tdRGS2o9hStJscRKwsHi/6Inqove0shht5nRgDpAAPwV+nafxUGtLkoKix3RH4HetrkVS+zFMSeoO5cEeYIdJ9vYDR9Vu+EP1BUfWnLvT5p9V+dMMV6epLQR+lKfxp1tdiDSB0R7TO1tahaS2ZJiS1C2OA64Ahgm38423WU/HPqUHNy6855r+RXddUeXdPFCzawC/G5ttAWGyVKkd7Vks72tpFZLakhcMkrrFlcBi4HzC6FulqQ5eWFq3YLvffIlw3Py5xeYRYH0jixRESbYH8FeE4NpD6JlylDQ1XJRkJeAlhADfR2h/d+ZpPFWY3xPYBDzS+AoldRrDlKTuUK5UgYsoD95EmBdmZ0Komq51wGPAmxtQnTZ3FbA/IbiOAM8Av2ppRZotPg6Ux237CjDVM5R7Arkj+UmaiGFKUncpV+6mPLgE+DzwbmDeNM5aC1wKfJByxZ6pxusBbszTeFmrC9GsczThR5P3E24JPgM4KkqyUjG32Q7AocASwqATOxKet7y1RfVKanOGKUndp1xZB5xCefBa4DKmCFQjVUaAd1GuXN2s8sTPgJOiJOvN03iY8uAgYSj011KuOPy5GukJYCBP46sBoiR7FXAk8P0oyXZlbLAJgJXA48BdwEXNLlRSZzBMSepm1xJ+fZ5KFciaUIvG3AycRvj1/w7gGMJFbEwIv1KjrAAOipKsDLwBWFpsfyPwQ+AbwI3AbcWkvZI0JcOUpG526DSPOwS4oZGFaDOPF8vReb9OrlkaptRIjwLbE56dupXwrN6rgWV5Gt/YysIkdaaeVhcgSQ10IjC/Zn3DuCU9JXqL49Q8/cVyE+XBBcBhxfoyyoPzJzlHmgnnA8cCO+RpvBT4ZrH9rtaVJKmTGaYkdacwie9xjH3PrQGuA3Yvlmtqjj6+OF7NMVAsf75sw2evu2joTeQjO0EIuUe3rix1uzyNn8zT+Jo8jZ8sNu0NPFGzLklbxYsHSd1qKdBLGHp7LXA6cCzlSk74ZfqMYvAJiuMOakWRs9Q2xXLd2uqc/f5h6L9tc/KmDwMsAk5qXVmahU4Bto+SrFq8bo6SbGCLZ0lSwWemJHWrEwgTc94PvJlyZfmf94Q5qS7+zbmlU5bs2Pvy4rgTgFtaUWhXCz1+O9RuOqvv2Du+PPTGc87v/6evHtp75+/O3HgqPxtZMrr79ZQHdyFMkjqRFZQrI5Psk56rgwlhf2OrC5HUGQxTkrrVUuBi4EOTzR21eiNrb310+JfAcsJFlGbeccAVhFEVNwKc1nc1p/VdDfC/gKGeUpUeqqPHDwH3TvA5A4QexHcAVza4Zs0eHwY+DSzK03hVq4uR1HkMU5K6U7ly4HQOG6kyQrlyaqPLmcWuBBYTHvyfC5TGHzBc7aE0FqYWTPAZI8B6wgSrVzWmTM1STxXLvaIkuwOIgL48jZdPfookjTFMSZIaJ9xSeRHlwZuAa4CdCaHqz6qU6C1NeufeOuAxwq2adzeyVM1KjxXLXxLmnCsRbg3es2UVSeooDkAhSWq8EISWAJcSBgT5s6T/m3yt/9MTnbUWuATY1yClBvkhYcLeU4Fzi223ta4cSZ3GnilJUnOUK+uAUygPXkuYnHcewE6lpye4+Y+1wLsoV65uao2aHcqDg8DN+RxeS7nyI4AoyfYHPgJ8u6W1Seoo9kxJkprtWsKAFFMZBrIm1KLZ6RhgHyCu2Tb6nOWtzS9HUqcyTEmSmu3QaR53SEOr0Gx28rglhDD1BJA3vRpJHcvb/CRJzXYiML9mfQNhbp/RJYRbAE8Ebmhuaep65cEFwGHF2jLKg/Oj9ZftCrwNuCFP4+rkJ0vS5uyZkiQ1T5jE9zjG/v9ZA1wH7F4s1xTbe4Hji+OlmXQ0xZxnI9XShvM3vf1vCQNRrAdOb2VhkjqP/0lJkpppKSEojRAGmTgdOJZyJQeOJcwltbbY3wsc1Joy1cVOunfkhQs/vemdHLLhgkUXDL+9DCwE4jyN89aWJqnTeJufJKmZTiBMzHs/Ye6osclRw5xUF9fMSbVHcfwtLahTnSr0Zu4w0a7TN562zz3Vc+Pl1d3oZZjDen7LWb1XbFrEmoOO6L39GcrsNMFpKyhXJp0ITdLsZpiSJDXTUuBi4EOUK+snPKJcWU55cAlwAXBwE2tTdzgOuIIwIuTG2h179zzU992hQ0of7fsGb+n9GduXnoHwrN4dE3zOAKF39B3AlY0tWVKnMkxJkpqnXDlwywdBEbRObWwx6lJXAouB84G51MxiVqnOp58h3tv7A0pjc5stmOAzRgjPUJ0BXNXQaiV1NMOUJEnqHuF20YtqbhfdmRCqWMlCFrOqNkhNZB3wGOE21LsbXK2kDucAFJIkqfuEILQEuJQwqAkrqwvYtrR6qrPWApcA+xqkJE2HPVOSJKk7lSvrgFMoD14LXLayunDetqya7Oi1wLsoV65uWn2SOp5hSpIkdbtrgeHLBs5hPQOTHTMMZM0rSVI38DY/SZLU7Q4F6CuNsKA08SCShUOaU46kbmGYkiRJ3e5EYH7N+oZxS4B5xXGSNG2GKUmS1L3CJL7HMXbNswa4Dti9WK4ptvcCxxfHS9K0+IUhSZK62VJCUBohDDJxOnAs5UoOHEuYS2ptsb8XOKg1ZUrqRA5AIUmSutkJhIl57yfMHbX8z3vCnFQX18xJtUdx/C0tqFNSB7JnSpIkdbOlwMXAks2CVK2wfUlx3MHNK01Sp7NnSpIkda9y5cBpHrceOLWxxUjqNvZMSZIkSVIdDFOSJEmSVAfDlCRJkiTVwTAlSZIkSXUwTEmSJElSHbZ2NL9Xlkql6xtRiLrCK4HbW13EOLZZTcU2q05jm1Unard2a5vVlky7zZaq1eq0P9WGpy2pVquHt7qGWrZZbYltVp3GNqtO1E7t1jar6Zhum92qMCVJkiRJCnxmSpIkSZLqYJiSJEmSpDoYpiRJkiSpDoYpSZIkSaqDYUqSJEmS6mCYkiRJkqQ6GKYkSZIkqQ6GKUmSJEmqg2FKkiRJkurw/wFfnRdN32w5/gAAAABJRU5ErkJggg==\n",
            "text/plain": [
              "<Figure size 1080x216 with 5 Axes>"
            ]
          },
          "metadata": {
            "tags": [],
            "needs_background": "light"
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z-wm1ClwR-ES",
        "colab_type": "text"
      },
      "source": [
        "When you run the previous cell, you should see the GCSL agent learning how to reach all the states in the room! And that's it!\n",
        "\n",
        "If you're interested in learning more about how *why* this simple procedure works, check out our [paper](https://arxiv.org/abs/1912.06088). If you're looking to apply the algorithm on larger domains, check out our [Github repository](https://github.com/notdibya/gcsl)\n",
        "\n",
        "\n"
      ]
    }
  ]
}