Section_8_deep_sarsa_complete.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%%\n"
        },
        "id": "LwYF7tHdI_3V"
      },
      "source": [
        "<div style=\"text-align:center\">\n",
        "    <h1>\n",
        "        Deep SARSA\n",
        "    </h1>\n",
        "</div>\n",
        "\n",
        "<br><br>\n",
        "\n",
        "<div style=\"text-align:center\">\n",
        "\n",
        "In this notebook, we extend the SARSA algorithm to use function approximators (Neural Networks). The resulting algorithm is known as Deep SARSA.\n",
        "</div>\n",
        "\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",
        "import matplotlib.patches as mpatches\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 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 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_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",
        "\n",
        "\n",
        "def plot_cost_to_go(env, q_network, xlabel=None, ylabel=None):\n",
        "    highx, highy = env.observation_space.high\n",
        "    lowx, lowy = env.observation_space.low\n",
        "    X = torch.linspace(lowx, highx, 100)\n",
        "    Y = torch.linspace(lowy, highy, 100)\n",
        "    X, Y = torch.meshgrid(X, Y)\n",
        "\n",
        "    q_net_input = torch.stack([X.flatten(), Y.flatten()], dim=-1)\n",
        "    Z = - q_network(q_net_input).max(dim=-1, keepdim=True)[0]\n",
        "    Z = Z.reshape(100, 100).detach().numpy()\n",
        "    X = X.numpy()\n",
        "    Y = Y.numpy()\n",
        "\n",
        "    fig = plt.figure(figsize=(8, 8))\n",
        "    ax = fig.add_subplot(111, projection='3d')\n",
        "    surf = ax.plot_surface(X, Y, Z, cmap='jet', linewidth=0, antialiased=False)\n",
        "    fig.colorbar(surf, shrink=0.5, aspect=5)\n",
        "    ax.set_xlabel(xlabel, size=14)\n",
        "    ax.set_ylabel(ylabel, size=14)\n",
        "    ax.set_title(\"Estimated cost-to-go\", size=18)\n",
        "    plt.tight_layout()\n",
        "    plt.show()\n",
        "\n",
        "\n",
        "def plot_max_q(env, q_network, xlabel=None, ylabel=None, action_labels=[]):\n",
        "    highx, highy = env.observation_space.high\n",
        "    lowx, lowy = env.observation_space.low\n",
        "    X = torch.linspace(lowx, highx, 100)\n",
        "    Y = torch.linspace(lowy, highy, 100)\n",
        "    X, Y = torch.meshgrid(X, Y)\n",
        "    q_net_input = torch.stack([X.flatten(), Y.flatten()], dim=-1)\n",
        "    Z = q_network(q_net_input).argmax(dim=-1, keepdim=True)\n",
        "    Z = Z.reshape(100, 100).T.detach().numpy()\n",
        "    values = np.unique(Z.ravel())\n",
        "    values.sort()\n",
        "\n",
        "    plt.figure(figsize=(5, 5))\n",
        "    plt.xlabel(xlabel, size=14)\n",
        "    plt.ylabel(ylabel, size=14)\n",
        "    plt.title(\"Optimal action\", size=18)\n",
        "\n",
        "    im = plt.imshow(Z, cmap='jet')\n",
        "    colors = [im.cmap(im.norm(value)) for value in values]\n",
        "    patches = [mpatches.Patch(color=color, label=label) for color, label in zip(colors, action_labels)]\n",
        "    plt.legend(handles=patches, bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)\n",
        "    plt.tight_layout()\n",
        "\n"
      ],
      "metadata": {
        "cellView": "form",
        "id": "3ghFo0wmJW6w"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ctg9zxuEI_3Z"
      },
      "source": [
        "## Import the necessary software libraries:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OUKiIm-JI_3Z"
      },
      "outputs": [],
      "source": [
        "import random\n",
        "import copy\n",
        "import gym\n",
        "import torch\n",
        "import torch.nn.functional as F\n",
        "import matplotlib.pyplot as plt\n",
        "from torch import nn as nn\n",
        "from torch.optim import AdamW\n",
        "from tqdm import tqdm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7wCg5jRUI_3a"
      },
      "source": [
        "## Create and prepare the environment"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Apbob5-dI_3a"
      },
      "source": [
        "### Create the environment"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AkRU4Rm8I_3b"
      },
      "outputs": [],
      "source": [
        "env = gym.make('MountainCar-v0')\n",
        "seed_everything(env)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "G-XpsiMYI_3b"
      },
      "outputs": [],
      "source": [
        "state_dims = env.observation_space.shape[0]\n",
        "num_actions = env.action_space.n\n",
        "print(f\"MountainCar env: State dimensions: {state_dims}, Number of actions: {num_actions}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZnfKvF-uI_3b"
      },
      "source": [
        "### Prepare the environment to work with PyTorch"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8Tc3FyCcI_3c"
      },
      "outputs": [],
      "source": [
        "class PreprocessEnv(gym.Wrapper):\n",
        "\n",
        "    def __init__(self, env):\n",
        "        gym.Wrapper.__init__(self, env)\n",
        "\n",
        "    def reset(self):\n",
        "        obs = self.env.reset()\n",
        "        return torch.from_numpy(obs).unsqueeze(dim=0).float()\n",
        "\n",
        "    def step(self, action):\n",
        "        action = action.item()\n",
        "        next_state, reward, done, info = self.env.step(action)\n",
        "        next_state = torch.from_numpy(next_state).unsqueeze(dim=0).float()\n",
        "        reward = torch.tensor(reward).view(1, -1).float()\n",
        "        done = torch.tensor(done).view(1, -1)\n",
        "        return next_state, reward, done, info"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "K24l42NqI_3c"
      },
      "outputs": [],
      "source": [
        "env = PreprocessEnv(env)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qdaT05K8I_3c"
      },
      "outputs": [],
      "source": [
        "state = env.reset()\n",
        "action = torch.tensor(0)\n",
        "next_state, reward, done, _ = env.step(action)\n",
        "print(f\"Sample state: {state}\")\n",
        "print(f\"Next state: {next_state}, Reward: {reward}, Done: {done}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZXihRkbGI_3c"
      },
      "source": [
        "## Create the Q-Network and policy\n",
        "\n",
        "<br><br>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qLKZe-pmI_3d"
      },
      "source": [
        "### Create the Q-Network: $\\hat q(s,a| \\theta)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lBVdTXmqI_3d"
      },
      "outputs": [],
      "source": [
        "q_network = nn.Sequential(\n",
        "    nn.Linear(state_dims, 128),\n",
        "    nn.ReLU(),\n",
        "    nn.Linear(128, 64),\n",
        "    nn.ReLU(),\n",
        "    nn.Linear(64, num_actions))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A06QbeZ4I_3d"
      },
      "source": [
        "### Create the target Q-Network: $\\hat q(s, a|\\theta_{targ})$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "giiJkUxmI_3d"
      },
      "outputs": [],
      "source": [
        "target_q_network = copy.deepcopy(q_network).eval()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "F-bq_Gg8I_3d"
      },
      "source": [
        "### Create the $\\epsilon$-greedy policy: $\\pi(s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cFV2Gsa4I_3e"
      },
      "outputs": [],
      "source": [
        "def policy(state, epsilon=0.):\n",
        "    if torch.rand(1) < epsilon:\n",
        "        return torch.randint(num_actions, (1, 1))\n",
        "    else:\n",
        "        av = q_network(state).detach()\n",
        "        return torch.argmax(av, dim=-1, keepdim=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wnCuL0reI_3e"
      },
      "source": [
        "### Plot the cost to go: $ - \\max_a \\hat q(s,a|\\theta)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jWHW3CxPI_3e"
      },
      "outputs": [],
      "source": [
        "plot_cost_to_go(env, q_network, xlabel='Car Position', ylabel='Velocity')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QHXYiYt-I_3e"
      },
      "source": [
        "## Create the Experience Replay buffer\n",
        "\n",
        "<br>\n",
        "<div style=\"text-align:center\">\n",
        "    <p>A simple buffer that stores transitions of arbitrary values, adapted from\n",
        "    <a href=\"https://pytorch.org/tutorials/intermediate/reinforcement_q_learning.html#training\">this source.</a></p>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v-xAxpkLI_3e"
      },
      "outputs": [],
      "source": [
        "class ReplayMemory:\n",
        "\n",
        "    def __init__(self, capacity=1000000):\n",
        "        self.capacity = capacity\n",
        "        self.memory = []\n",
        "        self.position = 0\n",
        "\n",
        "    def insert(self, transition):\n",
        "        if len(self.memory) < self.capacity:\n",
        "            self.memory.append(None)\n",
        "        self.memory[self.position] = transition\n",
        "        self.position = (self.position + 1) % self.capacity\n",
        "\n",
        "    def sample(self, batch_size):\n",
        "        assert self.can_sample(batch_size)\n",
        "\n",
        "        batch = random.sample(self.memory, batch_size)\n",
        "        batch = zip(*batch)\n",
        "        return [torch.cat(items) for items in batch]\n",
        "\n",
        "    def can_sample(self, batch_size):\n",
        "        return len(self.memory) >= batch_size * 10\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.memory)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z-I6QfkqI_3e"
      },
      "source": [
        "## Implement the algorithm\n",
        "\n",
        "</br></br>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0nSDqmbaI_3e"
      },
      "outputs": [],
      "source": [
        "def deep_sarsa(q_network, policy, episodes, alpha=0.001,\n",
        "               batch_size=32, gamma=0.99, epsilon=0.05):\n",
        "    optim = AdamW(q_network.parameters(), lr=alpha)\n",
        "    memory = ReplayMemory()\n",
        "    stats = {'MSE Loss': [], 'Returns': []}\n",
        "\n",
        "    for episode in tqdm(range(1, episodes + 1)):\n",
        "        state = env.reset()\n",
        "        done = False\n",
        "        ep_return = 0\n",
        "        while not done:\n",
        "            action = policy(state, epsilon)\n",
        "            next_state, reward, done, _ = env.step(action)\n",
        "            memory.insert([state, action, reward, done, next_state])\n",
        "\n",
        "            if memory.can_sample(batch_size):\n",
        "                state_b, action_b, reward_b, done_b, next_state_b = memory.sample(batch_size)\n",
        "                qsa_b = q_network(state_b).gather(1, action_b)\n",
        "                next_action_b = policy(next_state_b, epsilon)\n",
        "                next_qsa_b = target_q_network(next_state_b).gather(1, next_action_b)\n",
        "                target_b = reward_b + ~done_b * gamma * next_qsa_b\n",
        "                loss = F.mse_loss(qsa_b, target_b)\n",
        "                q_network.zero_grad()\n",
        "                loss.backward()\n",
        "                optim.step()\n",
        "\n",
        "                stats['MSE Loss'].append(loss.item())\n",
        "\n",
        "            state = next_state\n",
        "            ep_return += reward.item()\n",
        "\n",
        "        stats['Returns'].append(ep_return)\n",
        "\n",
        "        if episode % 10 == 0:\n",
        "            target_q_network.load_state_dict(q_network.state_dict())\n",
        "\n",
        "    return stats"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "scrolled": true,
        "id": "55v03cjFI_3e"
      },
      "outputs": [],
      "source": [
        "stats = deep_sarsa(q_network, policy, 2500, epsilon=0.01)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hPe6EfdhI_3e"
      },
      "source": [
        "## Show results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vhTAl32hI_3e"
      },
      "source": [
        "### Plot execution stats"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m7wOzLBTI_3e"
      },
      "outputs": [],
      "source": [
        "plot_stats(stats)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tV-guw7II_3e"
      },
      "source": [
        "### Plot the cost to go: $ - \\max_a \\hat q(s,a|\\theta)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DKp-7t4SI_3f"
      },
      "outputs": [],
      "source": [
        "plot_cost_to_go(env, q_network, xlabel='Car Position', ylabel='Velocity')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bcC52BmbI_3f"
      },
      "source": [
        "### Show resulting policy: $\\pi(s)$"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "F9sYjDFeI_3f"
      },
      "outputs": [],
      "source": [
        "plot_max_q(env, q_network, xlabel='Car Position', ylabel='Velocity',\n",
        "           action_labels=['Back', 'Do nothing', 'Forward'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9lnUdHHTI_3f"
      },
      "source": [
        "### Test the resulting agent"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lhCD24CMI_3f"
      },
      "outputs": [],
      "source": [
        "test_agent(env, policy, episodes=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dFVeLHTSI_3f"
      },
      "source": [
        "## Resources"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FiLkSo28I_3g"
      },
      "source": [
        "[[1] Deep Reinforcement Learning with Experience Replay Based on SARSA](https://www.researchgate.net/publication/313803199_Deep_reinforcement_learning_with_experience_replay_based_on_SARSA)"
      ]
    }
  ],
  "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
}