distributional_rl.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "distributional-rl.ipynb",
      "provenance": [],
      "authorship_tag": "ABX9TyOazqp9aRRkxusKbem8lCZy",
      "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/distributional_rl.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dk6lSq2Mje1V",
        "colab_type": "text"
      },
      "source": [
        "## A Distributional Perspective on Reinforcement Learning (Bellemare, 2017)\n",
        "\n",
        "Stems from the realisation cthat under the RL problem, the distribution over returns is very complex, e.g. may be multimodal. If you look at the mean, you're often looking at a reward you will *never* see in reality. The value (expected return) therefore hides the intrinsic randomness, and so we are not modelling an aspect of reality. This work is basically about *removing the expectations from the Bellman equation*.\n",
        "\n",
        "### Concepts: Return, The Value, and CNC\n",
        "* Value is expected sum of future rewards: $$V^\\pi(x) = \\mathbb{E}[Z^\\pi(x)]  = \\mathbb{E}\\bigg[\\sum_{t=0}^\\infty R(x_t)\\bigg]$$\n",
        "* The return (sum of future reward) from a state $x$ is a random variable $Z^\\pi$. \n",
        "* That means we can estimate it using Monte Carlo estimation\n",
        "* CNC: Let's assume the reward is undiscounted and the horizon is finite. Then we have a tractable supervised learning problem. We have a number of possible outcomes (return amounts) and we can just predict all of them (learn a distribution over them)\n",
        "\n",
        "####Bellman equations for different moments\n",
        "* Bellman (1957): Bellman equation for mean (describes the relation between the mean return at one state and the mean return at another state)\n",
        "$$ V^\\pi(x) = \\mathbb{E}\\big[Z^\\pi(x)\\big] = \\mathbb{E}R(x) + \\gamma \\mathbb{E}\\big[Z^\\pi(X')\\big]$$\n",
        "* But there are other Bellman equations for higher moments:\n",
        "\t* Sobel (1982): Bellman equation for variance\n",
        "\t* Engel (2003): Bellman equation for Bayesian uncertainty - this describes *our* uncertainty about the system, not the actual stochasticity of the return.\n",
        "\t* Azar et al. (2011), Lattimore & Hutter (2012): Bellman equation for higher moments.\n",
        "* Bellemare is trying to quantify the *whole* distribution and not just certain moments, and hopefully that gives a more unified picture.\n",
        "\n",
        "### Fixed Point of the Distributional Bellman Equation:\n",
        "$$ Z^\\pi(x) \\overset{D}{=} R(x) + \\gamma Z^\\pi(X')\\qquad X' \\sim P^\\pi(\\cdot\\mid x)$$\n",
        "* This is an instance of a *recursive distributional equation*. They're also called *stochastic fixed point equations*. Here's a recursive distributional equation whose solution is the standard normal distribution $\\mathcal{N}(0,1)$ (Rosler, 1992):\n",
        "\t$$ X \\overset{D}{=}  \\frac{1}{\\sqrt{2}}X_1 + \\frac{1}{\\sqrt{2}}X_2 $$\n",
        "\t$$ X \\sim X_1 \\sim X_2$$\n",
        "\n",
        "#### Look for an operator\n",
        "Whenever B writes an RL paper he asks \"Is there an **operator** I can look at?\"\n",
        "* An operator is basically a way of looking at the expected (average) behaviour of a learning algorithm\n",
        "* When we do RL we sample states from the environment and we update. The operator tells you what happens to your current prediction after one update.\n",
        "* So we have this value distribution which satisfies this recursive equation. So we have a fixed point...? He says this like it's obvious\n",
        "* Now we ask: *is there an operator we could apply and use to study the learning behaviour*\n",
        "* Let's call the operator $\\mathcal{T}^\\pi$ \n",
        "\n",
        "#### Applying the operator\n",
        "The result of applying the operator to a given value distribution is going to be distributed like: $$ \\mathcal{T}^\\pi Z(x,a) :\\overset{D}{=} R(x,a) + \\gamma Z(X', A') $$\n",
        "$$ X' \\sim P(x\\mid x,a)\\qquad A' \\sim \\pi(\\cdot\\mid X') $$\n",
        "So we have three sources of intrinsic randomness that define the compound distribution $\\mathcal{T}^\\pi Z$. This makes it fundamentally different to the usual Bellman Operator:\n",
        "1. Reward\n",
        "2. Next state\n",
        "3. Next state return\n",
        "\n",
        "## Theoretical Results\n",
        "### The Distributional Bellman Operator $\\mathcal{T}^\\pi$ is a contraction in $\\bar{d}_p$\n",
        "##### Aside: The inverse c.d.f transform\n",
        " * The probability integral transform states that if $X$ is a continuous random variable with cumulative distribution function $F_{X}$ then the random variable $Y=F_{X}(X)$ has a uniform distribution on $[0, 1]$. \n",
        " * The inverse probability integral transform is just the inverse of this: specifically, if  $Y$ has a uniform distribution on $[0, 1]$ and if $X$ has a cumulative distribution $F_{X}$ then the random variable $F_{X}^{-1}(Y)$ has the same distribution as $X$.\n",
        "\n",
        "#### The Wasserstein Metric (Earth Mover Distance)\n",
        "The Wasserstein metric measures the distance between two c.d.f's $F$ and $G$.\n",
        "$$d_p(F,G) := \\inf_{U,V}\\left\\Vert U - V\\right\\Vert $$\n",
        "where the infimum is taken over all pairs of random variables $(U,V)$ with respective cumulative distributions F and G. The infimum is attained by the inverse c.d.f. transform of a random varible $\\mathcal{U}$ that is uniformly distributed on $[0,1]$:\n",
        "\\begin{aligned}\n",
        " d_p(F,G)\\\n",
        " :=& \\left\\Vert F^{-1}(\\mathcal{U})-G^{-1}(\\mathcal{U}) \\right\\Vert_p \\\\\n",
        " :=& \\bigg( \\int_0^1\\left\\vert F^{-1}(u) - G^{-1}(u)\\right\\vert^p du\\bigg)^{\\frac{1}{p}}\n",
        " \\end{aligned}\n",
        "* The last equation requires that $p<\\infty$. I'm pretty sure the above is a result that allows you to compute the Wasserstein metric more easily (as opposed to somehow taking the infimum over \"all pairs of RVs with cdfs $F$ and $G$\").\n",
        "* Note: Though inaccurate, they use $d_p(F,G)$  (W distance between cdfs) and $d_p(U,V)$ (where $U$ and $V$ are random variables) interchangeably, even though the W distance involves taking the inf over (all pairs of RVs with said c.d.f's) $U$ and $V$.\n",
        "\n",
        "#### Maximal form of the Wasserstein metric, applied to returns\n",
        "* Note that a value distribution $Z^\\pi$ is a vector of random variables, with each element $Z^\\pi(x,a)$ being a random variable corresponding to the distribution over returns for one state-action pair.\n",
        "\t* In other words it's a mapping from state-action pairs to distributions over returns.\n",
        "* As before, let $\\mathcal{U}$ be a random variable uniformly distributed on $[0,1]$. A maximal (over $\\mathcal{X}\\times\\mathcal{A}$) form of the W metric is:\n",
        "\\begin{aligned}\n",
        "\\bar{d}_p(Z_1,Z_2) \n",
        ":=& \\sup_{x,a} d_p(Z_1(x,a),Z_2(x,a)) \\\\\n",
        ":=& \\sup_{x,a}\\inf_{Z_1(x,a),\\ Z_2(x,a)} \\left\\Vert Z_1(x,a) - Z_2(x,a) \\right\\Vert \\\\\n",
        ":=& \\sup_{x,a}\\left\\Vert F_{Z_1(x,a)}^{-1}(\\mathcal{U})-F_{Z_2(x,a)}^{-1}(\\mathcal{U}) \\right\\Vert_p \\\\\n",
        " :=& \\sup_{x,a}\\bigg( \\int_0^1\\left\\vert F_{Z_1(x,a)}^{-1}(u)-F_{Z_2(x,a)}^{-1}(u)\\right\\vert^p du\\bigg)^{\\frac{1}{p}}\n",
        "\\end{aligned}\n",
        "\n",
        "\n",
        "#### Lemma 3: $\\mathcal{T}^\\pi$ is a contraction in $\\bar{d}_p$\n",
        "* Let $\\mathcal{Z}$ denote the space of value distributions with bounded moments. For two value distributions $Z_1, Z_2\\in\\mathcal{Z}$, we make use of a maximal form of the Wasserstein metric (above). Now, consider the process $Z_{k+1} := \\mathcal{T}^\\pi Z_k$, starting with some $Z_0\\in\\mathcal{Z}$.\n",
        "* We can expect the limiting expectation of $\\{Z_k\\}$ to converge exponentially quickly, as usual, to $Q^\\pi$. But we want stronger convergence.\n",
        "\n",
        "#### The argument for *distributional* convergence:\n",
        "> 1. $\\mathcal{T}^\\pi$ is a contraction in $\\bar{d}_p$, so all moments converge exponentially quickly.\n",
        "2. By Banach's fixed point theorem, $\\mathcal{T^\\pi}$ has a unique fixed point.\n",
        "3. By inspection (of the standard Q function def?) this fixed point must be $Z^\\pi$.\n",
        "4. Since we assumed all moments were bounded, this is sufficient to conclude that the sequence $\\{Z_k\\}$ converges to $Z^\\pi$ in $\\bar{d}_p$ for $1\\leq p \\leq \\infty$.\n",
        "\n",
        "Note that it is not a contraction in total variation distance, KL-divergence, or Kolmogorov distance.\n",
        "\n",
        "## Algorithms\n",
        "### The categorical algorithm and Categorical DQN\n",
        "They propose an algorithm based on the distributional Bellman operator.\n",
        "#### Model: Parametric approximation to the value distribution\n",
        "We model the value distribution $Z$ it using a discrete distribution $Z_\\theta$ parametrized by\n",
        "* $N\\in\\mathbb{N}$ the number of canonical returns (return buckets)\n",
        "* $V_{min},V_{max}\\in\\mathbb{R}$\n",
        "\n",
        "Its support is the set of atoms $\\{z_i = V_{min} + i\\Delta z : 0 \\leq i < N\\}, \\Delta z := \\frac{V_{max}-V_{min}}{N-1}$. This discrete distribution is computationally friendly and expressive (apparently says so in the PixelRNN paper).\n",
        "#### Loss\n",
        "* The analysis in section 3 suggests that the Wasserstein metric would be a natural distance to minimize as a loss between $\\mathcal{T}Z_\\theta$ and $Z_\\theta$.\n",
        "* It is also robust to discrepancies in support, which we need. BUT it doesn't work for some other reason, because we're learning from sample transitions, which isn't possible under the Wasserstein loss.\n",
        "* Instead they project the sample Bellman update $\\hat{\\mathcal{T}}Z_\\theta$ onto the support of $Z_\\theta$\n",
        "* The sample loss is then the KL-divergence between the approximate value distribution and the projected sample Bellman update (equation 7)\n",
        "#### Learning algorithm\n",
        "A modified DQN architecture that outputs the atom probabilities $p_i(x,a)$ instead of action-values. \n",
        "* It's pretty natural to use DQN, instead of outputting a 1D vector of values for each action (given a state), you output a distribution over return for each action, given a state: i.e. a 2D tensor of return-bucket-probabilities (the elements of the value distribution) $\\times$ actions.\n",
        "* **Exploration**: Simple $\\epsilon$-greedy policy. THey leave to future work the many ways in which an agent could select actions on the basis of the full distribution.\n",
        "\n",
        "## Results\n",
        "Why does it work better?\n",
        "* Results suggest that value distributions are better able to propagate rarely occurring events. \n",
        "\t* Makes sense because rarely occuring events with high reward might have that high reward averaged out if the mean of the return is always taken.\n",
        "* Also works much better in the stochastic setting (random action rejection by environment)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fST_huwXjlV-",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}