3a_reinforcement_learning.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Khipu Practical 3A: Reinforcement Learning",
      "provenance": [],
      "collapsed_sections": [
        "MIcGGuCC4bDu",
        "S-KrdnsONUKh",
        "iAnootcVLZ-x",
        "P7BiER06QQqm"
      ]
    },
    "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.7.5rc1"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XgowXSLLVU4c",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/khipu-ai/practicals-2019/blob/master/3a_reinforcement_learning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "y__eyOeyzbp1"
      },
      "source": [
        "# Practical 3A: Reinforcement Learning\n",
        "\n",
        "© Deep Learning Indaba. Apache License 2.0.\n",
        "\n",
        "Adapted for Khipu."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1lDTQg9HNF84"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "In this practical, we will cover the basics of reinforcement learning, which has successfully been used to [control robotic hands](https://openai.com/blog/learning-dexterity/); play [Chess](https://deepmind.com/blog/article/alphazero-shedding-new-light-grand-games-chess-shogi-and-go), [Go](https://en.wikipedia.org/wiki/AlphaGo), and [StarCraft](https://deepmind.com/blog/article/alphastar-mastering-real-time-strategy-game-starcraft-ii). It can even be used for supervised learning problems. For example, training a neural net to optimize a non-differentiable objective (such as accuracy, or the [BLEU score in machine translation](https://arxiv.org/abs/1609.08144)) or [finding a good neural net architecture](https://arxiv.org/abs/1611.01578) for a particular supervised learning problem.\n",
        "\n",
        "## Learning Objectives\n",
        "\n",
        "* Understand the problem reinforcement learning tries to solve.\n",
        "* Understand the terminology used in reinforcement learning: **environment**, **action**, **observation**, **reward**, **agent**,  **policy**, **episode**.\n",
        "* Learn how to use the [OpenAI Gym](https://gym.openai.com/) reinforcement learning environments.\n",
        "* Learn how to solve the classic reinforcement learning problem of balancing a pole on a moving cart using Random Search and Policy Gradients. Understand some of the limitations of these methods."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ROYVaW7ZOrPD"
      },
      "source": [
        "**IMPORTANT: Please fill out the exit ticket form before you leave the practical: https://forms.gle/FLQ78eAEQprsWXp38**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "rExe21JpsoG-"
      },
      "source": [
        "## Imports"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "LAFmDu7NzuhZ",
        "colab": {}
      },
      "source": [
        "#@title Imports (RUN ME!)\n",
        "\n",
        "\n",
        "# Note: enviroments like CartPole-v0 require a display to render. We need to install pyvirtualdisplay etc \n",
        "# in order to render from these environments\n",
        "\n",
        "!pip install pyglet~=1.3.2 > /dev/null 2>&1\n",
        "!pip install 'gym[atari]' > /dev/null 2>&1\n",
        "!apt-get install python-opengl -y > /dev/null 2>&1\n",
        "!apt install xvfb -y > /dev/null 2>&1\n",
        "!pip install pyvirtualdisplay > /dev/null 2>&1\n",
        "!pip install tensorflow==2.0.0-beta0 > /dev/null 2>&1\n",
        "\n",
        "from IPython.display import HTML\n",
        "\n",
        "# Start virtual display\n",
        "from pprint import pprint\n",
        "import logging\n",
        "from pyvirtualdisplay import Display\n",
        "logging.getLogger(\"pyvirtualdisplay\").setLevel(logging.ERROR)\n",
        "\n",
        "display = Display(visible=0, size=(1024, 768))\n",
        "display.start()\n",
        "import os\n",
        "os.environ[\"DISPLAY\"] = \":\" + str(display.display) + \".\" + str(display.screen)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Y8W2Gee9m2G3"
      },
      "source": [
        "## The Reinforcement Learning Problem\n",
        "\n",
        "So far we have encountered **supervised learning**, where we have an input and a target value or class that we want to predict. We have also encountered **unsupervised learning**, where we are only given an input and look for patterns in that input. In this practical, we look into **reinforcement learning**, which can loosely be defined as training an **agent** to maximise the total **reward** it obtains through many interactions with an **environment**.\n",
        "\n",
        "At timestep $t$, the agent can make an **observation** of the environment $o_t$. For example, if the environment is a computer game, the observation could be the pixel values of the current screen.\n",
        "\n",
        "The environment defines a set of **actions** that an agent can take.  The agent performs an action $a_t$ informed by the observations it has made, and will receive a **reward** $r_t$ from the environment after every action. The *reinforcement learning problem* is to find an agent whose actions maximize the total rewards obtained from the environment over many actions.\n",
        "\n",
        "The following diagram illustrates the interaction between the agent and environment. We will explore each of the terms in more detail throughout this practical.\n",
        "\n",
        "<!-- ![Interaction of Agent and Environment](https://github.com/sbodenstein/Tutorials/blob/master/indaba_2019/Images/RL_Environment.png?raw=true) -->\n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/sbodenstein/Tutorials/master/indaba_2019/Images/RL_Environment.png\" alt=\"drawing\" width=600/>\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Xuz5T5jEpE9t"
      },
      "source": [
        "**Optional Recommended Reading**: \n",
        "- [*OpenAI Spinning Up: Key Concepts in RL*](https://spinningup.openai.com/en/latest/spinningup/rl_intro.html) is a great summary of the terminology used in reinforcement learning. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1ORYJinC0INe"
      },
      "source": [
        "### The Environment"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Anhtx1DVoWBI"
      },
      "source": [
        "We will focus on the cartpole environment for this practical. This environment consists of a pole attached to a cart via a hinge, with the pole initially balanced close to upright. The agent needs to move the cart to the left or to the right in order to prevent the pole from falling over.\n",
        "\n",
        "![CartPole Illustration](https://user-images.githubusercontent.com/10624937/42135683-dde5c6f0-7d13-11e8-90b1-8770df3e40cf.gif)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zbb_4Ezl6NpZ"
      },
      "source": [
        "We will use the OpenAI Gym implementation of the cartpole environment. [OpenAI Gym](https://gym.openai.com/) is probably the most popular set of reinforcement learning environments (the available environments in Gym can be seen [here](https://gym.openai.com/envs)). Every Gym environment has the same interface, allowing code written for one environment to work for all of them. Popular reinforcement learning frameworks, such as [Ray](https://ray.readthedocs.io/en/latest/index.html), often use the Gym interface as [their default interface](https://ray.readthedocs.io/en/latest/rllib-env.html#openai-gym) for reinforcement learning environments.\n",
        "\n",
        "Let us import Gym and open a cartpole environment:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "qvInNnDj5a7i",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import gym\n",
        "env_cartpole = gym.make('CartPole-v1')"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "RXRdB9f-pp5N"
      },
      "source": [
        "The first step to using a Gym environment is to initialize the environment to some initial configuration using the `reset` method. The `reset` method also returns the first observation of the environment:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "PlxBug7fpsKv",
        "colab": {}
      },
      "source": [
        "env_cartpole.reset()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Vu8gPKvOpw7h"
      },
      "source": [
        "These four numbers represent the position and velocity of the cart and pole, `(cart position, cart velocity, pole angle, velocity of the top of the pole)`. We will want our agent to use this observation when deciding to move left or right. \n",
        "\n",
        "We also want to see an image of the system that we can interpret. Gym provides the `render` method to do this:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "BAPiiBUZ1on0",
        "colab": {}
      },
      "source": [
        "import matplotlib.pyplot as plt\n",
        "plt.imshow(env_cartpole.render(mode='rgb_array'));"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4mxir4Ti1r_-"
      },
      "source": [
        "Note that reinitializing this environment with `reset` will randomly change the starting angle of the pole. By running this multiple times, see if you can notice this:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "mxQu--tI1vJ1",
        "colab": {}
      },
      "source": [
        "env_cartpole.reset()\n",
        "plt.imshow(env_cartpole.render(mode='rgb_array'));"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "LmlS44FO13zy"
      },
      "source": [
        "Now we want to take an action using the `step` method. But which actions are allowed, and how are they represented? We can find this out using the `action_space` property of Gym environments:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "m58afkbh110t",
        "colab": {}
      },
      "source": [
        "env_cartpole.action_space"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1eV7U47o17JR"
      },
      "source": [
        "In Gym, `Discrete(n)` means actions are represented by the integers `0,1,...,n-1`. Action spaces always provide a way to take a random action sampled from the space:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "d2yymSqg1-2p",
        "colab": {}
      },
      "source": [
        "[env_cartpole.action_space.sample() for _ in range(10)]"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "w2rn0FHq18Rd"
      },
      "source": [
        "Now let us take an action using the `step` method, which takes an action and returns a tuple `(observation, reward, done, info)`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "v8opAXnH2EOT",
        "colab": {}
      },
      "source": [
        "action = 1\n",
        "env_cartpole.step(action)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "4CwkWYnD2HyX"
      },
      "source": [
        "**Question**: Does the action `0` move the cart to the left or to the right? Use the `step` and `render` methods to figure this out.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PQVrk9Qw2Mu5"
      },
      "source": [
        "What is `done`? Many environments are **episodic**, which means they have a natural end-point, after which more actions cannot be taken. For example in chess, the game ends after a checkmate. In the case of cart-pole, this end-point is when the pole falls down too far. Then no matter the actions taken, the pole will keep falling. Gym will set `done` as `True` in this case, and will print a warning if `step` is called again. The `reset` method needs to be called to reinitialize the environment when `done` is `True`.  We won't use `info` (see the [official Gym docs](https://gym.openai.com/docs/#observations) for more information).\n",
        "\n",
        "In addition, many Gym environments have a maximum number of steps they can take, after which `done` is set to `True`. This can be seen using the `_max_episode_steps` property:\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "zV0ueJJiANb4",
        "colab": {}
      },
      "source": [
        "env_cartpole._max_episode_steps"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xdXd4kGe2ZtX"
      },
      "source": [
        "**Question**: What is an appropriate reward for a game of chess? Suppose the environment is one where your agent plays against an existing chess-playing computer. Remember that there are three outcomes in chess, a win, loss and draw.\n",
        "\n",
        "#### Question\n",
        "\n",
        "Write a function `actions_till_done(env, act)` that takes an environment `env` and action `act`, reinitializes `env` with `reset`, and returns the number of actions taken before `done` is `True`. If you always take the `0` action, how many steps can you take on average before `done` is `True`?\n",
        "\n",
        "Reveal the cell below by double-clicking and running it, to check your answer when you're done or you get stuck!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "D3gvOFdv2cvp",
        "colab": {}
      },
      "source": [
        "#@title Answer { display-mode: \"form\" }\n",
        "def actions_till_done(env, action):\n",
        "    env.reset()\n",
        "    done = False\n",
        "    count = 0\n",
        "    while done == False:\n",
        "        _, _, done, _, = env.step(action)\n",
        "        count += 1\n",
        "    return count\n",
        "  \n",
        "# Get an estimate for how many 0 actions you can take on average before the \n",
        "# pole falls over\n",
        "np.mean([actions_till_done(env_cartpole, 0) for _ in range(100)])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bDtZsDmc2xoU"
      },
      "source": [
        "## Agents and Policies"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MlfNAnx623yN"
      },
      "source": [
        "A **policy** is a function used by an agent to decide what actions to take. This function takes as input an observation and returns either an action (a **deterministic policy**) or a probability distribution over possible actions (a **stochastic policy**).\n",
        "\n",
        "**Note**: the terms **policy** and **agent** are often used interchangeably:\n",
        "> Because the policy is essentially the agent’s brain, it’s not uncommon to substitute the word “policy” for “agent”, e.g. saying “The policy is trying to maximize reward.” ~ [OpenAI Spinning Up](https://spinningup.openai.com/en/latest/spinningup/rl_intro.html#policies)\n",
        "\n",
        "We will consider an **agent** to be a function whose input is the observation, and it returns a dictionary which always contains the `\"Action\"` key. It might also contain other keys, such as the action probabilities, the value function, etc. \n",
        "\n",
        "The simplest possible policy for a cartpole agent is one that always takes the `0` action, no matter what the observation is:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "7HvWk6Q52_Qh",
        "colab": {}
      },
      "source": [
        "def agent_left(observation):\n",
        "    return {\"Action\":0}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "lmP-cN3Y3Bg6",
        "colab": {}
      },
      "source": [
        "agent_left(None)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ErZWjRDd3Ere"
      },
      "source": [
        "## Rewards and Returns"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "JglPsEwj3Hip"
      },
      "source": [
        "Recall that in Reinforcement Learning, the **reward** is a scalar signal that the agent observes upon acting on the environment and which it will try to maximize. In the cartpole example, a reward of **+1** is provided for every timestep that the pole remains upright (less than 15 degrees from the vertical). \n",
        "\n",
        "Is our agent `agent_left` any good? To answer that, we need some measure of 'goodness', which is usually related to the total return obtained during an episode. This is called the **return**, $R$, and there are different ways of defining it. The most obvious definition for return is to simply add up all the rewards the agent received during an episode. If the episode took $T$ steps, and got a reward $r_t$ at each step, then this return is $R = \\sum_{t=1}^{T} r_{t}$, where $r_t$ is the reward received at timestep $t$.\n",
        "\n",
        "One issue is that some environments will have different initial states, agents may take actions stochastically and the rewards themselves could be stochastic. The solution is to average the returns over multiple starting points. Here is a function that does this:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "qG7mfsIj3Jts",
        "colab": {}
      },
      "source": [
        "def average_episodic_return(env, agent, episodes=10, max_steps_per_episode=500):\n",
        "    episode_rewards = []\n",
        "    for episode in range(episodes):\n",
        "        total_rewards = 0\n",
        "        obs = env.reset()\n",
        "        for t in range(max_steps_per_episode):\n",
        "            out = agent(obs) # we don't care about the probabilities here\n",
        "            assert (\"Action\" in out), \"The key 'Action' was missing from the agents output.\"\n",
        "            obs, rew, done, _ = env.step(out[\"Action\"])\n",
        "            total_rewards += rew\n",
        "            # check if we are done, if so, exit loop\n",
        "            if done:\n",
        "                break\n",
        "                \n",
        "        episode_rewards.append(total_rewards)\n",
        "        \n",
        "    return {\n",
        "        \"AverageEpisodicReturn\": np.mean(episode_rewards), \n",
        "        \"StandardDeviation\":np.sqrt(np.var(episode_rewards))\n",
        "    }"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZZpMRVbG3MHI"
      },
      "source": [
        "Now we can see how good our agent `agent_left` is:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "qYcWYEvk3OJH",
        "colab": {}
      },
      "source": [
        "average_episodic_return(env_cartpole, agent_left)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "evLyMPk-22Qk"
      },
      "source": [
        "## Animating Agents"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "kSzeeRPu3VzJ"
      },
      "source": [
        "This code cell defines a function `animate_agent(environment, agent)` that produces a human understandable animation of the `agent` controlling the `environment`. **Note:** It is not necessary to understand these functions to follow this practical."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "DC0JN7FC4OQ0",
        "colab": {}
      },
      "source": [
        "from collections import deque\n",
        "import matplotlib.animation\n",
        "import matplotlib.pyplot as plt\n",
        "%matplotlib inline\n",
        "\n",
        "def animate_agent(env, agent, max_steps = 400):\n",
        "    obs = env.reset()\n",
        "    frames = deque()\n",
        "    frames.append(env.render(mode='rgb_array'))\n",
        "    for _ in range(max_steps):\n",
        "        out = agent(obs)\n",
        "        assert (\"Action\" in out), \"The key 'Action' was missing from the agents output.\"\n",
        "        action = out[\"Action\"]\n",
        "        obs, _ , done, _ = env.step(action)\n",
        "        frames.append(env.render(mode='rgb_array'))\n",
        "        if done:\n",
        "            break\n",
        "            \n",
        "    return animate_frames(frames)\n",
        "\n",
        "def animate_frames(frames):\n",
        "    new_height = 2.2\n",
        "    original_height = frames[0].shape[0]\n",
        "    original_width = frames[0].shape[1]\n",
        "    new_width = (new_height / original_height) * original_width\n",
        "    fig = plt.figure(figsize=(new_width, new_height), dpi = 120)\n",
        "    \n",
        "    ax = plt.Axes(fig, [0., 0., 1., 1.])\n",
        "    ax.set_axis_off()\n",
        "    fig.add_axes(ax)\n",
        "    patch = ax.imshow(frames[0], aspect='auto', animated=True, interpolation='bilinear')\n",
        "    animate = lambda i: patch.set_data(frames[i])\n",
        "    \n",
        "    ani = matplotlib.animation.FuncAnimation(fig, animate, frames=len(frames), interval = 50)\n",
        "    \n",
        "    plt.close()\n",
        "    return HTML(ani.to_jshtml())"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Trd4e0OH39Eq"
      },
      "source": [
        "Now we can easily get an animation of the cartpole environment controlled by `policy_left`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "9K7GMmSN4AsL",
        "colab": {}
      },
      "source": [
        "def agent_left(observation):\n",
        "    return {\"Action\":0}\n",
        "  \n",
        "animate_agent(env_cartpole, agent_left)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "7sUB5pCH4WaA"
      },
      "source": [
        "## Agent 1: Random Action Agent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Z58gnHdR4Y5W"
      },
      "source": [
        "Rather than having an agent whose policy is to always take the same action, we could also randomly choose between possible actions\n",
        "\n",
        "**Question**: would you expect the agent taking random actions to get a higher or lower return than the `agent_left` agent for cartpole?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "WFRvbzjS4I87",
        "colab": {}
      },
      "source": [
        "def agent_random(observation):\n",
        "    return {\"Action\":np.random.choice([0, 1])}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uYQVVCZ04qGU"
      },
      "source": [
        "**Note**: that we can easily write a general random agent for any gym environment using `env.action_space.sample()`. \n",
        "\n",
        "Computing the average return of `agent_random`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "2bPpR1kS4g15",
        "colab": {}
      },
      "source": [
        "average_episodic_return(env_cartpole, agent_random)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Uy3QD9YP4tc_"
      },
      "source": [
        "And animating it:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "urP2_C324rkq",
        "colab": {}
      },
      "source": [
        "animate_agent(env_cartpole, agent_random)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "mvptf3-q5BEJ"
      },
      "source": [
        "### Random Agent Tasks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "20QHSARZ4ypU"
      },
      "source": [
        "1. [**ALL**] Create a random agent for \"MsPacman-v0\" environment using `env.action_space.sample()` and animate the agent using `animate_agent` for 300 steps.\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Xz3gGrY95BuV"
      },
      "source": [
        "## Agent 2: Neural Net Policy Trained with Random Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pREW2PbF5HE8"
      },
      "source": [
        "To improve on the random agent `agent_random`, we need a policy that takes the observation of the environment into account. In this section, we will show how to create an agent that solves the cartpole problem without any fancy reinforcement learning methods.\n",
        "\n",
        "First, define an agent with a single hidden-layer multi-layer perceptron (MLP) policy implemented using NumPy. Its input is a vector of length 4 and its output is either `0` or `1`:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "IrYznq-Y42y6",
        "colab": {}
      },
      "source": [
        "class AgentMLPNumpy:\n",
        "    def __init__(self, num_hidden):\n",
        "        self.W1 = np.random.randn(num_hidden, 4)\n",
        "        self.W2 = np.random.randn(2, num_hidden)\n",
        "        \n",
        "    def __call__(self, x):\n",
        "        y = self.W1.dot(x)\n",
        "        y = np.tanh(y)\n",
        "        y = self.W2.dot(y)\n",
        "        return {\"Action\":np.argmax(y)}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0bFz32i05VC_"
      },
      "source": [
        "The reate an instance of `AgentMLPNumpy` (with random weight initializations), and compute its return. If the return is better than previous agents, then keep this agent. Otherwise, throw it away and try a new randomly initialized agent `AgentMLPNumpy`. Repeat this process `n` times."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DOKViXLs4Ptc"
      },
      "source": [
        "### Random Search Tasks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "pJYyhg8Y1Wt3"
      },
      "source": [
        "1. [**ALL**] Write a function `find_good_mlp_agent(env, num_tries, mlp_size)` that tries `num_tries` instances of `AgentMLPNumpy`, each with different weights. It should return both the agent that got the best return, along with that return.\n",
        "2. [**All**] Try out different values of `num_tries` and `mlp_size` until you find values that produces an agent getting a return of 500. **Hint**: `num_tries` shouldn't need to be larger than 300.\n",
        "3. [**ALL**] Visualize your best agent controlling cartpole using `animate_agent`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "MIcGGuCC4bDu"
      },
      "source": [
        "### Extra Reading on Random Search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0AZf9ABt6MSH"
      },
      "source": [
        "Variations of random search have proven to be very succesful for training agents to solve the reinforcement learning problem. Here are some examples:\n",
        "- [Evolution Strategies as a Scalable Alternative to Reinforcement Learning](https://openai.com/blog/evolution-strategies/), OpenAI, 2017\n",
        "- [Simple random search provides a competitive approach to reinforcement learning](https://arxiv.org/abs/1803.07055), H. Mania *et al*, 2018\n",
        "- [Welcoming the Era of Deep Neuroevolution](https://eng.uber.com/deep-neuroevolution/), Uber AI Labs, 2017\n",
        "\n",
        "Note that random search used to solve the *reinforcement learning problem* is not generally considered to be a *reinforcement learning method* (as implied by the titles of some of the readings above, eg. \"Evolution Strategies as a Scalable Alternative to Reinforcement Learning\"). The reason for this:\n",
        "\n",
        "> Reinforcement learning, like many topics whose names end with “ing,” such as machine learning and mountaineering, is simultaneously a problem, a class of solution methods that work well on the problem, and the field that studies this problem and its solution methods. It is convenient to use a single name for all three things, but at the same time essential to keep the three conceptually separate. In particular, the distinction between problems and solution methods is very important in reinforcement learning; failing to make this distinction is the source of many confusions. ~ *Reinforcement Learning*, Sutton and Barto, **2nd Edition**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "sPaQdlYa6TT6"
      },
      "source": [
        "## Agent 3: Neural Net Agent Trained with Policy Gradients"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "yRvNPjeY6VmJ"
      },
      "source": [
        "Unlike the random policy, the MLP policy trained with random search *does* take environment observations into account when choosing to take an action. However, the training procedure is clearly not optimal: each policy in the search makes no use of observations and rewards seen by previous policies. For environments where it takes significant time and money to take actions and make observations (for example, a real robot picking strawberries), this is incredibly wasteful. In this section, we present a method that does take previous observations into account when learning.\n",
        "\n",
        "**Warning**: This section requires more mathematics than previous sections, and hence more precise notation. It will thus be significantly more abstract."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "S-KrdnsONUKh"
      },
      "source": [
        "### Policy Gradient Theory"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "tVXz9Z3P6YFG"
      },
      "source": [
        "As we saw earlier, the reinforcement learning problem is to find an agent/policy that maximizes the return (rewards). We will now use a neural net to represent a stochastic policy  $\\pi_\\mathbf{\\theta}(a_t|o_t)$, where $\\mathbf{\\theta}$ is the set of learnable parameters in the neural net, and $\\pi_\\mathbf{\\theta}(a_t|o_t)$ is a probability distribution over the possible actions conditioned on the observation.\n",
        "\n",
        "A **trajectory**  $\\tau$ is a sequence of observations of the environment and the corresponding actions an agent takes. For an episodic task that ends after $T$ steps:\n",
        "$$\\tau = (o_0, a_0, o_1, a_1, \\ldots , o_T, a_T)$$\n",
        "The return (or total rewards) is a function of the trajectory, $R(\\tau)$.\n",
        "\n",
        "We want to find the values of $\\theta$ that maximize the **expected return** over trajectories:\n",
        "$$J(\\pi_\\theta)= \\mathop{\\mathbb{E}} _{\\tau \\sim \\pi_{\\theta}}[R(\\tau)]$$\n",
        "\n",
        "The usual way of efficiently finding good values of $\\theta$ is to differentiate the objective with respect to $\\theta$ and then use gradient descent. Luckily, we can obtain an *estimate* of this gradient. The key result is the *Policy Gradient Theorem*, which says that you can evaluate the gradient using the gradients of the policy itself:\n",
        "$$\\nabla_\\theta J(\\pi_\\theta) = \n",
        "\\mathop{\\mathbb{E}} _{\\tau \\sim \\pi_{\\theta}}\\left[ \\sum_{t=0}^{T}\\Phi_t\\nabla_\\theta\\log \\pi_\\theta(a_t|o_t) \\right] \n",
        "$$\n",
        "There are many possibilities for $\\Phi_t$ (see page 2 [here](https://arxiv.org/pdf/1506.02438.pdf) for a list of these possibilities). We  will use $\\Phi_t = \\hat{R}_t \\equiv \\sum_{t'=t}^{T} r_{t'}$, sometimes known as the [reward-to-go](https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html#don-t-let-the-past-distract-you). Note that $\\mathop{\\mathbb{E}} _{\\tau \\sim \\pi_{\\theta}}$ is an expectation over trajectories $\\tau$ which can be estimated by using a sample mean. Let us call this estimate of the gradient $\\hat{g}$:\n",
        "\n",
        "$$\\hat{g} = \\frac{1}{|D|}\\sum_{\\tau \\in \\mathcal{D}}\n",
        "\\sum_{t=0}^{T}\\hat{R}_t\\nabla_\\theta\\log \\pi_\\theta(a_t|o_t) \n",
        "$$\n",
        "where $\\mathcal{D}$ is a set of trajectories of size $|\\mathcal{D}|$. The more trajectories we sample from, the lower the variance of the estimator $\\hat{g}$. But to simplify things for this practical, we will only be sampling from only one trajectory at a time to get our gradient estimate. Then:\n",
        "$$\\hat{g} =\n",
        "\\sum_{t=0}^{T}\\hat{R}_t\\nabla_\\theta\\log \\pi_\\theta(a_t|o_t) \n",
        "$$\n",
        "Translating this into a loss to be *maximized*:\n",
        "$$\\mathcal{L}=\\sum_{t=0}^{T}\\hat{R}_t\\log \\pi_\\theta(a_t|o_t) $$\n",
        "\n",
        "This looks very similar to the usual supervised learning cross-entropy loss, $\\sum_{i}\\log p(y_i|x_i) $, where $y_i$ is the label and $x_i$ is the feature. Indeed, learning with policy gradients is almost the same doing supervised learning with the following small changes:\n",
        "\n",
        "- we use the action $a_t$ taken when we saw $o_t$ as the label ($y_i$)\n",
        "- the loss of each example gets multiplied (or *weighted*) by $\\hat{R}_t$. This increases the log probability for good actions and decrease it by bad actions.\n",
        "- we are doing gradient *ascent* rather than the usual gradient *descent*\n",
        "- the observations $o_t$ are correlated with each other, and hence are not independent and identically distributed (i.i.d) as the features $x_i$ are assumed to be. \n",
        "\n",
        "In summary:\n",
        "\n",
        "| Reinforcement Learning         \t|  Supervised Learning \t|\n",
        "|-------------------------------------\t|-----------------------------------------------------\t\t|\n",
        "|  $\\sum_{t}\\log \\pi_\\theta(a_t|o_t) \\hat{R}_t$     \t| $\\sum_{i}\\log p(y_i|x_i) $    |\n",
        "|  action $a_t$    | label $y_i$    |\n",
        "|  observation $o_t$    | feature $x_i$    |\n",
        "| $\\pi_\\theta(a_t|o_t)$   |   $p(y_i|x_i)$  |\n",
        "|  gradient ascent     \t|  gradient descent    |\n",
        "| non-i.i.d observations $o_t$ | i.i.d features $x_i$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "C9Pqx_jkmVSU"
      },
      "source": [
        "For those interested in the derivation, there are two excellent resources:\n",
        "- [Part 3: Intro to Policy Optimization](https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html) of OpenAI Spinning Up. Note that they often use *state* where we used *observation*.\n",
        "- Chapter 13 of Sutton and Barto, *Reinforcement Learning: and Introduction*, 2nd Edition. A free copy can be found [here](http://incompleteideas.net/book/RLbook2018.pdf).\n",
        "\n",
        "**Historical Note**: The name \"policy gradient\" comes from the fact that we're directly taking the gradient of the policy. The particular flavour of policy gradient which uses the loss function above, along with the Monte-carlo approximation of the objective is known as the **REINFORCE** algorithm ([Williams 1992](https://link.springer.com/article/10.1007/BF00992696)).\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "77BPPeblNOjf"
      },
      "source": [
        "#### Optional extra reading: More on Policy Gradients"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "2JNQrTtT6mpD"
      },
      "source": [
        "This section is for the mathematically inclined, and gives more insight into the Policy Gradient Theorem.\n",
        "\n",
        "Outside of the reinforcement learning setting, the Policy Gradient Theorem is known as a [score function estimator](http://blog.shakirm.com/2015/11/machine-learning-trick-of-the-day-5-log-derivative-trick/):\n",
        "$$\n",
        "\\frac{\\partial }{\\partial \\theta}\\mathop{\\mathbb{E}} _{x \\sim p_{\\theta}}[f(x)]=\\mathop{\\mathbb{E}} _{x \\sim p_{\\theta}}[f(x)\\frac{\\partial }{\\partial \\theta}\\log p(x;\\theta)]\n",
        "$$\n",
        "\n",
        "The power of this theorem comes from the ability to estimate the RHS, even if $f(x)$ is a *black-box function* (i.e. one for which we don't have the analytical form for). We can do this by replacing the mean by the sample mean:\n",
        "$$\\mathop{\\mathbb{E}} _{x \\sim p_{\\theta}}[f(x)\\frac{\\partial }{\\partial \\theta}\\log p(x;\\theta)] \n",
        "\\approx \\frac{1}{N}\\sum_{n=1}^{N}  f(\\hat{x}^{(n)})\\frac{\\partial }{\\partial \\theta}\\log p(\\hat{x}^{(n)};\\theta)$$\n",
        "\n",
        "where $\\hat{x}^{(n)}$ is sampled from $p(x;\\theta)$. In the reinforcement learning setting, $x$ is the action sampled from the parameterized policy $p(x;\\theta)$, and $f(x)$ is the reward function.\n",
        "\n",
        "The derivation of the score function estimator:\n",
        "$$\n",
        "\\begin{align}\n",
        "\\frac{\\partial }{\\partial \\theta}\\mathop{\\mathbb{E}} _{x \\sim p_{\\theta}}[f(x)] \n",
        "&= \\frac{\\partial }{\\partial \\theta} \\int p(x; \\theta)f(x) \\, dx \\\\ \n",
        "&= \\int  \\frac{\\partial }{\\partial \\theta} p(x; \\theta)f(x) \\, dx \\\\ \n",
        "&= \\int  p(x; \\theta) \\frac{\\partial }{\\partial \\theta} \\log p(x; \\theta)f(x) \\, dx \\ \\ \\ \\text{ (Log-Derivative trick)}\\\\ \n",
        "&= \\mathop{\\mathbb{E}} _{x \\sim p_{\\theta}}[f(x)\\frac{\\partial }{\\partial \\theta}\\log p(x;\\theta)]\n",
        "\\end{align}\n",
        "$$\n",
        "\n",
        "The Log-Derivative trick used in the third line is the property, for any function $f(x)$:\n",
        "$$\\frac{d f(x)}{dx}=\\frac{f(x)}{f(x)}\\frac{d f(x)}{dx}=f(x)\\frac{d\\log(f(x))}{dx}\n",
        "$$ \n",
        "The right-hand side was obtained using the chain rule and the fact that $d \\log(x)/dx=1/x$. \n",
        "\n",
        "The score function estimator is valid under **very** general conditions: $p(x;\\theta)$ needs to be a continuous function of $\\theta$, but does not need to be a continuous function of $x$. Furthermore, $f(x)$ can be discontinuous and $x$ can be either a discrete or continuous random variable (a reference for these conditions is [here](https://arxiv.org/abs/1506.05254)). This should reassure us when using this theorem for reinforcement learning, as it means we don't need to worry about showing that our environments produce a smooth or even continuous return function (which might not be true, or might be very hard to show for most environments)!\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5icJDwjW6sI5"
      },
      "source": [
        "**Example: A Toy Model**\n",
        "\n",
        "To get some intuition for how the score function estimator works, lets consider a toy model: Let $f(x)=-(x-k)^2$ (where $k$ is some real number) and let $p(x;\\mu,\\sigma)$ be normally distributed $\\mathcal{N}(\\mu, \\sigma)$ with mean $\\mu$ and variance $\\sigma$:\n",
        "$$\n",
        "p(x;\\mu,\\sigma) = \\frac{e^{-\\frac{1}{2\\sigma} (x-\\mu)^2}}{\\sqrt{2 \\pi } \\sigma}\n",
        "$$\n",
        "\n",
        "In reinforcement learning, the task is to find the values of $\\mu$ and $\\sigma$ that maximize $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$. Here is a summary as to how this toy problem fits into the reinforcement learning framework:\n",
        "\n",
        "|     Reinforcement Learning Concept    \t|  Toy Example \t|\n",
        "|-------------------------------------\t|-----------------------------------------------------\t\t|\n",
        "|  Action     \t| choice of $x$    |\n",
        "|  Policy    | $p(x;\\mu,\\sigma)$    |\n",
        "|  Observation $o_t$     | None    |\n",
        "| Reward $r_t$   |   $-(x-k)^2$  |\n",
        "\n",
        "\n",
        "Unlike the general reinforcement learning case where $f(x)$ is a black-box,  we can analytically evaluate $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$ with our example:\n",
        "$$\n",
        "\\begin{align}\n",
        "\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]\n",
        "&= \\int^{\\infty}_{-\\infty} f(x) p(x;\\mu,\\sigma) dx \\\\\n",
        "&=\\int^{\\infty}_{-\\infty} \\left[-(x-k)^2\\right] \\left[\\frac{e^{-\\frac{1}{2\\sigma} (x-\\mu)^2}}{\\sqrt{2 \\pi } \\sigma}\\right]dx \\\\\n",
        "&=-(k-\\mu)^2 - \\sigma^2\n",
        "\\end{align}\n",
        "$$\n",
        "\n",
        "The maximimum of $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$ is when the distribution $p(x;\\mu,\\sigma)$ has mean $\\mu=k$ (which corresponds to the position of the maximum of $f(x)=-(x-k)^2$) and variance $\\sigma=0$ (so it is infinitely 'spiked', i.e. a [Dirac Delta function](https://en.wikipedia.org/wiki/Dirac_delta_function)). This makes sense: to maximize $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$, we want all of the probability mass of $p(x;\\mu,\\sigma)$ to be concentrated at the maximum of $f(x)=-(x-k)^2$, i.e. $x=k$.\n",
        "\n",
        "We can now also easily verify that using the score function estimator is equivalent to differentiating $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$ directly: as we have computed $\n",
        "\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$, we immediately have that \n",
        "$$\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)] = \\frac{\\partial }{\\partial \\mu}\\left[-((\\mu-k)^2 + \\sigma^2)\\right]=2(k-\\mu)\n",
        "$$\n",
        "But using the score function estimator instead:\n",
        "$$\\begin{align}\n",
        "\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)] \n",
        "&= \\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)\\frac{\\partial }{\\partial \\mu}\\log p(x;\\mu,\\sigma)] \\\\\n",
        "&= \\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[-(x-k)^2 \\frac{\\partial }{\\partial \\mu}\\log \\left(\\frac{e^{-\\frac{1}{2\\sigma} (x-\\mu)^2}}{\\sqrt{2 \\pi } \\sigma}\\right)] \\\\\n",
        "&= \\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[\\frac{(x-k)^2(\\mu-x)}{\\sigma^2}] \\\\\n",
        "&=2(k-\\mu)\n",
        "\\end{align}\n",
        "$$\n",
        "\n",
        "Note: If you have trouble following that last step see [moments of the Gaussian](https://en.wikipedia.org/wiki/Normal_distribution#Moments)\n",
        "\n",
        "It is left as an exercise to verify that the same holds for $\\frac{\\partial }{\\partial \\sigma}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)] $. \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aYPN5xG3f8i0"
      },
      "source": [
        "**Sampling using the Toy Model**\n",
        "\n",
        "So far, the score function estimator has not given us anything extra: we can \n",
        "directly compute the derivatives of $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$, and we don't even need derivatives to find its minimum. But lets pretend that $f(x)$ is a black box: we can only evaluate $f(x)$ at a finite number of values of $x$. We can no longer compute $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$ directly. How do we find the values of $\\mu$ and $\\sigma$ that maximize $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$? \n",
        "\n",
        "**Further simplification**: set $\\sigma=1$ and $k=0$, so $f(x)=-x^2$ and the task is to find the optimal value of $\\mu$.\n",
        "\n",
        "**Approach 1: Direct Search.** We can replace the mean $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}$ with the sample mean: \n",
        "\n",
        "$$\n",
        "\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x)]\\approx \\frac{1}{N} \\sum_{n=1}^{N}f(\\hat{x}^{(n)})\n",
        "$$\n",
        "where $\\hat{x}^{(n)} \\sim p(x;\\mu,1)$ is the $n^{th}$ sample from the distribution $p(x;\\mu,1)$. The sample mean is equivalent to the expectation in the limit $N\\to\\infty$. Then we simply try out lots of different values of $\\mu$ and $\\sigma$ and take the ones that maximize the sample mean. The random search method explored in the previous section is an example of this approach.\n",
        "\n",
        "**Approach 2: Gradient Ascent**. If we can get an estimate of the derivative of $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x)]$ w.r.t $\\mu$, we can modify our initial values of $\\mu$ by gradient ascent rather than just trying out random values. Using the score function estimator and then taking the sample mean:\n",
        "\n",
        "$$\n",
        "\\begin{align}\n",
        "\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)] \n",
        "&= \\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x)\\frac{\\partial }{\\partial \\mu}\\log p(x;\\mu,1)] \\\\\n",
        "&= \\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x) (x-\\mu)] \\\\\n",
        "&\\approx \\frac{1}{N} \\sum_{n=1}^{N}f(\\hat{x}^{(n)}) (\\hat{x}^{(n)}-\\mu) \\\\\n",
        "&\\equiv \\eta_N \n",
        "\\end{align}$$\n",
        "Our estimator $\\eta_N \\equiv \\sum_{n=1}^{N}f(\\hat{x}^{(n)}) (\\hat{x}^{(n)}-\\mu)$ allows us to estimate the gradient $\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu,\\sigma}}[f(x)]$ using a finite number of samples of $p(x;\\mu,1)$. In this case, we were able to explicitly compute $\\frac{\\partial }{\\partial \\mu}\\log p(x;\\mu,1)$ as we have an analytical expression for $p(x;\\mu,1)$, which is usually not the case for reinforcement learning (where $p(x;\\mu,1)$ is often a neural net).\n",
        "Let us implement this estimator $\\eta_N$:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "ibHTfIWg6uvD",
        "colab": {}
      },
      "source": [
        "def f_estimator(x, mean):\n",
        "    return (-x**2) * (x - mean)\n",
        "\n",
        "def sf_estimator(mean, sample_size, f):\n",
        "    x = np.random.normal(mean, 1, sample_size)\n",
        "    return np.mean(f(x, mean))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "OCk94F6f6wgN"
      },
      "source": [
        "If our starting value is $\\mu=3$, then the exact derivate is $2\\times (-3)=-6$, and our estimated derivative is:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "MEqsttrE6yYR",
        "colab": {}
      },
      "source": [
        "sample_size=200\n",
        "mean=3\n",
        "sf_estimator(mean, sample_size, f_estimator)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "qjdP2MAY62GL"
      },
      "source": [
        "**Baselines and Variance Reduction**\n",
        "\n",
        "We can also get an idea of the **variance** of our derivative estimates, which tells us how noisy these estimates are:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "0nczQsIT67Uz",
        "colab": {}
      },
      "source": [
        "def sf_estimator_variance(mean, mean_sample_size, grad_f, variance_samples = 200):\n",
        "    grads = [grad_f(mean, mean_sample_size, f_estimator) for _ in range(variance_samples)];\n",
        "    return np.var(grads)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UTzM4Kli69hP"
      },
      "source": [
        "With a sample size of 200, the variance of our derivative estimator is:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "both",
        "colab_type": "code",
        "id": "6hFFR6eU6_MP",
        "colab": {}
      },
      "source": [
        "var1 = sf_estimator_variance(mean, sample_size, sf_estimator)\n",
        "var1"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "tOj7sIvvDpTF"
      },
      "source": [
        "One obvious way of reducing the variance is to increase the sample size:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "M-fXQeHzDuEV",
        "colab": {}
      },
      "source": [
        "sample_size2 = 500\n",
        "sf_estimator_variance(mean, sample_size2, sf_estimator)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "0OVRUJ_L7UGW"
      },
      "source": [
        "In reinforcement learning, taking extra samples is often expensive. It turns out that we can reduce the variance *without* taking any more samples! This idea is very important in reinforcement learning, where it is known as the [baseline method](http://jmlr.csail.mit.edu/papers/volume5/greensmith04a/greensmith04a.pdf). The general procedure is also known as the [control variate method](https://arxiv.org/abs/1906.10652).\n",
        "\n",
        "To see how this works, let us derive an expression for the variance of our estimator:\n",
        "\n",
        "$$\n",
        "\\begin{align}\n",
        "\\text{Var}[\\eta_N] &= \\text{Var}\\left[\\frac{1}{N}\\sum_{n=1}^{N}f(x) (x-\\mu) \\right] \\ \\ \\  \\ (x \\sim \\mathcal{N}(\\mu, 1))  \\\\\n",
        "&=\\frac{1}{N^2}\\text{Var}\\left[\\sum_{n=1}^{N}f(x) (x-\\mu) \\right] \\ \\ \\ \\ (\\text{Var}[aX]=a^2\\text{Var}[X]) \\\\\n",
        "&= \\frac{1}{N^2}\\sum_{n=1}^{N}\\text{Var}\\left[f(x) (x-\\mu) \\right] \\ \\ \\ \\ (\\text{Var}[X+Y]=\\text{Var}[X]+\\text{Var}[Y] \\ \\text{when} \\ X,Y \\ \\text{independent}) \\\\\n",
        "&= \\frac{1}{N^2}\\cdot N\\cdot\\text{Var}\\left[f(x) (x-\\mu) \\right] \\\\\n",
        "&= \\frac{1}{N}\\text{Var}\\left[f(x) (x-\\mu) \\right] \n",
        "\\end{align}\n",
        "$$\n",
        "As we already saw, the obvious way of reducing the variance $\\text{Var}[\\eta_N]$ is to increase $N$, but this comes at the cost of needing to take more samples. We have learnt though that the variance decreases as $\\mathcal{O}(1/N)$ with sample size.\n",
        "\n",
        "Another way is with the following trick: replace the $f(x)$ with a new $\\tilde{f}(x)$ that leaves the gradient $\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[(f(x))]$ unchanged, but reduces the variance of our gradient estimator $\\eta_N$, i.e reduces $\\text{Var}[\\eta_N]$. A simple choice is $\\tilde{f}(x)=f(x)-b(x)$ for any function $b(x)$ that does not depend on $\\mu$. It is easy to check that this leaves $\\eta_N$ unchanged:\n",
        "$$\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[(\\tilde{f}(x))]\n",
        "=\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[(f(x)-b(x))]\n",
        "=\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x)]+\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[b(x)]\n",
        "=\\frac{\\partial }{\\partial \\mu}\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x)]$$\n",
        "Now [note that](https://en.wikipedia.org/wiki/Variance#Sum_of_variables)\n",
        "$$\n",
        "\\text{Var}[f-b]= \\text{Var}[f]-2\\text{Cov}[f,b]+\\text{Var}[b]\n",
        "$$\n",
        "So if $f$ and $b$ are strongly correlated, then the covariance $2\\text{Cov}[f,b]$ can be greater than the extra variance term $\\text{Var}[b]$, reducing the total variance $\\text{Var}[f-b]$.\n",
        "\n",
        "A common choice for $b$ is $b=\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x)]$. This choice of $b$ is clearly correlated with $f(x)$, as increases in $f(x)$ will lead to increases in the mean $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x)]$ (and decreases will lead to decreases in the mean).\n",
        "\n",
        "Let us check for our toy example that this does actually reduce the variance. First, note that $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x)] = -(1+\\mu^2)$ and $\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x) (x-\\mu) ] = -2\\mu$ (leaving out the details of this calculation). Then (skipping the details of the computation: this is what computer algebra systems are there for!):\n",
        "$$\n",
        "\\begin{align}\n",
        "\\text{Var}[\\eta_N] \n",
        "&=  \\frac{1}{N}\\text{Var}\\left[f(x) (x-\\mu)\\right] \\\\\n",
        "&=  \\frac{1}{N}\\int^{\\infty}_{-\\infty}\\left(f(x)(x-\\mu) -\\mathop{\\mathbb{E}} _{x \\sim p_{\\mu}}[f(x) (x-\\mu) ] \\right)^2 p(x;\\mu)dx \\\\\n",
        "&= \\frac{1}{N}\\left( \\mu^4 +14\\mu^2+15\\right) \\\\\n",
        "\\end{align}\n",
        "$$\n",
        "And similarly for $\\tilde{f}(x)$:\n",
        "$$\n",
        "\\begin{align}\n",
        "\\text{Var}[\\tilde{\\eta}_N] &= \\frac{1}{N}\\text{Var}\\left[\\tilde{f}(x) (x-\\mu)\\right] \\\\\n",
        "&=\\frac{1}{N}\\left( 8\\mu^2+10\\right)\n",
        "\\end{align}\n",
        "$$\n",
        "The variance is indeed reduced for all values of $\\mu$, as the difference in variances $\\text{Var}[\\eta_N] -\\text{Var}[\\tilde{\\eta}_N]$ is a positive quantity:\n",
        "$$\n",
        "\\text{Var}[\\eta_N] -\\text{Var}[\\tilde{\\eta}_N]=\\frac{1}{N}(\\mu^4+6 \\mu^2+5)\n",
        "$$\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "C_9Lms7PMdvg"
      },
      "source": [
        "To get more insight into how the addition of the baseline changes the distribution of $f(x) (x-\\mu)$ (where $x$ is normally distributed $\\mathcal{N}(\\mu, 1)$), we could analytically compute what this distribution is and plot it, and hopefully see how the baseline changes its shape. This is tricky as $f(x) (x-\\mu)$ is a cubic polynomial, which is non-bijective in a non-trivial way, complicating the use of the standard probability density [change-of-variables formula](https://en.wikipedia.org/wiki/Probability_density_function#Function_of_random_variables_and_change_of_variables_in_the_probability_density_function). Rather than do this analytically, it is much easier to sample values of $x$ and plot the histogram of values of $f(x) (x-\\mu)$. The Seaborn plotting package makes this very easy:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "kz7F0AZeMJvZ",
        "colab": {}
      },
      "source": [
        "def f_estimator_baseline(x, mean):\n",
        "  return (-x**2 - (-1 - mean**2)) * (x-mean)\n",
        "\n",
        "mean2 = 7\n",
        "sample_size = 10000\n",
        "x_rand=np.random.normal(mean2, 1, sample_size)\n",
        "rands_f = f_estimator(x_rand, mean2)\n",
        "rands_f_baseline = f_estimator_baseline(x_rand, mean2)\n",
        "\n",
        "import seaborn as sns\n",
        "sns.set_style('darkgrid')\n",
        "sns.kdeplot(rands_f, shade=True, label=\"Without Baseline\", legend=True)\n",
        "sns.kdeplot(rands_f_baseline, shade=True, label=\"With Baseline\", legend=True)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Ej0IiB6RnYyu"
      },
      "source": [
        "This shows how the addition of the baseline changes the distribution to be a lot *sharper*, which means that it has less variance. The sampling we did above allows us to check our analytic calculations as well: with $\\mu=7$, $\\text{Var}\\left[f(x) (x-\\mu)\\right] = 3102$ and $\\text{Var}\\left[\\tilde{f}(x) (x-\\mu)\\right] = 402$. This is similar to what we see with our sampled data:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "i3mOC7dOoQT5",
        "colab": {}
      },
      "source": [
        "print(\"Variance without Baseline: \", np.var(rands_f))\n",
        "print(\"Variance with Baseline: \", np.var(rands_f_baseline))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gjwhgz53nOUZ"
      },
      "source": [
        "The sample means should be similar (but only the same in the limit $N\\to\\infty$):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "uIPbvbjCnAAJ",
        "colab": {}
      },
      "source": [
        "print(\"Mean without Baseline: \", np.mean(rands_f))\n",
        "print(\"Mean with Baseline: \", np.mean(rands_f_baseline))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "hHZxhJn8L2_Y"
      },
      "source": [
        "In practice, we don't know the exact mean of $f(x)$ to use as a baseline, but we can just take the sample mean.\n",
        "\n",
        "**Question**: verify that using `f_estimator_baseline` instead of `f_estimator` in `sf_estimator_variance` reduces the variance"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xMeP6P_9h86-"
      },
      "source": [
        "**Extra Reading**\n",
        "- A great review of score function estimators:\n",
        "    * [Monte Carlo Gradient Estimation in Machine Learning](https://arxiv.org/abs/1906.10652), Mohamed *et al*, 2019"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "cxMIN-4N7pom"
      },
      "source": [
        "### Episode Trajectory Collector"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "y5cUCGZ5hGIr"
      },
      "source": [
        "Let us implement a function `get_episode_trajectory(env, agent)` that controls the environment `env` with `agent` and returns lists of rewards, observations and actions taken at each timestep:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "UQJlw2MU7lqL",
        "colab": {}
      },
      "source": [
        "def get_episode_trajectory(env, agent, max_steps=1000):\n",
        "    observation_list = []\n",
        "    reward_list = []\n",
        "    action_list = []\n",
        "    value_list = []\n",
        "\n",
        "    done = False\n",
        "    obs = env.reset()\n",
        "    for _ in range(max_steps):\n",
        "        observation_list.append(obs)\n",
        "        out = agent(obs)\n",
        "        assert (\"Action\" in out), \"The key 'Action' was missing from the agents output.\"\n",
        "        action = out[\"Action\"]        \n",
        "        obs, rew, done, _, = env.step(action)\n",
        "        reward_list.append(rew)\n",
        "        action_list.append(action)\n",
        "        if \"Value\" in out:\n",
        "            value_list.append(out[\"Value\"])\n",
        "            \n",
        "        if done:\n",
        "            break\n",
        "        \n",
        "    ret = {\n",
        "        \"Observations\": observation_list, \n",
        "        \"Actions\": action_list, \n",
        "        \"Rewards\": np.array(reward_list, dtype=np.float32)\n",
        "    }\n",
        "    if len(value_list) > 0:\n",
        "        ret[\"Values\"] = value_list\n",
        "        \n",
        "    return ret"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TbXEoE0RlM8h"
      },
      "source": [
        "Let us try out `get_episode_trajectory` on cartpole using the random agent we implemented earlier:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "kSs6AHhnkb93",
        "colab": {}
      },
      "source": [
        "traj = get_episode_trajectory(env_cartpole, agent_random, max_steps=3)\n",
        "pprint(traj)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jO-hBdOSumNA"
      },
      "source": [
        "### Reward-to-go $\\hat{R}_t$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_mhIQMcBvOta"
      },
      "source": [
        "We need to compute the reward-to-go $\\hat{R}_t= \\sum_{t'=t}^{T} r_{t'}$. Here is a simple NumPy implementation:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "55QdJMQs8Ia-",
        "colab": {}
      },
      "source": [
        "def reward_to_go(rewards):\n",
        "    return np.flip(np.cumsum(np.flip(rewards)))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "EDeiES4Qvgq0"
      },
      "source": [
        "Checking that this works as expected:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "rZAFQxCWvg9R",
        "colab": {}
      },
      "source": [
        "reward_to_go([1, 1.2, 1])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "iAnootcVLZ-x"
      },
      "source": [
        "#### Optional extra reading: estimator issues"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TY9W1M27wV7j"
      },
      "source": [
        "The estimator $\\Phi_t \\equiv \\hat{R}_t= \\sum_{t'=t}^{T} r_{t'}$ we used has a number of problems. For example, if we don't get to the end of an episode (as it takes too long and we stop sampling), or our task is not episodic, we will have incorrect estimates of the advantages of our states. One popular alternative that works around this difficulty is the GAE-Lamba estimator: \n",
        " - [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://arxiv.org/abs/1506.02438), Schulman *et al*, 2015\n",
        " - A good blog post on this: https://danieltakeshi.github.io/2017/04/02/notes-on-the-generalized-advantage-estimation-paper/"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "zqAM505qxmUz"
      },
      "source": [
        "### TensorFlow Policy Net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "_OfnxJD7x3FC"
      },
      "source": [
        "Let us use a neural net policy with parameters $\\theta$, so we have $\\pi_\\theta(a|s) = NN(s; \\theta)$, where $NN(s; \\theta)$ is some potentially complex function represented by a neural network with parameters $\\theta$:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "loX4BH07xiz9",
        "colab": {}
      },
      "source": [
        "import tensorflow as tf\n",
        "from tensorflow.keras.layers import Dense, Flatten, Conv2D\n",
        "from tensorflow.keras import Model"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "En5f2yILwN7p",
        "colab": {}
      },
      "source": [
        "class AgentMLPTF(Model):\n",
        "  def __init__(self):\n",
        "    super(AgentMLPTF, self).__init__()\n",
        "    self.d1 = Dense(15, activation='tanh')\n",
        "    self.d2 = Dense(2)\n",
        "\n",
        "  def call(self, x):\n",
        "    # 1. Define Policy\n",
        "    batch = True\n",
        "    if x.ndim == 1:\n",
        "        batch = False\n",
        "        x = np.expand_dims(x, axis=0)\n",
        "    x = self.d1(x)\n",
        "    action_logits = self.d2(x)\n",
        "    \n",
        "    # 2. Sample policy to get action\n",
        "    action = tf.random.categorical(action_logits, 1)\n",
        "    action = action.numpy().flatten()\n",
        "    if not batch:\n",
        "        action = np.asscalar(action)\n",
        "        \n",
        "    return {\"Action\":action, \"LogProbability\":action_logits}"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "tEHv9UzDxpwB",
        "colab": {}
      },
      "source": [
        "agent_mlp_tf = AgentMLPTF()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "TsSCeL-nySKj",
        "colab": {}
      },
      "source": [
        "obs = np.array([0.1,0.2,0.3,0.4])\n",
        "agent_mlp_tf(obs)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "uNjVoPGRO0wV"
      },
      "source": [
        "We can also verify that this works on a batch of observations:\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "IDHVarkCyefP",
        "colab": {}
      },
      "source": [
        "obs_batch = np.array([[0.1,0.2,0.3,0.4], [0.5,0.3,0.2,0.1]])\n",
        "agent_mlp_tf(obs_batch)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ZwENPlLtjPmZ"
      },
      "source": [
        "### Loss Function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "DPJBU0tljjGN"
      },
      "source": [
        "**Exercise [ALL]**: implement the loss function $\\mathcal{L}=\\sum_{t=0}^{T}\\hat{R}_t\\log \\pi_\\theta(a_t|s_t) $ using TensorFlow operations. Call this function `loss_pg(actions, log_probs, returns)`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "IFkB4aDjjUFr",
        "colab": {}
      },
      "source": [
        "#@title Answer { display-mode: \"form\" }\n",
        "\n",
        "def loss_pg(actions, log_probs, returns):\n",
        "    action_masks = tf.one_hot(actions, 2, dtype=np.float64)\n",
        "    log_probs = tf.reduce_sum(action_masks * tf.nn.log_softmax(log_probs), axis=1)\n",
        "    return -tf.reduce_sum(returns * log_probs)\n",
        "  \n",
        "# Verify that this works on some example data:\n",
        "\n",
        "actions = [1,0,0]\n",
        "logits = np.array([[0.2,0.8],[0.2,0.8],[0.6,0.4]])\n",
        "weights = [2.3, 4.3, 2.1]\n",
        "loss_pg(actions, logits, weights)\n",
        "\n",
        "# Note: this is equivalent to:\n",
        "def loss_pg2(actions, log_probs, returns):\n",
        "    loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)\n",
        "    return tf.reduce_mean(returns * loss(actions, log_probs))"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "n5IOFz5O9leQ"
      },
      "source": [
        "### Train the Policy Net"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "PVA-sEny9lAC"
      },
      "source": [
        "Now we can put all the pieces together in the training loop. It simply collects a single trajectory and uses it to do a single gradient update of the agent:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "JrG1V6rd9kfZ",
        "colab": {}
      },
      "source": [
        "def train_policy_grad(env, agent, num_epochs=300):\n",
        "    optimizer = tf.keras.optimizers.Adam(learning_rate=1e-2)\n",
        "    log_reward = 0\n",
        "    log_reward_list = []\n",
        "    logging_period = 20\n",
        "    \n",
        "    for epoch in range(num_epochs):\n",
        "        # get the training data\n",
        "        traj = get_episode_trajectory(env, agent)\n",
        "        obs = np.stack(traj[\"Observations\"])\n",
        "        rew = traj[\"Rewards\"]\n",
        "        actions = traj[\"Actions\"]\n",
        "        \n",
        "        # compute 'reward-to-go'\n",
        "        rew_2_go = reward_to_go(rew)\n",
        "        \n",
        "        # compute gradients + update weights\n",
        "        with tf.GradientTape() as tape:\n",
        "            logits = agent(obs)[\"LogProbability\"]\n",
        "            loss = loss_pg(actions, logits, rew_2_go)\n",
        "            \n",
        "        gradients = tape.gradient(loss, agent.trainable_variables)\n",
        "        optimizer.apply_gradients(zip(gradients, agent.trainable_variables))\n",
        "        \n",
        "        # log the reward\n",
        "        log_reward += np.sum(rew)\n",
        "        if (epoch % logging_period) == 0:\n",
        "            template = 'Training Epoch {}, Averaged Return: {}'\n",
        "            print(template.format(epoch, log_reward / logging_period))\n",
        "            log_reward_list.append(log_reward / logging_period)\n",
        "            log_reward = 0\n",
        "       \n",
        "    return (range(0, num_epochs, logging_period), log_reward_list)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "tGY1zD869i7h"
      },
      "source": [
        "Note that `train_policy_grad(env, agent)` takes a TensorFlow agent as input and mutates it during training. Create and train a `AgentMLPTF` agent:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "-eiTdaZ--MaS",
        "colab": {}
      },
      "source": [
        "agent_mlp_tf = AgentMLPTF()\n",
        "(episodes, rewards) = train_policy_grad(env_cartpole, agent_mlp_tf)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "f1Eto6aH9ivg"
      },
      "source": [
        "We can plot the average rewards obtained versus the number of episodes trained for:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "PGbPMo4u9idL",
        "colab": {}
      },
      "source": [
        "plt.plot(episodes, rewards, 'bo')\n",
        "plt.xlabel('Episode Number')\n",
        "plt.ylabel('Average Episode Reward')\n",
        "plt.show()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "bAcp0ffR9iNl"
      },
      "source": [
        "Let us animate this agent controlling the cartpole environment:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "t6T-X6eU9hzh",
        "colab": {}
      },
      "source": [
        "animate_agent(env_cartpole, agent_mlp_tf)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "jNTwVhJuxL6B"
      },
      "source": [
        "### Policy Gradient Tasks"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "1Lut8O5hxI_T"
      },
      "source": [
        "1. [**ALL**] Try using TensorFlows default initial learning rate of the ADAM optimizer in `train_policy_grad`. What happens? Try come up with an explanation for what is going on.\n",
        "\n",
        "**Variance Reduction**: We are trying to estimate the policy gradient using $\\hat{g} = \\frac{1}{|D|}\\sum_{\\tau \\in \\mathcal{D}}\n",
        "\\sum_{t=0}^{T}\\hat{R}_t\\nabla_\\theta\\log \\pi_\\theta(a_t|o_t) \n",
        "$. This gradient is only exact if we have an infinite number of samples. Otherwise it is *noisy*: we will get a different estimate of $\\hat{g}$ for every sample we take. \n",
        "\n",
        "2. [**ALL**] The simplest way of reducing the variance of $\\hat{g}$ is to use a *baseline*: rather than using $\\hat{R}_t$, use $\\hat{R}_t-b$ where $b$ is the mean of $\\hat{R}$, i.e. $b=\\frac{1}{T}\\sum_{t=1}^{T}\\hat{R}$. Modify `train_policy_grad` to use this. Does it speed up training? An in-depth discussion of why this works is in the section *Optional extra reading: More on Policy Gradients*.\n",
        "\n",
        "3. [**Advanced**] The policy used to solve the Cartpole environment only required the most recent observation to inform its actions. Observations taken before the most recent one give no extra information that the policy could use to improve its decision making. This is called the [Markov Property](http://www.cs.cmu.edu/~10601b/slides/MDP_RL.pdf). Most real-world environments are not like this. Consider the problem of using a reinforcement learning agent to control a self-driving car, where an observation is a frame from its onboard video cameras. List various ways in which the Markov Property is violated in this environment (this is also known as a *partially-observed environment*). How might you modify the architecture of your policy net to operate in a non-Markovian environment like this?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "I9870JT0SiZf"
      },
      "source": [
        "**IMPORTANT: Please fill out the exit ticket form before you leave the practical: https://forms.gle/FLQ78eAEQprsWXp38**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "P7BiER06QQqm"
      },
      "source": [
        "### Optional extra reading: Policy Gradients"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "utNUx7nBhg9p"
      },
      "source": [
        "\n",
        "\n",
        "- The simple REINFORCE algorithm shown here is very unstable. There are a number of improvements to this method, the most popular of which is probably PPO:\n",
        "    * [Proximal Policy Optimization Algorithms](https://arxiv.org/abs/1707.06347), Schulman *et al*, 2017\n",
        "    * OpenAI SpinningUp: https://spinningup.openai.com/en/latest/algorithms/ppo.html\n",
        "    * OpenAI Blog: https://openai.com/blog/openai-baselines-ppo/\n",
        "- PPO borrows many ideas from TRPO. [Depth First Learning](https://www.depthfirstlearning.com/2018/TRPO) has a course with a lot of good resources on TRPO.\n",
        "- The idea of being able to differentiate through stochastic nodes in computation graphs is a very powerful idea, and it is used in a number of other areas in deep learning, such as *variational autoencoders*. \n",
        "    * For the general setting, see [Gradient Estimation Using Stochastic Computation Graphs](https://arxiv.org/abs/1506.05254), Schulman *et al*, 2015.\n",
        "- For a skeptical take on policy gradients, see Ben Recht's post [The Policy of Truth](https://www.argmin.net/2018/02/20/reinforce/).\n",
        "- Karpathy has [good post](http://karpathy.github.io/2016/05/31/rl/) on reinforcement learning using policy gradients, including an implementation of policy gradients that solves Pong using only NumPy in [131 lines of Python](https://gist.github.com/karpathy/a4166c7fe253700972fcbc77e4ea32c5)! "
      ]
    }
  ]
}