Section_4_on_policy_control_complete.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "FegQfsZN8Pap"
      },
      "source": [
        "<div style=\"text-align:center\">\n",
        "    <h1>\n",
        "        On-policy Monte Carlo Control\n",
        "    </h1>\n",
        "</div>\n",
        "<br>\n",
        "\n",
        "<div style=\"text-align:center\">\n",
        "    <p>\n",
        "        In this notebook we are going to implement one of the two major strategies that exist when learning a policy by interacting with the environment, called on-policy learning. The agent will perform the task from start to finish and based on the sample experience generated, update its estimates of the q-values of each state-action pair $Q(s,a)$.\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",
        "\n",
        "def plot_action_values(action_values):\n",
        "\n",
        "    text_positions = [\n",
        "        [(0.35, 4.75), (1.35, 4.75), (2.35, 4.75), (3.35, 4.75), (4.35, 4.75),\n",
        "         (0.35, 3.75), (1.35, 3.75), (2.35, 3.75), (3.35, 3.75), (4.35, 3.75),\n",
        "         (0.35, 2.75), (1.35, 2.75), (2.35, 2.75), (3.35, 2.75), (4.35, 2.75),\n",
        "         (0.35, 1.75), (1.35, 1.75), (2.35, 1.75), (3.35, 1.75), (4.35, 1.75),\n",
        "         (0.35, 0.75), (1.35, 0.75), (2.35, 0.75), (3.35, 0.75), (4.35, 0.75)],\n",
        "        [(0.6, 4.45), (1.6, 4.45), (2.6, 4.45), (3.6, 4.45), (4.6, 4.45),\n",
        "         (0.6, 3.45), (1.6, 3.45), (2.6, 3.45), (3.6, 3.45), (4.6, 3.45),\n",
        "         (0.6, 2.45), (1.6, 2.45), (2.6, 2.45), (3.6, 2.45), (4.6, 2.45),\n",
        "         (0.6, 1.45), (1.6, 1.45), (2.6, 1.45), (3.6, 1.45), (4.6, 1.45),\n",
        "         (0.6, 0.45), (1.6, 0.45), (2.6, 0.45), (3.6, 0.45), (4.6, 0.45)],\n",
        "        [(0.35, 4.15), (1.35, 4.15), (2.35, 4.15), (3.35, 4.15), (4.35, 4.15),\n",
        "         (0.35, 3.15), (1.35, 3.15), (2.35, 3.15), (3.35, 3.15), (4.35, 3.15),\n",
        "         (0.35, 2.15), (1.35, 2.15), (2.35, 2.15), (3.35, 2.15), (4.35, 2.15),\n",
        "         (0.35, 1.15), (1.35, 1.15), (2.35, 1.15), (3.35, 1.15), (4.35, 1.15),\n",
        "         (0.35, 0.15), (1.35, 0.15), (2.35, 0.15), (3.35, 0.15), (4.35, 0.15)],\n",
        "        [(0.05, 4.45), (1.05, 4.45), (2.05, 4.45), (3.05, 4.45), (4.05, 4.45),\n",
        "         (0.05, 3.45), (1.05, 3.45), (2.05, 3.45), (3.05, 3.45), (4.05, 3.45),\n",
        "         (0.05, 2.45), (1.05, 2.45), (2.05, 2.45), (3.05, 2.45), (4.05, 2.45),\n",
        "         (0.05, 1.45), (1.05, 1.45), (2.05, 1.45), (3.05, 1.45), (4.05, 1.45),\n",
        "         (0.05, 0.45), (1.05, 0.45), (2.05, 0.45), (3.05, 0.45), (4.05, 0.45)]]\n",
        "\n",
        "    fig, ax = plt.subplots(figsize=(7, 7))\n",
        "    tripcolor = quatromatrix(action_values, ax=ax,\n",
        "                             triplotkw={\"color\": \"k\", \"lw\": 1}, tripcolorkw={\"cmap\": \"coolwarm\"})\n",
        "    ax.margins(0)\n",
        "    ax.set_aspect(\"equal\")\n",
        "    fig.colorbar(tripcolor)\n",
        "\n",
        "    for j, av in enumerate(text_positions):\n",
        "        for i, (xi, yi) in enumerate(av):\n",
        "            plt.text(xi, yi, round(action_values[:, :, j].flatten()[i], 2), size=8, color=\"w\", weight=\"bold\")\n",
        "\n",
        "    plt.title(\"Action values Q(s,a)\", size=18)\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def quatromatrix(action_values, ax=None, triplotkw=None, tripcolorkw=None):\n",
        "    action_values = np.flipud(action_values)\n",
        "    n = 5\n",
        "    m = 5\n",
        "    a = np.array([[0, 0], [0, 1], [.5, .5], [1, 0], [1, 1]])\n",
        "    tr = np.array([[0, 1, 2], [0, 2, 3], [2, 3, 4], [1, 2, 4]])\n",
        "    A = np.zeros((n * m * 5, 2))\n",
        "    Tr = np.zeros((n * m * 4, 3))\n",
        "    for i in range(n):\n",
        "        for j in range(m):\n",
        "            k = i * m + j\n",
        "            A[k * 5:(k + 1) * 5, :] = np.c_[a[:, 0] + j, a[:, 1] + i]\n",
        "            Tr[k * 4:(k + 1) * 4, :] = tr + k * 5\n",
        "    C = np.c_[action_values[:, :, 3].flatten(), action_values[:, :, 2].flatten(),\n",
        "              action_values[:, :, 1].flatten(), action_values[:, :, 0].flatten()].flatten()\n",
        "\n",
        "    ax.triplot(A[:, 0], A[:, 1], Tr, **triplotkw)\n",
        "    tripcolor = ax.tripcolor(A[:, 0], A[:, 1], Tr, facecolors=C, **tripcolorkw)\n",
        "    return tripcolor\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "WQ3CjcEb9Fkx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "i1QHHcC38Pas"
      },
      "source": [
        "## Import the necessary software libraries:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ddSL9RtK8Pas"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aujcXnZg8Pat"
      },
      "source": [
        "## Initialize the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aPtijbjT8Pat"
      },
      "outputs": [],
      "source": [
        "env = Maze()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bEYRDn1Q8Pau"
      },
      "outputs": [],
      "source": [
        "frame = env.render(mode='rgb_array')\n",
        "plt.axis('off')\n",
        "plt.imshow(frame)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tWRgQvYn8Pau"
      },
      "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": "nYSxH6vb8Pau"
      },
      "source": [
        "## Define value table $Q(s, a)$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "thgvYa_S8Pav"
      },
      "source": [
        "#### Create the $Q(s, a)$ table"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fehyO_5p8Pav"
      },
      "outputs": [],
      "source": [
        "action_values = np.zeros(shape=(5, 5, 4))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aVNIyhiH8Pav"
      },
      "source": [
        "#### Plot Q(s, a)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZSvw1wgF8Pav"
      },
      "outputs": [],
      "source": [
        "plot_action_values(action_values)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ysJQJGXm8Pav"
      },
      "source": [
        "## Define the policy $\\pi(s)$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0lRYXyx68Pav"
      },
      "source": [
        "#### Create the policy $\\pi(s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zc2nSOTb8Pav"
      },
      "outputs": [],
      "source": [
        "def policy(state, epsilon=0.):\n",
        "    if np.random.random() < epsilon:\n",
        "        return np.random.randint(4)\n",
        "    else:\n",
        "        av = action_values[state]\n",
        "        return np.random.choice(np.flatnonzero(av == av.max()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V6y0E6nx8Pav"
      },
      "source": [
        "#### Test the policy with state (0, 0)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ck167K4M8Pav"
      },
      "outputs": [],
      "source": [
        "action = policy((0,0))\n",
        "print(f\"Action taken in state (0,0): {action}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "g2xH7qSZ8Pav"
      },
      "source": [
        "#### Plot the policy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": false,
        "id": "KzYsJsnO8Pav"
      },
      "outputs": [],
      "source": [
        "plot_policy(action_values, frame)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PqAqHkpk8Paw"
      },
      "source": [
        "## Implement the algorithm\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Os79GeR88Paw"
      },
      "outputs": [],
      "source": [
        "def on_policy_mc_control(policy, action_values, episodes, gamma=0.99, epsilon=0.2):\n",
        "\n",
        "    sa_returns = {}\n",
        "\n",
        "    for episode in range(1, episodes+1):\n",
        "        state = env.reset()\n",
        "        done = False\n",
        "        transitions = []\n",
        "\n",
        "        while not done:\n",
        "            action = policy(state, epsilon)\n",
        "            next_state, reward, done, _ = env.step(action)\n",
        "            transitions.append([state, action, reward])\n",
        "            state = next_state\n",
        "\n",
        "        G = 0\n",
        "        for state_t, action_t, reward_t in reversed(transitions):\n",
        "            G = reward_t + gamma * G\n",
        "\n",
        "            if not (state_t, action_t) in sa_returns:\n",
        "                sa_returns[(state_t, action_t)] = []\n",
        "            sa_returns[(state_t, action_t)].append(G)\n",
        "            action_values[state_t][action_t] = np.mean(sa_returns[(state_t, action_t)])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GMHevlFT8Paw"
      },
      "outputs": [],
      "source": [
        "on_policy_mc_control(policy, action_values, episodes=10000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VMNGJYQ38Paw"
      },
      "source": [
        "## Show results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rvgNcDkX8Paw"
      },
      "source": [
        "#### Show resulting value table $Q(s, a)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yIyUuFZE8Paw"
      },
      "outputs": [],
      "source": [
        "plot_action_values(action_values)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61A_1lhs8Paw"
      },
      "source": [
        "#### Show resulting policy $\\pi(\\cdot|s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lmmo4uMN8Paw"
      },
      "outputs": [],
      "source": [
        "plot_policy(action_values, frame)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qaQWXzb28Paw"
      },
      "source": [
        "#### Test the resulting agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aMm5uC9x8Paw"
      },
      "outputs": [],
      "source": [
        "test_agent(env, policy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MNTpDkGO8Paw"
      },
      "source": [
        "## Resources"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NDaGZUNh8Paw"
      },
      "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
}