ml-finance-python

01_gridworld_dynamic_programming.ipynb

(33250B)

 {
  "cells": [
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "# Dynamic programming: Value and Policy Iteration"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "In this section, we will apply value and policy iteration to a toy environment that consists of a 3 x 4 grid that's depicted in the following diagram with the following features:\n",
     "\n",
     "- **States**: 11 states represented as two-dimensional coordinates. One field is not accessible and the top two states in the rightmost column are terminal, that is, they end the episode.\n",
     "- **Actions**: Movements on each step, that is, up, down, left, and right. The environment is randomized so that actions can have unintended outcomes. For each action, there is an 80% probability to move to the expected state, and 10% probability to move in an adjacent direction (for example, right or left instead of up or up or down instead of right).\n",
     "- **Rewards**: As depicted in the right-hand side panel, each state results in -.02, except for the +1/-1 rewards in the terminal states:"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "<img src=\"img/03_mdp.png\" width=\"500\">"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "The right panel of the preceding GridWorld diagram shows the optimal value estimate that's produced by Value Iteration and the corresponding greedy policy. The negative rewards, combined with the uncertainty in the environment, produce an optimal policy that involves moving away from the negative terminal state.\n",
     "\n",
     "The results are sensitive to both the rewards and the discount factor. The cost of the negative state affects the policy in the surrounding fields, and you should modify the example in the corresponding notebook to identify threshold levels that alter the optimal action selection."
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Imports & Settings"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
     "%matplotlib inline\n",
     "\n",
     "from pathlib import Path\n",
     "from time import process_time\n",
     "import numpy as np\n",
     "import pandas as pd\n",
     "from mdptoolbox import mdp\n",
     "from itertools import product\n",
     "import gym"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Set up Gridworld"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### States, Actions and Rewards"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "We will begin by defining the environment parameters:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
     "grid_size = (3, 4)\n",
     "blocked_cell = (1, 1)\n",
     "baseline_reward = -0.02\n",
     "absorbing_cells = {(0, 3): 1, (1, 3): -1}"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
     "actions = ['L', 'U', 'R', 'D']\n",
     "num_actions = len(actions)\n",
     "probs = [.1, .8, .1, 0]"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "We will frequently need to convert between one-dimensional and two-dimensional representations, so we will define two helper functions for this purpose; states are one-dimensional and cells are the corresponding two-dimensional coordinates:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 4,
    "metadata": {},
    "outputs": [],
    "source": [
     "to_1d = lambda x: np.ravel_multi_index(x, grid_size)\n",
     "to_2d = lambda x: np.unravel_index(x, grid_size)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "Furthermore, we will precompute some data points to make the code more concise:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 5,
    "metadata": {},
    "outputs": [],
    "source": [
     "num_states = np.product(grid_size)\n",
     "cells = list(np.ndindex(grid_size))\n",
     "states = list(range(len(cells)))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
     "cell_state = dict(zip(cells, states))\n",
     "state_cell= dict(zip(states, cells))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
     "absorbing_states = {to_1d(s):r for s, r in absorbing_cells.items()}\n",
     "blocked_state = to_1d(blocked_cell)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "We store the rewards for each state:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 8,
    "metadata": {},
    "outputs": [],
    "source": [
     "state_rewards = np.full(num_states, baseline_reward)\n",
     "state_rewards[blocked_state] = 0\n",
     "for state, reward in absorbing_states.items():\n",
     "    state_rewards[state] = reward"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 9,
    "metadata": {},
    "outputs": [],
    "source": [
     "action_outcomes = {}\n",
     "for i, action in enumerate(actions):\n",
     "    probs_ = dict(zip([actions[j % 4] for j in range(i, num_actions + i)], probs))\n",
     "    action_outcomes[actions[(i + 1) % 4]] = probs_"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "To account for the probabilistic environment, we also need to compute the probability distribution over the actual move for a given action:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 10,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "{'U': {'L': 0.1, 'U': 0.8, 'R': 0.1, 'D': 0},\n",
        " 'R': {'U': 0.1, 'R': 0.8, 'D': 0.1, 'L': 0},\n",
        " 'D': {'R': 0.1, 'D': 0.8, 'L': 0.1, 'U': 0},\n",
        " 'L': {'D': 0.1, 'L': 0.8, 'U': 0.1, 'R': 0}}"
       ]
      },
      "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "action_outcomes"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "Now, we are ready to compute the transition matrix, which is the key input to the MDP."
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Transition Matrix"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "The transition matrix defines the probability to end up in a certain state, S, for each previous state and action, A, $P(s^\\prime \\mid s, a)$. We will demonstrate `pymdptoolbox`, and use one of the formats that's available to us to specify transitions and rewards. For both transition probabilities, we will create a `NumPy` array with dimensions of $A \\times S \\times S$.\n",
     "\n",
     "First, we compute the target cell for each starting cell and move:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 11,
    "metadata": {},
    "outputs": [],
    "source": [
     "def get_new_cell(state, move):\n",
     "    cell = to_2d(state)\n",
     "    if actions[move] == 'U':\n",
     "        return cell[0] - 1, cell[1]\n",
     "    elif actions[move] == 'D':\n",
     "        return cell[0] + 1, cell[1]\n",
     "    elif actions[move] == 'R':\n",
     "        return cell[0], cell[1] + 1\n",
     "    elif actions[move] == 'L':\n",
     "        return cell[0], cell[1] - 1"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 12,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "array([-0.02, -0.02, -0.02,  1.  , -0.02,  0.  , -0.02, -1.  , -0.02,\n",
        "       -0.02, -0.02, -0.02])"
       ]
      },
      "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "state_rewards"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "The following function uses the argument's starting `state`, `action`, and `outcome` to fill in the transition probabilities and rewards:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 13,
    "metadata": {},
    "outputs": [],
    "source": [
     "def update_transitions_and_rewards(state, action, outcome):\n",
     "    if state in absorbing_states.keys() or state == blocked_state:\n",
     "        transitions[action, state, state] = 1\n",
     "    else:\n",
     "        new_cell = get_new_cell(state, outcome)\n",
     "        p = action_outcomes[actions[action]][actions[outcome]]\n",
     "        if new_cell not in cells or new_cell == blocked_cell:\n",
     "            transitions[action, state, state] += p\n",
     "            rewards[action, state, state] = baseline_reward\n",
     "        else:\n",
     "            new_state= to_1d(new_cell)\n",
     "            transitions[action, state, new_state] = p\n",
     "            rewards[action, state, new_state] = state_rewards[new_state]"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "We generate the transition and reward values by creating placeholder data structures and iterating over the Cartesian product of $A \\times S \\times S$, as follows:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 14,
    "metadata": {},
    "outputs": [],
    "source": [
     "rewards = np.zeros(shape=(num_actions, num_states, num_states))\n",
     "transitions = np.zeros((num_actions, num_states, num_states))\n",
     "actions_ = list(range(num_actions))\n",
     "for action, outcome, state in product(actions_, actions_, states):\n",
     "    update_transitions_and_rewards(state, action, outcome)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 15,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "((4, 12, 12), (4, 12, 12))"
       ]
      },
      "execution_count": 15,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "rewards.shape, transitions.shape"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## PyMDPToolbox"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "We can also solve MDPs using the [pymdptoolbox](https://pymdptoolbox.readthedocs.io/en/latest/api/mdptoolbox.html) Python library, which includes a few more algorithms, including Q-learning."
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Value Iteration"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 16,
    "metadata": {},
    "outputs": [],
    "source": [
     "gamma = .99\n",
     "epsilon = 1e-5"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "To run `ValueIteration`, just instantiate the corresponding object with the desired configuration options and the rewards and transition matrices before calling the `.run()` method:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 17,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "'# Iterations: 31 | Time: 0.0006'"
       ]
      },
      "execution_count": 17,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "vi = mdp.ValueIteration(transitions=transitions,\n",
     "                        reward=rewards,\n",
     "                        discount=gamma,\n",
     "                        epsilon=epsilon)\n",
     "\n",
     "vi.run()\n",
     "f'# Iterations: {vi.iter:,d} | Time: {vi.time:.4f}'"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 18,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
        "    .dataframe tbody tr th:only-of-type {\n",
        "        vertical-align: middle;\n",
        "    }\n",
        "\n",
        "    .dataframe tbody tr th {\n",
        "        vertical-align: top;\n",
        "    }\n",
        "\n",
        "    .dataframe thead th {\n",
        "        text-align: right;\n",
        "    }\n",
        "</style>\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>0</th>\n",
        "      <th>1</th>\n",
        "      <th>2</th>\n",
        "      <th>3</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>R</td>\n",
        "      <td>R</td>\n",
        "      <td>R</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
        "   0  1  2  3\n",
        "0  R  R  R  L\n",
        "1  U  L  U  L\n",
        "2  U  L  L  L"
       ]
      },
      "execution_count": 18,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "policy = np.asarray([actions[i] for i in vi.policy])\n",
     "pd.DataFrame(policy.reshape(grid_size))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 19,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
        "    .dataframe tbody tr th:only-of-type {\n",
        "        vertical-align: middle;\n",
        "    }\n",
        "\n",
        "    .dataframe tbody tr th {\n",
        "        vertical-align: top;\n",
        "    }\n",
        "\n",
        "    .dataframe thead th {\n",
        "        text-align: right;\n",
        "    }\n",
        "</style>\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>0</th>\n",
        "      <th>1</th>\n",
        "      <th>2</th>\n",
        "      <th>3</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>0.884143</td>\n",
        "      <td>0.925054</td>\n",
        "      <td>0.961986</td>\n",
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>0.848181</td>\n",
        "      <td>0.000000</td>\n",
        "      <td>0.714643</td>\n",
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>0.808345</td>\n",
        "      <td>0.773328</td>\n",
        "      <td>0.736099</td>\n",
        "      <td>0.516083</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
        "          0         1         2         3\n",
        "0  0.884143  0.925054  0.961986  0.000000\n",
        "1  0.848181  0.000000  0.714643  0.000000\n",
        "2  0.808345  0.773328  0.736099  0.516083"
       ]
      },
      "execution_count": 19,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "value = np.asarray(vi.V).reshape(grid_size)\n",
     "pd.DataFrame(value)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Policy Iteration"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "The `PolicyIteration` function works similarly:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 20,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "'# Iterations: 7 | Time: 0.0087'"
       ]
      },
      "execution_count": 20,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "pi = mdp.PolicyIteration(transitions=transitions,\n",
     "                        reward=rewards,\n",
     "                        discount=gamma,\n",
     "                        max_iter=1000)\n",
     "\n",
     "pi.run()\n",
     "f'# Iterations: {pi.iter:,d} | Time: {pi.time:.4f}'"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "It also yields the same policy, but the value function varies by run and does not need to achieve the optimal value before the policy converges."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 21,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
        "    .dataframe tbody tr th:only-of-type {\n",
        "        vertical-align: middle;\n",
        "    }\n",
        "\n",
        "    .dataframe tbody tr th {\n",
        "        vertical-align: top;\n",
        "    }\n",
        "\n",
        "    .dataframe thead th {\n",
        "        text-align: right;\n",
        "    }\n",
        "</style>\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>0</th>\n",
        "      <th>1</th>\n",
        "      <th>2</th>\n",
        "      <th>3</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>R</td>\n",
        "      <td>R</td>\n",
        "      <td>R</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
        "   0  1  2  3\n",
        "0  R  R  R  L\n",
        "1  U  L  U  L\n",
        "2  U  L  L  L"
       ]
      },
      "execution_count": 21,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "policy = np.asarray([actions[i] for i in pi.policy])\n",
     "pd.DataFrame(policy.reshape(grid_size))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 22,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
        "    .dataframe tbody tr th:only-of-type {\n",
        "        vertical-align: middle;\n",
        "    }\n",
        "\n",
        "    .dataframe tbody tr th {\n",
        "        vertical-align: top;\n",
        "    }\n",
        "\n",
        "    .dataframe thead th {\n",
        "        text-align: right;\n",
        "    }\n",
        "</style>\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>0</th>\n",
        "      <th>1</th>\n",
        "      <th>2</th>\n",
        "      <th>3</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>0.884143</td>\n",
        "      <td>0.925054</td>\n",
        "      <td>0.961986</td>\n",
        "      <td>1.594721e-16</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>0.848181</td>\n",
        "      <td>0.000000</td>\n",
        "      <td>0.714643</td>\n",
        "      <td>-0.000000e+00</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>0.808345</td>\n",
        "      <td>0.773328</td>\n",
        "      <td>0.736099</td>\n",
        "      <td>5.160828e-01</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
        "          0         1         2             3\n",
        "0  0.884143  0.925054  0.961986  1.594721e-16\n",
        "1  0.848181  0.000000  0.714643 -0.000000e+00\n",
        "2  0.808345  0.773328  0.736099  5.160828e-01"
       ]
      },
      "execution_count": 22,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "value = np.asarray(pi.V).reshape(grid_size)\n",
     "pd.DataFrame(value)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Value Iteration"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 26,
    "metadata": {},
    "outputs": [],
    "source": [
     "skip_states = list(absorbing_states.keys())+[blocked_state]\n",
     "states_to_update = [s for s in states if s not in skip_states]"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "Then, we initialize the value function and set the discount factor gamma and the convergence threshold epsilon:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 27,
    "metadata": {},
    "outputs": [],
    "source": [
     "V = np.random.rand(num_states)\n",
     "V[skip_states] = 0"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 28,
    "metadata": {},
    "outputs": [],
    "source": [
     "gamma = .99\n",
     "epsilon = 1e-5"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "The algorithm updates the value function using the Bellman optimality equation, and terminates when the L1 norm of V changes less than epsilon in absolute terms:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 29,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "'# Iterations 17 | Time 0.0037'"
       ]
      },
      "execution_count": 29,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "iterations = 0\n",
     "start = process_time()\n",
     "converged = False\n",
     "while not converged:\n",
     "    V_ = np.copy(V)\n",
     "    for state in states_to_update:\n",
     "        q_sa = np.sum(transitions[:, state] * (rewards[:, state] + gamma* V), axis=1)\n",
     "        V[state] = np.max(q_sa)\n",
     "    if np.sum(np.fabs(V - V_)) < epsilon:\n",
     "        converged = True\n",
     "\n",
     "    iterations += 1\n",
     "    if iterations % 1000 == 0:\n",
     "        print(np.sum(np.fabs(V - V_)))\n",
     "\n",
     "f'# Iterations {iterations} | Time {process_time() - start:.4f}'"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Value Function"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 30,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
       "          0         1         2         3\n",
       "0  0.884143  0.925054  0.961986  0.000000\n",
       "1  0.848181  0.000000  0.714643  0.000000\n",
       "2  0.808344  0.773327  0.736099  0.516082\n"
      ]
     }
    ],
    "source": [
     "print(pd.DataFrame(V.reshape(grid_size)))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 31,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "True"
       ]
      },
      "execution_count": 31,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "np.allclose(V.reshape(grid_size), np.asarray(vi.V).reshape(grid_size))"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "### Optimal Policy"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 35,
    "metadata": {
     "scrolled": true
    },
    "outputs": [
     {
      "data": {
       "text/plain": [
        "array([2, 2, 2, 0, 1, 0, 0, 0, 1, 0, 0, 0])"
       ]
      },
      "execution_count": 35,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "for state, reward in absorbing_states.items():\n",
     "    V[state] = reward\n",
     "\n",
     "policy = np.argmax(np.sum(transitions * V, 2),0)\n",
     "policy"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 36,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
        "    .dataframe tbody tr th:only-of-type {\n",
        "        vertical-align: middle;\n",
        "    }\n",
        "\n",
        "    .dataframe tbody tr th {\n",
        "        vertical-align: top;\n",
        "    }\n",
        "\n",
        "    .dataframe thead th {\n",
        "        text-align: right;\n",
        "    }\n",
        "</style>\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>0</th>\n",
        "      <th>1</th>\n",
        "      <th>2</th>\n",
        "      <th>3</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>R</td>\n",
        "      <td>R</td>\n",
        "      <td>R</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
        "   0  1  2  3\n",
        "0  R  R  R  L\n",
        "1  U  L  L  L\n",
        "2  U  L  L  L"
       ]
      },
      "execution_count": 36,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "pd.DataFrame(policy.reshape(grid_size)).replace(dict(enumerate(actions)))"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "## Policy Iteration"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "Policy iterations involves separate evaluation and improvement steps. We define the improvement part by selecting the action that maximizes the sum of expected reward and next-state value. Note that we temporarily fill in the rewards for the terminal states to avoid ignoring actions that would lead us there:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 41,
    "metadata": {},
    "outputs": [],
    "source": [
     "def policy_improvement(value, transitions):\n",
     "    for state, reward in absorbing_states.items():\n",
     "        value[state] = reward\n",
     "    return np.argmax(np.sum(transitions * value, 2),0)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "V = np.random.rand(num_states)\n",
     "V[skip_states] = 0\n",
     "pi = np.random.choice(list(range(num_actions)), size=num_states)"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "The algorithm alternates between policy evaluation for a greedily selected action and policy improvement until the policy stabilizes:"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 42,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
        "'# Iterations 3 | Time 0.0059'"
       ]
      },
      "execution_count": 42,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "iterations = 0\n",
     "start = process_time()\n",
     "converged = False\n",
     "while not converged:\n",
     "    pi_ = np.copy(pi)\n",
     "    for state in states_to_update:\n",
     "        action = policy[state]\n",
     "        V[state] = np.dot(transitions[action, state], (rewards[action, state] + gamma* V))\n",
     "        pi = policy_improvement(V.copy(), transitions)\n",
     "    if np.array_equal(pi_, pi):\n",
     "        converged = True\n",
     "    iterations += 1\n",
     "\n",
     "f'# Iterations {iterations} | Time {process_time() - start:.4f}'"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
     "Policy iteration converges after only three iterations. The policy stabilizes before the algorithm finds the optimal value function, and the optimal policy differs slightly, most notably by suggesting up instead of the safer left for the field next to the negative terminal state. This can be avoided by tightening the convergence criteria, for example, by requiring a stable policy of several rounds or adding a threshold for the value function."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 43,
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
        "    .dataframe tbody tr th:only-of-type {\n",
        "        vertical-align: middle;\n",
        "    }\n",
        "\n",
        "    .dataframe tbody tr th {\n",
        "        vertical-align: top;\n",
        "    }\n",
        "\n",
        "    .dataframe thead th {\n",
        "        text-align: right;\n",
        "    }\n",
        "</style>\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>0</th>\n",
        "      <th>1</th>\n",
        "      <th>2</th>\n",
        "      <th>3</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>R</td>\n",
        "      <td>R</td>\n",
        "      <td>R</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>U</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "      <td>L</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
        "   0  1  2  3\n",
        "0  R  R  R  L\n",
        "1  U  L  U  L\n",
        "2  U  L  L  L"
       ]
      },
      "execution_count": 43,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "pd.DataFrame(pi.reshape(grid_size)).replace(dict(enumerate(actions)))"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 44,
    "metadata": {
     "scrolled": true
    },
    "outputs": [
     {
      "data": {
       "text/html": [
        "<div>\n",
        "<style scoped>\n",
        "    .dataframe tbody tr th:only-of-type {\n",
        "        vertical-align: middle;\n",
        "    }\n",
        "\n",
        "    .dataframe tbody tr th {\n",
        "        vertical-align: top;\n",
        "    }\n",
        "\n",
        "    .dataframe thead th {\n",
        "        text-align: right;\n",
        "    }\n",
        "</style>\n",
        "<table border=\"1\" class=\"dataframe\">\n",
        "  <thead>\n",
        "    <tr style=\"text-align: right;\">\n",
        "      <th></th>\n",
        "      <th>0</th>\n",
        "      <th>1</th>\n",
        "      <th>2</th>\n",
        "      <th>3</th>\n",
        "    </tr>\n",
        "  </thead>\n",
        "  <tbody>\n",
        "    <tr>\n",
        "      <th>0</th>\n",
        "      <td>0.765729</td>\n",
        "      <td>0.874981</td>\n",
        "      <td>0.923802</td>\n",
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>1</th>\n",
        "      <td>0.697293</td>\n",
        "      <td>0.000000</td>\n",
        "      <td>0.405803</td>\n",
        "      <td>0.000000</td>\n",
        "    </tr>\n",
        "    <tr>\n",
        "      <th>2</th>\n",
        "      <td>0.625634</td>\n",
        "      <td>0.563798</td>\n",
        "      <td>0.506690</td>\n",
        "      <td>0.305024</td>\n",
        "    </tr>\n",
        "  </tbody>\n",
        "</table>\n",
        "</div>"
       ],
       "text/plain": [
        "          0         1         2         3\n",
        "0  0.765729  0.874981  0.923802  0.000000\n",
        "1  0.697293  0.000000  0.405803  0.000000\n",
        "2  0.625634  0.563798  0.506690  0.305024"
       ]
      },
      "execution_count": 44,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "pd.DataFrame(V.reshape(grid_size))"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
    "display_name": "Python 3",
    "language": "python",
    "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
     "name": "ipython",
     "version": 3
    },
    "file_extension": ".py",
    "mimetype": "text/x-python",
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
    "version": "3.6.8"
   },
   "toc": {
    "base_numbering": 1,
    "nav_menu": {},
    "number_sections": true,
    "sideBar": true,
    "skip_h1_title": true,
    "title_cell": "Table of Contents",
    "title_sidebar": "Contents",
    "toc_cell": false,
    "toc_position": {},
    "toc_section_display": true,
    "toc_window_display": true
   }
  },
  "nbformat": 4,
  "nbformat_minor": 2
 }