Section_7_continuous_observation_spaces.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "9vN8l4-UvgqJ"
      },
      "source": [
        "<div style=\"text-align:center\">\n",
        "    <h1>\n",
        "        Continuous state spaces\n",
        "    </h1>\n",
        "</div>\n",
        "\n",
        "<br>\n",
        "\n",
        "<div style=\"text-align:center\">\n",
        "    In this notebook we will learn how to adapt tabular methods to continuous state spaces. We will do it with two methods:\n",
        "    state aggregation and tile coding.\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",
        "import torch\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(env, 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",
        "def seed_everything(env: gym.Env, seed: int = 42) -> None:\n",
        "    env.seed(seed)\n",
        "    env.action_space.seed(seed)\n",
        "    env.observation_space.seed(seed)\n",
        "    np.random.seed(seed)\n",
        "    torch.manual_seed(seed)\n",
        "    torch.use_deterministic_algorithms(True)\n",
        "\n",
        "\n",
        "def plot_tabular_cost_to_go(action_values, xlabel, ylabel):\n",
        "    plt.figure(figsize=(8, 8))\n",
        "    cost_to_go = -action_values.max(axis=-1)\n",
        "    plt.imshow(cost_to_go, cmap='jet')\n",
        "    plt.title(\"Estimated cost-to-go\", size=24)\n",
        "    plt.xlabel(xlabel, size=18)\n",
        "    plt.ylabel(ylabel, size=18)\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    plt.xticks()\n",
        "    plt.colorbar()\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def plot_stats(stats):\n",
        "    rows = len(stats)\n",
        "    cols = 1\n",
        "\n",
        "    fig, ax = plt.subplots(rows, cols, figsize=(12, 6))\n",
        "\n",
        "    for i, key in enumerate(stats):\n",
        "        vals = stats[key]\n",
        "        vals = [np.mean(vals[i-10:i+10]) for i in range(10, len(vals)-10)]\n",
        "        if len(stats) > 1:\n",
        "            ax[i].plot(range(len(vals)), vals)\n",
        "            ax[i].set_title(key, size=18)\n",
        "        else:\n",
        "            ax.plot(range(len(vals)), vals)\n",
        "            ax.set_title(key, size=18)\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n"
      ],
      "metadata": {
        "id": "RH8HXBTBvwBs",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1a40Bl59vgqL"
      },
      "source": [
        "## Import the necessary software libraries:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ITOgJ74WvgqL"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import gym\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from tqdm import tqdm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AYUqQBl1vgqM"
      },
      "source": [
        "## Implement state aggregation\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-3COnudvvgqM"
      },
      "source": [
        "### Create the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Zpo8BQEAvgqM"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qo3ssx5cvgqM"
      },
      "source": [
        "### Create the state aggregation wrapper"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IVjaqy9FvgqM"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GNAXqTtlvgqN"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J4cAudZsvgqN"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iGyKnzn6vgqN"
      },
      "source": [
        "### Compare the original environment to the one with aggregated states"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rd_a-ui9vgqN"
      },
      "outputs": [],
      "source": [
        "print(f\"Modified observation space: {saenv.observation_space}, \\n\\\n",
        "Sample state: {saenv.observation_space.sample()}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9csRyU1SvgqN"
      },
      "outputs": [],
      "source": [
        "print(f\"Original observation space: {env.observation_space}, \\n\\\n",
        "Sample state: {env.observation_space.sample()}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a_igeDTqvgqN"
      },
      "source": [
        "### Create the $Q(s,a)$ value table"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XmmC_IwfvgqN"
      },
      "outputs": [],
      "source": [
        "action_values = np.zeros((20,20, 3))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dxEX6ZvGvgqN"
      },
      "source": [
        "### Create the $\\epsilon$-greedy policy: $\\pi(s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "U2CSJjIzvgqN"
      },
      "outputs": [],
      "source": [
        "def policy(state, epsilon=0.):\n",
        "    if np.random.random() < epsilon:\n",
        "        return np.random.randint(3)\n",
        "    else:\n",
        "        av = action_values[state]\n",
        "        return np.random.choice(np.flatnonzero(av == av.max()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rnZ2AUVmvgqN"
      },
      "source": [
        "### Test the SARSA algorithm on the modified environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jitSZTRRvgqO"
      },
      "outputs": [],
      "source": [
        "def sarsa(action_values, policy, episodes, alpha=0.1, gamma=0.99, epsilon=0.2):\n",
        "    stats = {'Returns': []}\n",
        "    for episode in tqdm(range(1, episodes + 1)):\n",
        "        state = saenv.reset()\n",
        "        action = policy(state, epsilon)\n",
        "        done = False\n",
        "        ep_return = 0\n",
        "        while not done:\n",
        "            next_state, reward, done, _ = saenv.step(action)\n",
        "            next_action = policy(next_state, epsilon)\n",
        "\n",
        "            qsa = action_values[state][action]\n",
        "            next_qsa = action_values[next_state][next_action]\n",
        "            action_values[state][action] = qsa + alpha * (reward + gamma * next_qsa - qsa)\n",
        "            state = next_state\n",
        "            action = next_action\n",
        "            ep_return += reward\n",
        "        stats['Returns'].append(ep_return)\n",
        "    return stats"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pRoh3U-QvgqO"
      },
      "outputs": [],
      "source": [
        "stats = sarsa(action_values, policy, 20000, alpha=0.1, epsilon=0.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RacbxBmZvgqO"
      },
      "outputs": [],
      "source": [
        "plot_stats(stats)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZZmOJxuUvgqO"
      },
      "source": [
        "### Plot the learned policy: $\\pi(s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "57NP8GQRvgqO"
      },
      "outputs": [],
      "source": [
        "plot_policy(action_values, env.render(mode='rgb_array'), \\\n",
        "            action_meanings={0: 'B', 1: 'N', 2: 'F'})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1nvBZFSgvgqO"
      },
      "source": [
        "### Plot the cost to go: $ - \\max_a \\hat q(s,a|\\theta)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": false,
        "id": "SUsd_F73vgqO"
      },
      "outputs": [],
      "source": [
        "plot_tabular_cost_to_go(action_values, xlabel=\"Car Position\", ylabel=\"Velocity\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B1P7kUmXvgqO"
      },
      "source": [
        "### Test the resulting policy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lZX89VryvgqO"
      },
      "outputs": [],
      "source": [
        "test_agent(saenv, policy, 2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W47pYQdCvgqO"
      },
      "source": [
        "<br><br><br><br>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OVg1SvxOvgqO"
      },
      "source": [
        "## Implement Tile Coding\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZOggOxgHvgqO"
      },
      "source": [
        "### Create the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YD3C41lgvgqO"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dbbJcOtCvgqO"
      },
      "source": [
        "### Create the Tile Coding wrapper"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HZls8wqNvgqO"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zo286x38vgqO"
      },
      "outputs": [],
      "source": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iJ4JBdNTvgqO"
      },
      "source": [
        "### Compare the original environment to the one with aggregated states"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0Bs_FL8jvgqO"
      },
      "outputs": [],
      "source": [
        "print(f\"Modified observation space: {tcenv.observation_space}, \\n\\\n",
        "Sample state: {tcenv.reset()}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mg7B8yH1vgqO"
      },
      "outputs": [],
      "source": [
        "print(f\"Original observation space: {env.observation_space}, \\n\\\n",
        "Sample state: {env.reset()}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Hxhy2cXmvgqO"
      },
      "source": [
        "### Create the $Q(s,a)$ value table"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y46HsY2TvgqS"
      },
      "outputs": [],
      "source": [
        "action_values = np.zeros((4, 20, 20, 3))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1bI9pWIevgqS"
      },
      "source": [
        "### Create the $\\epsilon$-greedy policy: $\\pi(s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9K2iDRbIvgqS"
      },
      "outputs": [],
      "source": [
        "def policy(state, epsilon=0.):\n",
        "    if np.random.random() < epsilon:\n",
        "        return np.random.randint(3)\n",
        "    else:\n",
        "        av_list = []\n",
        "        for i, idx in enumerate(state):\n",
        "            av = action_values[i][idx]\n",
        "            av_list.append(av)\n",
        "\n",
        "        av = np.mean(av_list, axis=0)\n",
        "        return np.random.choice(np.flatnonzero(av==av.max()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mRGnsmKlvgqS"
      },
      "source": [
        "### Test the SARSA algorithm on the modified environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7UtvHUrxvgqS"
      },
      "outputs": [],
      "source": [
        "def sarsa(action_values, policy, episodes, alpha=0.1, gamma=0.99, epsilon=0.2):\n",
        "    stats = {'Returns': []}\n",
        "    for episode in tqdm(range(1, episodes + 1)):\n",
        "        state = tcenv.reset()\n",
        "        action = policy(state, epsilon)\n",
        "        done = False\n",
        "        ep_return = 0\n",
        "        while not done:\n",
        "            next_state, reward, done, _ = tcenv.step(action)\n",
        "            next_action = policy(next_state, epsilon)\n",
        "\n",
        "            for i, (idx, next_idx) in enumerate(zip(state, next_state)):\n",
        "                qsa = action_values[i][idx][action]\n",
        "                next_qsa = action_values[i][next_idx][next_action]\n",
        "                action_values[i][idx][action] = qsa + alpha * (reward + gamma * next_qsa - qsa)\n",
        "\n",
        "            state = next_state\n",
        "            action = next_action\n",
        "            ep_return += reward\n",
        "        stats['Returns'].append(ep_return)\n",
        "    return stats"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MW_884NBvgqS"
      },
      "outputs": [],
      "source": [
        "stats = sarsa(action_values, policy, 20000, alpha=0.1, epsilon=0.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SGdpV9TNvgqS"
      },
      "outputs": [],
      "source": [
        "plot_stats(stats)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u_MzBqr2vgqS"
      },
      "source": [
        "### Plot the learned policy: $\\pi(s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jV_XpCylvgqS"
      },
      "outputs": [],
      "source": [
        "plot_policy(action_values.mean(axis=0), env.render(mode='rgb_array'), \\\n",
        "            action_meanings={0: 'B', 1: 'N', 2: 'F'})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z_0TY7oovgqT"
      },
      "source": [
        "### Plot the cost to go: $ - \\max_a \\hat q(s,a|\\theta)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p3G9nCyMvgqT"
      },
      "outputs": [],
      "source": [
        "plot_tabular_cost_to_go(action_values.mean(axis=0), \\\n",
        "                        xlabel=\"Car Position\", ylabel=\"Velocity\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K8z4PLeQvgqT"
      },
      "source": [
        "### Test the resulting policy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fmopjBSTvgqT"
      },
      "outputs": [],
      "source": [
        "test_agent(tcenv, policy, 2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J_qLNpc0vgqT"
      },
      "source": [
        "## Resources"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R8lQEV8uvgqT"
      },
      "source": [
        "[[1] Reinforcement Learning: An Introduction. Section 9.5.4: Tile Coding](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
}