Policy_Value_New_Solutions.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Policy_Value_New_Solutions.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/Chris-Fourie/rl_at_ammi/blob/master/Policy_Value_New_Solutions.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZO_beNp3sWOZ",
        "colab_type": "text"
      },
      "source": [
        "## Policy Evaluation\n",
        "$$\n",
        "\\def\\E{\\mathbb{E}}\n",
        "\\def\\given{\\mid}\n",
        "\\def\\states{\\mathcal{S}}\n",
        "\\def\\argmax{\\text{argmax}}\n",
        "$$\n",
        "The first step to improving a policy is to evaluate the state-value function $v_\\pi$ for an arbitraty policy (see Sutton & Barto 4.1 for details). The state-value according to a policy is computed as:\n",
        "\n",
        "$$\n",
        "\\begin{aligned}\n",
        "v_\\pi (s) &= \\E_\\pi \\left[ G_t \\given S_t = s\\right]\\\\\n",
        "&= \\E_\\pi \\left[ R_{t+1} + \\gamma G_{t+1} \\given S_t = s\\right]\\\\\n",
        "&= \\E_\\pi \\left[ R_{t+1} + \\gamma v_\\pi ( S_{t+1}) \\mid S_t = s \\right]\\\\\n",
        "&= \\sum_a \\pi(a|s) \\sum_{s^\\prime, r} p(s^\\prime, r \\mid s,a) \\left[r + \\gamma v_\\pi (s^\\prime)\\right]\n",
        "\\end{aligned}\n",
        "$$\n",
        "\n",
        "The value of $v_\\pi$ can be updated iteratively for all $s$:\n",
        "\n",
        "$$\n",
        "\\begin{aligned}\n",
        "v_{k+1}(s) &= \\E_{\\pi} \\left[ R_{t+1} + \\gamma v_k (S_{t_1} \\given S_t = s)\\right] \\\\\n",
        "&= \\sum_a \\pi(a|s) \\sum_{s^\\prime, r} p(s^\\prime, r \\mid s,a) \\left[r + \\gamma v_\\pi (s^\\prime)\\right]\n",
        "\\end{aligned}\n",
        "$$\n",
        "\n",
        "For any policy $\\pi$, your task to to implement a policy evaluation function based on the following pseudocode:\n",
        "\n",
        "\n",
        "![policy-evaluation](https://i.ibb.co/j4zZ9Xw/policy-evaluation.png)\n",
        "\n",
        "(source: S&B Section 4.1, page 75)\n",
        "\n",
        "Recall $p$ represents the transition probabilities, which we will get from a simple GridWorld below, $\\gamma$ is the discount factor, whereas $\\theta$ is a  small threshold which indicates when to stop updating our value function. A simple GridWorld looks something like the following:\n",
        "\n",
        "![gridworld](https://i.ibb.co/Lk166s4/gridworld.png)\n",
        "\n",
        "(source: S&B Section 4.1, page 76)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kcyiLAxzsWOe",
        "colab_type": "text"
      },
      "source": [
        "### Exercise 1\n",
        "Note the code for this and the following exercises is based on the reinforcement-learning [repo](https://github.com/dennybritz/reinforcement-learning) from Denny Britz.\n",
        "\n",
        "Complete the `policy_eval` function above so that the following evaluates without error:\n",
        "```python\n",
        "v = policy_eval(random_policy, env)\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IlHK529IyDIG",
        "colab_type": "code",
        "outputId": "089602a6-8cbe-49a9-8934-35fc1739192a",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 103
        }
      },
      "source": [
        "!git clone https://github.com/dennybritz/reinforcement-learning"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Cloning into 'reinforcement-learning'...\n",
            "remote: Enumerating objects: 1213, done.\u001b[K\n",
            "Receiving objects:   0% (1/1213)   \rReceiving objects:   1% (13/1213)   \rReceiving objects:   2% (25/1213)   \rReceiving objects:   3% (37/1213)   \rReceiving objects:   4% (49/1213)   \rReceiving objects:   5% (61/1213)   \rReceiving objects:   6% (73/1213)   \rReceiving objects:   7% (85/1213)   \rReceiving objects:   8% (98/1213)   \rReceiving objects:   9% (110/1213)   \rReceiving objects:  10% (122/1213)   \rReceiving objects:  11% (134/1213)   \rReceiving objects:  12% (146/1213)   \rReceiving objects:  13% (158/1213)   \rReceiving objects:  14% (170/1213)   \rReceiving objects:  15% (182/1213)   \rReceiving objects:  16% (195/1213)   \rReceiving objects:  17% (207/1213)   \rReceiving objects:  18% (219/1213)   \rReceiving objects:  19% (231/1213)   \rReceiving objects:  20% (243/1213)   \rReceiving objects:  21% (255/1213)   \rReceiving objects:  22% (267/1213)   \rReceiving objects:  23% (279/1213)   \rReceiving objects:  24% (292/1213)   \rReceiving objects:  25% (304/1213)   \rReceiving objects:  26% (316/1213)   \rReceiving objects:  27% (328/1213)   \rReceiving objects:  28% (340/1213)   \rReceiving objects:  29% (352/1213)   \rReceiving objects:  30% (364/1213)   \rReceiving objects:  31% (377/1213)   \rReceiving objects:  32% (389/1213)   \rReceiving objects:  33% (401/1213)   \rReceiving objects:  34% (413/1213)   \rReceiving objects:  35% (425/1213)   \rReceiving objects:  36% (437/1213)   \rReceiving objects:  37% (449/1213)   \rReceiving objects:  38% (461/1213)   \rReceiving objects:  39% (474/1213)   \rReceiving objects:  40% (486/1213)   \rReceiving objects:  41% (498/1213)   \rReceiving objects:  42% (510/1213)   \rReceiving objects:  43% (522/1213)   \rReceiving objects:  44% (534/1213)   \rReceiving objects:  45% (546/1213)   \rReceiving objects:  46% (558/1213)   \rReceiving objects:  47% (571/1213)   \rReceiving objects:  48% (583/1213)   \rReceiving objects:  49% (595/1213)   \rReceiving objects:  50% (607/1213)   \rReceiving objects:  51% (619/1213)   \rReceiving objects:  52% (631/1213)   \rReceiving objects:  53% (643/1213)   \rReceiving objects:  54% (656/1213)   \rReceiving objects:  55% (668/1213)   \rReceiving objects:  56% (680/1213)   \rReceiving objects:  57% (692/1213)   \rReceiving objects:  58% (704/1213)   \rReceiving objects:  59% (716/1213)   \rReceiving objects:  60% (728/1213)   \rReceiving objects:  61% (740/1213)   \rReceiving objects:  62% (753/1213)   \rReceiving objects:  63% (765/1213)   \rReceiving objects:  64% (777/1213)   \rReceiving objects:  65% (789/1213)   \rReceiving objects:  66% (801/1213)   \rReceiving objects:  67% (813/1213)   \rReceiving objects:  68% (825/1213)   \rReceiving objects:  69% (837/1213)   \rReceiving objects:  70% (850/1213)   \rReceiving objects:  71% (862/1213)   \rReceiving objects:  72% (874/1213)   \rReceiving objects:  73% (886/1213)   \rReceiving objects:  74% (898/1213)   \rReceiving objects:  75% (910/1213)   \rReceiving objects:  76% (922/1213)   \rReceiving objects:  77% (935/1213)   \rReceiving objects:  78% (947/1213)   \rReceiving objects:  79% (959/1213)   \rReceiving objects:  80% (971/1213)   \rReceiving objects:  81% (983/1213)   \rReceiving objects:  82% (995/1213)   \rReceiving objects:  83% (1007/1213)   \rReceiving objects:  84% (1019/1213)   \rReceiving objects:  85% (1032/1213)   \rReceiving objects:  86% (1044/1213)   \rReceiving objects:  87% (1056/1213)   \rReceiving objects:  88% (1068/1213)   \rReceiving objects:  89% (1080/1213)   \rReceiving objects:  90% (1092/1213)   \rReceiving objects:  91% (1104/1213)   \rReceiving objects:  92% (1116/1213)   \rReceiving objects:  93% (1129/1213)   \rReceiving objects:  94% (1141/1213)   \rReceiving objects:  95% (1153/1213)   \rReceiving objects:  96% (1165/1213)   \rReceiving objects:  97% (1177/1213)   \rReceiving objects:  98% (1189/1213)   \rremote: Total 1213 (delta 0), reused 0 (delta 0), pack-reused 1213\u001b[K\n",
            "Receiving objects:  99% (1201/1213)   \rReceiving objects: 100% (1213/1213)   \rReceiving objects: 100% (1213/1213), 5.34 MiB | 29.88 MiB/s, done.\n",
            "Resolving deltas:   0% (0/785)   \rResolving deltas:   2% (17/785)   \rResolving deltas:   4% (33/785)   \rResolving deltas:   9% (78/785)   \rResolving deltas:  12% (95/785)   \rResolving deltas:  13% (105/785)   \rResolving deltas:  14% (114/785)   \rResolving deltas:  16% (127/785)   \rResolving deltas:  18% (145/785)   \rResolving deltas:  19% (152/785)   \rResolving deltas:  20% (158/785)   \rResolving deltas:  21% (165/785)   \rResolving deltas:  22% (176/785)   \rResolving deltas:  23% (182/785)   \rResolving deltas:  24% (192/785)   \rResolving deltas:  25% (198/785)   \rResolving deltas:  27% (215/785)   \rResolving deltas:  31% (248/785)   \rResolving deltas:  33% (263/785)   \rResolving deltas:  34% (268/785)   \rResolving deltas:  37% (293/785)   \rResolving deltas:  38% (299/785)   \rResolving deltas:  39% (309/785)   \rResolving deltas:  40% (315/785)   \rResolving deltas:  41% (324/785)   \rResolving deltas:  42% (332/785)   \rResolving deltas:  43% (339/785)   \rResolving deltas:  47% (371/785)   \rResolving deltas:  48% (377/785)   \rResolving deltas:  51% (401/785)   \rResolving deltas:  52% (411/785)   \rResolving deltas:  55% (433/785)   \rResolving deltas:  57% (450/785)   \rResolving deltas:  58% (456/785)   \rResolving deltas:  59% (465/785)   \rResolving deltas:  60% (476/785)   \rResolving deltas:  66% (520/785)   \rResolving deltas:  67% (527/785)   \rResolving deltas:  69% (544/785)   \rResolving deltas:  71% (558/785)   \rResolving deltas:  72% (566/785)   \rResolving deltas:  73% (576/785)   \rResolving deltas:  74% (587/785)   \rResolving deltas:  76% (601/785)   \rResolving deltas:  77% (612/785)   \rResolving deltas:  81% (636/785)   \rResolving deltas:  82% (651/785)   \rResolving deltas:  83% (655/785)   \rResolving deltas:  85% (673/785)   \rResolving deltas:  89% (699/785)   \rResolving deltas:  90% (714/785)   \rResolving deltas:  91% (715/785)   \rResolving deltas:  94% (744/785)   \rResolving deltas:  95% (747/785)   \rResolving deltas:  96% (754/785)   \rResolving deltas:  99% (779/785)   \rResolving deltas: 100% (785/785)   \rResolving deltas: 100% (785/785), done.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l1hb2MCusWOh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import numpy as np\n",
        "import sys\n",
        "if \"/content/reinforcement-learning\" not in sys.path:\n",
        "  sys.path.append(\"/content/reinforcement-learning\") \n",
        "from lib.envs.gridworld import GridworldEnv"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l-uVxEF5sWOv",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "env = GridworldEnv()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aJpUFdu-sWO5",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "def policy_eval(policy, env, discount_factor=1.0, theta=0.00001):\n",
        "    \"\"\"\n",
        "    Evaluate a policy given an environment and a full description of the environment's dynamics.\n",
        "    \n",
        "    Args:\n",
        "        policy: [S, A] shaped matrix representing the policy.\n",
        "        env: OpenAI env. env.P represents the transition probabilities of the environment.\n",
        "            env.P[s][a] is a list of transition tuples (prob, next_state, reward, done).\n",
        "            env.nS is a number of states in the environment. \n",
        "            env.nA is a number of actions in the environment.\n",
        "        theta: We stop evaluation once our value function change is less than theta for all states.\n",
        "        discount_factor: Gamma discount factor.\n",
        "    \n",
        "    Returns:\n",
        "        Vector of length env.nS representing the value function.\n",
        "    \"\"\"\n",
        "    # Start with a random (all 0) value function\n",
        "    V = np.zeros(env.nS)\n",
        "    while True:\n",
        "      \n",
        "        delta = 0\n",
        "        # For each state, perform a \"full backup\"\n",
        "        for s in range(env.nS):\n",
        "            v = 0\n",
        "            \n",
        "            # Look at the possible next actions\n",
        "            for a, action_prob in enumerate(policy[s]):\n",
        "              \n",
        "                # For each action, look at the possible next states...\n",
        "                for  prob, next_state, reward, done in env.P[s][a]:\n",
        "                  \n",
        "                    # Calculate the expected value. Ref: Sutton book eq. 4.6.\n",
        "                    v += action_prob * prob * (reward + discount_factor * V[next_state])\n",
        "                    \n",
        "            # How much our value function changed (across any states)\n",
        "            delta = max(delta, np.abs(v - V[s]))\n",
        "            V[s] = v\n",
        "              \n",
        "        # Stop evaluating once our value function change is below a threshold\n",
        "        if delta < theta:\n",
        "            break\n",
        "    return np.array(V)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "nu4YWgt6sWPB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "random_policy = np.ones([env.nS, env.nA]) / env.nA\n",
        "v = policy_eval(random_policy, env)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eahtKZidsWPK",
        "colab_type": "text"
      },
      "source": [
        "Run the cell below to verify your $v_\\pi$ has converged to the proper value."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FVB6jrAzsWPO",
        "colab_type": "code",
        "outputId": "8f30d62b-f0bd-48d6-af89-5af6b9428ef0",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        }
      },
      "source": [
        "expected_v = np.array([0, -14, -20, -22, -14, -18, -20, -20, -20, -20, -18, -14, -22, -20, -14, 0])\n",
        "result = np.testing.assert_array_almost_equal(v, expected_v, decimal=2)\n",
        "print(result)\n",
        "print(\"If `result` prints None, then your algorithm was successfully implemented\")"
      ],
      "execution_count": 0,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "None\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lR3Erk1GsWPX",
        "colab_type": "text"
      },
      "source": [
        "# Policy Iteration\n",
        "\n",
        "Now that we can evaluate the state value function for any given policy, we can use the values to improve our policy. \n",
        "\n",
        "First consider the state-action value function (q-function), which consists of selecting $a$ in $s$ and then behaving according to $\\pi$:\n",
        "\n",
        "$$\n",
        "\\begin{aligned}\n",
        "q_\\pi(s,a) &= \\E \\left[ R_{t+1} + \\gamma v_\\pi (S_{s+1}) \\given S_t = s, A_t = a \\right] \\\\\n",
        "&= \\sum_{s^\\prime, r} p(s^\\prime, r \\given, s, a) \\left[ r + \\gamma v_\\pi (s^\\prime) \\right]\n",
        "\\end{aligned}\n",
        "$$\n",
        "\n",
        "Let $\\pi^\\prime$ be a policy such that, for all $s \\in \\states$:\n",
        "$$\n",
        "q(s, \\pi^\\prime (s)) \\geq v_\\pi (s)\n",
        "$$\n",
        "\n",
        "Then the following holds for all $s \\in \\states$:\n",
        "$$\n",
        "v_\\pi^\\prime (s) \\geq v_\\pi (s)\n",
        "$$\n",
        "(See S&B page 78 for proof)\n",
        "\n",
        "Consider a simple policy improvement which consists of $\\pi$ selecting an action according to the maximum state-action value:\n",
        "\n",
        "$$\n",
        "\\begin{aligned}\n",
        "\\pi^\\prime (s) &= \\argmax_{a} q_\\pi (s,a) \\\\\n",
        "&= \\argmax_a \\E \\left[ R_{t+1} + \\gamma v_\\pi (S_{t+1}) \\given S_t = s, A_t = a \\right] \\\\\n",
        "&= \\argmax_a \\sum_{s^\\prime, r} p(s^\\prime, s \\given s, a) [r + \\gamma v_\\pi (s^\\prime)]\n",
        "\\end{aligned}\n",
        "$$\n",
        "\n",
        "Alternating between policy evaluation and policy improvement is known as **policy iteration**. Implement the algorithm based on the following pseudocode:\n",
        "\n",
        "![gridworld.png](https://i.ibb.co/GHpV8kV/policy-iteration.png)\n",
        "\n",
        "(source: S&B page 80)\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "R1B_VsImsWPZ",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "pp = pprint.PrettyPrinter(indent=2)\n",
        "env = GridworldEnv()"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XGyLfXKOsWPl",
        "colab_type": "text"
      },
      "source": [
        "### Exercise 2\n",
        "Complete the `policy_improvement` function, based on the above pseudocode, so that the following runs without error:\n",
        "```\n",
        "policy, v = policy_improvement(env)\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_JuXGYh_sWPq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "def policy_improvement(env, policy_eval_fn=policy_eval, discount_factor=1.0):\n",
        "    \"\"\"\n",
        "    Policy Improvement Algorithm. Iteratively evaluates and improves a policy\n",
        "    until an optimal policy is found.\n",
        "    \n",
        "    Args:\n",
        "        env: The OpenAI envrionment.\n",
        "        policy_eval_fn: Policy Evaluation function that takes 3 arguments:\n",
        "            policy, env, discount_factor.\n",
        "        discount_factor: gamma discount factor.\n",
        "        \n",
        "    Returns:\n",
        "        A tuple (policy, V). \n",
        "        policy is the optimal policy, a matrix of shape [S, A] where each state s\n",
        "        contains a valid probability distribution over actions.\n",
        "        V is the value function for the optimal policy.\n",
        "        \n",
        "    \"\"\"\n",
        "\n",
        "    def one_step_lookahead(state, V):\n",
        "        \"\"\"\n",
        "        Helper function to calculate the value for all action in a given state.\n",
        "        \n",
        "        Args:\n",
        "            state: The state to consider (int)\n",
        "            V: The value to use as an estimator, Vector of length env.nS\n",
        "        \n",
        "        Returns:\n",
        "            A vector of length env.nA containing the expected value of each action.\n",
        "        \"\"\"\n",
        "        A = np.zeros(env.nA)\n",
        "        for a in range(env.nA):\n",
        "            for prob, next_state, reward, done in env.P[state][a]:\n",
        "                A[a] += prob * (reward + discount_factor * V[next_state])\n",
        "        return A\n",
        "      \n",
        "    # Start with a random policy\n",
        "    policy = np.ones([env.nS, env.nA]) / env.nA\n",
        "    \n",
        "    while True:\n",
        "        # Evaluate the current policy\n",
        "        V = policy_eval_fn(policy, env, discount_factor)\n",
        "        \n",
        "        # Will be set to false if we make any changes to the policy\n",
        "        policy_stable = True\n",
        "        \n",
        "        # For each state...\n",
        "        for s in range(env.nS):\n",
        "            # The best action we would take under the currect policy\n",
        "            chosen_a = np.argmax(policy[s])\n",
        "            \n",
        "            # Find the best action by one-step lookahead\n",
        "            # Ties are resolved arbitarily\n",
        "            action_values = one_step_lookahead(s, V)\n",
        "            best_a = np.argmax(action_values)\n",
        "            \n",
        "            # Greedily update the policy\n",
        "            if chosen_a != best_a:\n",
        "                policy_stable = False\n",
        "            policy[s] = np.eye(env.nA)[best_a]\n",
        "        \n",
        "        # If the policy is stable we've found an optimal policy. Return it\n",
        "        if policy_stable:\n",
        "            return policy, V"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hFFqHVzxsWP2",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "policy, v = policy_improvement(env)\n",
        "print(\"Policy Probability Distribution:\")\n",
        "print(policy)\n",
        "print(\"\")\n",
        "\n",
        "print(\"Reshaped Grid Policy (0=up, 1=right, 2=down, 3=left):\")\n",
        "print(np.reshape(np.argmax(policy, axis=1), env.shape))\n",
        "print(\"\")\n",
        "\n",
        "print(\"Value Function:\")\n",
        "print(v)\n",
        "print(\"\")\n",
        "\n",
        "print(\"Reshaped Grid Value Function:\")\n",
        "print(v.reshape(env.shape))\n",
        "print(\"\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eAD93o78sWQB",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Test the value function\n",
        "expected_v = np.array([ 0, -1, -2, -3, -1, -2, -3, -2, -2, -3, -2, -1, -3, -2, -1,  0])\n",
        "np.testing.assert_array_almost_equal(v, expected_v, decimal=2)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I0rpOroesWQJ",
        "colab_type": "text"
      },
      "source": [
        "## Value Iteration\n",
        "The isue with policy iteration is that every time we improve our policy, we must do a full sweep through all states again to evaluate the new policy! This is extremely inefficient. Value iteration combines both evaluation and improvement in the same step. Implement value iteration based on the following pseudocode:\n",
        "\n",
        "\n",
        "![gridworld.png](https://i.ibb.co/SVTmBFQ/value-iteration.png)\n",
        "\n",
        "(source: S&B Section 4.4, page 83)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LtMKX6-QsWQL",
        "colab_type": "text"
      },
      "source": [
        "### Exercise 3\n",
        "Complete the `value_iteration` function below, based on the above pseudocode, so that the following evaluates without error\n",
        "```\n",
        "policy, v = value_iteration(env)\n",
        "```"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "v5yBKZKOsWQN",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "\n",
        "def value_iteration(env, theta=0.0001, discount_factor=1.0):\n",
        "    \"\"\"\n",
        "    Value Iteration Algorithm.\n",
        "    \n",
        "    Args:\n",
        "        env: OpenAI env. env.P represents the transition probabilities of the environment.\n",
        "            env.P[s][a] is a list of transition tuples (prob, next_state, reward, done).\n",
        "            env.nS is a number of states in the environment. \n",
        "            env.nA is a number of actions in the environment.\n",
        "        theta: We stop evaluation once our value function change is less than theta for all states.\n",
        "        discount_factor: Gamma discount factor.\n",
        "        \n",
        "    Returns:\n",
        "        A tuple (policy, V) of the optimal policy and the optimal value function.\n",
        "    \"\"\"\n",
        "    \n",
        "    def one_step_lookahead(state, V):\n",
        "        \"\"\"\n",
        "        Helper function to calculate the value for all action in a given state.\n",
        "        \n",
        "        Args:\n",
        "            state: The state to consider (int)\n",
        "            V: The value to use as an estimator, Vector of length env.nS\n",
        "        \n",
        "        Returns:\n",
        "            A vector of length env.nA containing the expected value of each action.\n",
        "        \"\"\"\n",
        "        A = np.zeros(env.nA)\n",
        "        for a in range(env.nA):\n",
        "            for prob, next_state, reward, done in env.P[state][a]:\n",
        "                A[a] += prob * (reward + discount_factor * V[next_state])\n",
        "        return A\n",
        "\n",
        "    V = np.zeros(env.nS)\n",
        "    while True:\n",
        "        # Stopping condition\n",
        "        delta = 0\n",
        "        # Update each state...\n",
        "        for s in range(env.nS):\n",
        "            # Do a one-step lookahead to find the best action\n",
        "            A = one_step_lookahead(s, V)\n",
        "            best_action_value = np.max(A)\n",
        "            # Calculate delta across all states seen so far\n",
        "            delta = max(delta, np.abs(best_action_value - V[s]))\n",
        "            # Update the value function. Ref: Sutton book eq. 4.10. \n",
        "            V[s] = best_action_value        \n",
        "        # Check if we can stop \n",
        "        if delta < theta:\n",
        "            break\n",
        "    \n",
        "    # Create a deterministic policy using the optimal value function\n",
        "    policy = np.zeros([env.nS, env.nA])\n",
        "    for s in range(env.nS):\n",
        "        # One step lookahead to find the best action for this state\n",
        "        A = one_step_lookahead(s, V)\n",
        "        best_action = np.argmax(A)\n",
        "        # Always take the best action\n",
        "        policy[s, best_action] = 1.0\n",
        "    \n",
        "    return policy, V"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fapSyrHAsWQW",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "policy, v = value_iteration(env)\n",
        "\n",
        "print(\"Policy Probability Distribution:\")\n",
        "print(policy)\n",
        "print(\"\")\n",
        "\n",
        "print(\"Reshaped Grid Policy (0=up, 1=right, 2=down, 3=left):\")\n",
        "print(np.reshape(np.argmax(policy, axis=1), env.shape))\n",
        "print(\"\")\n",
        "\n",
        "print(\"Value Function:\")\n",
        "print(v)\n",
        "print(\"\")\n",
        "\n",
        "print(\"Reshaped Grid Value Function:\")\n",
        "print(v.reshape(env.shape))\n",
        "print(\"\")"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LXCsMRo-sWQf",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Test the value function\n",
        "expected_v = np.array([ 0, -1, -2, -3, -1, -2, -3, -2, -2, -3, -2, -1, -3, -2, -1,  0])\n",
        "np.testing.assert_array_almost_equal(v, expected_v, decimal=2)"
      ],
      "execution_count": 0,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "buewi_IUsWQm",
        "colab_type": "text"
      },
      "source": [
        "## Exercise 4\n",
        "Now that you have implemented Policy Iteration and Value Iteration, plot the average running time for both algorithms by varying the discount rate $\\gamma$ between 0 and 1. What do you observe?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAUy7uTfsWQp",
        "colab_type": "text"
      },
      "source": [
        "## Alternative implementation\n",
        "The algorithms you've seen in this notebook can alternatively be solved using matrix notation and in closed form in some cases. See [this](https://drive.google.com/file/d/1UR20JtQRjFyrvCseusVuPBmQIpB3XFAH/view?usp=sharing) notebook."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jD7MwFnEsWQs",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 0,
      "outputs": []
    }
  ]
}