labs/DRL.01.REINFORCE+A2C.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Xz5VM03vyFeb"
   },
   "source": [
    "# Reinforce & Advantage Actor Critic (A2C)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mjm8RBoQyFed"
   },
   "source": [
    "[You can find the original paper here](https://arxiv.org/pdf/1602.01783.pdf)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CWjie7h0zGdD"
   },
   "source": [
    "## Install, import and utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VrO05StZyFee"
   },
   "outputs": [],
   "source": [
    "!pip install torch==1.1.0 torchvision pyvirtualdisplay matplotlib seaborn pandas numpy pathlib gym\n",
    "!sudo apt-get install xvfb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Yx0xvpJeyFel"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F \n",
    "from torch import optim\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "from pyvirtualdisplay import Display\n",
    "from IPython import display as ipythondisplay\n",
    "from IPython.display import clear_output\n",
    "from pathlib import Path\n",
    "\n",
    "import random, os.path, math, glob, csv, base64, itertools, sys\n",
    "import gym\n",
    "from gym.wrappers import Monitor\n",
    "from pprint import pprint\n",
    "\n",
    "# The following code is will be used to visualize the environments.\n",
    "\n",
    "def show_video(directory):\n",
    "    html = []\n",
    "    for mp4 in Path(directory).glob(\"*.mp4\"):\n",
    "        video_b64 = base64.b64encode(mp4.read_bytes())\n",
    "        html.append('''<video alt=\"{}\" autoplay \n",
    "                      loop controls style=\"height: 400px;\">\n",
    "                      <source src=\"data:video/mp4;base64,{}\" type=\"video/mp4\" />\n",
    "                 </video>'''.format(mp4, video_b64.decode('ascii')))\n",
    "    ipythondisplay.display(ipythondisplay.HTML(data=\"<br>\".join(html)))\n",
    "    \n",
    "display = Display(visible=0, size=(1400, 900))\n",
    "display.start();\n",
    "\n",
    "def make_seed(seed):\n",
    "    np.random.seed(seed=seed)\n",
    "    torch.manual_seed(seed=seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pMy7W6NkyFeo"
   },
   "source": [
    "The following code is will be used to visualize the environments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "sL3r2QLayFeu"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "In this tutorial we will focus on Deep Reinforcement Learning with **Reinforce** and the **Actor-Advantage Critic** algorithm. This tutorial is composed of:\n",
    "* An introduction to the deep learning framework: **PyTorch**, \n",
    "* A quick reminder of the RL setting,\n",
    "* A theoritical and coding approch of Reinforce\n",
    "* A theoritical and coding approch of A2C."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "XsY-KC1myFeu"
   },
   "source": [
    "## Introduction to PyTorch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Uu2UA5-uyFev"
   },
   "source": [
    "*If you already know PyTorch you can skip this part. From this part on we assume that you have some experience with Python and Numpy. This part is extracted from PyTorch docs.*\n",
    "\n",
    "PyTorch is a Python package that provides two high-level features:\n",
    "- Tensor computation (like NumPy) with strong GPU acceleration\n",
    "- Deep neural networks built on a tape-based autograd system\n",
    "\n",
    "At a granular level, PyTorch is a library that consists of the following components:\n",
    "\n",
    "| Component | Description |\n",
    "| ---- | --- |\n",
    "| [**torch**](https://pytorch.org/docs/stable/torch.html) | a Tensor library like NumPy, with strong GPU support |\n",
    "| [**torch.autograd**](https://pytorch.org/docs/stable/autograd.html) | a tape-based automatic differentiation library that supports all differentiable Tensor operations in torch |\n",
    "| [**torch.jit**](https://pytorch.org/docs/stable/jit.html) | a compilation stack (TorchScript) to create serializable and optimizable models from PyTorch code  |\n",
    "| [**torch.nn**](https://pytorch.org/docs/stable/nn.html) | a neural networks library deeply integrated with autograd designed for maximum flexibility |\n",
    "| [**torch.multiprocessing**](https://pytorch.org/docs/stable/multiprocessing.html) | Python multiprocessing, but with magical memory sharing of torch Tensors across processes. Useful for data loading and Hogwild training |\n",
    "| [**torch.utils**](https://pytorch.org/docs/stable/data.html) | DataLoader and other utility functions for convenience |\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A5rsYAuvyFew"
   },
   "source": [
    "PyTorch works in a very similar way as Numpy and PyTorch's Tensors are the equivalent of Numpy's Arrays."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "HKlz3-vLyFex"
   },
   "source": [
    "You can initialize an zero filled tensor just like in numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UY-u1dchyFex"
   },
   "outputs": [],
   "source": [
    "torch.zeros(5,3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Nr6M39BHyFe1"
   },
   "outputs": [],
   "source": [
    "torch.eye(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "i8NsRchhyFe6"
   },
   "source": [
    "You can also convert an array to a tensor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "eYKCvjXqyFe7"
   },
   "outputs": [],
   "source": [
    "torch.tensor(np.eye(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WaABqYttyFfA"
   },
   "source": [
    "And you can transform a tensor to an array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F0XefR9dyFfB"
   },
   "outputs": [],
   "source": [
    "torch.tensor(np.eye(3)).numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZVG0BE0RyFfH"
   },
   "source": [
    "You can sum, substract, multiply arrays just like in numpy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "VQOUnDhByFfI"
   },
   "outputs": [],
   "source": [
    "a = torch.randint(0,10,(2,3))\n",
    "print(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Vb6ZbeicyFfN"
   },
   "outputs": [],
   "source": [
    "b = torch.randint(0,10,(2,3))\n",
    "print(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mQJcMF8ayFfe"
   },
   "outputs": [],
   "source": [
    "print(f'a + b = {a + b}')\n",
    "print(f'a * b = {a * b}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Eigd9ij0yFfj"
   },
   "source": [
    "You can make matrix products."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YQwH7cFjyFfk"
   },
   "outputs": [],
   "source": [
    "a @ b.t()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "G1c2c2A4yFfs"
   },
   "source": [
    "### AUTOGRAD: automatic differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GndhwBQ0yFft"
   },
   "source": [
    "The autograd package provides automatic differentiation for all operations on Tensors. It is a define-by-run framework, which means that your backprop is defined by how your code is run, and that every single iteration can be different.\n",
    "\n",
    "``torch.Tensor`` is the central class of the package. If you set its attribute\n",
    "``.requires_grad`` as ``True``, it starts to track all operations on it. When\n",
    "you finish your computation you can call ``.backward()`` and have all the\n",
    "gradients computed automatically. The gradient for this tensor will be\n",
    "accumulated into ``.grad`` attribute.\n",
    "\n",
    "To stop a tensor from tracking history, you can call ``.detach()`` to detach\n",
    "it from the computation history, and to prevent future computation from being\n",
    "tracked.\n",
    "\n",
    "To prevent tracking history (and using memory), you can also wrap the code block\n",
    "in ``with torch.no_grad():``. This can be particularly helpful when evaluating a\n",
    "model because the model may have trainable parameters with\n",
    "``requires_grad=True``, but for which we don't need the gradients.\n",
    "\n",
    "There’s one more class which is very important for autograd\n",
    "implementation - a ``Function``.\n",
    "\n",
    "``Tensor`` and ``Function`` are interconnected and build up an acyclic\n",
    "graph, that encodes a complete history of computation. Each tensor has\n",
    "a ``.grad_fn`` attribute that references a ``Function`` that has created\n",
    "the ``Tensor`` (except for Tensors created by the user - their\n",
    "``grad_fn is None``).\n",
    "\n",
    "If you want to compute the derivatives, you can call ``.backward()`` on\n",
    "a ``Tensor``. If ``Tensor`` is a scalar (i.e. it holds a one element\n",
    "data), you don’t need to specify any arguments to ``backward()``,\n",
    "however if it has more elements, you need to specify a ``gradient``\n",
    "argument that is a tensor of matching shape."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NkcsoISWyFfu"
   },
   "source": [
    "## Reminder of the RL setting\n",
    "\n",
    "As always we will consider a MDP $M = (\\mathcal{S}, \\mathcal{A}, p, r, \\gamma)$ with:\n",
    "* $\\mathcal{S}$ the state space,\n",
    "* $\\mathcal{A}$ the action space,\n",
    "* $p(x^\\prime \\mid x, a)$ the transition probability,\n",
    "* $r(x, a, x^\\prime)$ the reward of the transition $(x, a, x^\\prime)$,\n",
    "* $\\gamma \\in [0,1)$ is the discount factor.\n",
    "\n",
    "A policy $\\pi$ is a mapping from the state space $\\mathcal{S}$ to the probability of selecting each action.\n",
    "\n",
    "The action value function of a policy is the overall expected reward from a state action. $Q^\\pi(s, a) = \\mathbb{E}_{\\tau \\sim \\pi}\\big[R(\\tau) \\mid s_0=s, a_0=a\\big]$ where $\\tau$ is an episode $(s_0, a_0, r_0, s_1, a_1, r_1, s_2, ..., s_T, a_T, r_T)$ with the actions drawn from $\\pi(s)$; $R(\\tau)$ is the random variable defined as the cumulative sum of the discounted reward.\n",
    "\n",
    "The goal is to maximize the agent's reward.\n",
    "\n",
    "$$ J(\\pi) = \\mathbb{E}_{\\tau \\sim \\pi}\\big[R(\\tau) \\big]$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5OJnuROCyFfv"
   },
   "source": [
    "## Gym Environment\n",
    "\n",
    "In this lab and also the next one we are going to use the [OpenAI's Gym library](https://gym.openai.com/envs/). This library provides a large number of environments to test RL algorithm.\n",
    "\n",
    "We will focus on the **CartPole-v1** environment in this lab but we encourage you to also test your code on:\n",
    "* **Acrobot-v1**\n",
    "* **MountainCar-v0**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "5W7VaPHSyFfw"
   },
   "source": [
    "| Env Info          \t| CartPole-v1 \t| Acrobot-v1                \t| MountainCar-v0 \t|\n",
    "|-------------------\t|-------------\t|---------------------------\t|----------------\t|\n",
    "| **Observation Space** \t| Box(4)      \t| Box(6)                    \t| Box(2)         \t|\n",
    "| **Action Space**      \t| Discrete(2) \t| Discrete(3)               \t| Discrete(3)    \t|\n",
    "| **Rewards**           \t| 1 per step  \t| -1 if not terminal else 0 \t| -1 per step    \t|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Hthuc34vyFfy"
   },
   "source": [
    "A gym environment is loaded with the command `env = gym.make(env_id)`. Once the environment is created, you need to reset it with `observation = env.reset()` and then you can interact with it using the method step: `observation, reward, done, info = env.step(action)`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "WTEGZ1nayFfz"
   },
   "source": [
    "### Carpole"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "mnrKdIhsyFf0"
   },
   "outputs": [],
   "source": [
    "# We load CartPole-v1\n",
    "env = gym.make('CartPole-v1')\n",
    "# We wrap it in order to save our experiment on a file.\n",
    "env = Monitor(env, \"./gym-results\", force=True, video_callable=lambda episode: True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "1Eq_p3yKyFf3"
   },
   "outputs": [],
   "source": [
    "done = False\n",
    "obs = env.reset()\n",
    "while not done:\n",
    "    action = env.action_space.sample()\n",
    "    obs, reward, done, info = env.step(action)\n",
    "env.close()\n",
    "show_video(\"./gym-results\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "L4wSjOadyFf8"
   },
   "source": [
    "### Acrobot-v1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "FKGFYo3zyFf9"
   },
   "outputs": [],
   "source": [
    "# We load Acrobot-v1\n",
    "env = gym.make('Acrobot-v1')\n",
    "# We wrap it in order to save our experiment on a file.\n",
    "env = Monitor(env, \"./gym-results\", force=True, video_callable=lambda episode: True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Pue8QgfFyFf_"
   },
   "outputs": [],
   "source": [
    "done = False\n",
    "obs = env.reset()\n",
    "while not done:\n",
    "    action = env.action_space.sample()\n",
    "    obs, reward, done, info = env.step(action)\n",
    "env.close()\n",
    "show_video(\"./gym-results\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZwRK_7ZtyFgC"
   },
   "source": [
    "### MountainCar-v0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uZ9qetK-yFgC"
   },
   "outputs": [],
   "source": [
    "# We load Acrobot-v1\n",
    "env = gym.make('MountainCar-v0')\n",
    "# We wrap it in order to save our experiment on a file.\n",
    "env = Monitor(env, \"./gym-results\", force=True, video_callable=lambda episode: True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SbGbUebiyFgE"
   },
   "outputs": [],
   "source": [
    "done = False\n",
    "obs = env.reset()\n",
    "while not done:\n",
    "    action = env.action_space.sample()\n",
    "    obs, reward, done, info = env.step(action)\n",
    "env.close()\n",
    "show_video(\"./gym-results\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9B6HnSCOyFgH"
   },
   "source": [
    "## REINFORCE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IMTSmNrMyFgH"
   },
   "source": [
    "### Introduction\n",
    "\n",
    "Reinforce is an actor-based **on policy** method. The policy $\\pi_{\\theta}$ is parametrized by a function approximator (e.g. a neural network).\n",
    "\n",
    "Recall: $$ J(\\pi) = \\mathbb{E}_{\\tau \\sim \\pi}\\big[ \\sum_{t} \\gamma^t R_t \\mid x_0, \\pi \\big].$$\n",
    "\n",
    "To update the parameters $\\theta$ of the policy, one has to do gradient ascent: $\\theta_{k+1} = \\theta_{k} + \\alpha \\nabla_{\\theta}J(\\pi_{\\theta})|_{\\theta_{k}}$.\n",
    "\n",
    "**Advantages of this approach:**\n",
    "- Compared to a Q-learning approach, here the policy is directly parametrized so a small change of the parameters will not dramatically change the policy whereas this is not the case for Q-learning approaches.\n",
    "- The stochasticity of the policy allows exploration. In off policy learning, one has to deal with both a behaviour policy and an exploration policy."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "P2M2e8h-yFgI"
   },
   "source": [
    "### Policy Gradient Theorem\n",
    "\n",
    "**Q.1: Prove the Policy Gradient Theorem:** $$ \\nabla_{\\theta} J(\\pi_{\\theta}) = \\mathbb{E}_{\\tau \\sim \\pi_{\\theta}}\\left[{\\sum_{t=0}^{T} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t |s_t) R(\\tau)}\\right]$$\n",
    "\n",
    "\n",
    "The policy gradient can be approximated with:\n",
    "$$ \\hat{g} = \\frac{1}{|\\mathcal{D}|} \\sum_{\\tau \\in \\mathcal{D}} \\sum_{t=0}^{T} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t |s_t) R(\\tau) $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2YCvTlKo2NeG"
   },
   "source": [
    "#### Hint 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "9Oyts6k72_wC"
   },
   "source": [
    "The probability of a trajectory $\\tau = (s_{0}, a_{0},\\dots, s_{T+1}$) with action chosen from $\\displaystyle \\pi_{\\theta}$ is $P(\\tau|\\theta) = \\rho_{0}(s_{0})\\prod_{t=0}^{T}P\\left(s_{t+1}|s_{t}, a_{t}\\right) \\pi_{\\theta}(a_{t}|s_{t})$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "t7qjxPVB3LF9"
   },
   "source": [
    "#### Hint 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "KrGQVsqI3Noh"
   },
   "source": [
    "Gradient-log trick: $  \\nabla_{\\theta}P(\\tau|\\theta)= P(\\tau|\\theta)\\nabla_{\\theta}\\log P(\\tau|\\theta). $\n",
    "\n",
    "The policy gradient can therefore be approximated with:\n",
    "$$ \\hat{g} = \\frac{1}{|\\mathcal{D}|} \\sum_{\\tau \\in \\mathcal{D}} \\sum_{t=0}^{T} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t |s_t) R(\\tau) $$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aGxnJVUkyFgI"
   },
   "source": [
    "### Implementation of vanilla REINFORCE\n",
    "\n",
    "**Q.2: Implement the REINFORCE algorithm**\n",
    "\n",
    "The code is splitted in two parts:\n",
    "* The Model class defines the architecture of our neural network which takes as input the current state and returns the policy,\n",
    "* The Agent class is responsible for the training and evaluation procedure. You will need to code the method `optimize_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "aNH3udIuyFgK"
   },
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, dim_observation, n_actions):\n",
    "        super(Model, self).__init__()\n",
    "        \n",
    "        self.n_actions = n_actions\n",
    "        self.dim_observation = dim_observation\n",
    "        \n",
    "        self.net = nn.Sequential(\n",
    "            nn.Linear(in_features=self.dim_observation, out_features=16),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(in_features=16, out_features=8),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(in_features=8, out_features=self.n_actions),\n",
    "            nn.Softmax(dim=0)\n",
    "        )\n",
    "        \n",
    "    def forward(self, state):\n",
    "        return self.net(state)\n",
    "    \n",
    "    def select_action(self, state):\n",
    "        action = torch.multinomial(self.forward(state), 1)\n",
    "        return action"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aKGCDD4KyFgN"
   },
   "source": [
    "It is always nice to visualize the differents layers of our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UzCLFgLEyFgN"
   },
   "outputs": [],
   "source": [
    "env_id = 'CartPole-v1'\n",
    "env = gym.make(env_id)\n",
    "model = Model(env.observation_space.shape[0], env.action_space.n)\n",
    "print(f'The model we created correspond to:\\n{model}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PpMv0QFsyFgR"
   },
   "source": [
    "We provide a base agent that you will need to extend in the next cell with your implementation of `optimize_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sBWGRN6OyFgR"
   },
   "outputs": [],
   "source": [
    "class BaseAgent:\n",
    "    \n",
    "    def __init__(self, config):\n",
    "        self.config = config\n",
    "        self.env = gym.make(config['env_id'])\n",
    "        make_seed(config['seed'])\n",
    "        self.env.seed(config['seed'])\n",
    "        self.model = Model(self.env.observation_space.shape[0], self.env.action_space.n)\n",
    "        self.gamma = config['gamma']\n",
    "        self.optimizer = torch.optim.Adam(self.model.net.parameters(), lr=config['learning_rate'])\n",
    "        self.monitor_env = Monitor(env, \"./gym-results\", force=True, video_callable=lambda episode: True)\n",
    "    \n",
    "    def _make_returns(self, rewards):\n",
    "        \"\"\"Returns the cumulative discounted rewards at each time step\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        rewards : array\n",
    "            The array of rewards of one episode\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        array\n",
    "            The cumulative discounted rewards at each time step\n",
    "            \n",
    "        Example\n",
    "        -------\n",
    "        for rewards=[1, 2, 3] this method outputs [1 + 2 * gamma + 3 * gamma**2, 2 + 3 * gamma, 3] \n",
    "        \"\"\"\n",
    "        \n",
    "        returns = np.zeros_like(rewards)\n",
    "        returns[-1] = rewards[-1]\n",
    "        for t in reversed(range(len(rewards) - 1)):\n",
    "            returns[t] = rewards[t] + self.gamma * returns[t + 1]\n",
    "        return returns\n",
    "    \n",
    "    # Method to implement\n",
    "    def optimize_model(self, n_trajectories):\n",
    "        \"\"\"Perform a gradient update using n_trajectories\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_trajectories : int\n",
    "            The number of trajectories used to approximate the expectation card(D) in the formula above\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        array\n",
    "            The cumulative discounted rewards of each trajectory\n",
    "        \"\"\"\n",
    "        \n",
    "        raise NotImplementedError\n",
    "    \n",
    "    def train(self, n_trajectories, n_update):\n",
    "        \"\"\"Training method\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_trajectories : int\n",
    "            The number of trajectories used to approximate the expectation card(D) in the formula above\n",
    "        n_update : int\n",
    "            The number of gradient updates\n",
    "            \n",
    "        \"\"\"\n",
    "        \n",
    "        rewards = []\n",
    "        for episode in range(n_update):\n",
    "            rewards.append(self.optimize_model(n_trajectories))\n",
    "            print(f'Episode {episode + 1}/{n_update}: rewards {round(rewards[-1].mean(), 2)} +/- {round(rewards[-1].std(), 2)}')\n",
    "        \n",
    "        # Plotting\n",
    "        r = pd.DataFrame((itertools.chain(*(itertools.product([i], rewards[i]) for i in range(len(rewards))))), columns=['Epoch', 'Reward'])\n",
    "        sns.lineplot(x=\"Epoch\", y=\"Reward\", data=r, ci='sd');\n",
    "        \n",
    "    def evaluate(self):\n",
    "        \"\"\"Evaluate the agent on a single trajectory            \n",
    "        \"\"\"\n",
    "        \n",
    "        observation = self.monitor_env.reset()\n",
    "        observation = torch.tensor(observation, dtype=torch.float)\n",
    "        reward_episode = 0\n",
    "        done = False\n",
    "            \n",
    "        while not done:\n",
    "            action = self.model.select_action(observation)\n",
    "            observation, reward, done, info = self.monitor_env.step(int(action))\n",
    "            observation = torch.tensor(observation, dtype=torch.float)\n",
    "            reward_episode += reward\n",
    "        \n",
    "        self.monitor_env.close()\n",
    "        show_video(\"./gym-results\")\n",
    "        print(f'Reward: {reward_episode}')\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-lS-k_EWyFgT"
   },
   "outputs": [],
   "source": [
    "class SimpleAgent(BaseAgent):\n",
    "    \n",
    "    def optimize_model(self, n_trajectories):\n",
    "        \"\"\"Perform a gradient update using n_trajectories\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_trajectories : int\n",
    "            The number of trajectories used to approximate the expectation card(D) in the formula above\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        array\n",
    "            The cumulative discounted rewards of each trajectory\n",
    "        \"\"\"\n",
    "        \n",
    "        ###\n",
    "        # Your code here\n",
    "        ###\n",
    "        \n",
    "        reward_trajectories = None\n",
    "        loss = None\n",
    "        \n",
    "        # The following lines take care of the gradient descent step for the variable loss\n",
    "        # that you need to compute.\n",
    "        \n",
    "        # Discard previous gradients\n",
    "        self.optimizer.zero_grad()\n",
    "        # Compute the gradient \n",
    "        loss.backward()\n",
    "        # Do the gradient descent step\n",
    "        self.optimizer.step()\n",
    "        return reward_trajectories\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the cell bellow are listed the parameters you should play with. Try out different configurations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "KMLAy5fY-1V4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Current config is:\n",
      "{'env_id': 'CartPole-v1', 'gamma': 1, 'learning_rate': 0.01, 'seed': 1235}\n"
     ]
    }
   ],
   "source": [
    "#@title Config {display-mode: \"form\", run: \"auto\"}\n",
    "\n",
    "env_id = 'CartPole-v1'  #@param [\"CartPole-v1\", \"Acrobot-v1\", \"MountainCar-v0\"]\n",
    "learning_rate = 0.01  #@param {type: \"number\"}\n",
    "gamma = 1  #@param {type: \"number\"}\n",
    "seed = 1235  #@param {type: \"integer\"}\n",
    "#@markdown ---\n",
    "\n",
    "config = {\n",
    "    'env_id': env_id,\n",
    "    'learning_rate': learning_rate,\n",
    "    'seed': seed,\n",
    "    'gamma': gamma\n",
    "}\n",
    "\n",
    "print(\"Current config is:\")\n",
    "pprint(config)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's train the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MHmntrq4yFgW",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "agent = SimpleAgent(config)\n",
    "agent.train(n_trajectories=50, n_update=50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Lu3QdaqhyFgZ"
   },
   "source": [
    "Let's evaluate the quality of the learned policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "G3phyHpgyFga"
   },
   "outputs": [],
   "source": [
    "agent.evaluate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rkKkc2OxyFgf"
   },
   "source": [
    "**Q.3: What are the strengths and drawbacks of this algorithm? How would you improve it?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "IvfQ8GSPyFgg"
   },
   "source": [
    "*Type your answer here*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Don't let the past distract you"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rkYWtsuuyFgj"
   },
   "source": [
    "- The sum of rewards during one episode has a high variance which affects the performance of this version of **REINFORCE**.\n",
    "- To assess the quality of an action, it make more sense to take into consideration only the rewards obtained after taking this action.\n",
    "- It can be proven that $$  \\nabla_{\\theta} J(\\pi_{\\theta}) = \\mathbb{E}_{\\tau \\sim \\pi_{\\theta}}\\left[{\\sum_{t=0}^{T} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t |s_t) \\sum_{t'=t}^T \\gamma^{t'-t} R(s_{t'}, a_{t'}, s_{t'+1})}\\right].$$\n",
    "- **Bonus**: proof of this claim.\n",
    "- This has for effect to reduce the variance. Past rewards have zero mean but nonzero variance so they just add noise.  \n",
    "\n",
    "**Q4: Implement this enhanced version of REINFORCE**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "LuwbPfLfyFgj"
   },
   "outputs": [],
   "source": [
    "class EnhancedAgent(BaseAgent):\n",
    "    \n",
    "    def optimize_model(self, n_trajectories):\n",
    "        \"\"\"Perform a gradient update using n_trajectories\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        n_trajectories : int\n",
    "            The number of trajectories used to approximate the expectation card(D) in the formula above\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        array\n",
    "            The cumulative discounted rewards of each trajectory\n",
    "        \"\"\"\n",
    "        \n",
    "        ###\n",
    "        # Your code here\n",
    "        ###\n",
    "        \n",
    "        reward_trajectories = None\n",
    "        loss = None\n",
    "        \n",
    "        self.optimizer.zero_grad()\n",
    "        # Compute the gradient \n",
    "        loss.backward()\n",
    "        # Do the gradient descent step\n",
    "        self.optimizer.step()\n",
    "        return reward_trajectories\n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bQ-fznp1A9w7"
   },
   "outputs": [],
   "source": [
    "#@title Config Enhanced {display-mode: \"form\", run: \"auto\"}\n",
    "\n",
    "env_id = 'CartPole-v1'  #@param [\"CartPole-v1\", \"Acrobot-v1\", \"MountainCar-v0\"]\n",
    "learning_rate = 0.001  #@param {type: \"number\"}\n",
    "gamma = 1  #@param {type: \"number\"}\n",
    "seed = 1  #@param {type: \"integer\"}\n",
    "#@markdown ---\n",
    "\n",
    "config_enhanced = {\n",
    "    'env_id': env_id,\n",
    "    'learning_rate': learning_rate,\n",
    "    'seed': seed,\n",
    "    'gamma': gamma\n",
    "}\n",
    "\n",
    "print(\"Current config_enhanced is:\")\n",
    "pprint(config_enhanced)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "HVGu8w84yFgn",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "agent = EnhancedAgent(config_enhanced)\n",
    "agent.train(n_trajectories=50, n_update=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XSIb6UuMyFgr"
   },
   "outputs": [],
   "source": [
    "agent.evaluate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "n653mbx-yFgx"
   },
   "source": [
    "**Q.5: Did this method improve over vanilla REINFORCE?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "3HuYUaXwyFgx"
   },
   "source": [
    "*Type your answer here*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "I9kdsiBQ31mF"
   },
   "source": [
    "## A2C"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e-dB7kWq4AOu"
   },
   "source": [
    "### Theory\n",
    "The cumulative discounted reward has a high variance and therefore **REINFORCE** needs lots of trajectories to converge.\n",
    "In order to reduce it we can substract a baseline $b(s_t)$.\n",
    "\n",
    "This is possible because:\n",
    "\n",
    "$$\\mathbb{E}_{a_t \\sim \\pi_{\\theta}}{\\nabla_{\\theta} \\log \\pi_{\\theta}(a_t|s_t) b(s_t)} = 0.$$\n",
    "\n",
    "**Proof**\n",
    "\n",
    "Let $P_\\theta$ be the parameterized probability distribution over a random variable x.\n",
    "$$\\int_x P_\\theta (x) = 1$$\n",
    "Taking the gradient we get: $\\nabla_\\theta \\int_x P_\\theta (x) = \\nabla_\\theta 1 = 0$\n",
    "\\begin{align*}\n",
    "0 &= \\nabla_\\theta \\int_x P_\\theta (x)\\\\\n",
    "&= \\int_x \\nabla_\\theta P_\\theta(x)\\\\\n",
    "&= \\int_x P_\\theta (x) \\nabla_\\theta \\log P_\\theta (x) \\\\\n",
    "&= \\mathbb{E}_{x \\sim P_\\theta} \\nabla_\\theta \\log P_\\theta (x)\n",
    "\\end{align*}\n",
    "\n",
    "Which leads to the following formula:\n",
    "\n",
    "$$ \\nabla_{\\theta} J(\\pi_{\\theta}) = \\mathbb{E}_{\\tau \\sim \\pi_{\\theta}}{\\sum_{t=0}^{T} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t |s_t) \\left(\\sum_{t'=t}^T \\gamma^{t' - t} R(s_{t'}, a_{t'}, s_{t'+1}) - b(s_t)\\right)}$$\n",
    "\n",
    "The most common choice of baseline is the on-policy value function $V^{\\pi}(s_t)$. This choice has the desirable effect of reducing variance in the sample estimate for the policy gradient. This results in faster and more stable policy learning. It is also appealing from a conceptual angle: it encodes the intuition that if an agent gets what it expected, it should “feel” neutral about it.\n",
    "\n",
    "However $V^\\pi$ is unknown and therefore we need to learn it. We will use a neural network to approximate $V^\\pi$ with Mean Square Error as the loss function.\n",
    "\n",
    "$$ \\arg \\min_{\\phi} \\mathbb{E}_{s_t, \\hat{R}_t \\sim \\pi_k}{\\left( V_{\\phi}(s_t) - \\hat{R}_t \\right)^2}, $$\n",
    "\n",
    "\n",
    "We can show that we can replace $\\sum_{t'=t}^T \\gamma^{t' - t} R(s_{t'}, a_{t'}, s_{t'+1})$ in the formula above by $Q^{\\pi_\\theta}(s_t, a_t)$ and by doing so we finally get $$ \\nabla_{\\theta} J(\\pi_{\\theta}) = \\mathbb{E}_{\\tau \\sim \\pi_{\\theta}}{\\sum_{t=0}^{T} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t |s_t) \\left(Q^{\\pi_\\theta}(s_t, a_t) - V^{\\pi_\\theta}(s_t)\\right)}$$\n",
    "\n",
    "**Bonus: Prove the statement above**\n",
    "\n",
    "$A(s_t, a_t) = Q(s_{t}, a_{t}) - V^\\pi(s_t)$ is called the Advantage function which gives the name **Advantage Actor Critic**.\n",
    "\n",
    "$Q^{\\pi_\\theta}(s_t, a_t)$ is approximated as the cumulative sum of rewards.\n",
    "\n",
    "Minh et al. in their paper explained that they added an entropy term to the loss in order to encourage exploration.\n",
    "\n",
    "$$ - \\sum_{a} \\pi(a | s) \\log \\pi(a | s) $$\n",
    "\n",
    "**Q.6: Explain why adding the entropy term encourage exploration.**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xDifFS9I4X7A"
   },
   "source": [
    "### Coding\n",
    "\n",
    "In the first part of this lab we had to wait for the end of an episode in order to compute the cumulative discounted rewards. Here we can use the critic to estimate the cumulative discounted reward and therefore we no longer need to wait for the episode termination.\n",
    "\n",
    "**Example**: For a trajectory $\\tau = (s_0, a_0, r_0, s_1, a_1, r_1, ..., s_t, a_t, r_t, ..., s_T, a_T, r_T)$ the cumulative discounted rewards $R(\\tau) = \\sum_{i=0}^T \\gamma^i r_i$ can be approximated by $ \\sum_{i=0}^t \\gamma^i r_i + \\gamma^{t+1} V_\\phi(s_{t+1})$ where $V_\\phi$ is our critic.\n",
    "\n",
    "This is allows us to train our model using batch data, which is more efficient for:\n",
    "* Computing: deep learning libraries (PyTorch, TensorFlow, ...) are optimized for batched data;\n",
    "* We don't have to wait for the end of a long episode in order to perform the update;\n",
    "* It is more sample efficient;\n",
    "* ...\n",
    "\n",
    "\n",
    "**Q.7: Implement `optimize_model` using batches**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5DlNKriHyFg3"
   },
   "outputs": [],
   "source": [
    "class ActorNetwork(nn.Module):\n",
    "\n",
    "    def __init__(self, input_size, hidden_size, action_size):\n",
    "        super(ActorNetwork, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_size, hidden_size)\n",
    "        self.fc2 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.fc3 = nn.Linear(hidden_size, action_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.fc1(x))\n",
    "        out = F.relu(self.fc2(out))\n",
    "        out = F.softmax(self.fc3(out), dim=-1)\n",
    "        return out\n",
    "    \n",
    "    def select_action(self, x):\n",
    "        return torch.multinomial(self(x), 1).detach().numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SgnRZXEayFhA"
   },
   "outputs": [],
   "source": [
    "class ValueNetwork(nn.Module):\n",
    "\n",
    "    def __init__(self, input_size, hidden_size, output_size):\n",
    "        super(ValueNetwork, self).__init__()\n",
    "        self.fc1 = nn.Linear(input_size, hidden_size)\n",
    "        self.fc2 = nn.Linear(hidden_size, hidden_size)\n",
    "        self.fc3 = nn.Linear(hidden_size, output_size)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = F.relu(self.fc1(x))\n",
    "        out = F.relu(self.fc2(out))\n",
    "        out = self.fc3(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_5pS7cRyyFhC"
   },
   "outputs": [],
   "source": [
    "class A2CAgent:\n",
    "\n",
    "    def __init__(self, config):\n",
    "        self.config = config\n",
    "        self.env = gym.make(config['env_id'])\n",
    "        make_seed(config['seed'])\n",
    "        self.env.seed(config['seed'])\n",
    "        self.monitor_env = Monitor(self.env, \"./gym-results\", force=True, video_callable=lambda episode: True)\n",
    "        self.gamma = config['gamma']\n",
    "        \n",
    "        # Our two networks\n",
    "        self.value_network = ValueNetwork(self.env.observation_space.shape[0], 16, 1)\n",
    "        self.actor_network = ActorNetwork(self.env.observation_space.shape[0], 16, self.env.action_space.n)\n",
    "        \n",
    "        # Their optimizers\n",
    "        self.value_network_optimizer = optim.RMSprop(self.value_network.parameters(), lr=config['value_network']['learning_rate'])\n",
    "        self.actor_network_optimizer = optim.RMSprop(self.actor_network.parameters(), lr=config['actor_network']['learning_rate'])\n",
    "        \n",
    "    # Hint: use it during training_batch\n",
    "    def _returns_advantages(self, rewards, dones, values, next_value):\n",
    "        \"\"\"Returns the cumulative discounted rewards at each time step\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        rewards : array\n",
    "            An array of shape (batch_size,) containing the rewards given by the env\n",
    "        dones : array\n",
    "            An array of shape (batch_size,) containing the done bool indicator given by the env\n",
    "        values : array\n",
    "            An array of shape (batch_size,) containing the values given by the value network\n",
    "        next_value : float\n",
    "            The value of the next state given by the value network\n",
    "        \n",
    "        Returns\n",
    "        -------\n",
    "        returns : array\n",
    "            The cumulative discounted rewards\n",
    "        advantages : array\n",
    "            The advantages\n",
    "        \"\"\"\n",
    "        \n",
    "        returns = np.append(np.zeros_like(rewards), [next_value], axis=0)\n",
    "        \n",
    "        for t in reversed(range(rewards.shape[0])):\n",
    "            returns[t] = rewards[t] + self.gamma * returns[t + 1] * (1 - dones[t])\n",
    "            \n",
    "        returns = returns[:-1]\n",
    "        advantages = returns - values\n",
    "        return returns, advantages\n",
    "\n",
    "    def training_batch(self, epochs, batch_size):\n",
    "        \"\"\"Perform a training by batch\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        epochs : int\n",
    "            Number of epochs\n",
    "        batch_size : int\n",
    "            The size of a batch\n",
    "        \"\"\"\n",
    "        episode_count = 0\n",
    "        actions = np.empty((batch_size,), dtype=np.int)\n",
    "        dones = np.empty((batch_size,), dtype=np.bool)\n",
    "        rewards, values = np.empty((2, batch_size), dtype=np.float)\n",
    "        observations = np.empty((batch_size,) + self.env.observation_space.shape, dtype=np.float)\n",
    "        observation = self.env.reset()\n",
    "        rewards_test = []\n",
    "\n",
    "        for epoch in range(epochs):\n",
    "            # Lets collect one batch\n",
    "            for i in range(batch_size):\n",
    "                observations[i] = observation\n",
    "                values[i] = self.value_network(torch.tensor(observation, dtype=torch.float)).detach().numpy()\n",
    "                policy = self.actor_network(torch.tensor(observation, dtype=torch.float))\n",
    "                actions[i] = torch.multinomial(policy, 1).detach().numpy()\n",
    "                observation, rewards[i], dones[i], _ = self.env.step(actions[i])\n",
    "\n",
    "                if dones[i]:\n",
    "                    observation = self.env.reset()\n",
    "\n",
    "            # If our epiosde didn't end on the last step we need to compute the value for the last state\n",
    "            if dones[-1]:\n",
    "                next_value = 0\n",
    "            else:\n",
    "                next_value = self.value_network(torch.tensor(observation, dtype=torch.float)).detach().numpy()[0]\n",
    "            \n",
    "            # Update episode_count\n",
    "            episode_count += sum(dones)\n",
    "\n",
    "            # Compute returns and advantages\n",
    "            returns, advantages = self._returns_advantages(rewards, dones, values, next_value)\n",
    "\n",
    "            # Learning step !\n",
    "            self.optimize_model(observations, actions, returns, advantages)\n",
    "\n",
    "            # Test it every 50 epochs\n",
    "            if epoch % 50 == 0 or epoch == epochs - 1:\n",
    "                rewards_test.append(np.array([self.evaluate() for _ in range(50)]))\n",
    "                print(f'Epoch {epoch}/{epochs}: Mean rewards: {round(rewards_test[-1].mean(), 2)}, Std: {round(rewards_test[-1].std(), 2)}')\n",
    "\n",
    "                # Early stopping\n",
    "                if rewards_test[-1].mean() > 490 and epoch != epochs -1:\n",
    "                    print('Early stopping !')\n",
    "                    break\n",
    "                observation = self.env.reset()\n",
    "                    \n",
    "        # Plotting\n",
    "        r = pd.DataFrame((itertools.chain(*(itertools.product([i], rewards_test[i]) for i in range(len(rewards_test))))), columns=['Epoch', 'Reward'])\n",
    "        sns.lineplot(x=\"Epoch\", y=\"Reward\", data=r, ci='sd');\n",
    "        \n",
    "        print(f'The trainnig was done over a total of {episode_count} episodes')\n",
    "\n",
    "    def optimize_model(self, observations, actions, returns, advantages):\n",
    "        actions = F.one_hot(torch.tensor(actions), self.env.action_space.n)\n",
    "        returns = torch.tensor(returns[:, None], dtype=torch.float)\n",
    "        advantages = torch.tensor(advantages, dtype=torch.float)\n",
    "        observations = torch.tensor(observations, dtype=torch.float)\n",
    "\n",
    "        # MSE for the values\n",
    "        # Actor & Entropy loss\n",
    "        \n",
    "        raise NotImplementedError\n",
    "\n",
    "    def evaluate(self, render=False):\n",
    "        env = self.monitor_env if render else self.env\n",
    "        observation = env.reset()\n",
    "        observation = torch.tensor(observation, dtype=torch.float)\n",
    "        reward_episode = 0\n",
    "        done = False\n",
    "\n",
    "        while not done:\n",
    "            policy = self.actor_network(observation)\n",
    "            action = torch.multinomial(policy, 1)\n",
    "            observation, reward, done, info = env.step(int(action))\n",
    "            observation = torch.tensor(observation, dtype=torch.float)\n",
    "            reward_episode += reward\n",
    "            \n",
    "        env.close()\n",
    "        if render:\n",
    "            show_video(\"./gym-results\")\n",
    "            print(f'Reward: {reward_episode}')\n",
    "        return reward_episode"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Qm9Rj2s5yFhF"
   },
   "source": [
    "**Q.8: Try out different hyperparameters (batch size, learning rate, optimizer, gamma) and identify how each one of them influence the learning.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5e734MaaDsal"
   },
   "outputs": [],
   "source": [
    "#@title Config A2C {display-mode: \"form\", run: \"auto\"}\n",
    "\n",
    "env_id = 'CartPole-v1'  #@param [\"CartPole-v1\", \"Acrobot-v1\", \"MountainCar-v0\"]\n",
    "value_learning_rate = 0.001  #@param {type: \"number\"}\n",
    "actor_learning_rate = 0.001  #@param {type: \"number\"}\n",
    "gamma = 1  #@param {type: \"number\"}\n",
    "entropy = 1  #@param {type: \"number\"}\n",
    "seed = 1  #@param {type: \"integer\"}\n",
    "#@markdown ---\n",
    "\n",
    "config_a2c = {\n",
    "    'env_id': env_id,\n",
    "    'gamma': gamma,\n",
    "    'seed': seed,\n",
    "    'value_network': {'learning_rate': value_learning_rate},\n",
    "    'actor_network': {'learning_rate': actor_learning_rate},\n",
    "    'entropy': entropy\n",
    "}\n",
    "\n",
    "print(\"Current config_a2c is:\")\n",
    "pprint(config_a2c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fG52hSe6yFhG"
   },
   "outputs": [],
   "source": [
    "agent = A2CAgent(config_a2c)\n",
    "rewards = agent.training_batch(1000, 256)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gXmqUtuuyFhJ"
   },
   "outputs": [],
   "source": [
    "agent.evaluate(True);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "iz4YuO-dyFhO"
   },
   "source": [
    "**Q.9: What are the strengths and drawbacks of this algorithm? How would you improve it?**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "G7AFxgDv4orr"
   },
   "source": [
    "*Type your answer here*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "CYidT03TyFhP"
   },
   "source": [
    "**Q.10: Compare the three algorithms (sample efficiency, stability, ...)**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rIKL7MUX4pq6"
   },
   "source": [
    "*Type your answer here*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NZULYBWcyFhQ"
   },
   "source": [
    "Make sure to try the algorithms with other environments!"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "name": "01_REINFORCE+A2C.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}