Policy_Gradient_Methods_Theory.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Policy Gradient Methods Theory",
      "provenance": [],
      "collapsed_sections": [
        "kAcDOeMNe8kn"
      ],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/jarrydmartinx/deep-rl/blob/master/Policy_Gradient_Methods_Theory.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6-5MBrUhlNt6",
        "colab_type": "text"
      },
      "source": [
        "# Policy Gradient Methods Theory\n",
        "\n",
        "* This review of the theory of Policy Gradient methods for reinforcment learning is based largely on David Silver's fantastic UCL course on reinforcement learning.\n",
        "* At the bottom there's some very old broken code that should serve as an anti-model, if anything. Don't write your agents like this in future.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "uuT9s4JrlZg0",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "! pip install -q plotnine"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-3lYRPTIlMq5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#@title Imports\n",
        "\n",
        "from abc import ABCMeta, abstractmethod\n",
        "import collections\n",
        "import enum\n",
        "import itertools\n",
        "\n",
        "import pdb\n",
        "\n",
        "import numpy as np\n",
        "import plotnine as gg\n",
        "import pandas as pd \n",
        "\n",
        "# Type stuff\n",
        "from typing import Dict, List, Tuple"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dJJIML9TdKRU",
        "colab_type": "text"
      },
      "source": [
        "# Theory Review"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ErzVoUgodZ-U",
        "colab_type": "text"
      },
      "source": [
        "## Policy-Based Methods vs. Value-Based Methods\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2CkhmgHmM8_",
        "colab_type": "text"
      },
      "source": [
        "* Value Based: Learnt Value Function but implicit policy\n",
        "* Policy Based: No value function and learnt policy\n",
        "* Actor-Critic: Learnt Value function + learnt policy\n",
        "\n",
        "\n",
        "In some environments, it might be much more compact to represent the policy rather than the value-function (e.g. Breakout)\n",
        "\n",
        "**Advantages:**\n",
        "* Better convergence properties: you are just smoothly updating your policy, and you don't get the dramatic swings in the policy can arise from updating the value-function.\n",
        "  * if you just follow the gradient of your policy wrt the return, you're guaranteed to converge at least to a local optimum\n",
        "* Effective in high-dimensional or continuous action-spaces: You don't have to compute the max at every step. This maximisation can be prohibitively expensive (say you have 10^10 actions, or continuous actions). \n",
        "  * With PG you just incrementally adjust the parameters of your policy in order to incrementally learn 'what the max will be', without solving a maximisation problem at every step.\n",
        "* They can learn stochastic policies: necessary in e.g. rock paper scissors, aliased states. Value-based methods learn deterministic (or near deterministic) implicit policies.\n",
        "\n",
        "**Disadvantages**\n",
        "* Typically converges to a local optimum rather than a global optimum\n",
        "* Naive policy based reinforcement learning can be slower, and higher variance, and less efficient than value based methods.\n",
        "  * The max in value-based actions (even in Sarsa you often have a near greedy policy with a max), is extremely aggressive (you push your policy as far as you can toward what you currently estimate the optimal policy to be).\n",
        "  * PG methods update smoothly in policy space, taking small steps in the *direction* of the estimated optimal policy, which makes them\n",
        "    * more stable\n",
        "    * less efficient\n",
        "  * Also computing the gradient naively can be very slow and high-variance (can mitigate this, e.g. in Actor Critic)\n",
        " \n",
        " \n",
        "#### Convergence properties of Policy Gradients vs. Value-based\n",
        "* Will we find a unique global optimum?\n",
        "* In the Tabular TD case, you are essentially applying the bellman operator, which is a contraction in value function space, so you're guaranteed to arrive at a global optimum\n",
        "* If you follow the softmax policy, if you just follow the gradient, you also find the global optimum, with separate softmax parameters for each state\n",
        "* If you have a more general value/policy approximator, like a neural network, you can get stuck in a local optimum in either case.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AkGy-vqIe3TL",
        "colab_type": "text"
      },
      "source": [
        "![alt text](https://spinningup.openai.com/en/latest/_images/math/262538f3077a7be8ce89066abbab523575132996.svg)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qa25olwudWdd",
        "colab_type": "text"
      },
      "source": [
        "## PG Objective Functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AAYIBJcqrjpX",
        "colab_type": "text"
      },
      "source": [
        "* Goal: given policy $\\pi_{\\theta}(s, a)$ with parameters $\\theta$, find the best $\\theta$\n",
        "* How do we measure the quality of a policy $\\pi$? What is the objective function for PG methods?\n",
        "* Three possible objective functions:\n",
        "  * Start Value: only for episodic environments. You have a defined distribution over start states s_1\n",
        "  $$ J_1(\\theta) = V^{\\pi_{\\theta}} (s_1) = \\mathbb{E}_{\\pi_{\\theta}}\\big[v_1\\big] $$\n",
        "  * Average Value: for continuing environments\n",
        "  $$ J_{avV}(\\theta) = \\sum_s d^{\\pi_\\theta}(s) V^{\\pi_\\theta} (s) $$\n",
        "  * Average reward per time-step\n",
        "  $$ J_{avR}(\\theta) = \\sum_s d^{\\pi_\\theta}(s) \\sum_a \\pi_\\theta (s, a) \\mathcal{R}_s^a       $$\n",
        "  where $d^{\\pi_\\theta}(s)$ is the stationary distribution of the Markov Chain induced by $\\pi_\\theta$\n",
        "* But don't worry because *they're just rescalings of each other*, and the **policy gradient direction is the same** for all three\n",
        "* Hence all methods that maximise one maximise the others"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FZ6BnD2JdT1c",
        "colab_type": "text"
      },
      "source": [
        "## Policy Optimisation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s0PZFcZx4CwK",
        "colab_type": "text"
      },
      "source": [
        "* Policy based reinforcment learning is an optimisation problem. Find $\\theta$ that maximises $J(\\theta)$\n",
        "* Some approaches do not use gradient\n",
        "  * Hill climbing\n",
        "  * Simplex / amoeba / Nelder Mead\n",
        "  * Genetic algorithms\n",
        "* Greater efficiency often possible using gradient\n",
        "  * Gradient descent\n",
        "  * Conjugate gradient\n",
        "  * Quasi-newton\n",
        "* We focus on gradient descent, many extensions possible\n",
        "* And on methods that exploit sequential structure\n",
        "  * We're not going to do blind optimisation like a genetic algorithm\n",
        "    * GA would just let the agent run around, die, get a return value at the end and then we'd change some parameters\n",
        "    * PG methods still allow us to break open the trajectory, and make use of the sequence of states of rewards to do better **by learning within the agent's lifetime**\n",
        " "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "elX7kIe55FRF",
        "colab_type": "text"
      },
      "source": [
        "### Gradient Ascent\n",
        "* The objective is something like 'reward gotten out of the system', we want to maximise that by adjusting the policy parameters, just a little, in the direction of the gradient of the object wrt the policy parameters\n",
        "$$ \\Delta\\theta = \\alpha\\nabla_\\theta J(\\theta)    $$\n",
        "* $\\nabla_\\theta J(\\theta)$ is the **policy gradient**\n",
        "\\begin{align}\n",
        "\\nabla_\\theta J(\\theta) &= \n",
        "\\begin{pmatrix}\n",
        "\\frac{\\partial J(\\theta)}{\\partial \\theta_1}\\\\\n",
        "\\vdots \\\\\n",
        "\\frac{\\partial J(\\theta)}{\\partial \\theta_n}\n",
        "\\end{pmatrix}\n",
        "\\end{align}\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UFcZm8ruwTPj",
        "colab_type": "text"
      },
      "source": [
        "## Computing the gradient of the policy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yaEPYXN-5FNT",
        "colab_type": "text"
      },
      "source": [
        "### Numerically Computing the gradient of the policy (not the same as the PG)\n",
        "* Numerical method of finite differences, works even if your objective isn't differentiable. \n",
        "  * You just do small perturbations to the policy in each parameter dimension, and then look at the value of the objective before and after the perturbation. That yields a numerical estimate of the gradient in each direction.\n",
        "  * Very slow, especially in high-dimensional policy space, but it works sometimes (AIBO Soccer, UT, Peter Stone)\n",
        "\n",
        "#### Computing the Gradient of the Policy Analytically \n",
        "* Here we still have no value function, these are still monte carlo approaches\n",
        "* Assume that the policy $\\pi_\\theta$ is differentiable whenever it is non-zero\n",
        "  * Technically only has to be differentiable wherever you are picking actions???\n",
        "* and assume we know the gradient \\nabla_theta \\pi_\\theta(s, a)\n",
        "  * e.g. our policy is represented as a neural network we have created\n",
        "* Likelihood ratios exploit the following identity\n",
        "\\begin{align}\n",
        "  \\nabla_\\theta\\pi_\\theta(s, a) \n",
        "  &= \\pi_\\theta(s, a)\\frac{\\nabla_\\theta \\pi_\\theta (s, a)}{\\pi_\\theta(s, a)} \\\\\n",
        "  &= \\pi_\\theta(s, a)\\nabla_\\theta\\log\\pi_\\theta (s, a)\n",
        "\\end{align}\n",
        "  * We can multiply and divide by our policy without changing the objective\n",
        "  * the gradient of the policy divided by the policy is equal to the gradient of the log of the policy\n",
        "* This term $\\nabla_\\theta\\log\\pi_\\theta(s,a)$ is called the score function\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WmyBi4CAAV6Q",
        "colab_type": "text"
      },
      "source": [
        "### What does the score look like for common policy choices?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1g4MGk4nT5k8",
        "colab_type": "text"
      },
      "source": [
        "#### Softmax Policy\n",
        "* For discrete actions, a very simple way to parameterise the policy\n",
        "  * A softmax policy is any policy that makes the probability of an action proportional to some exponentiated 'value'\n",
        "* It's a smoothly parameterised policy that tells us how often we should take a particular action\n",
        "* We choose some features, and some linear weights, and the linear combination of state-action features becomes something like a value\n",
        "* This is called the linear softmax policy: the probability of taking an action a from state s is proportional to the exponentiated value\n",
        "$$ \\pi_\\theta(s, a) \\propto e^{\\phi(s, a)^\\top\\theta}  $$\n",
        "* Now we want to find the gradient of this policy, i.e. the score function:\n",
        "$$ \\nabla_\\theta\\log\\pi_\\theta(s, a) = \\phi(s, a) - \\mathbb{E}_{\\pi_\\theta} \\big[\\phi(s, \\cdot)\\big] $$\n",
        "  * Intuitively, it's scoring 'how much more of this feature do I have than usual'. If this feature occurs more than usual with this action, and it gets a good reward, then we want to adjust the policy to do more of that action."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uaT9YkFCskY8",
        "colab_type": "text"
      },
      "source": [
        "#### Gaussian Policy\n",
        "* For continuous action spaces.\n",
        "* Just a different parameterisation of the policy, here we usually just parameterise the mean $\\mu$ but could also parameterise the variance $\\sigma^2$\n",
        "* Policy is a Gaussian over some support, $a \\sim \\mathcal{N}(\\mu(s), \\sigma^2)$\n",
        "* The score function is:\n",
        "$$ \\nabla_\\theta\\log\\pi_\\theta(s,a) = \\frac{(a-\\mu(s))\\phi(s)}{\\sigma^2} $$\n",
        "* Intuitively, this again measures how much 'more' we're taking this action than usual (than the mean), multiplied by the features that were observed. If observing those features resulted in more reward, we adjust the policy parameters to increase the probability of taking that action again in that state."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_XpCHeM0wyp0",
        "colab_type": "text"
      },
      "source": [
        "## Computing the Policy Gradient $\\nabla_\\theta J(\\theta)$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_z3Eb7jbw-Lt",
        "colab_type": "text"
      },
      "source": [
        "### One-step MDP (a.k.a contextual bandit)\n",
        "* Agents starts in state $s\\sim d(s)$\n",
        "* Environment terminates after one time-step with reward $r = \\mathcal{R_{s,a}}$\n",
        "* We use the likelihood ratio trick to compute the PG:\n",
        "\\begin{align}\n",
        "J(\\theta) &= \\mathbb{E}_{\\pi_\\theta}[r] \\\\\n",
        "&= \\sum_{s\\in \\mathcal{S}}d(s)\\sum_{a\\in\\mathcal{A}}\\pi_\\theta(s, a)\\mathcal{R_{s,a}}\\\\\n",
        "\\nabla_{\\theta} J(\\theta) &=  \n",
        "\\sum_{s\\in \\mathcal{S}} d(s) \\sum_{a\\in\\mathcal{A}} \\pi_\\theta (s, a) \\nabla_{\\theta} \\log \\pi_\\theta (s, a) \\mathcal{R_{s,a}}  \\\\\n",
        "&= \\mathbb{E}\\big[\\nabla_\\theta \\log \\pi_\\theta (s, a) r \\big]\n",
        "\\end{align}\n",
        "* Note that the whole point of using the likelihood ratio trick is to derive an expression for the gradient that is itself an expectation wrt the policy (i.e. wrt the behaviour distribution)\n",
        "  * We can sample by taking actions in the environment in order to estimate this expectation\n",
        "  * This expectation tells us the direction in which to shift the policy parame"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZTa0LJxs2gQd",
        "colab_type": "text"
      },
      "source": [
        "### Policy Gradient Theorem\n",
        "* We want to know which direction to shift the parameters in a multi-step MDP\n",
        "* Beautifully, all we need to do is replace the immediate reward with the value function. That turns out to give us the true gradient of the policy objective function.\n",
        "* For any differentiable policy $\\pi_\\theta(s, a)$, \\\\\n",
        "for any of the policy objective functions $J = J_1, J_{avR}$, or $\\frac{1}{1-\\gamma} J_{avV}$\n",
        "\n",
        "\\begin{align}\n",
        "  \\nabla_\\theta J(\\theta) = \\mathbb{E}_{\\pi_\\theta} \\big[\\nabla_\\theta \\log \\pi_\\theta(s, a)Q^{\\pi_\\theta}(s, a)\\big] \n",
        "\\end{align}\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HIpZo4HX80Y-",
        "colab_type": "text"
      },
      "source": [
        "## Algorithms\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sBa5mAoT87no",
        "colab_type": "text"
      },
      "source": [
        "### Monte-Carlo Policy Gradient (REINFORCE)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H4Nw6hvCATxJ",
        "colab_type": "text"
      },
      "source": [
        "* Update parameters by stochastic gradient ascent, using PG theorem\n",
        "* Using policy gradient theorem\n",
        "* Using return $v_t$ as an unbiased sample of $Q^{\\pi_\\theta}(s_t,a_t)$\n",
        "* Gradient step: \n",
        "$$ \\Delta\\theta_t = \\alpha \\nabla_\\theta\\log\\pi_\\theta(s_t, a_t) v_t    $$\n",
        "* You get a very smooth learning curve (in contrast to value-based, which chatters and sometimes collapses)\n",
        "* Problem: Slow. Very high variance, like any Monte Carlo estimate.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IMAey7sT9AoL",
        "colab_type": "text"
      },
      "source": [
        "### Actor-Critic"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2hhem4fyYdLH",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "\n",
        "Desiderata: \n",
        "* Smooth learning with good convergence properties\n",
        "* Be faster and lower variance than Monte-Carlo methods\n",
        "\n",
        "To reduce variance we introduce a critic:\n",
        "* The critic estimates the action-value function,\n",
        "$$ Q_w(s, a) \\approx Q^{\\pi_\\theta} $$\n",
        "* This is just the familiar problem of *policy evaluation*: how good is policy $\\pi_\\theta$\n",
        "  * Could do it with MC policy evalutation, TD learning, TD(\\lambda)\n",
        "* we will substitute these in place of the return in our policy gradient algorithm\n",
        "* Actor-critic algorithms maintain *two* sets of parameters\n",
        "  * Critic: Updates action-value function parameters $w$\n",
        "  * Actor: Updates policy parameters $\\theta$, in direction suggested by critic\n",
        "  \n",
        "Actor-critic algorithms follow an approximate policy gradient\n",
        "\\begin{align}\n",
        "\\nabla_\\theta J(\\theta) &\\approx \\mathbb{E}_{\\pi_\\theta}\\big[\\nabla_\\theta\\log\\pi_\\theta(s, a) Q_w(s, a)\\big] \\\\\n",
        "\\Delta\\theta &= \\alpha\\nabla_\\theta\\log\\pi_\\theta(s, a)Q_w(s, a)\n",
        "\\end{align}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R2agexyLwu0E",
        "colab_type": "text"
      },
      "source": [
        "#### Czepesvari on Actor-Critic\n",
        "* **Actor-Critic methods implement Generalised Policy Iteration**\n",
        "  * Remember that policy-iteration works by alternating between a complete policy evaluation and a complete policy improvement step. \n",
        "  * When using sample-vased methods or function approximation, exactly evaluation of the policies may require infinitely many samples or might be impossible due to the restrictions of the function-approximation technique\n",
        "  * Hence RL algs simulating policy iteration must change the policy based on incomplete knowledge of the value function.\n",
        "  * This is GPI: there are two closely interacting processes, actor and critic, policy improvement and policy evaluation.\n",
        "    * Note that the policy used to generate samples (the behaviour policy), can be different from the one that is evaluated by the critic and improved by the actor\n",
        "    * This can be very useful because the critic can learn about actions not preferred by the current target policy so that the critic can improve the target policy\n",
        "    \n",
        " #### Problems\n",
        " * Note that unlike perfect policy iteration, a GPI method may generate a policy taht is substantially worse than the previous one.\n",
        "  * The quality of the sequence of generated policies may oscillate or even diverge when the PE step is incomplete, regardless of whether the PI step is exact or approximate (B&T 1996)\n",
        "  * In practice, they often improve in the beginning and oscillate later\n",
        "  "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fQETRdh5e62h",
        "colab_type": "text"
      },
      "source": [
        "#### Algorithm: Action-Value Actor-Critic\n",
        "* Simple action-value critic using linear value fn approimation Q_w(s, a) = \\phi(s, a)^\\top_w\n",
        "  * Critic: Updates $w$ by linear TD(0)\n",
        "  * Actor: Updates $\\theta$ by policy gradient"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8k-M_1lUes_E",
        "colab_type": "text"
      },
      "source": [
        "### The Bias Variance Tradeoff in RL\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wuU3Wa94Y4Ve",
        "colab_type": "text"
      },
      "source": [
        "* Monte Carlo methods sample the actual end-of-episode return before doing an update. This return is an unbiased estimators of the true value\n",
        "* TD methods do backups; that is, they estimate the value of a particular state using estimates of the value of other states. These estimates are much lower variance, but this also introduces bias.\n",
        "* Methods like TD($\\lambda$) attempt to balance this tradeoff\n",
        "\n",
        "#### Bias/Variance in Policy Gradients\n",
        "* Approximating the policy gradient introduces bias\n",
        "* A biased policy gradient may not find the right solution\n",
        "  * e.g. if Q_w(s, a) uses aliased features, can we solve the gridworld example\n",
        "* Amazingly, if we choose value function approximation carefully, we can avoid introducing bias and follow the **exact** policy gradient\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tg8f5mkjfEjy",
        "colab_type": "text"
      },
      "source": [
        "### Compatible Function Approximation Theorem\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9AjDMV5RgagA",
        "colab_type": "text"
      },
      "source": [
        "If the following two conditions are satisfied:\n",
        "1.  Value function approximator is **compatible** to the policy:\n",
        "$$ \\nabla_w Q_w(s, a) = \\nabla_\\theta\\log \\pi_\\theta(s, a) $$\n",
        "2. Value function parameters $w$ minimise the mean-squared error\n",
        "$$ \\mathcal{L}(w) = \\mathbb{E}_{\\pi_\\theta} \\big[(Q^{\\pi_\\theta}(s, a) - Q_w(s, a))^2\\big] $$\n",
        "\n",
        "Then the policy gradient is exact,\n",
        "$$ \\nabla_\\theta J(\\theta) = \\mathbb{E}_{\\pi_\\theta}\\big[\\nabla_\\theta\\log \\pi_\\theta(s, a) Q_w(s, a)\\big]  $$\n",
        "\n",
        "\n",
        "**TODO(jarryd@google.com): Prove this theorem (slides have a proof)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VYFFRTX1fNXa",
        "colab_type": "text"
      },
      "source": [
        "### Advantage Actor-Critic"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LU9a9VHQqmst",
        "colab_type": "text"
      },
      "source": [
        "#### Reducing variance using a baseline\n",
        "* We substract a baseline function $B(s)$ from the policy gradient\n",
        "* A good baseline is the state value function $B(s) = V^{\\pi_\\theta}(s)$\n",
        "* So we can rewrite the policy gradient using the advantage function $A^{\\pi_\\theta}(s, a)$\n",
        "\\begin{align}\n",
        "A^{\\pi_\\theta}(s, a) &= Q^{\\pi_\\theta}(s, a) - V^{\\pi_\\theta}(s) \\\\ \n",
        "\\nabla_\\theta J(\\theta) &= \\mathbb{E}_{\\pi_\\theta}\\big[\\nabla_\\theta\\log\\pi_\\theta(s,a) A^{\\pi_\\theta}(s, a)\\big]\n",
        "\\end{align}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5cAbJkaAstYK",
        "colab_type": "text"
      },
      "source": [
        "#### Estimating the advantage function\n",
        "* You could learn both $Q^{\\pi_\\theta}$ and $V^{\\pi_\\theta}$ with two different function approximators and sets of parameters\n",
        "* Better way: **the td-error $\\delta^{\\pi_\\theta}$ is an unbiased estimator of the advantage function**\n",
        "\\begin{align}\n",
        "\\mathbb{E}_{\\pi_\\theta}\\big[\\delta^{ \\pi_\\theta} \\mid s, a\\big]\n",
        "&= \\mathbb{E}_{\\pi_\\theta}\\big[r + \\gamma V^{\\pi_\\theta}(s')\\mid s, a \\big]-V^{\\pi_\\theta}(s) \\\\\n",
        "&= Q^{\\pi_\\theta}(s,a) - V^{\\pi_theta}(s) \\\\\n",
        "&= A^{\\pi_\\theta}(s, a)\n",
        "\\end{align}\n",
        "* So we just use the TD error to compute the PG:\n",
        "$$ \\nabla_\\theta J(\\theta) = \\mathbb{E}_{\\pi_\\theta}\\big[\\nabla_\\theta \\log \\pi_\\theta(s, a)\\delta^{\\pi_\\theta}\\big] **$$**\n",
        "* In practice we can use an approximate TD error\n",
        "$$ \\delta_v = r + \\gamma V_v(s') - V_v(s) $$\n",
        "which only requires one set of critic parameters, **we don't need to estimate $ /Q^{\\pi_\\theta}$**\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "apM5iLKOhvxl",
        "colab_type": "text"
      },
      "source": [
        "## PG algorithms are on-policy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5KRkMR8uhzhr",
        "colab_type": "text"
      },
      "source": [
        "* On-policy MEANS that when you update your policy, you throw away your old samples and generate new samples\n",
        "* When we deal with very complicated policy classes represented by neural networks, which we only change a bit with each gradient step, this can make on-policy learning very inefficient.\n",
        "* What if we want to learn from samples generated under a (perhaps only slightly) different policy from the current one?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mNyAZyG0iSDs",
        "colab_type": "text"
      },
      "source": [
        "### Importance Sampling and the off-policy Policy Gradient\n",
        "* It's a way to estimate the expectation of some function wrt to some distribution, but when we can't sample from that distribution\n",
        "* Instead we sample from another proposal distribution\n",
        "\\begin{align}\n",
        "\\mathbb{E}_{x\\sim p(x)}[f(x)]\n",
        "&= \\int p(x)f(x)dx \\\\\n",
        "&= \\int \\frac{q(x)}{q(x)}p(x)f(x)dx \\\\\n",
        "&= \\mathbb{E}_{x\\sim q(x)} \\bigg[\\frac{p(x)}{q(x)}f(x)\\bigg]\n",
        "\\end{align}\n",
        "* In the RL setup the relevant expectation is our objective\n",
        "$$ \\mathbb{E}_{\\tau\\sim \\pi_{\\theta}(\\tau)}\\big[r(\\tau)\\big]         $$\n",
        "and we have samples from some $\\bar{\\pi}(\\tau)$ instead\n",
        "* Using importance sampling we can estimate this expectation:\n",
        "$$ J(\\theta) = \\mathbb{E}_{\\tau\\sim\\bar{\\pi}(\\tau)}\\bigg[\\frac{\\pi_\\theta(\\tau)}{\\bar{\\pi}(\\tau)}r(\\tau)\\bigg]  $$\n",
        "* Computing the gradient and using the likelihood ratio trick yields the off-policy PG:\n",
        "\\begin{align}\n",
        "\\nabla_{\\theta'} J(\\theta')\n",
        "&= \\mathbb{E}_{\\tau\\sim\\pi_\\theta(\\tau)}\\bigg[\\frac{\\pi_{\\theta'}(\\tau)}{\\pi_\\theta(\\tau)}\\nabla_{\\theta'}\\log \\pi_{\\theta'}(\\tau)r(\\tau) \\bigg] , \\qquad \\text{when } \\theta \\neq \\theta'\n",
        "\\end{align}\n",
        "* Substituting terms in for $\\pi_\\theta$ and $\\bar{\\pi}$, we can cancel out transition probability terms in the product, and all we are left with is the ratio of action probabilities under each policy\n",
        "* A first order approximation for IS:\n",
        "\\begin{align}\n",
        "J(\\theta') = \n",
        "\\sum_{t=1}^{T}\\mathbb{E}_{s_t\\sim p_\\theta(s_t)}\\big[\\frac{p_{\\theta'}(s_t)}{p_{\\theta}(s_t)}\\big]\n",
        "\\mathbb{E}_{a_t\\sim\\pi_\\theta(a_t \\mid s_t)} \\bigg[\\frac{\\pi_{\\theta'}(a_t\\mid s_t)}{\\pi_{\\theta'}(a_t\\mid s_t)}r(s_t,a_t)\\bigg]\n",
        "\\end{align}\n",
        "where we then ignore the importance ratios for the transition probabilities. Will see later in Levine course why this is reasonable.\n",
        "* Something about importance weights being exponential in $T$, and that's why we left off the aforementioned weights. Revisit."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kAcDOeMNe8kn",
        "colab_type": "text"
      },
      "source": [
        "# Implementation"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y1UklWVDMmVO",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "Action = int\n",
        "State = int\n",
        "Vector = np.ndarray\n",
        "TimeStep = collections.namedtuple('TimeStep', ['observation', 'reward', 'done'])"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kAxwAJmsKr0m",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class Agent(metaclass=ABCMeta):\n",
        "  \n",
        "  @abstractmethod\n",
        "  def policy(self, timestep: TimeStep) -> Action:\n",
        "    pass\n",
        "\n",
        "  @abstractmethod\n",
        "  def update(self,\n",
        "             timestep: TimeStep,\n",
        "             action: Action,\n",
        "             new_timetstep: TimeStep) -> None:\n",
        "    pass"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lJqcT5YiO3mh",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title TD-Learning Agent with Linear Function Approximation\n",
        "\n",
        "class TDAgentLFA:\n",
        "  \n",
        "  def __init__(self,\n",
        "               num_actions: int,\n",
        "               exploration_policy: Policy,\n",
        "               alpha: float = 0.01,\n",
        "               gamma: float = 0.9):\n",
        "    \"\"\"Builds a tabular Q-learning agent.\n",
        "    \n",
        "    Glossary:\n",
        "      A: Size of action space.\n",
        "      X: Size of feature space.\n",
        "    \n",
        "    Args:\n",
        "      TODO(jarryd): Describe what each of these means.\n",
        "    \"\"\"\n",
        "    \n",
        "    # Exploration Policy\n",
        "    self._exploration_policy = exploration_policy\n",
        "    \n",
        "    # Learning parameters\n",
        "    self._gamma = gamma\n",
        "    self._alpha = alpha\n",
        "\n",
        "    # Value function parameters/environment metadata\n",
        "    self._num_actions = num_actions\n",
        "    self._weights = None\n",
        "  \n",
        "  def policy(self, timestep: TimeStep) -> Action:\n",
        "    \"\"\"Compute the policy.\"\"\"\n",
        "    \n",
        "    q_values = self._get_q_values(timestep.observation.flatten())\n",
        "\n",
        "    return self._exploration_policy(q_values)\n",
        "  \n",
        "  def update(self, \n",
        "             timestep: TimeStep,\n",
        "             last_action: Action,\n",
        "             new_timestep: TimeStep) -> None:\n",
        "    \"\"\"\n",
        "    Performs a SGD step to improve the approximation of q(., .)\n",
        "    \"\"\"\n",
        "\n",
        "    x = timestep.observation.flatten()\n",
        "    a = last_action\n",
        "    r = new_timestep.reward\n",
        "    x_ = new_timestep.observation.flatten()\n",
        "    \n",
        "    old_q = self._get_q_values(x, a)\n",
        "    target_q = self._get_target_q_value(x_)\n",
        "\n",
        "    target = r\n",
        "    if not new_timestep.done:\n",
        "      target += self._gamma * target_q\n",
        "    \n",
        "    td_error = target - old_q\n",
        "    step = self._alpha * td_error * x\n",
        "    self._weights[a] += step\n",
        "\n",
        "  def _get_q_values(self, x: Vector, action: Action = None):\n",
        "    \"\"\"\n",
        "    TODO(jarryd@google.com)s\n",
        "    Compute the approximate q-values\n",
        "    \"\"\"\n",
        "    \n",
        "    if self._weights is None:\n",
        "      self._weights = np.random.rand(self._num_actions, x.shape[0])  # [A, X]\n",
        "    \n",
        "    if action is None:\n",
        "#       raise NotImplemented(\"You haven't dealt with the case for Q-Learning\")\n",
        "      return np.dot(self._weights, x)  # []\n",
        "    else:\n",
        "      return np.dot(self._weights[action], x)\n",
        "\n",
        "  @abstractmethod\n",
        "  def _get_target_q_value(self, new_observation: Vector) -> float:\n",
        "    \"\"\"\n",
        "    Concrete subclasses like SarsaAgent/QLearningAgent must implement this\n",
        "    method to compute the next action-value Q(s',a') in on/off policy manner.\n",
        "    \"\"\"\n",
        "    pass\n",
        "  \n",
        "  \n",
        "#@title Sarsa Agent with LFA\n",
        "\n",
        "class SarsaAgentLFA(TDAgentLFA):\n",
        "  \"\"\"A Sarsa agent that uses linear value function approximation\"\"\"\n",
        "  def __init__(self, *args, **kwargs):\n",
        "    super().__init__(*args, **kwargs)\n",
        "    self._action = None    \n",
        "  \n",
        "  def _get_target_q_value(self, new_observation: Vector) -> float:\n",
        "    \"\"\"Computes the (on-policy) ....# TODO\"\"\"\n",
        "    \n",
        "    q_target = np.dot(self._weights[self._action], \n",
        "                      new_observation)\n",
        "    \n",
        "    return q_target"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NMRwusDWM4VH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class ActorCriticAgentLFA(Agent):\n",
        "  \n",
        "  def __init__(num_actions: int,\n",
        "               critic: TDAgentLFA, \n",
        "               behaviour_policy: Policy,\n",
        "               alpha: float = 0.01\n",
        "               gamma: float = 0.9):\n",
        "    \n",
        "    # Value function parameters\n",
        "    self._num_actions = num_actions\n",
        "    \n",
        "    # Critic module, a TD-learning Agent that performs policy evaluation\n",
        "    self._critic = critic\n",
        "    # \n",
        "    self._behaviour_policy = behaviour_policy\n",
        "    \n",
        "    # Learning Parameters\n",
        "    self._gamma = gamma\n",
        "    self._alpha = alpha\n",
        "    \n",
        "  def policy(self, timestep:TimeStep) -> Action:\n",
        "    \n",
        "    q_values = self._critic._get_q_values(timestep.observation.flatten())\n",
        "    \n",
        "    return self._behaviour_policy(q_values)\n",
        "  \n",
        "  def update(self,\n",
        "            timestep: TimeStep,\n",
        "            action: Action,\n",
        "            new_timestep: TimeStep) -> None:\n",
        "    \n",
        "    \n",
        "    self._critic.update(timestep, action, new_timestep)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lOv7wNgMe_vZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Initialise states s, policy weights theta\n",
        "# Sample a ~ pi_theta\n",
        "# for each step in the trajectories\n",
        "  # r = R(s,a), transition s' ~ P(s,a)\n",
        "\n",
        "  # Critic computes td_error\n",
        "  # Sample action a' ~ \\pi_theta(s', a')\n",
        "  # td_error = r + gamma Q_w(s', a') - Q_w(s, a)\n",
        "\n",
        "  # Update policy weights\n",
        "  # theta = theta + alpha * score * Q_w(s, a)\n",
        "\n",
        "  # Update critic weights\n",
        "  # w = w + beta * td_error * phi(s, a)\n",
        "\n",
        "  # action = a'\n",
        "  # state = s'"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rOyipUkBlcQ9",
        "colab_type": "code",
        "cellView": "code",
        "colab": {}
      },
      "source": [
        "#@title Four Rooms.\n",
        "\n",
        "class Tile(enum.Enum):\n",
        "  EMPTY = 0\n",
        "  WALL = 1\n",
        "  AGENT = 2\n",
        "  GOAL = 3\n",
        "  \n",
        "  \n",
        "# class Action(enum.Enum):\n",
        "#   UP = (-1, 0)\n",
        "#   DOWN = (1, 0)\n",
        "#   LEFT = (0, -1)\n",
        "#   RIGHT = (0, 1)\n",
        "\n",
        "\n",
        "MAP = \"\"\"\n",
        "#########\n",
        "#       #\n",
        "#   #   #\n",
        "#   #   #\n",
        "### ## ##\n",
        "#   #   #\n",
        "#       #\n",
        "#   #   #\n",
        "#########\n",
        "\"\"\"\n",
        "\n",
        "\n",
        "class FourRooms(object):\n",
        "  \n",
        "  def __init__(self, \n",
        "               map_string=MAP, \n",
        "               initial_pos=(1, 1), \n",
        "               goal_pos=(7, 7)):\n",
        "    # Convert map to binary\n",
        "    map_string = map_string.replace('#', str(Tile.WALL.value))\n",
        "    binary_map = map_string.replace(' ', str(Tile.EMPTY.value))\n",
        "\n",
        "    # Convert to list\n",
        "    map_list = [list(x) for x in binary_map.split('\\n')[1:-1]]\n",
        "    \n",
        "    # Convert to numpy array\n",
        "    self._grid = np.array(map_list, dtype=np.int32)\n",
        "    \n",
        "    self._size = self._grid.shape[0]\n",
        "    \n",
        "    # Make sure that the grid is square.\n",
        "    if self._grid.shape[1] != self._grid.shape[0]:\n",
        "      raise ValueError('Map must be square, soz.')\n",
        "      \n",
        "    self._actions = {\n",
        "        0: np.array((-1, 0)),\n",
        "        1: np.array((1, 0)),\n",
        "        2: np.array((0, -1)),\n",
        "        3: np.array((0, 1)),\n",
        "    }\n",
        "    \n",
        "    self._initial_pos = np.array(initial_pos)\n",
        "    self._pos = self._initial_pos\n",
        "    \n",
        "    self._goal_pos = np.array(goal_pos)\n",
        "    self._grid[goal_pos[0], goal_pos[1]] = Tile.GOAL.value\n",
        "  \n",
        "  def step(self, action: int) -> TimeStep:\n",
        "    # TODO(jaslanides): Noisy actions?\n",
        "    action_vec = self._actions[action]\n",
        "    next_x, next_y = self._pos + action_vec\n",
        "    next_tile = self._grid[next_x, next_y]\n",
        "    \n",
        "    reward = 0\n",
        "    done = False\n",
        "    \n",
        "    if next_tile != Tile.WALL.value:  # Not a wall\n",
        "      self._pos = (next_x, next_y)  # Move there :)\n",
        "\n",
        "    if next_tile == Tile.GOAL.value:  # If it's the goal, get money get paid\n",
        "      reward = 1\n",
        "      done = True  # Episode ends when u get paid  \n",
        "    timestep = TimeStep(observation=self._get_observation(), \n",
        "                        reward=reward,\n",
        "                        done=done)\n",
        "\n",
        "    return timestep\n",
        "  \n",
        "  def reset(self) -> TimeStep:\n",
        "    self._pos = self._initial_pos\n",
        "    # TODO(jaslanides): Randomize goals?\n",
        "    return TimeStep(observation=self._get_observation(), \n",
        "                    reward=0, \n",
        "                    done=False)\n",
        "      \n",
        "  def _get_observation(self):\n",
        "    obs = self._grid.copy()\n",
        "    obs[self._pos[0], self._pos[1]] = Tile.AGENT.value\n",
        "    obs = np.float32(obs)\n",
        "    obs -= obs.mean()\n",
        "\n",
        "    return obs\n",
        "  \n",
        "  @property\n",
        "  def num_actions(self) -> int:\n",
        "    return 4  # L, U, R, D\n",
        "  \n",
        "  @property\n",
        "  def obs_shape(self) -> np.ndarray:\n",
        "    # return the obs shape\n",
        "    return self._grid.shape\n",
        "  "
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}