Section_3_policy_iteration_complete.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "b9jmpP4J6VbY"
      },
      "source": [
        "<div style=\"text-align:center\">\n",
        "    <h1>\n",
        "        Policy Iteration\n",
        "    </h1>\n",
        "</div>\n",
        "<br>\n",
        "\n",
        "<div style=\"text-align:center\">\n",
        "    <p>\n",
        "        In this notebook we are going to look at a dynamic programming algorithm called policy iteration. In it, we will iteratively interleave two processes: policy evaluation and policy improvement, until the optimal policy and state values are found.\n",
        "    </p>\n",
        "</div>\n",
        "\n",
        "<br>"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# @title Setup code (not important) - Run this cell by pressing \"Shift + Enter\"\n",
        "\n",
        "\n",
        "\n",
        "!pip install -qq gym==0.23.0\n",
        "\n",
        "\n",
        "from typing import Tuple, Dict, Optional, Iterable, Callable\n",
        "\n",
        "import numpy as np\n",
        "import seaborn as sns\n",
        "import matplotlib\n",
        "from matplotlib import animation\n",
        "\n",
        "from IPython.display import HTML\n",
        "\n",
        "import gym\n",
        "from gym import spaces\n",
        "from gym.error import DependencyNotInstalled\n",
        "\n",
        "import pygame\n",
        "from pygame import gfxdraw\n",
        "\n",
        "\n",
        "class Maze(gym.Env):\n",
        "\n",
        "    def __init__(self, exploring_starts: bool = False,\n",
        "                 shaped_rewards: bool = False, size: int = 5) -> None:\n",
        "        super().__init__()\n",
        "        self.exploring_starts = exploring_starts\n",
        "        self.shaped_rewards = shaped_rewards\n",
        "        self.state = (size - 1, size - 1)\n",
        "        self.goal = (size - 1, size - 1)\n",
        "        self.maze = self._create_maze(size=size)\n",
        "        self.distances = self._compute_distances(self.goal, self.maze)\n",
        "        self.action_space = spaces.Discrete(n=4)\n",
        "        self.action_space.action_meanings = {0: 'UP', 1: 'RIGHT', 2: 'DOWN', 3: \"LEFT\"}\n",
        "        self.observation_space = spaces.MultiDiscrete([size, size])\n",
        "\n",
        "        self.screen = None\n",
        "        self.agent_transform = None\n",
        "\n",
        "    def step(self, action: int) -> Tuple[Tuple[int, int], float, bool, Dict]:\n",
        "        reward = self.compute_reward(self.state, action)\n",
        "        self.state = self._get_next_state(self.state, action)\n",
        "        done = self.state == self.goal\n",
        "        info = {}\n",
        "        return self.state, reward, done, info\n",
        "\n",
        "    def reset(self) -> Tuple[int, int]:\n",
        "        if self.exploring_starts:\n",
        "            while self.state == self.goal:\n",
        "                self.state = tuple(self.observation_space.sample())\n",
        "        else:\n",
        "            self.state = (0, 0)\n",
        "        return self.state\n",
        "\n",
        "    def render(self, mode: str = 'human') -> Optional[np.ndarray]:\n",
        "        assert mode in ['human', 'rgb_array']\n",
        "\n",
        "        screen_size = 600\n",
        "        scale = screen_size / 5\n",
        "\n",
        "        if self.screen is None:\n",
        "            pygame.init()\n",
        "            self.screen = pygame.Surface((screen_size, screen_size))\n",
        "\n",
        "        surf = pygame.Surface((screen_size, screen_size))\n",
        "        surf.fill((22, 36, 71))\n",
        "\n",
        "\n",
        "        for row in range(5):\n",
        "            for col in range(5):\n",
        "\n",
        "                state = (row, col)\n",
        "                for next_state in [(row + 1, col), (row - 1, col), (row, col + 1), (row, col - 1)]:\n",
        "                    if next_state not in self.maze[state]:\n",
        "\n",
        "                        # Add the geometry of the edges and walls (i.e. the boundaries between\n",
        "                        # adjacent squares that are not connected).\n",
        "                        row_diff, col_diff = np.subtract(next_state, state)\n",
        "                        left = (col + (col_diff > 0)) * scale - 2 * (col_diff != 0)\n",
        "                        right = ((col + 1) - (col_diff < 0)) * scale + 2 * (col_diff != 0)\n",
        "                        top = (5 - (row + (row_diff > 0))) * scale - 2 * (row_diff != 0)\n",
        "                        bottom = (5 - ((row + 1) - (row_diff < 0))) * scale + 2 * (row_diff != 0)\n",
        "\n",
        "                        gfxdraw.filled_polygon(surf, [(left, bottom), (left, top), (right, top), (right, bottom)], (255, 255, 255))\n",
        "\n",
        "        # Add the geometry of the goal square to the viewer.\n",
        "        left, right, top, bottom = scale * 4 + 10, scale * 5 - 10, scale - 10, 10\n",
        "        gfxdraw.filled_polygon(surf, [(left, bottom), (left, top), (right, top), (right, bottom)], (40, 199, 172))\n",
        "\n",
        "        # Add the geometry of the agent to the viewer.\n",
        "        agent_row = int(screen_size - scale * (self.state[0] + .5))\n",
        "        agent_col = int(scale * (self.state[1] + .5))\n",
        "        gfxdraw.filled_circle(surf, agent_col, agent_row, int(scale * .6 / 2), (228, 63, 90))\n",
        "\n",
        "        surf = pygame.transform.flip(surf, False, True)\n",
        "        self.screen.blit(surf, (0, 0))\n",
        "\n",
        "        return np.transpose(\n",
        "                np.array(pygame.surfarray.pixels3d(self.screen)), axes=(1, 0, 2)\n",
        "            )\n",
        "\n",
        "    def close(self) -> None:\n",
        "        if self.screen is not None:\n",
        "            pygame.display.quit()\n",
        "            pygame.quit()\n",
        "            self.screen = None\n",
        "\n",
        "    def compute_reward(self, state: Tuple[int, int], action: int) -> float:\n",
        "        next_state = self._get_next_state(state, action)\n",
        "        if self.shaped_rewards:\n",
        "            return - (self.distances[next_state] / self.distances.max())\n",
        "        return - float(state != self.goal)\n",
        "\n",
        "    def simulate_step(self, state: Tuple[int, int], action: int):\n",
        "        reward = self.compute_reward(state, action)\n",
        "        next_state = self._get_next_state(state, action)\n",
        "        done = next_state == self.goal\n",
        "        info = {}\n",
        "        return next_state, reward, done, info\n",
        "\n",
        "    def _get_next_state(self, state: Tuple[int, int], action: int) -> Tuple[int, int]:\n",
        "        if action == 0:\n",
        "            next_state = (state[0] - 1, state[1])\n",
        "        elif action == 1:\n",
        "            next_state = (state[0], state[1] + 1)\n",
        "        elif action == 2:\n",
        "            next_state = (state[0] + 1, state[1])\n",
        "        elif action == 3:\n",
        "            next_state = (state[0], state[1] - 1)\n",
        "        else:\n",
        "            raise ValueError(\"Action value not supported:\", action)\n",
        "        if next_state in self.maze[state]:\n",
        "            return next_state\n",
        "        return state\n",
        "\n",
        "    @staticmethod\n",
        "    def _create_maze(size: int) -> Dict[Tuple[int, int], Iterable[Tuple[int, int]]]:\n",
        "        maze = {(row, col): [(row - 1, col), (row + 1, col), (row, col - 1), (row, col + 1)]\n",
        "                for row in range(size) for col in range(size)}\n",
        "\n",
        "        left_edges = [[(row, 0), (row, -1)] for row in range(size)]\n",
        "        right_edges = [[(row, size - 1), (row, size)] for row in range(size)]\n",
        "        upper_edges = [[(0, col), (-1, col)] for col in range(size)]\n",
        "        lower_edges = [[(size - 1, col), (size, col)] for col in range(size)]\n",
        "        walls = [\n",
        "            [(1, 0), (1, 1)], [(2, 0), (2, 1)], [(3, 0), (3, 1)],\n",
        "            [(1, 1), (1, 2)], [(2, 1), (2, 2)], [(3, 1), (3, 2)],\n",
        "            [(3, 1), (4, 1)], [(0, 2), (1, 2)], [(1, 2), (1, 3)],\n",
        "            [(2, 2), (3, 2)], [(2, 3), (3, 3)], [(2, 4), (3, 4)],\n",
        "            [(4, 2), (4, 3)], [(1, 3), (1, 4)], [(2, 3), (2, 4)],\n",
        "        ]\n",
        "\n",
        "        obstacles = upper_edges + lower_edges + left_edges + right_edges + walls\n",
        "\n",
        "        for src, dst in obstacles:\n",
        "            maze[src].remove(dst)\n",
        "\n",
        "            if dst in maze:\n",
        "                maze[dst].remove(src)\n",
        "\n",
        "        return maze\n",
        "\n",
        "    @staticmethod\n",
        "    def _compute_distances(goal: Tuple[int, int],\n",
        "                           maze: Dict[Tuple[int, int], Iterable[Tuple[int, int]]]) -> np.ndarray:\n",
        "        distances = np.full((5, 5), np.inf)\n",
        "        visited = set()\n",
        "        distances[goal] = 0.\n",
        "\n",
        "        while visited != set(maze):\n",
        "            sorted_dst = [(v // 5, v % 5) for v in distances.argsort(axis=None)]\n",
        "            closest = next(x for x in sorted_dst if x not in visited)\n",
        "            visited.add(closest)\n",
        "\n",
        "            for neighbour in maze[closest]:\n",
        "                distances[neighbour] = min(distances[neighbour], distances[closest] + 1)\n",
        "        return distances\n",
        "\n",
        "\n",
        "def plot_policy(probs_or_qvals, frame, action_meanings=None):\n",
        "    if action_meanings is None:\n",
        "        action_meanings = {0: 'U', 1: 'R', 2: 'D', 3: 'L'}\n",
        "    fig, axes = plt.subplots(1, 2, figsize=(8, 4))\n",
        "    max_prob_actions = probs_or_qvals.argmax(axis=-1)\n",
        "    probs_copy = max_prob_actions.copy().astype(object)\n",
        "    for key in action_meanings:\n",
        "        probs_copy[probs_copy == key] = action_meanings[key]\n",
        "    sns.heatmap(max_prob_actions, annot=probs_copy, fmt='', cbar=False, cmap='coolwarm',\n",
        "                annot_kws={'weight': 'bold', 'size': 12}, linewidths=2, ax=axes[0])\n",
        "    axes[1].imshow(frame)\n",
        "    axes[0].axis('off')\n",
        "    axes[1].axis('off')\n",
        "    plt.suptitle(\"Policy\", size=18)\n",
        "    plt.tight_layout()\n",
        "\n",
        "\n",
        "def plot_values(state_values, frame):\n",
        "    f, axes = plt.subplots(1, 2, figsize=(10, 4))\n",
        "    sns.heatmap(state_values, annot=True, fmt=\".2f\", cmap='coolwarm',\n",
        "                annot_kws={'weight': 'bold', 'size': 12}, linewidths=2, ax=axes[0])\n",
        "    axes[1].imshow(frame)\n",
        "    axes[0].axis('off')\n",
        "    axes[1].axis('off')\n",
        "    plt.tight_layout()\n",
        "\n",
        "\n",
        "def display_video(frames):\n",
        "    # Copied from: https://colab.research.google.com/github/deepmind/dm_control/blob/master/tutorial.ipynb\n",
        "    orig_backend = matplotlib.get_backend()\n",
        "    matplotlib.use('Agg')\n",
        "    fig, ax = plt.subplots(1, 1, figsize=(5, 5))\n",
        "    matplotlib.use(orig_backend)\n",
        "    ax.set_axis_off()\n",
        "    ax.set_aspect('equal')\n",
        "    ax.set_position([0, 0, 1, 1])\n",
        "    im = ax.imshow(frames[0])\n",
        "    def update(frame):\n",
        "        im.set_data(frame)\n",
        "        return [im]\n",
        "    anim = animation.FuncAnimation(fig=fig, func=update, frames=frames,\n",
        "                                    interval=50, blit=True, repeat=False)\n",
        "    return HTML(anim.to_html5_video())\n",
        "\n",
        "\n",
        "def test_agent(environment, policy, episodes=10):\n",
        "    frames = []\n",
        "    for episode in range(episodes):\n",
        "        state = env.reset()\n",
        "        done = False\n",
        "        frames.append(env.render(mode=\"rgb_array\"))\n",
        "\n",
        "        while not done:\n",
        "            p = policy(state)\n",
        "            if isinstance(p, np.ndarray):\n",
        "                action = np.random.choice(4, p=p)\n",
        "            else:\n",
        "                action = p\n",
        "            next_state, reward, done, extra_info = env.step(action)\n",
        "            img = env.render(mode=\"rgb_array\")\n",
        "            frames.append(img)\n",
        "            state = next_state\n",
        "\n",
        "    return display_video(frames)\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "dnULhrB06fg1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j-pns0SL6Vbc"
      },
      "source": [
        "## Import the necessary software libraries:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MiP1siGs6Vbc"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NM6THrGl6Vbd"
      },
      "source": [
        "## Initialize the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "q2ytq7oR6Vbe"
      },
      "outputs": [],
      "source": [
        "env = Maze()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "glvnH2SX6Vbe"
      },
      "outputs": [],
      "source": [
        "frame = env.render(mode='rgb_array')\n",
        "plt.axis('off')\n",
        "plt.imshow(frame)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YrpWHqxu6Vbe"
      },
      "outputs": [],
      "source": [
        "print(f\"Observation space shape: {env.observation_space.nvec}\")\n",
        "print(f\"Number of actions: {env.action_space.n}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rx947Ii46Vbf"
      },
      "source": [
        "## Define the policy $\\pi(\\cdot|s)$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mj-DEzgD6Vbf"
      },
      "source": [
        "#### Create the policy $\\pi(\\cdot|s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kivjfd-t6Vbf"
      },
      "outputs": [],
      "source": [
        "policy_probs = np.full((5, 5, 4), 0.25)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sunXzggw6Vbf"
      },
      "outputs": [],
      "source": [
        "def policy(state):\n",
        "    return policy_probs[state]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PGrxXujC6Vbg"
      },
      "source": [
        "#### Test the policy with state (0, 0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VlXckfZ96Vbg"
      },
      "outputs": [],
      "source": [
        "action_probabilities = policy((0,0))\n",
        "for action, prob in zip(range(4), action_probabilities):\n",
        "    print(f\"Probability of taking action {action}: {prob}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "THcQNGHa6Vbg"
      },
      "source": [
        "#### See how the random policy does in the maze"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gsSYs5Ry6Vbg"
      },
      "outputs": [],
      "source": [
        "test_agent(env, policy, episodes=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tF2UjZft6Vbg"
      },
      "source": [
        "#### Plot the policy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HJdUvc1T6Vbg"
      },
      "outputs": [],
      "source": [
        "plot_policy(policy_probs, frame)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "btCVXrhJ6Vbg"
      },
      "source": [
        "## Define value table $V(s)$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "byOFEsaA6Vbg"
      },
      "source": [
        "#### Create the $V(s)$ table"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j2lX_33D6Vbg"
      },
      "outputs": [],
      "source": [
        "state_values = np.zeros(shape=(5,5))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D9_EfMnY6Vbh"
      },
      "source": [
        "#### Plot $V(s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TI5YZTSF6Vbh"
      },
      "outputs": [],
      "source": [
        "plot_values(state_values, frame)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qfiam4Ef6Vbh"
      },
      "source": [
        "## Implement the Policy Iteration algorithm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N_5XQChT6Vbh"
      },
      "outputs": [],
      "source": [
        "def policy_evaluation(policy_probs, state_values, theta=1e-6, gamma=0.99):\n",
        "    delta = float(\"inf\")\n",
        "\n",
        "    while delta > theta:\n",
        "        delta = 0\n",
        "\n",
        "        for row in range(5):\n",
        "            for col in range(5):\n",
        "                old_value = state_values[(row, col)]\n",
        "                new_value = 0\n",
        "                action_probabilities = policy_probs[(row, col)]\n",
        "\n",
        "                for action, prob in enumerate(action_probabilities):\n",
        "                    next_state, reward, _, _ = env.simulate_step((row, col), action)\n",
        "                    new_value += prob * (reward + gamma * state_values[next_state])\n",
        "\n",
        "                state_values[(row, col)] = new_value\n",
        "\n",
        "                delta = max(delta, abs(old_value - new_value))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JRj7ZxuO6Vbh"
      },
      "outputs": [],
      "source": [
        "def policy_improvement(policy_probs, state_values, gamma=0.99):\n",
        "\n",
        "    policy_stable = True\n",
        "    for row in range(5):\n",
        "        for col in range(5):\n",
        "            old_action = policy_probs[(row, col)].argmax()\n",
        "\n",
        "            new_action = None\n",
        "            max_qsa = float(\"-inf\")\n",
        "\n",
        "            for action in range(4):\n",
        "                next_state, reward, _, _ = env.simulate_step((row, col), action)\n",
        "                qsa = reward + gamma * state_values[next_state]\n",
        "                if qsa > max_qsa:\n",
        "                    max_qsa = qsa\n",
        "                    new_action = action\n",
        "\n",
        "            action_probs = np.zeros(4)\n",
        "            action_probs[new_action] = 1.\n",
        "            policy_probs[(row, col)] = action_probs\n",
        "\n",
        "            if new_action != old_action:\n",
        "                policy_stable = False\n",
        "\n",
        "    return policy_stable"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-0F0aB_f6Vbh"
      },
      "outputs": [],
      "source": [
        "def policy_iteration(policy_probs, state_values, theta=1e-6, gamma=0.99):\n",
        "    policy_stable = False\n",
        "\n",
        "    while not policy_stable:\n",
        "\n",
        "        policy_evaluation(policy_probs, state_values, theta, gamma)\n",
        "\n",
        "        policy_stable = policy_improvement(policy_probs, state_values, gamma)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0eZr-1ko6Vbh"
      },
      "outputs": [],
      "source": [
        "policy_iteration(policy_probs, state_values)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zHHT9iLY6Vbh"
      },
      "source": [
        "## Show results"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "Show resulting value table V(s)"
      ],
      "metadata": {
        "id": "4bVf7AMd7ajW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plot_values(state_values, frame)"
      ],
      "metadata": {
        "id": "Y1GDXzvt7atl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Show resulting policy $\\pi(\\cdot|s)$"
      ],
      "metadata": {
        "id": "1OdMtp1n7hBU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "plot_policy(policy_probs, frame)"
      ],
      "metadata": {
        "id": "y4u6alml7hHQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "onu9Yu8D6Vbh"
      },
      "source": [
        "#### Test the resulting agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U_8oUWxl6Vbh"
      },
      "outputs": [],
      "source": [
        "test_agent(env, policy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bixFf0kp6Vbh"
      },
      "source": [
        "## Resources"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NIU61hO36Vbh"
      },
      "source": [
        "[[1] Reinforcement Learning: An Introduction. Ch. 4: Dynamic Programming](https://web.stanford.edu/class/psych209/Readings/SuttonBartoIPRLBook2ndEd.pdf)"
      ]
    }
  ],
  "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.8.5"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}