notebooks/chapter_01/chapter-01.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following is a function to print policies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_policy(pi, P):\n",
    "    arrs = {k:v for k,v in enumerate(('<', 'v', '>', '^'))}\n",
    "    for key, value in pi.items():\n",
    "        print(\"| \", end=\"\")\n",
    "        if P[key][0][0][0] == 1.0:\n",
    "            print(\"    \", end=\" \")\n",
    "        else:\n",
    "            print(str(key).zfill(2), arrs[value], end=\" \")\n",
    "        if (key + 1) % np.sqrt(len(pi)) == 0: print(\"|\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the MDP of the Frozen Lake environment as described in the book"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "P = {0: {0: [(0.3333333333333333, 0, 0.0, False),\n",
    "   (0.3333333333333333, 0, 0.0, False),\n",
    "   (0.3333333333333333, 4, 0.0, False)],\n",
    "  1: [(0.3333333333333333, 0, 0.0, False),\n",
    "   (0.3333333333333333, 4, 0.0, False),\n",
    "   (0.3333333333333333, 1, 0.0, False)],\n",
    "  2: [(0.3333333333333333, 4, 0.0, False),\n",
    "   (0.3333333333333333, 1, 0.0, False),\n",
    "   (0.3333333333333333, 0, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 1, 0.0, False),\n",
    "   (0.3333333333333333, 0, 0.0, False),\n",
    "   (0.3333333333333333, 0, 0.0, False)]},\n",
    " 1: {0: [(0.3333333333333333, 1, 0.0, False),\n",
    "   (0.3333333333333333, 0, 0.0, False),\n",
    "   (0.3333333333333333, 5, 0.0, True)],\n",
    "  1: [(0.3333333333333333, 0, 0.0, False),\n",
    "   (0.3333333333333333, 5, 0.0, True),\n",
    "   (0.3333333333333333, 2, 0.0, False)],\n",
    "  2: [(0.3333333333333333, 5, 0.0, True),\n",
    "   (0.3333333333333333, 2, 0.0, False),\n",
    "   (0.3333333333333333, 1, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 2, 0.0, False),\n",
    "   (0.3333333333333333, 1, 0.0, False),\n",
    "   (0.3333333333333333, 0, 0.0, False)]},\n",
    " 2: {0: [(0.3333333333333333, 2, 0.0, False),\n",
    "   (0.3333333333333333, 1, 0.0, False),\n",
    "   (0.3333333333333333, 6, 0.0, False)],\n",
    "  1: [(0.3333333333333333, 1, 0.0, False),\n",
    "   (0.3333333333333333, 6, 0.0, False),\n",
    "   (0.3333333333333333, 3, 0.0, False)],\n",
    "  2: [(0.3333333333333333, 6, 0.0, False),\n",
    "   (0.3333333333333333, 3, 0.0, False),\n",
    "   (0.3333333333333333, 2, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 3, 0.0, False),\n",
    "   (0.3333333333333333, 2, 0.0, False),\n",
    "   (0.3333333333333333, 1, 0.0, False)]},\n",
    " 3: {0: [(0.3333333333333333, 3, 0.0, False),\n",
    "   (0.3333333333333333, 2, 0.0, False),\n",
    "   (0.3333333333333333, 7, 0.0, True)],\n",
    "  1: [(0.3333333333333333, 2, 0.0, False),\n",
    "   (0.3333333333333333, 7, 0.0, True),\n",
    "   (0.3333333333333333, 3, 0.0, False)],\n",
    "  2: [(0.3333333333333333, 7, 0.0, True),\n",
    "   (0.3333333333333333, 3, 0.0, False),\n",
    "   (0.3333333333333333, 3, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 3, 0.0, False),\n",
    "   (0.3333333333333333, 3, 0.0, False),\n",
    "   (0.3333333333333333, 2, 0.0, False)]},\n",
    " 4: {0: [(0.3333333333333333, 0, 0.0, False),\n",
    "   (0.3333333333333333, 4, 0.0, False),\n",
    "   (0.3333333333333333, 8, 0.0, False)],\n",
    "  1: [(0.3333333333333333, 4, 0.0, False),\n",
    "   (0.3333333333333333, 8, 0.0, False),\n",
    "   (0.3333333333333333, 5, 0.0, True)],\n",
    "  2: [(0.3333333333333333, 8, 0.0, False),\n",
    "   (0.3333333333333333, 5, 0.0, True),\n",
    "   (0.3333333333333333, 0, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 5, 0.0, True),\n",
    "   (0.3333333333333333, 0, 0.0, False),\n",
    "   (0.3333333333333333, 4, 0.0, False)]},\n",
    " 5: {0: [(1.0, 5, 0, True)],\n",
    "  1: [(1.0, 5, 0, True)],\n",
    "  2: [(1.0, 5, 0, True)],\n",
    "  3: [(1.0, 5, 0, True)]},\n",
    " 6: {0: [(0.3333333333333333, 2, 0.0, False),\n",
    "   (0.3333333333333333, 5, 0.0, True),\n",
    "   (0.3333333333333333, 10, 0.0, False)],\n",
    "  1: [(0.3333333333333333, 5, 0.0, True),\n",
    "   (0.3333333333333333, 10, 0.0, False),\n",
    "   (0.3333333333333333, 7, 0.0, True)],\n",
    "  2: [(0.3333333333333333, 10, 0.0, False),\n",
    "   (0.3333333333333333, 7, 0.0, True),\n",
    "   (0.3333333333333333, 2, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 7, 0.0, True),\n",
    "   (0.3333333333333333, 2, 0.0, False),\n",
    "   (0.3333333333333333, 5, 0.0, True)]},\n",
    " 7: {0: [(1.0, 7, 0, True)],\n",
    "  1: [(1.0, 7, 0, True)],\n",
    "  2: [(1.0, 7, 0, True)],\n",
    "  3: [(1.0, 7, 0, True)]},\n",
    " 8: {0: [(0.3333333333333333, 4, 0.0, False),\n",
    "   (0.3333333333333333, 8, 0.0, False),\n",
    "   (0.3333333333333333, 12, 0.0, True)],\n",
    "  1: [(0.3333333333333333, 8, 0.0, False),\n",
    "   (0.3333333333333333, 12, 0.0, True),\n",
    "   (0.3333333333333333, 9, 0.0, False)],\n",
    "  2: [(0.3333333333333333, 12, 0.0, True),\n",
    "   (0.3333333333333333, 9, 0.0, False),\n",
    "   (0.3333333333333333, 4, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 9, 0.0, False),\n",
    "   (0.3333333333333333, 4, 0.0, False),\n",
    "   (0.3333333333333333, 8, 0.0, False)]},\n",
    " 9: {0: [(0.3333333333333333, 5, 0.0, True),\n",
    "   (0.3333333333333333, 8, 0.0, False),\n",
    "   (0.3333333333333333, 13, 0.0, False)],\n",
    "  1: [(0.3333333333333333, 8, 0.0, False),\n",
    "   (0.3333333333333333, 13, 0.0, False),\n",
    "   (0.3333333333333333, 10, 0.0, False)],\n",
    "  2: [(0.3333333333333333, 13, 0.0, False),\n",
    "   (0.3333333333333333, 10, 0.0, False),\n",
    "   (0.3333333333333333, 5, 0.0, True)],\n",
    "  3: [(0.3333333333333333, 10, 0.0, False),\n",
    "   (0.3333333333333333, 5, 0.0, True),\n",
    "   (0.3333333333333333, 8, 0.0, False)]},\n",
    " 10: {0: [(0.3333333333333333, 6, 0.0, False),\n",
    "   (0.3333333333333333, 9, 0.0, False),\n",
    "   (0.3333333333333333, 14, 0.0, False)],\n",
    "  1: [(0.3333333333333333, 9, 0.0, False),\n",
    "   (0.3333333333333333, 14, 0.0, False),\n",
    "   (0.3333333333333333, 11, 0.0, True)],\n",
    "  2: [(0.3333333333333333, 14, 0.0, False),\n",
    "   (0.3333333333333333, 11, 0.0, True),\n",
    "   (0.3333333333333333, 6, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 11, 0.0, True),\n",
    "   (0.3333333333333333, 6, 0.0, False),\n",
    "   (0.3333333333333333, 9, 0.0, False)]},\n",
    " 11: {0: [(1.0, 11, 0, True)],\n",
    "  1: [(1.0, 11, 0, True)],\n",
    "  2: [(1.0, 11, 0, True)],\n",
    "  3: [(1.0, 11, 0, True)]},\n",
    " 12: {0: [(1.0, 12, 0, True)],\n",
    "  1: [(1.0, 12, 0, True)],\n",
    "  2: [(1.0, 12, 0, True)],\n",
    "  3: [(1.0, 12, 0, True)]},\n",
    " 13: {0: [(0.3333333333333333, 9, 0.0, False),\n",
    "   (0.3333333333333333, 12, 0.0, True),\n",
    "   (0.3333333333333333, 13, 0.0, False)],\n",
    "  1: [(0.3333333333333333, 12, 0.0, True),\n",
    "   (0.3333333333333333, 13, 0.0, False),\n",
    "   (0.3333333333333333, 14, 0.0, False)],\n",
    "  2: [(0.3333333333333333, 13, 0.0, False),\n",
    "   (0.3333333333333333, 14, 0.0, False),\n",
    "   (0.3333333333333333, 9, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 14, 0.0, False),\n",
    "   (0.3333333333333333, 9, 0.0, False),\n",
    "   (0.3333333333333333, 12, 0.0, True)]},\n",
    " 14: {0: [(0.3333333333333333, 10, 0.0, False),\n",
    "   (0.3333333333333333, 13, 0.0, False),\n",
    "   (0.3333333333333333, 14, 0.0, False)],\n",
    "  1: [(0.3333333333333333, 13, 0.0, False),\n",
    "   (0.3333333333333333, 14, 0.0, False),\n",
    "   (0.3333333333333333, 15, 1.0, True)],\n",
    "  2: [(0.3333333333333333, 14, 0.0, False),\n",
    "   (0.3333333333333333, 15, 1.0, True),\n",
    "   (0.3333333333333333, 10, 0.0, False)],\n",
    "  3: [(0.3333333333333333, 15, 1.0, True),\n",
    "   (0.3333333333333333, 10, 0.0, False),\n",
    "   (0.3333333333333333, 13, 0.0, False)]},\n",
    " 15: {0: [(1.0, 15, 0, True)],\n",
    "  1: [(1.0, 15, 0, True)],\n",
    "  2: [(1.0, 15, 0, True)],\n",
    "  3: [(1.0, 15, 0, True)]}}\n",
    "\n",
    "# import gym\n",
    "# P = gym.make('FrozenLake-v0').env.P"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This code snippet introduces and executes policy evaluation. Remember, policy evaluation takes in a policy, an MDP and it calculates the true state-value function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([0.02134932, 0.00914971, 0.00980529, 0.00735397, 0.04066537,\n",
       "       0.        , 0.01618064, 0.        , 0.07353654, 0.13091989,\n",
       "       0.04413016, 0.        , 0.        , 0.31873294, 0.61279031,\n",
       "       0.        ])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "# we first define the policy we will evaluate\n",
    "LEFT, DOWN, RIGHT, UP = range(4)\n",
    "pi = {\n",
    "    0:RIGHT, 1:LEFT, 2:DOWN, 3:UP,\n",
    "    4:LEFT, 5:LEFT, 6:RIGHT, 7:LEFT,\n",
    "    8:UP, 9:DOWN, 10:UP, 11:LEFT,\n",
    "    12:LEFT, 13:RIGHT, 14:DOWN, 15:LEFT\n",
    "}\n",
    "\n",
    "# Now we define the policy evaluation method\n",
    "def policy_evaluation(pi, P, gamma=0.9, theta=1e-10):\n",
    "    V = np.zeros(len(pi))\n",
    "    \n",
    "    while True:\n",
    "        max_delta = 0\n",
    "        old_V = V.copy()\n",
    "\n",
    "        for s in range(len(P)):\n",
    "            V[s] = 0\n",
    "            for prob, new_state, reward, done in P[s][pi[s]]:\n",
    "                if done:\n",
    "                    value = reward\n",
    "                else:\n",
    "                    value = reward + gamma * old_V[new_state]\n",
    "                V[s] += prob * value\n",
    "            max_delta = max(max_delta, abs(old_V[s] - V[s]))\n",
    "        if max_delta < theta:\n",
    "            break\n",
    "    return V.copy()\n",
    "\n",
    "# Run it and display the results\n",
    "V = policy_evaluation(pi, P)\n",
    "V"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we introduce the policy improvement algorithm. Remember, policy improvement takes in a policy, a state-value function and an MDP and it returns a better, improved policy (unless it is an optimal policy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{0: 0,\n",
       " 1: 3,\n",
       " 2: 0,\n",
       " 3: 3,\n",
       " 4: 0,\n",
       " 5: 0,\n",
       " 6: 0,\n",
       " 7: 0,\n",
       " 8: 3,\n",
       " 9: 1,\n",
       " 10: 0,\n",
       " 11: 0,\n",
       " 12: 0,\n",
       " 13: 2,\n",
       " 14: 1,\n",
       " 15: 0}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# define the algorithm\n",
    "def policy_improvement(pi, V, P, gamma=0.9):\n",
    "\n",
    "    for s in range(len(V)):\n",
    "        Qs = np.zeros(len(P[0]), dtype=np.float64)\n",
    "        for a in range(len(P[s])):\n",
    "            for prob, new_state, reward, done in P[s][a]:\n",
    "                if done:\n",
    "                    value = reward\n",
    "                else:\n",
    "                    value = reward + gamma * V[new_state]\n",
    "                Qs[a] += prob * value\n",
    "        pi[s] = np.argmax(Qs)\n",
    "    return pi.copy()\n",
    "\n",
    "# run and return the results\n",
    "new_pi = policy_improvement(pi, V, P)\n",
    "new_pi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now let's print the improved policy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| 00 < | 01 ^ | 02 < | 03 ^ |\n",
      "| 04 < |      | 06 < |      |\n",
      "| 08 ^ | 09 v | 10 < |      |\n",
      "|      | 13 > | 14 v |      |\n"
     ]
    }
   ],
   "source": [
    "print_policy(new_pi, P)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Having the policy evaluation and policy improvement functions, we can combine these methods to create an algorithm that goes from a random policy and iterates between those two methods to obtain an optimal policy (and optimal state-value function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0.0688909 , 0.06141457, 0.07440976, 0.05580732, 0.09185454,\n",
       "        0.        , 0.11220821, 0.        , 0.14543635, 0.24749695,\n",
       "        0.29961759, 0.        , 0.        , 0.3799359 , 0.63902015,\n",
       "        0.        ]),\n",
       " {0: 0,\n",
       "  1: 3,\n",
       "  2: 0,\n",
       "  3: 3,\n",
       "  4: 0,\n",
       "  5: 0,\n",
       "  6: 0,\n",
       "  7: 0,\n",
       "  8: 3,\n",
       "  9: 1,\n",
       "  10: 0,\n",
       "  11: 0,\n",
       "  12: 0,\n",
       "  13: 2,\n",
       "  14: 1,\n",
       "  15: 0})"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# lets define the algorithm\n",
    "def policy_iteration(P, gamma=0.9):\n",
    "\n",
    "    random_actions = np.random.choice(tuple(P[0].keys()), len(P))\n",
    "    pi = {s:a for s, a in enumerate(random_actions)}\n",
    "\n",
    "    while True:\n",
    "        old_pi = pi.copy()\n",
    "\n",
    "        V = policy_evaluation(pi, P, gamma)\n",
    "        pi = policy_improvement(pi, V, P, gamma)\n",
    "\n",
    "        if old_pi == pi:\n",
    "            break\n",
    "\n",
    "    return V, pi\n",
    "\n",
    "# we call it passing an MDP, and return the \n",
    "# optimal policy and state-value function\n",
    "V_best, pi_best = policy_iteration(P)\n",
    "V_best, pi_best"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's not print the optimal policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| 00 < | 01 ^ | 02 < | 03 ^ |\n",
      "| 04 < |      | 06 < |      |\n",
      "| 08 ^ | 09 v | 10 < |      |\n",
      "|      | 13 > | 14 v |      |\n"
     ]
    }
   ],
   "source": [
    "print_policy(pi_best, P)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now implement value iteration. This algorithm is similar to policy iteration but as if the policy evaluation step was stopped before termination. This is often called truncated policy evaluation. We improve the policy just as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([0.06889091, 0.06141457, 0.07440976, 0.05580732, 0.09185454,\n",
       "        0.        , 0.11220821, 0.        , 0.14543636, 0.24749695,\n",
       "        0.29961759, 0.        , 0.        , 0.3799359 , 0.63902015,\n",
       "        0.        ]),\n",
       " {0: 0,\n",
       "  1: 3,\n",
       "  2: 0,\n",
       "  3: 3,\n",
       "  4: 0,\n",
       "  5: 0,\n",
       "  6: 0,\n",
       "  7: 0,\n",
       "  8: 3,\n",
       "  9: 1,\n",
       "  10: 0,\n",
       "  11: 0,\n",
       "  12: 0,\n",
       "  13: 2,\n",
       "  14: 1,\n",
       "  15: 0})"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# define the value iteration algorithm\n",
    "def value_iteration(P, gamma=0.9, theta = 1e-10):\n",
    "\n",
    "    V = np.random.random(len(P))\n",
    "    while True:\n",
    "        max_delta = 0\n",
    "        Q = np.zeros((len(P), len(P[0])), dtype=np.float64)\n",
    "\n",
    "        for s in range(len(P)):\n",
    "            v = V[s]\n",
    "            for a in range(len(P[s])):\n",
    "                for prob, new_state, reward, done in P[s][a]:\n",
    "                    if done:\n",
    "                        value = reward\n",
    "                    else:\n",
    "                        value = reward + gamma * V[new_state]\n",
    "                    Q[s][a] += prob * value\n",
    "            V[s] = np.max(Q[s])\n",
    "            max_delta = max(max_delta, abs(v - V[s]))\n",
    "        if max_delta < theta:\n",
    "            break\n",
    "    pi = {s:a for s, a in enumerate(np.argmax(Q, axis=1))}\n",
    "    return V, pi\n",
    "\n",
    "# run it and return the optimal policy and state-value function\n",
    "V_best, pi_best = value_iteration(P, gamma=0.9)\n",
    "V_best, pi_best"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Print the policy and compare to the one above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "| 00 < | 01 ^ | 02 < | 03 ^ |\n",
      "| 04 < |      | 06 < |      |\n",
      "| 08 ^ | 09 v | 10 < |      |\n",
      "|      | 13 > | 14 v |      |\n"
     ]
    }
   ],
   "source": [
    "print_policy(pi_best, P)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This snippet will change the MDP. I recommend you change the values and run the algorithms."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change reward function\n",
    "reward_goal, reward_holes, reward_others = 1, -1, -0.01\n",
    "goal, hole = 15, [5, 7, 11, 12]\n",
    "for s in range(len(P)):\n",
    "    for a in range(len(P[s])):\n",
    "        for t in range(len(P[s][a])):\n",
    "            values = list(P[s][a][t])\n",
    "            if values[1] == goal:\n",
    "                values[2] = 1\n",
    "                values[3] = False\n",
    "            elif values[1] in hole:\n",
    "                values[2] = -1\n",
    "                values[3] = False\n",
    "            else:\n",
    "                values[2] = -0.01\n",
    "                values[3] = False\n",
    "            if s in hole or s == goal:\n",
    "                values[2] = 0\n",
    "                values[3] = True\n",
    "            P[s][a][t] = tuple(values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same with the snippet below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# change transition function\n",
    "prob_action, drift_right, drift_left = 0.8, 0.1, 0.1\n",
    "for s in range(len(P)):\n",
    "    for a in range(len(P[s])):\n",
    "        for t in range(len(P[s][a])):\n",
    "            if P[s][a][t][0] == 1.0:\n",
    "                continue\n",
    "            values = list(P[s][a][t])\n",
    "            if t == 0:\n",
    "                values[0] = drift_left\n",
    "            elif t == 1:\n",
    "                values[0] = prob_action\n",
    "            elif t == 2:\n",
    "                values[0] = drift_right\n",
    "            P[s][a][t] = tuple(values)"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}