lab4_rl1_dynamic_programming.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P5gNz6b0lpCB"
      },
      "source": [
        "# INF581 Lab4: Dynamic Programming - Value Iteration and Policy Iteration\n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/jeremiedecock/polytechnique-inf581-2023/master/logo.jpg\" style=\"float: left; width: 15%\" />\n",
        "\n",
        "[INF581-2023](https://moodle.polytechnique.fr/course/view.php?id=14259) Lab session #4\n",
        "\n",
        "2019-2023 Jérémie Decock"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zlL77UvTlpCC"
      },
      "source": [
        "[![Open in Google Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/jeremiedecock/polytechnique-inf581-2023/blob/master/lab4_rl1_dynamic_programming.ipynb)\n",
        "\n",
        "[![My Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/jeremiedecock/polytechnique-inf581-2023/master?filepath=lab4_rl1_dynamic_programming.ipynb)\n",
        "\n",
        "[![NbViewer](https://raw.githubusercontent.com/jupyter/design/master/logos/Badges/nbviewer_badge.svg)](https://nbviewer.jupyter.org/github/jeremiedecock/polytechnique-inf581-2023/blob/master/lab4_rl1_dynamic_programming.ipynb)\n",
        "\n",
        "[![Local](https://img.shields.io/badge/Local-Save%20As...-blue)](https://github.com/jeremiedecock/polytechnique-inf581-2023/raw/master/lab4_rl1_dynamic_programming.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "tags": [],
        "id": "EY5DJdfplpCD"
      },
      "source": [
        "## Introduction\n",
        "\n",
        "The purpose of this lab is to introduce some classic concepts used\n",
        "in reinforcement learning like *Dynamic Programming*, *Bellman's Principle of Optimality* and *Bellman equations*.\n",
        "\n",
        "You will implement and test the two main dynamic programming algorithms (*Value Iteration* and *Policy Iteration*) in this Python notebook.\n",
        "\n",
        "You can either:\n",
        "- open, edit and execute the notebook in *Google Colab* following this link: http://www.jdhp.org/inf581/lab4 (or https://colab.research.google.com/github/jeremiedecock/polytechnique-inf581-2023/blob/master/lab4_rl1_dynamic_programming.ipynb); this is the **recommended** choice as you have nothing to install on your computer\n",
        "- open, edit and execute the notebook in *MyBinder* (if for any reason the Google Colab solution doesn't work): https://mybinder.org/v2/gh/jeremiedecock/polytechnique-inf581-2023/master?filepath=lab4_rl1_dynamic_programming.ipynb\n",
        "- download, edit and execute the notebook on your computer if Python3 and JypyterLab are already installed: https://github.com/jeremiedecock/polytechnique-inf581-2023/raw/master/lab4_rl1_dynamic_programming.ipynb\n",
        "\n",
        "If you work with Google Colab or MyBinder, remember to save or download your work regularly or you may lose it!\n",
        "\n",
        "$\n",
        "\\newcommand{\\vs}[1]{\\mathbf{#1}} % vector symbol (\\boldsymbol, \\textbf or \\vec)\n",
        "\\newcommand{\\ms}[1]{\\mathbf{#1}} % matrix symbol (\\boldsymbol, \\textbf)\n",
        "\\def\\U{V}\n",
        "\\def\\action{\\vs{a}}       % action\n",
        "\\def\\A{\\mathcal{A}}        % TODO\n",
        "\\def\\actionset{\\mathcal{A}} %%%\n",
        "\\def\\discount{\\gamma}  % discount factor\n",
        "\\def\\state{\\vs{s}}         % state\n",
        "\\def\\S{\\mathcal{S}}         % TODO\n",
        "\\def\\stateset{\\mathcal{S}}  %%%\n",
        "%\n",
        "\\def\\E{\\mathbb{E}}\n",
        "%\\newcommand{transition}{T(s,a,s')}\n",
        "%\\newcommand{transitionfunc}{\\mathcal{T}^a_{ss'}}\n",
        "\\newcommand{transitionfunc}{P}\n",
        "\\newcommand{transitionfuncinst}{P(\\nextstate|\\state,\\action)}\n",
        "\\newcommand{transitionfuncpi}{\\mathcal{T}^{\\pi_i(s)}_{ss'}}\n",
        "\\newcommand{rewardfunc}{r}\n",
        "\\newcommand{rewardfuncinst}{r(\\state,\\action,\\nextstate)}\n",
        "\\newcommand{rewardfuncpi}{r(s,\\pi_i(s),s')}\n",
        "\\newcommand{statespace}{\\mathcal{S}}\n",
        "\\newcommand{statespaceterm}{\\mathcal{S}^F}\n",
        "\\newcommand{statespacefull}{\\mathcal{S^+}}\n",
        "\\newcommand{actionspace}{\\mathcal{A}}\n",
        "\\newcommand{reward}{R}\n",
        "\\newcommand{statet}{S}\n",
        "\\newcommand{actiont}{A}\n",
        "%\\newcommand{newstatet}{S'}\n",
        "\\newcommand{nextstate}{\\state'}\n",
        "\\newcommand{newactiont}{A'}\n",
        "\\newcommand{stepsize}{\\alpha}\n",
        "\\newcommand{discount}{\\gamma}\n",
        "\\newcommand{qtable}{q_*}\n",
        "\\newcommand{finalstate}{\\state_F}\n",
        "%\n",
        "\\newcommand{\\vs}[1]{\\boldsymbol{#1}} % vector symbol (\\boldsymbol, \\textbf or \\vec)\n",
        "\\newcommand{\\ms}[1]{\\boldsymbol{#1}} % matrix symbol (\\boldsymbol, \\textbf)\n",
        "\\def\\vit{Value Iteration}\n",
        "\\def\\pit{Policy Iteration}\n",
        "\\def\\discount{\\gamma}  % discount factor\n",
        "\\def\\state{\\vs{s}}         % state\n",
        "\\def\\S{\\mathcal{S}}         % TODO\n",
        "\\def\\stateset{\\mathcal{S}}  %%%\n",
        "\\def\\cstateset{\\mathcal{X}} %%%\n",
        "\\def\\x{\\vs{x}}                    % TODO cstate\n",
        "\\def\\cstate{\\vs{x}}               %%%\n",
        "\\def\\policy{\\pi}\n",
        "\\def\\piparam{\\vs{\\theta}}         % TODO pparam\n",
        "\\def\\action{\\vs{a}}       % action\n",
        "\\def\\A{\\mathcal{A}}        % TODO\n",
        "\\def\\actionset{\\mathcal{A}} %%%\n",
        "\\def\\caction{\\vs{u}}       % action\n",
        "\\def\\cactionset{\\mathcal{U}} %%%\n",
        "\\def\\decision{\\vs{d}}       % decision\n",
        "\\def\\randvar{\\vs{\\omega}}       %%%\n",
        "\\def\\randset{\\Omega}       %%%\n",
        "\\def\\transition{T}       %%%\n",
        "\\def\\immediatereward{r}    %%%\n",
        "\\def\\strategichorizon{s}    %%% % TODO\n",
        "\\def\\tacticalhorizon{k}    %%%  % TODO\n",
        "\\def\\operationalhorizon{h}    %%%\n",
        "\\def\\constalpha{a}    %%%\n",
        "\\def\\U{V}              % utility function\n",
        "\\def\\valuefunc{V}\n",
        "\\def\\X{\\mathcal{X}}\n",
        "\\def\\meu{Maximum Expected Utility}\n",
        "\\def\\finaltime{T}\n",
        "\\def\\timeindex{t}\n",
        "\\def\\iterationindex{i}\n",
        "\\def\\decisionfunc{d}       % action\n",
        "\\def\\mdp{\\text{MDP}}\n",
        "$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HPdnxJ-7lpCF"
      },
      "source": [
        "**Notice**: Here we assume that the reward only depends on the state: $r(\\state) \\equiv \\mathcal{R}(\\state, \\action, \\state')$."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zTNzM-UXlpCF"
      },
      "source": [
        "## Python requirements"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LuQ-pNiWlpCG"
      },
      "source": [
        "This notebook requires the following Python libraries: *Gym*, NumPy, Pandas, Seaborn and Imageio.\n",
        "\n",
        "### If you use Google Colab\n",
        "\n",
        "Execute the next cell to install required libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "RmMZvg3KlpCG",
        "outputId": "045d224b-aa72-40d4-8067-90d211af9ef1",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: gym in /usr/local/lib/python3.8/dist-packages (0.25.2)\n",
            "Requirement already satisfied: gym-notices>=0.0.4 in /usr/local/lib/python3.8/dist-packages (from gym) (0.0.8)\n",
            "Requirement already satisfied: importlib-metadata>=4.8.0 in /usr/local/lib/python3.8/dist-packages (from gym) (6.0.0)\n",
            "Requirement already satisfied: cloudpickle>=1.2.0 in /usr/local/lib/python3.8/dist-packages (from gym) (2.2.0)\n",
            "Requirement already satisfied: numpy>=1.18.0 in /usr/local/lib/python3.8/dist-packages (from gym) (1.21.6)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.8/dist-packages (from importlib-metadata>=4.8.0->gym) (3.11.0)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.8/dist-packages (1.21.6)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.8/dist-packages (1.3.5)\n",
            "Requirement already satisfied: numpy>=1.17.3 in /usr/local/lib/python3.8/dist-packages (from pandas) (1.21.6)\n",
            "Requirement already satisfied: python-dateutil>=2.7.3 in /usr/local/lib/python3.8/dist-packages (from pandas) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.8/dist-packages (from pandas) (2022.7)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil>=2.7.3->pandas) (1.15.0)\n",
            "Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
            "Requirement already satisfied: seaborn in /usr/local/lib/python3.8/dist-packages (0.11.2)\n",
            "Requirement already satisfied: matplotlib>=2.2 in /usr/local/lib/python3.8/dist-packages (from seaborn) (3.2.2)\n",
            "Requirement already satisfied: pandas>=0.23 in /usr/local/lib/python3.8/dist-packages (from seaborn) (1.3.5)\n",
            "Requirement already satisfied: scipy>=1.0 in /usr/local/lib/python3.8/dist-packages (from seaborn) (1.7.3)\n",
            "Requirement already satisfied: numpy>=1.15 in /usr/local/lib/python3.8/dist-packages (from seaborn) (1.21.6)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.8/dist-packages (from matplotlib>=2.2->seaborn) (0.11.0)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib>=2.2->seaborn) (2.8.2)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib>=2.2->seaborn) (3.0.9)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.8/dist-packages (from matplotlib>=2.2->seaborn) (1.4.4)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.8/dist-packages (from pandas>=0.23->seaborn) (2022.7)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.8/dist-packages (from python-dateutil>=2.1->matplotlib>=2.2->seaborn) (1.15.0)\n"
          ]
        }
      ],
      "source": [
        "colab_requirements = [\n",
        "    \"gym\",\n",
        "    \"numpy\",\n",
        "    \"pandas\",\n",
        "    \"seaborn\"\n",
        "]\n",
        "import sys, subprocess\n",
        "def run_subprocess_command(cmd):\n",
        "    # run the command\n",
        "    process = subprocess.Popen(cmd.split(), stdout=subprocess.PIPE)\n",
        "    # print the output\n",
        "    for line in process.stdout:\n",
        "        print(line.decode().strip())\n",
        "        \n",
        "if \"google.colab\" in sys.modules:\n",
        "    for i in colab_requirements:\n",
        "        run_subprocess_command(\"pip install \" + i)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E1RgfYoYlpCH"
      },
      "source": [
        "### If you use MyBinder\n",
        "\n",
        "Required libraries are already installed, you have nothing to do.\n",
        "\n",
        "### If you have downloaded the notebook on your computer and execute it in your own Python environment\n",
        "\n",
        "Uncomment and execute the following cell to install required packages in your local environment (remove only the `#` not the `!`)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BQToQlSVlpCH"
      },
      "outputs": [],
      "source": [
        "#!pip install gym imageio numpy pandas seaborn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b6zWv0S6lpCI"
      },
      "source": [
        "## Import required libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "0SC0sy5SlpCI"
      },
      "outputs": [],
      "source": [
        "%matplotlib inline\n",
        "\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import math\n",
        "import gym\n",
        "import numpy as np\n",
        "import copy\n",
        "import pandas as pd\n",
        "import seaborn as sns"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "id": "AtKKE6XplpCJ",
        "outputId": "ea09ffd3-c251-43ca-8714-4b48062a5607",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        }
      },
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "'0.25.2'"
            ],
            "application/vnd.google.colaboratory.intrinsic+json": {
              "type": "string"
            }
          },
          "metadata": {},
          "execution_count": 3
        }
      ],
      "source": [
        "gym.__version__"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "T9E6hwVdlpCJ"
      },
      "outputs": [],
      "source": [
        "sns.set_context(\"talk\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SWKMPbrRlpCK"
      },
      "outputs": [],
      "source": [
        "#matplotlib.rcParams['figure.figsize'] = (20.0, 10.0)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XSkOxVMflpCK"
      },
      "source": [
        "## Backward Induction"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8lt43DF2lpCK"
      },
      "source": [
        "*Backward Induction* is a basic *Dynamic Programming* method **[BELLMAN57]**.\n",
        "Like other Dynamic Programming algorithms, it uses the *Bellman's\n",
        "Principle of Optimality* **[BELLMAN57]** for accelerating computation (compared\n",
        "to an exhaustive search). It can be applied to problems that exhibit a compatible structure, i.e., a problem that has *overlapping subproblems* or a problem having an *optimal substructure* **[BELLMAN57]**.\n",
        "Actually, this acceleration is obtained by breaking problems down into simpler subproblems in such a manner\n",
        "that redundant computations are avoided by storing results.\n",
        "When applicable, the method takes far less time than naïve methods that don't take advantage of the subproblem overlap (like depth-first search)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iX92QfmslpCL"
      },
      "source": [
        "*Backward Induction* computes non-stationary policies: a new policy is computed for each time step.\n",
        "Thus the number of time steps used to solve the problem is set in advance.\n",
        "*Backward Induction* algorithms solve Sequential Decision Making problems defined with\n",
        "discrete actions and state spaces."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EsQzQd31lpCL"
      },
      "source": [
        "The *value* (or *utility*) $\\U^*$ for each state $\\state$ at the latest time step $T$ is\n",
        "$$\n",
        "\\U^*_T(\\state) = r(\\state) \\label{eq:backward-induction-last-value} \\tag{1}\n",
        "$$\n",
        "where $r$ is the immediate reward function.\n",
        "\n",
        "The best expected value $\\U^*$ for each state $\\state$ at the $t^{\\text{th}}$ time step is\n",
        "$$\n",
        "\\U^*_t(\\state) = r(\\state) + \\max_{\\action \\in \\actionset} \\left[ \\sum_{\\state' \\in \\stateset} P(\\state' | \\state, \\action) \\U^*_{t+1}(\\state') \\right]  \\label{eq:backward-induction-tth-value} \\tag{2}\n",
        "$$\n",
        "and the $t^{\\text{th}}$ optimal action (or decision) $d^*_t(\\state)$ among the set of\n",
        "possible actions $\\actionset$ is\n",
        "$$\n",
        "d^*_t(\\state) = \\arg\\max_{\\action \\in \\actionset} \\left[ \\sum_{\\state' \\in \\stateset} P(\\state' | \\state, \\action) \\U^*_{t+1}(\\state') \\right]  \\label{eq:backward-induction-tth-decision} \\tag{3}\n",
        "$$\n",
        "where $T$ is the transition function.\n",
        "\n",
        "The main idea is to compute the expected value of each state\n",
        "(Eq. \\ref{eq:backward-induction-tth-value}) and then to use it to select the\n",
        "best action for any given state (Eq. \\ref{eq:backward-induction-tth-decision}).\n",
        "\n",
        "Eq. $\\ref{eq:backward-induction-tth-value}$ cannot be solved analytically because\n",
        "the system of equations to compute $V$ contains non-linear terms (due to the\n",
        "\"max\" operator).\n",
        "As an alternative, Eq. $\\ref{eq:backward-induction-tth-value}$\n",
        "is usually computed using Dynamic Programming method, as described in algorithm 1."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CQ9XvF0GlpCM"
      },
      "source": [
        "___\n",
        "### Algorithm 1: Backward Induction\n",
        "\n",
        "**Input**:<br>\n",
        "$\\quad$ $mdp = \\langle \\stateset, \\actionset, T, r \\rangle$, a Markov Decision Process <br>\n",
        "$\\quad$ $T$, the resolution horizon (i.e. the number of time steps) <br>\n",
        "**Local variables**: <br>\n",
        "$\\quad$ $\\U^*_t ~~ \\forall t \\in \\{1, ..., T\\}$, value array (expected global reward following the optimal policy for states in $\\stateset$) <br>\n",
        "<br>\n",
        "$\\U^*_T[\\state] \\leftarrow r(\\state) ~~ \\forall \\state \\in \\stateset$ <br>\n",
        "**for all** $t \\in \\{T-1, T-2, ..., 1\\}$ **do** <br>\n",
        "$\\quad$ **for all** $\\state \\in \\stateset$ **do** <br>\n",
        "$\\quad\\quad$ **if** $\\state$ is a final state **then** <br>\n",
        "$\\quad\\quad\\quad$ $\\displaystyle \\U^*_t[\\state] \\leftarrow r(\\state)$ <br>\n",
        "$\\quad\\quad$ **else** <br>\n",
        "$\\quad\\quad\\quad$ $\\displaystyle \\U^*_t[\\state] \\leftarrow r(\\state) + \\max_{\\action \\in \\actionset} \\left[ \\sum_{\\state' \\in \\stateset} P(\\state' | \\state, \\action) \\U^*_{t+1}[\\state'] \\right]$ <br>\n",
        "$\\quad\\quad$ **end if** <br>\n",
        "$\\quad$ **end for** <br>\n",
        "**end for** <br>\n",
        "<br>\n",
        "**return** $\\U^*_t ~~ \\forall t \\in \\{1, ..., T\\}$\n",
        "___"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "68_coHVQlpCM"
      },
      "source": [
        "## Value Iteration\n",
        "\n",
        "*Value Iteration* **[BELLMAN57]** is one of the most famous Dynamic Programming algorithm to compute the optimal policy for a Markov Decision Process (MDP).\n",
        "Similarly to Backward Induction, the\n",
        "main idea implemented by Value Iteration is to compute the best expected value of each state and then to use\n",
        "these values to select the best action from any given state.\n",
        "\n",
        "The main difference with the Backward Induction algorithm is that Value Iteration\n",
        "is used to compute stationary policies.\n",
        "Indeed, the same resulting policy is used for each time step and thus there is\n",
        "no assumption about the number of time steps to consider for the solution.\n",
        "\n",
        "The expected value $\\U^{\\pi}$ for each state $\\state$ when the agent follows a\n",
        "given (stationary) policy $\\pi$ is \n",
        "$$\n",
        "\\U^{\\pi}(\\state) = E \\left[ \\sum^{\\infty}_{t=0} \\discount^t r(\\state_t) | \\pi, \\state_0 = \\state \\right] \\label{eq:vi-value-of-s-for-pi} \\tag{4}\n",
        "$$\n",
        "\n",
        "The optimal (stationary) policy $\\pi^*$ is defined using the best expected value $\\U^{\\pi^*}$ and using the principle of *Maximum Expected Utility* as follows\n",
        "$$\n",
        "\\pi^*(\\state) = \\arg\\max_{\\action \\in \\actionset} \\left[ \\sum_{\\state' \\in \\stateset} P(\\state' | \\state, \\action) \\U^{\\pi^*}(\\state') \\right]  \\label{eq:vi-optimal-policy} \\tag{5}\n",
        "$$\n",
        "\n",
        "Eq. \\ref{eq:vi-bellman-eq} is commonly called *Bellman equation*; it gives the best\n",
        "value we can expect for any given\n",
        "state (assuming the optimal policy $\\pi^*$ is\n",
        "followed). There are $|\\stateset|$ Bellman equations, one for each state.\n",
        "As for the Backward Induction method,\n",
        "this system of equations cannot be solved analytically because\n",
        "Bellman equations contain non-linear terms (due to the\n",
        "\"max\" operator).  As an alternative, Eq. \\ref{eq:vi-bellman-eq}\n",
        "can be computed iteratively using Value Iteration, a Dynamic Programming method\n",
        "described in Algorithm 2.\n",
        "\n",
        "\\begin{equation}\n",
        "    \\U(\\state) := \\U^{\\pi^*}(\\state) = \\left\\{\n",
        "    \\begin{array}{l l}\n",
        "        r(\\state)                                                                                                                                 & \\quad \\text{if $\\state$ is a final state} \\\\\n",
        "        \\displaystyle r(\\state) + \\discount \\max_{\\action \\in \\actionset} \\left[ \\sum_{\\state' \\in \\stateset} P(\\state' | \\state, \\action) \\U(\\state') \\right]    & \\quad \\text{otherwise}\\\\\n",
        "    \\end{array} \\right.\n",
        "    \\label{eq:vi-bellman-eq} \\tag{6}\n",
        "\\end{equation}\n",
        "\n",
        "Equation \\ref{eq:vi-bellman-update} -- called *Bellman update* -- is\n",
        "used in the iterative method described in Algorithm 2, to update $\\U$ at each iteration.\n",
        "\n",
        "\\begin{equation}\n",
        "    \\U_{i+1}(\\state) \\leftarrow \\left\\{\n",
        "    \\begin{array}{l l}\n",
        "        r(\\state)                                                                                                                                   & \\quad \\text{if $\\state$ is a final state} \\\\\n",
        "        \\displaystyle r(\\state) + \\discount \\max_{\\action \\in \\actionset} \\left[ \\sum_{\\state' \\in \\stateset} P(\\state' | \\state, \\action) \\U_i(\\state') \\right]    & \\quad \\text{otherwise}\\\\\n",
        "    \\end{array} \\right.\n",
        "    \\label{eq:vi-bellman-update} \\tag{7}\n",
        "\\end{equation}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bF7CDGNulpCM"
      },
      "source": [
        "___\n",
        "### Algorithm 2: Value Iteration\n",
        "\n",
        "**Input**:<br>\n",
        "$\\quad$ $mdp = \\langle \\stateset, \\actionset, P, r \\rangle$, a Markov Decision Process <br>\n",
        "$\\quad$ $\\discount$, the discount factor <br>\n",
        "$\\quad$ $\\epsilon$, the stopping criteria: the algorithm is stopped if the largest update in an iteration is lower than $\\epsilon$ <br>\n",
        "**Local variables**: <br>\n",
        "$\\quad$ $\\U, \\U'$, old and new estimated value array (estimation of the expected global reward following the optimal policy for all states in $\\stateset$), initially zero <br>\n",
        "$\\quad$ $\\delta$, the largest change in the value array in an iteration <br>\n",
        "<br>\n",
        "**repeat** <br>\n",
        "$\\quad$ $\\U \\leftarrow \\U'$ <br>\n",
        "$\\quad$ $\\delta \\leftarrow 0$ <br>\n",
        "$\\quad$ **for all** $\\state \\in \\stateset$ **do** <br>\n",
        "$\\quad\\quad$ **if** $\\state$ is a final state **then** <br>\n",
        "$\\quad\\quad\\quad$ $\\displaystyle \\U'[\\state] \\leftarrow r[\\state]$ <br>\n",
        "$\\quad\\quad$ **else** <br>\n",
        "$\\quad\\quad\\quad$ $\\displaystyle \\U'[\\state] \\leftarrow r[\\state] + \\discount \\max_{\\action \\in \\actionset} \\left[ \\sum_{\\state' \\in \\stateset} P(\\state' | \\state, \\action) \\U[\\state'] \\right]$ <br>\n",
        "$\\quad\\quad$ **end if** <br>\n",
        "$\\quad\\quad$ **if** $|\\U'[\\state] - \\U[\\state]| > \\delta$ **then** <br>\n",
        "$\\quad\\quad\\quad$ $\\delta \\leftarrow |\\U'[\\state] - \\U[\\state]|$ <br>\n",
        "$\\quad\\quad$ **end if** <br>\n",
        "$\\quad$ **end for** <br>\n",
        "**until** $\\delta < \\epsilon$ <br>\n",
        "<br>\n",
        "**return** $\\U$\n",
        "___"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vWx6kwcTlpCM"
      },
      "source": [
        "### Convergence\n",
        "\n",
        "The convergence of Value Iteration has been proved, but this convergence is asymptotic **[BELLMAN57]**.\n",
        "However, each iteration is easy and fast to compute."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CdnrpvCKlpCM"
      },
      "source": [
        "## Hands on Gym and the FrozenLake toy problem"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iiS-lhYolpCN"
      },
      "source": [
        "For the purpose of focusing on the algorithms, we will use standard environments provided by OpenAI Gym framework.\n",
        "OpenAI Gym provides controllable environments (https://www.gymlibrary.dev/environments/classic_control/) for research in Reinforcement Learning.\n",
        "We will use a simple toy problem to illustrate Dynamic Programming algorithms properties."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WbHsilWalpCN"
      },
      "source": [
        "**Task:** read https://www.gymlibrary.dev/content/basic_usage/ to discover Gym and get familiar with its main concepts.\n",
        "\n",
        "In this lab, we will try to solve the FrozenLake-v1 environment (https://www.gymlibrary.dev/environments/toy_text/frozen_lake/).\n",
        "Additional information is available [here](https://github.com/openai/gym/wiki/FrozenLake-v0) and [here](https://github.com/openai/gym/wiki/FrozenLake-v0)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3yY02W3llpCN"
      },
      "source": [
        "**Notice**: this environment is *fully observable*, thus here the terms (environment) *state* and (agent) *observation* are equivalent.\n",
        "This is not always the case for example in poker, the agent doesn't know the opponent's cards."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uSBLxsWYlpCN"
      },
      "source": [
        "### Get the FrozenLake state space and action space"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "grwBedHrlpCN"
      },
      "outputs": [],
      "source": [
        "try:\n",
        "    env = gym.make('FrozenLake-v1')\n",
        "except:\n",
        "    # Fallback for older Gym versions\n",
        "    env = gym.make('FrozenLake-v0')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cm-dTXUjlpCO"
      },
      "source": [
        "Possible states in FrozenLake are:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RhkkbfBzlpCO"
      },
      "outputs": [],
      "source": [
        "states = list(range(env.observation_space.n))\n",
        "states"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BwKK7p-RlpCO"
      },
      "source": [
        "Possible actions are:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0pIZaEgblpCO"
      },
      "outputs": [],
      "source": [
        "actions = list(range(env.action_space.n))\n",
        "actions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6ebR1H_AlpCP"
      },
      "source": [
        "The following dictionary may be used to understand actions:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kge2gJRylpCP"
      },
      "outputs": [],
      "source": [
        "action_labels = {\n",
        "    0: \"Move Left\",\n",
        "    1: \"Move Down\",\n",
        "    2: \"Move Right\",\n",
        "    3: \"Move Up\"\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dBvQhOYklpCP"
      },
      "source": [
        "### Display functions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-06grXkglpCP"
      },
      "source": [
        "The next cells contain functions that can be used to display states, transitions and policies with the FrozenLake environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kBVLeKSolpCQ"
      },
      "outputs": [],
      "source": [
        "def states_display(state_seq, title=None, figsize=(5,5), annot=True, fmt=\"0.1f\", linewidths=.5, square=True, cbar=False, cmap=\"Reds\"):\n",
        "    size = int(math.sqrt(len(state_seq)))\n",
        "    state_array = np.array(state_seq)\n",
        "    state_array = state_array.reshape(size, size)\n",
        "    \n",
        "    fig, ax = plt.subplots(figsize=figsize)         # Sample figsize in inches\n",
        "    sns.heatmap(state_array, annot=annot, fmt=fmt, linewidths=linewidths, square=square, cbar=cbar, cmap=cmap)\n",
        "    plt.title(title)\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vwP2MNrelpCQ"
      },
      "outputs": [],
      "source": [
        "def transition_display(state, action):\n",
        "    states_display(transition_array[state,action], title=\"Transition probabilities for action {} ({}) in state {}\".format(action, action_labels[action], state))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vbXXPQK_lpCQ"
      },
      "outputs": [],
      "source": [
        "def display_policy(policy):\n",
        "    actions_src = [\"{}={}\".format(action, action_labels[action].replace(\"Move \", \"\")) for action in actions]\n",
        "    title = \"Policy (\" + \", \".join(actions_src) + \")\"\n",
        "    states_display(policy, title=title, fmt=\"d\", cbar=False, cmap=\"Reds\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ggdku5gdlpCR"
      },
      "source": [
        "### Make the `is_final_array`, `reward_array` and `transition_array`"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DRxuonMrlpCR"
      },
      "source": [
        "To implement Dynamic Programming algorithms, we need the transition probability (or transition function) and the reward function, both defined in `env.P`.\n",
        "\n",
        "`env.P[S][A]` gives the list of reachable states from state S executing action A.\n",
        "\n",
        "These reachable states are coded in a tuple defined like this: `(probability, next state, reward, is_final_state)`.\n",
        "\n",
        "You will not need to use `env.P` to solve exercises.\n",
        "In the following cell, `is_final_array`, `reward_array` and `transition_array` are defined for convenience."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L5Hm3apGlpCR"
      },
      "outputs": [],
      "source": [
        "is_final_array = np.full(shape=len(states), fill_value=np.nan, dtype=bool)\n",
        "reward_array = np.full(shape=len(states), fill_value=np.NINF)                # np.NINF = negative infinity\n",
        "transition_array = np.zeros(shape=(len(states), len(actions), len(states)))\n",
        "\n",
        "for state in states:\n",
        "    for action in actions:\n",
        "        for next_state_tuple in env.P[state][action]:              # env.P[state][action] contains the next states list (a list of tuples)\n",
        "            transition_probability, next_state, next_state_reward, next_state_is_final = next_state_tuple\n",
        "\n",
        "            is_final_array[next_state] = next_state_is_final\n",
        "            reward_array[next_state] = max(reward_array[next_state], next_state_reward)   # workaround: when we already are in state 15, reward is 0 if we stay in state 15 (in practice this never append as the simulation stop when we arrive in state 15 as any other terminal state)\n",
        "            transition_array[state, action, next_state] += transition_probability"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1_mUUTf9lpCR"
      },
      "outputs": [],
      "source": [
        "def reachable_states(state, action):\n",
        "    return np.nonzero(transition_array[state, action])[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D0c2agThlpCS"
      },
      "source": [
        "The following plot shows the state corresponding to square of the FrozenLake grid."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f1Hf9MfelpCS"
      },
      "outputs": [],
      "source": [
        "states_display(states, fmt=\"d\", title=\"States ID\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2k7evyl9lpCS"
      },
      "source": [
        "The following plot shows the reward obtained in each square of the FrozenLake grid."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xZktjYzIlpCT"
      },
      "outputs": [],
      "source": [
        "states_display(reward_array, title=\"Rewards\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ZkkB08QlpCT"
      },
      "source": [
        "The following plot shows whether a square is a final state or not (i.e. whether it ends the simulation or not)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FXVIrhvSlpCT"
      },
      "outputs": [],
      "source": [
        "states_display(is_final_array, fmt=\"d\", title=\"Final states\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WhJd11uelpCU"
      },
      "source": [
        "The following cells show how to display transitions with the provided `transition_display` function. Figures displayed in squares are the probability to reach these squares from the given (`state`, `action`) pair. Colored squares are the states that may be reached from this pair (a non-zero probability)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "daXD8GD_lpCV"
      },
      "outputs": [],
      "source": [
        "transition_display(state=0, action=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b2WuRRn-lpCV"
      },
      "outputs": [],
      "source": [
        "transition_display(state=6, action=0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Iz23V6vlpCV"
      },
      "outputs": [],
      "source": [
        "transition_display(state=6, action=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jaI9Wyy1lpCV"
      },
      "source": [
        "## Exercise 1: Implement the Value Iteration algorithm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CVrb3F_elpCW"
      },
      "source": [
        "To solve the FrozenLake-v1 problem with Dynamic Programming, we will first use the Value Iteration algorithm described in Algorithm 2.\n",
        "\n",
        "Notice that the FrozenLake-v1 environment is non-deterministic.\n",
        "To implement Value Iteration, you will need the transition probability (or the transition function) defined in `transition_array`.\n",
        "- Use `reachable_states(S, A)` to get the list of reachable states from state `S` executing action `A`.\n",
        "- Use `transition_array[S, A]` to get the probability of reaching each state from state `S` executing action `A`.\n",
        "- Use `transition_array[S, A, S']` to get the probability of reaching state `S'` from state `S` executing action `A`.\n",
        "\n",
        "You will also need the previously defined `is_final_array` matrix.\n",
        "- Use `is_final_array[S]` to know whether `S` is a final state (`True`) or not (`False`).\n",
        "\n",
        "Finally, you will need the previously defined `reward_array` matrix.\n",
        "- Use `reward_array[S]` to get the reward obtained by the agent each time it reaches state `S`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2CGz1M4_lpCW"
      },
      "source": [
        "In the following cell, we define `expected_value` and `expected_values` functions for convenience.\n",
        "The first one returns the expected reward\n",
        "$$\\sum P(\\state' | \\state, \\action) \\U(\\state')$$\n",
        "for a given pair $(\\state, \\action)$ and a given V-table (value function) $\\U$.\n",
        "The second one computes the expected reward for all the actions in $\\state$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gE6R2sVTlpCW"
      },
      "outputs": [],
      "source": [
        "def expected_value(state, action, v_array):\n",
        "    return (transition_array[state, action] * v_array).sum() # compute sum(P(s'|a,s).V(s'))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DKnMwcfXlpCW"
      },
      "outputs": [],
      "source": [
        "def expected_values(state, v_array):\n",
        "    return (transition_array[state] * v_array).sum(axis=1)   # compute sum(P(s'|a,s).V(s')) for all the actions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n1aJjaSJlpCW"
      },
      "source": [
        "### Question 1: Implement the Value Iteration algorithm (compute the *value function* `v_array`)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0t0ojHC7lpCX"
      },
      "source": [
        "**Note**: here we use the `state_display` function to show the evolution of the value function `v_array` over iterations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QlZ-EQQelpCX"
      },
      "outputs": [],
      "source": [
        "stop = False\n",
        "\n",
        "value_function_history = []\n",
        "delta_history = []\n",
        "\n",
        "def value_iteration(gamma=0.95, epsilon=0.001, display=False):\n",
        "    v_array = np.zeros(len(states))   # Initial value function\n",
        "    stop = False\n",
        "\n",
        "    while not stop:\n",
        "        if display:\n",
        "            states_display(v_array, title=\"Value function\", cbar=True, cmap=\"Reds\")\n",
        "        else:\n",
        "            print('.', end=\"\")\n",
        "        value_function_history.append(v_array)\n",
        "        \n",
        "        delta = 0.\n",
        "        \n",
        "        # TODO...\n",
        "\n",
        "        delta_history.append(delta)\n",
        "        \n",
        "        if delta < epsilon:\n",
        "            stop = True\n",
        "    \n",
        "    return v_array\n",
        "        \n",
        "v_array = value_iteration(display=True)\n",
        "states_display(v_array, title=\"Value function\", cbar=True, cmap=\"Reds\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8BPvxEsulpCX"
      },
      "source": [
        "### Display the evolution of the value function over iterations"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mRhlEsW9lpCX"
      },
      "outputs": [],
      "source": [
        "df_v_hist = pd.DataFrame(value_function_history)\n",
        "df_v_hist"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BjXH6LjYlpCX"
      },
      "source": [
        "Evolution of `v_array` (the estimated value of each state) over iterations (one curve per state):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9_Up36EVlpCY"
      },
      "outputs": [],
      "source": [
        "df_v_hist.plot(figsize=(14,8))\n",
        "plt.title(\"V(s) w.r.t iteration\")\n",
        "plt.ylabel(\"V(s)\")\n",
        "plt.xlabel(\"iteration\")\n",
        "plt.legend(loc='upper right');"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YucQuSkBlpCY"
      },
      "source": [
        "Evolution of `delta` over iterations:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I4W6q-A8lpCY"
      },
      "outputs": [],
      "source": [
        "plt.plot(delta_history)\n",
        "plt.yscale(\"log\")\n",
        "plt.title(r\"$\\max~\\delta$ w.r.t iteration\")\n",
        "plt.ylabel(r\"$\\max~\\delta$\")\n",
        "plt.xlabel(\"iteration\");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vGOyGth7lpCY"
      },
      "source": [
        "### Question 2: Define the greedy policy (Maximum Expected Utility)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v4s28oIKlpCY"
      },
      "outputs": [],
      "source": [
        "def greedy_policy(state, v_array):\n",
        "\n",
        "    # TODO...\n",
        "\n",
        "    return policy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PU3LFf1wlpCZ"
      },
      "source": [
        "### Display the opimized policy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BSDm9TrXlpCZ"
      },
      "source": [
        "Applying the `greedy_policy` on each state gives us the policy matrix:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "os_gjN1alpCZ"
      },
      "outputs": [],
      "source": [
        "policy = [greedy_policy(state, v_array) for state in states]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f3cQlHjflpCZ"
      },
      "source": [
        "The following cell gives us a graphical representation of the optimal policy we have computed. The figure in each square is the optimal action to execute in the corresponding state (0 = \"move left\", 1 = \"move down\", 2 = \"move right\", 3 = \"move up\")."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fS8ouEMplpCa"
      },
      "outputs": [],
      "source": [
        "display_policy(policy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tXScij3OlpCa"
      },
      "source": [
        "### Evaluate Value Iteration with Gym (single trial)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HKHcJIwflpCa"
      },
      "source": [
        "So far, we have computed the value function `v_array` for one *episode*.\n",
        "The environment is stochastic, thus if we apply the computed policy several times on the environment, we may have different results.\n",
        "To measure the performance of our value function `v_array`, we should assess it several times and count the number of successful trials.\n",
        "Gym considers an agent to successfully solve the FrozenLake problem if it reaches 76% success rate over the last 100 trials (or \"episodes\")."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rJLcwQqulpCa"
      },
      "outputs": [],
      "source": [
        "env._max_episode_steps = 1000"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h1GJ_k45lpCa"
      },
      "outputs": [],
      "source": [
        "reward_list = []\n",
        "\n",
        "NUM_EPISODES = 1000\n",
        "\n",
        "for episode_index in range(NUM_EPISODES):\n",
        "    state, info = env.reset()\n",
        "    done = False\n",
        "    #t = 0\n",
        "\n",
        "    while not done:\n",
        "        action = greedy_policy(state, v_array)\n",
        "        state, reward, done, truncated, info = env.step(action)\n",
        "        #t += 1\n",
        "\n",
        "    reward_list.append(reward)\n",
        "    #print(\"Episode finished after {} timesteps ; reward = {}\".format(t, reward))\n",
        "\n",
        "print(sum(reward_list) / NUM_EPISODES)            \n",
        "\n",
        "env.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nh79qjKXlpCc"
      },
      "source": [
        "### Question 3: What do you think the discount factor $\\gamma$ is for?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XASafXWclpCc"
      },
      "source": [
        "TODO..."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zdWVTyFglpCc"
      },
      "source": [
        "### Evaluate Value Iteration for different value of $\\gamma$ with confidence interval (bootstrap)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ICcTr9VrlpCc"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "NUM_EPISODES = 1000\n",
        "\n",
        "reward_list = []\n",
        "\n",
        "for gamma in (0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.84, 0.9, 0.95, 0.99):\n",
        "    v_array = value_iteration(gamma=gamma)\n",
        "    \n",
        "    for episode_index in range(NUM_EPISODES):\n",
        "        state, info = env.reset()\n",
        "        done = False\n",
        "\n",
        "        while not done:\n",
        "            action = greedy_policy(state, v_array)\n",
        "            state, reward, done, truncated, info = env.step(action)\n",
        "\n",
        "        reward_list.append({\"gamma\": gamma, \"reward\": reward})\n",
        "\n",
        "env.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wwOZZNlqlpCc"
      },
      "outputs": [],
      "source": [
        "df = pd.DataFrame(reward_list)\n",
        "df.tail()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b7ifsFaflpCd"
      },
      "outputs": [],
      "source": [
        "# Plot mean reward (with its 95% confidence interval)\n",
        "\n",
        "sns.relplot(x=\"gamma\", y=\"reward\", kind=\"line\", data=df, height=6, aspect=1.5)\n",
        "plt.axhline(0.76, color=\"red\", linestyle=\":\", label=\"76% success threshold\");   # 76% success threshold\n",
        "plt.legend();"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BzXfklRUlpCd"
      },
      "source": [
        "### Display the Value Iteration optimal policy with respect to $\\gamma$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3vCZRqXXlpCd"
      },
      "outputs": [],
      "source": [
        "for gamma in (0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.84, 0.9, 0.95, 0.99):\n",
        "    print()\n",
        "    print(\"=\" * 10, \"GAMMA = \", gamma, \"=\" * 10)\n",
        "    print()\n",
        "    \n",
        "    v_array = value_iteration(gamma=gamma)\n",
        "    \n",
        "    print()\n",
        "    print()\n",
        "    \n",
        "    policy = [greedy_policy(state, v_array) for state in states]\n",
        "    display_policy(policy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P8f-pjUOlpCe"
      },
      "source": [
        "## Exercise 2: Implement the Policy Iteration algorithm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eol5TEsrlpCe"
      },
      "source": [
        "## Policy Iteration\n",
        "\n",
        "*Policy Iteration* **[HOWARD60]** is another popular Dynamic Programming algorithm to\n",
        "compute MDP's optimal policy. In practice, it is often faster than Value Iteration.\n",
        "\n",
        "The Policy Iteration algorithm alternates the following two steps, starting with an initial policy $\\pi_0$:\n",
        "1. **Policy Evaluation**: given a policy $\\pi_i$, compute $\\U^{\\pi_i}(\\state) ~ \\forall \\state \\in \\stateset$, the expected value of each state when $\\pi_i$ is followed.\n",
        "2. **Policy Improvement**: compute a new policy $\\pi_{i+1}$, using one-step look-ahead based on $\\U^{\\pi_i}$ and using the principle of *Maximum Expected Utility* as follows\n",
        "$$\n",
        "\\policy_{\\iterationindex+1}(\\state) = \\arg \\max_{\\action \\in \\actionspace} \\sum_{\\nextstate  \\in \\statespace} \\transitionfuncinst \\left( \\rewardfuncinst + \\discount \\valuefunc^{\\policy_{\\iterationindex}}(\\nextstate) \\right)\n",
        "$$\n",
        "\n",
        "In the following exercise we will assume that the reward only depends on the state: $r(\\state) \\equiv \\mathcal{R}(\\state, \\action, \\state')$.\n",
        "Thus the *Policy Improvement* can be rewritten as follow:\n",
        "\n",
        "$$\n",
        "\\pi_{i+1}(\\state) = \\arg\\max_{\\action \\in \\actionset} \\sum_{\\state' \\in \\stateset} P(\\state' | \\state, \\action) \\U^{\\pi_i}(\\state')  \\label{eq:pi-policy-improvement} \\tag{8}\n",
        "$$\n",
        "\n",
        "($\\rewardfunc$ and $\\discount$ can disappear as they have no influence on the $\\arg\\max_{\\action}$ result).\n",
        "\n",
        "Algorithm 3 describes the two-step procedure.\n",
        "The algorithm terminates when the *Policy Improvement* step yields no change in the utilities."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nv2ioiC_lpCe"
      },
      "source": [
        "___\n",
        "### Algorithm 3: Policy Iteration\n",
        "\n",
        "**Input**:<br>\n",
        "$\\quad$ $MDP = \\langle \\stateset, \\actionset, T, r \\rangle$, a Markov Decision Process<br>\n",
        "**Local variables**: <br>\n",
        "$\\quad$ $\\U$, vector of utilities for states in $\\stateset$, initially zero <br>\n",
        "$\\quad$ $\\pi$, a policy vector indexed by state, initially random <br>\n",
        "<br>\n",
        "**repeat** <br>\n",
        "$\\quad$ $\\valuefunc \\leftarrow \\text{POLICY-EVALUATION}(\\policy, \\valuefunc, \\mdp)$ <br>\n",
        "$\\quad$ unchanged $\\leftarrow$ true <br>\n",
        "$\\quad$ **for all** $\\mbox{state} ~ \\state \\in \\stateset$ **do** <br>\n",
        "$\\quad\\quad$ **if** $\\displaystyle \\max_{\\action \\in \\actionset} \\left[ \\sum_{\\state' \\in \\stateset} \\transitionfuncinst \\U[\\state'] \\right] > \\sum_{\\state' \\in \\stateset} \\transitionfunc(\\nextstate | \\state, \\policy_{\\iterationindex}(\\state)) \\U[\\state']$ **then** <br>\n",
        "$\\quad\\quad\\quad$ $\\displaystyle \\pi[\\state] \\leftarrow \\arg\\max_{\\action \\in \\actionset} \\left[ \\sum_{\\state' \\in \\stateset} \\transitionfuncinst \\U[\\state'] \\right]$ <br>\n",
        "$\\quad\\quad\\quad$ unchanged $\\leftarrow$ false <br>\n",
        "$\\quad\\quad$ **end if** <br>\n",
        "$\\quad$ **end for** <br>\n",
        "**until** unchanged <br>\n",
        "<br>\n",
        "**return** $\\pi$\n",
        "___"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dIxdZB3xlpCe"
      },
      "source": [
        "Solving the POLICY-EVALUATION routine is much simpler than solving the standard\n",
        "Bellman equations (which is what Value Iteration does).  Indeed, the action in each\n",
        "state is fixed by the policy, thus the \"max\" operator disappears and Bellman\n",
        "equations become linear.\n",
        "As a result, $\\U^{\\pi_i}$ can be computed by solving the linear system of these\n",
        "*simplified Bellman equations* (Eq. \\ref{eq:pi-simplified-bellman}) for\n",
        "each state.\n",
        "\n",
        "\\begin{equation}\n",
        "    \\U^{\\pi_i}(\\state) = \\left\\{\n",
        "    \\begin{array}{l l}\n",
        "        r(\\state)               & \\quad \\text{if $\\state$ is a final state} \\\\\n",
        "        \\displaystyle r(\\state) + \\discount \\sum_{\\nextstate \\in \\statespace} \\transitionfunc(\\nextstate | \\state, \\policy_{\\iterationindex}(\\state)) ~ \\valuefunc^{\\policy_{\\iterationindex}}(\\nextstate)     & \\quad \\text{otherwise}\\\\\n",
        "    \\end{array} \\right.\n",
        "    \\label{eq:pi-simplified-bellman} \\tag{9}\n",
        "\\end{equation}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WeCc0UgflpCf"
      },
      "source": [
        "### Convergence\n",
        "\n",
        "As the number of states and policies is finite, and as the policy is improved\n",
        "at each iteration, Policy Iteration converges in a finite number of iterations (often\n",
        "small in practice).\n",
        "However, within each iteration, solving the\n",
        "POLICY-EVALUATION routine may cost a lot (its complexity is $O(|\\stateset|^3)$)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CQyWu1aUlpCf"
      },
      "source": [
        "An approach alternative to Value Iteration (Exercise 1) is Policy Iteration (described in Algorithm 3).\n",
        "\n",
        "**Task:** implement Iterative Policy Iteration (for the same environment). Note that as part of this task you should also implement iterative policy evaluation. Compare the policies obtained by both approaches (they should be the same)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9ak80di-lpCf"
      },
      "source": [
        "### Question 1: Define the (exact) Policy Evaluation function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GOyC6YL8lpCg"
      },
      "outputs": [],
      "source": [
        "def policy_evaluation(policy, gamma):\n",
        "\n",
        "    # TODO...\n",
        "\n",
        "    return x"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zfEgZmwrlpCg"
      },
      "source": [
        "### Question 2: Define the Policy Improvement function"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9y2xBB-qlpCg"
      },
      "outputs": [],
      "source": [
        "def policy_iteration(gamma, initial_policy=None, policy_evaluation_function=policy_evaluation):\n",
        "\n",
        "    # TODO...\n",
        "\n",
        "    return policy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v3yuyCe2lpCg"
      },
      "outputs": [],
      "source": [
        "gamma = 0.99\n",
        "\n",
        "policy = policy_iteration(gamma=gamma, policy_evaluation_function=policy_evaluation)\n",
        "\n",
        "display_policy(policy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "It1YHVY-lpCh"
      },
      "source": [
        "### Evaluate Policy Iteration with Gym"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lyy9fx7ulpCi"
      },
      "outputs": [],
      "source": [
        "env._max_episode_steps = 1000"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rR2LsgHOlpCj"
      },
      "outputs": [],
      "source": [
        "reward_list = []\n",
        "\n",
        "NUM_EPISODES = 1000\n",
        "\n",
        "for episode_index in range(NUM_EPISODES):\n",
        "    state, info = env.reset()\n",
        "    done = False\n",
        "    #t = 0\n",
        "\n",
        "    while not done:\n",
        "        action = policy[state]\n",
        "        state, reward, done, truncated, info = env.step(action)\n",
        "        #t += 1\n",
        "\n",
        "    reward_list.append(reward)\n",
        "    #print(\"Episode finished after {} timesteps ; reward = {}\".format(t, reward))\n",
        "\n",
        "print(sum(reward_list) / NUM_EPISODES)            \n",
        "\n",
        "env.close()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Jm9oKgvIlpCj"
      },
      "source": [
        "### Evaluate Policy Iteration for different $\\gamma$ with confidence interval (bootstrap)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5emHTu1NlpCj"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "NUM_EPISODES = 1000\n",
        "\n",
        "reward_list = []\n",
        "\n",
        "for gamma in (0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.84, 0.9, 0.99):\n",
        "    print(\"gamma:\", gamma)\n",
        "    policy = policy_iteration(gamma=gamma)\n",
        "    \n",
        "    for episode_index in range(NUM_EPISODES):\n",
        "        state, info = env.reset()\n",
        "        done = False\n",
        "\n",
        "        while not done:\n",
        "            action = policy[state]\n",
        "            state, reward, done, truncated, info = env.step(action)\n",
        "\n",
        "        reward_list.append({\"gamma\": gamma, \"reward\": reward})\n",
        "\n",
        "env.close()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NNQXultylpCk"
      },
      "outputs": [],
      "source": [
        "df = pd.DataFrame(reward_list)\n",
        "df.tail()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "p36PtIHnlpCk"
      },
      "source": [
        "### Plot mean reward (with its 95% confidence interval)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wdsbOymelpCk"
      },
      "outputs": [],
      "source": [
        "sns.relplot(x=\"gamma\", y=\"reward\", kind=\"line\", data=df, height=6, aspect=1.5)\n",
        "plt.axhline(0.76, color=\"red\", linestyle=\":\", label=\"76% success threshold\");   # 76% success threshold\n",
        "plt.legend();"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NOeDV-cLlpCk"
      },
      "source": [
        "## References"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x4Hp-WDylpCk"
      },
      "source": [
        "**[BELLMAN57]** Richard Ernest Bellman. *Dynamic Programming*. Princeton University Press, Princeton,\n",
        "New Jersey, USA, 1957."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "47zViAiOlpCk"
      },
      "source": [
        "**[HOWARD60]** R.A. Howard. Dynamic Programming and Markov Processes. MIT Press, Cambridge,\n",
        "Massachusetts, 1960."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "nbgrader": {
          "grade": false,
          "grade_id": "cell-b82aa116ac616c8d",
          "locked": true,
          "schema_version": 3,
          "solution": false,
          "task": false
        },
        "id": "QfD0wT2VlpCk"
      },
      "source": [
        "## Going further\n",
        "\n",
        "In this lab we have introduced Reinforcement Learning in a very specific case where the *agent* (the algorithm) has a perfect knowledge of the environment (transition and reward functions).\n",
        "\n",
        "This is convenient to introduce basic concepts but we cannot expect this assumption to be true in many practical problems.\n",
        "A lot of sophisticated algorithms have been developed recently and most of them have been implemented in [Stable Baselines 3](https://stable-baselines3.readthedocs.io/en/master/) library and can be used in [Gym](https://www.gymlibrary.dev/) benchmark library.\n",
        "\n",
        "Also, for those who want to go further, one of the best book in reinforcement learning is freely available on the web: http://incompleteideas.net/book/RLbook2018.pdf"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nFEagzwXlpCl"
      },
      "source": [
        "Example of what can be done in RL:\n",
        "- AlphaGo (movie) https://www.youtube.com/watch?v=WXuK6gekU1Y (this work had huge impact in the AI community)\n",
        "- AlphaGo https://deepmind.com/research/case-studies/alphago-the-story-so-far\n",
        "- AlphaZero https://deepmind.com/blog/article/alphazero-shedding-new-light-grand-games-chess-shogi-and-go\n",
        "- AlphaStar (StarCraft II) https://deepmind.com/blog/article/AlphaStar-Grandmaster-level-in-StarCraft-II-using-multi-agent-reinforcement-learning\n",
        "- DQN https://deepmind.com/blog/article/deep-reinforcement-learning\n",
        "- Dota 2 https://openai.com/blog/openai-five/\n",
        "- Robotics https://openai.com/blog/solving-rubiks-cube/\n",
        "- Robotics https://openai.com/blog/learning-dexterity/\n",
        "- Breakout https://www.youtube.com/watch?v=V1eYniJ0Rnk\n",
        "- Walker https://youtu.be/pgaEE27nsQw\n",
        "- Helicopter https://www.youtube.com/watch?v=VCdxqn0fcnE\n",
        "- Energy https://deepmind.com/blog/article/deepmind-ai-reduces-google-data-centre-cooling-bill-40\n",
        "- Self-driving cars"
      ]
    }
  ],
  "metadata": {
    "celltoolbar": "Create Assignment",
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.8.10"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}