rl/EEML2020_RL_Tutorial.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "EEML2020_RL_Tutorial",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ULdrhOaVbsdO"
      },
      "source": [
        "#RL Tutorial \n",
        "\n",
        "<a href=\"https://colab.research.google.com/github/eemlcommunity/PracticalSessions2020/blob/master/rl/EEML2020_RL_Tutorial.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
        "\n",
        "Contact us at feryal@google.com & gcomanici@google.com for any questions/comments :)\n",
        "\n",
        "Special thanks to Diana Borsa and Loic Matthey.\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dv-846KxIqPD",
        "colab_type": "text"
      },
      "source": [
        "The tutorial covers a number of important reinforcement learning (RL) algorithms, including policy iteration, Q-Learning, and Neural Fitted Q. In the first part, we will guide you through the general interaction between RL agents and environments, where the agents ought to take actions in order to maximize returns (i.e. cumulative reward). Next, we will implement Policy Iteration, SARSA, and Q-Learning for a simple tabular environment. The core ideas in the latter will be scaled to more complex MDPs through the use of function approximation. Lastly, we will provide a short introduction to deep reinforcement learning and the DQN algorithm.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ffeeXVm4AuZ6",
        "colab_type": "text"
      },
      "source": [
        "# Overview"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kT8rkxTUAZzL",
        "colab_type": "text"
      },
      "source": [
        "The agent interacts with the environment in a loop corresponding to the following diagram. The environment defines a set of <font color='blue'>**actions**</font>  that an agent can take.  The agent takes an action informed by the <font color='red'>**observations**</font> it recieves, and will get a <font color='green'>**reward**</font> from the environment after each action. The goal in RL is to find an agent whose actions maximize the total accumulation of rewards obtained from the environment. \n",
        "\n",
        "\n",
        "<center><img src=\"https://drive.google.com/uc?id=1nAj1yf0kmTb369dVkIgXaMSPEpTUslnC\" width=\"500\" /></center>\n",
        "\n",
        "\n",
        "The tutorial is mainly focused on **value based methods**: agents are maintaining a value for all state-action pairs and use those estimates to choose actions that maximize that value (instead of maintaining a policy directly, like in policy gradient methods). \n",
        "\n",
        "We represent the action-value function (otherwise known as Q-function) associated with following/employing a policy $\\pi$ in a given MDP as:\n",
        "\n",
        "$$ Q^{\\pi}(\\color{red}{s},\\color{blue}{a}) = \\mathbb{E}_{\\tau \\sim P^{\\pi}} \\left[ \\sum_t \\gamma^t \\color{green}{R_t}| s_0=\\color{red}s,a=\\color{blue}{a_0} \\right]$$\n",
        "\n",
        "where $\\tau = \\{\\color{red}{s_0}, \\color{blue}{a_0}, \\color{green}{r_0}, \\color{red}{s_1}, \\color{blue}{a_1}, \\color{green}{r_1}, \\cdots \\}$\n",
        "\n",
        "\n",
        "Recall that efficient value estimations are based on the famous **_Bellman Optimallity Equation_**:\n",
        "\n",
        "$$ Q^\\pi(\\color{red}{s},\\color{blue}{a}) =  \\color{green}{r}(\\color{red}{s},\\color{blue}{a}) + \\gamma  \\sum_{\\color{red}{s'}\\in \\color{red}{\\mathcal{S}}} P(\\color{red}{s'} |\\color{red}{s},\\color{blue}{a}) V^\\pi(\\color{red}{s'}) $$\n",
        "\n",
        "where $V^\\pi$ is the expected $Q^\\pi$ value for a particular state, i.e. $V^\\pi(\\color{red}{s}) = \\sum_{\\color{blue}{a} \\in \\color{blue}{\\mathcal{A}}} \\pi(\\color{blue}{a} |\\color{red}{s}) Q^\\pi(\\color{red}{s},\\color{blue}{a})$."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "xaJxoatMhJ71"
      },
      "source": [
        "## Installation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ovuCuHCC78Zu"
      },
      "source": [
        "### Install required libraries\n",
        "\n",
        "1. [Acme](https://github.com/deepmind/acme) is a library of reinforcement learning (RL) agents and agent building blocks. Acme strives to expose simple, efficient, and readable agents, that serve both as reference implementations of popular algorithms and as strong baselines, while still providing enough flexibility to do novel research. The design of Acme also attempts to provide multiple points of entry to the RL problem at differing levels of complexity.\n",
        "\n",
        "\n",
        "2. [Haiku](https://github.com/deepmind/dm-haiku) is a simple neural network library for JAX developed by some of the authors of Sonnet, a neural network library for TensorFlow.\n",
        "\n",
        "3. [dm_env](https://github.com/deepmind/dm_env): DeepMind Environment API, which will be covered in more details in the [Environment subsection](https://colab.research.google.com/drive/1oKyyhOFAFSBTpVnmuOm9HXh5D5ekqhh5#scrollTo=I6KuVGSk4uc9) below."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "KH3O0zcXUeun",
        "colab": {}
      },
      "source": [
        "#@title Installations  { form-width: \"30%\" }\n",
        "\n",
        "!pip install dm-acme\n",
        "!pip install dm-acme[reverb]\n",
        "!pip install dm-acme[jax]\n",
        "!pip install dm-acme[tf]\n",
        "!pip install dm-acme[envs]\n",
        "!pip install dm-env\n",
        "!pip install dm-haiku\n",
        "!sudo apt-get install -y xvfb ffmpeg\n",
        "!pip install imageio\n",
        "\n",
        "from IPython.display import clear_output\n",
        "clear_output()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "c-H2d6UZi7Sf"
      },
      "source": [
        "## Import Modules"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "cellView": "form",
        "colab_type": "code",
        "id": "HJ74Id-8MERq",
        "colab": {}
      },
      "source": [
        "#@title Imports  { form-width: \"30%\" }\n",
        "\n",
        "import IPython\n",
        "\n",
        "import acme\n",
        "from acme import environment_loop\n",
        "from acme import datasets\n",
        "from acme import specs\n",
        "from acme import wrappers\n",
        "from acme.wrappers import gym_wrapper\n",
        "from acme.agents.jax import dqn\n",
        "from acme.adders import reverb as adders\n",
        "from acme.utils import counting\n",
        "from acme.utils import loggers\n",
        "import base64\n",
        "import collections\n",
        "from collections import namedtuple\n",
        "import dm_env\n",
        "import enum\n",
        "import functools\n",
        "import gym\n",
        "import haiku as hk\n",
        "import io\n",
        "import imageio\n",
        "import itertools\n",
        "import jax\n",
        "from jax import tree_util\n",
        "from jax.experimental import optix\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import pandas as pd\n",
        "import random\n",
        "import reverb\n",
        "import rlax\n",
        "import time\n",
        "\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')\n",
        "\n",
        "np.set_printoptions(precision=3, suppress=1)\n",
        "\n",
        "plt.style.use('seaborn-notebook')\n",
        "plt.style.use('seaborn-whitegrid')\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HeGPIOMkUTEn",
        "colab_type": "text"
      },
      "source": [
        "# RL Lab - Part 0: Environment & Agent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "I6KuVGSk4uc9"
      },
      "source": [
        "## Environment\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UhZwB__DPcyM",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "We will focus on a simple grid world environment for this practical session. \n",
        "\n",
        "\n",
        "<img src=\"https://drive.google.com/uc?id=1qBjh_PRdZ4GWTDqB9pmjLEOlUAsOfrZi\" width=\"500\" />\n",
        "\n",
        "\n",
        "\n",
        "This environment consists of either walls and empty cells. The agent starts from an initial location and needs to navigate to reach a goal location. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "inIAhwLKuHKr",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Gridworld Implementation { form-width: \"30%\" }\n",
        "\n",
        "\n",
        "class ObservationType(enum.IntEnum):\n",
        "  STATE_INDEX = enum.auto()\n",
        "  AGENT_ONEHOT = enum.auto()\n",
        "  GRID = enum.auto()\n",
        "  AGENT_GOAL_POS = enum.auto()\n",
        "\n",
        "\n",
        "class GridWorld(dm_env.Environment):\n",
        "\n",
        "  def __init__(self,\n",
        "               layout,\n",
        "               start_state,\n",
        "               goal_state=None,\n",
        "               observation_type=ObservationType.STATE_INDEX,\n",
        "               discount=0.9,\n",
        "               penalty_for_walls=-5,\n",
        "               reward_goal=10,\n",
        "               max_episode_length=None,\n",
        "               randomize_goals=False):\n",
        "    \"\"\"Build a grid environment.\n",
        "\n",
        "    Simple gridworld defined by a map layout, a start and a goal state.\n",
        "\n",
        "    Layout should be a NxN grid, containing:\n",
        "      * 0: empty\n",
        "      * -1: wall\n",
        "      * Any other positive value: value indicates reward; episode will terminate\n",
        "\n",
        "    Args:\n",
        "      layout: NxN array of numbers, indicating the layout of the environment.\n",
        "      start_state: Tuple (y, x) of starting location.\n",
        "      goal_state: Optional tuple (y, x) of goal location. Will be randomly\n",
        "        sampled once if None.\n",
        "      observation_type: Enum observation type to use. One of:\n",
        "        * ObservationType.STATE_INDEX: int32 index of agent occupied tile.\n",
        "        * ObservationType.AGENT_ONEHOT: NxN float32 grid, with a 1 where the \n",
        "          agent is and 0 elsewhere.\n",
        "        * ObservationType.GRID: NxNx3 float32 grid of feature channels. \n",
        "          First channel contains walls (1 if wall, 0 otherwise), second the \n",
        "          agent position (1 if agent, 0 otherwise) and third goal position\n",
        "          (1 if goal, 0 otherwise)\n",
        "        * ObservationType.AGENT_GOAL_POS: float32 tuple with \n",
        "          (agent_y, agent_x, goal_y, goal_x)\n",
        "      discount: Discounting factor included in all Timesteps.\n",
        "      penalty_for_walls: Reward added when hitting a wall (should be negative).\n",
        "      reward_goal: Reward added when finding the goal (should be positive).\n",
        "      max_episode_length: If set, will terminate an episode after this many \n",
        "        steps.\n",
        "      randomize_goals: If true, randomize goal at every episode.\n",
        "    \"\"\"\n",
        "    if observation_type not in ObservationType:\n",
        "      raise ValueError('observation_type should be a ObservationType instace.')\n",
        "    self._layout = np.array(layout)\n",
        "    self._start_state = start_state\n",
        "    self._state = self._start_state\n",
        "    self._number_of_states = np.prod(np.shape(self._layout))\n",
        "    self._discount = discount\n",
        "    self._penalty_for_walls = penalty_for_walls\n",
        "    self._reward_goal = reward_goal\n",
        "    self._observation_type = observation_type\n",
        "    self._layout_dims = self._layout.shape\n",
        "    self._max_episode_length = max_episode_length\n",
        "    self._num_episode_steps = 0\n",
        "    self._randomize_goals = randomize_goals\n",
        "    if goal_state is None:\n",
        "      # Randomly sample goal_state if not provided\n",
        "      goal_state = self._sample_goal()\n",
        "    self.goal_state = goal_state\n",
        "\n",
        "  def _sample_goal(self):\n",
        "    \"\"\"Randomly sample reachable non-starting state.\"\"\"\n",
        "    # Sample a new goal\n",
        "    n = 0\n",
        "    max_tries = 1e5\n",
        "    while n < max_tries:\n",
        "      goal_state = tuple(np.random.randint(d) for d in self._layout_dims)\n",
        "      if goal_state != self._state and self._layout[goal_state] == 0:\n",
        "        # Reachable state found!\n",
        "        return goal_state\n",
        "      n += 1\n",
        "    raise ValueError('Failed to sample a goal state.')\n",
        "\n",
        "  @property\n",
        "  def number_of_states(self):\n",
        "    return self._number_of_states\n",
        "\n",
        "  @property\n",
        "  def goal_state(self):\n",
        "    return self._goal_state\n",
        "\n",
        "  def set_state(self, x, y):\n",
        "    self._state = (y, x)\n",
        "\n",
        "  @goal_state.setter\n",
        "  def goal_state(self, new_goal):\n",
        "    if new_goal == self._state or self._layout[new_goal] < 0:\n",
        "      raise ValueError('This is not a valid goal!')\n",
        "    # Zero out any other goal\n",
        "    self._layout[self._layout > 0] = 0\n",
        "    # Setup new goal location\n",
        "    self._layout[new_goal] = self._reward_goal\n",
        "    self._goal_state = new_goal\n",
        "\n",
        "  def observation_spec(self):\n",
        "    if self._observation_type is ObservationType.AGENT_ONEHOT:\n",
        "      return specs.Array(\n",
        "          shape=self._layout_dims,\n",
        "          dtype=np.float32,\n",
        "          name='observation_agent_onehot')\n",
        "    elif self._observation_type is ObservationType.GRID:\n",
        "      return specs.Array(\n",
        "          shape=self._layout_dims + (3,),\n",
        "          dtype=np.float32,\n",
        "          name='observation_grid')\n",
        "    elif self._observation_type is ObservationType.AGENT_GOAL_POS:\n",
        "      return specs.Array(\n",
        "          shape=(4,), dtype=np.float32, name='observation_agent_goal_pos')\n",
        "    elif self._observation_type is ObservationType.STATE_INDEX:\n",
        "      return specs.DiscreteArray(\n",
        "          self._number_of_states, dtype=int, name='observation_state_index')\n",
        "\n",
        "  def action_spec(self):\n",
        "    return specs.DiscreteArray(4, dtype=int, name='action')\n",
        "\n",
        "  def get_obs(self):\n",
        "    if self._observation_type is ObservationType.AGENT_ONEHOT:\n",
        "      obs = np.zeros(self._layout.shape, dtype=np.float32)\n",
        "      # Place agent\n",
        "      obs[self._state] = 1\n",
        "      return obs\n",
        "    elif self._observation_type is ObservationType.GRID:\n",
        "      obs = np.zeros(self._layout.shape + (3,), dtype=np.float32)\n",
        "      obs[..., 0] = self._layout < 0\n",
        "      obs[self._state[0], self._state[1], 1] = 1\n",
        "      obs[self._goal_state[0], self._goal_state[1], 2] = 1\n",
        "      return obs\n",
        "    elif self._observation_type is ObservationType.AGENT_GOAL_POS:\n",
        "      return np.array(self._state + self._goal_state, dtype=np.float32)\n",
        "    elif self._observation_type is ObservationType.STATE_INDEX:\n",
        "      y, x = self._state\n",
        "      return y * self._layout.shape[1] + x\n",
        "\n",
        "  def reset(self):\n",
        "    self._state = self._start_state\n",
        "    self._num_episode_steps = 0\n",
        "    if self._randomize_goals:\n",
        "      self.goal_state = self._sample_goal()\n",
        "    return dm_env.TimeStep(\n",
        "        step_type=dm_env.StepType.FIRST,\n",
        "        reward=None,\n",
        "        discount=None,\n",
        "        observation=self.get_obs())\n",
        "\n",
        "  def step(self, action):\n",
        "    y, x = self._state\n",
        "\n",
        "    if action == 0:  # up\n",
        "      new_state = (y - 1, x)\n",
        "    elif action == 1:  # right\n",
        "      new_state = (y, x + 1)\n",
        "    elif action == 2:  # down\n",
        "      new_state = (y + 1, x)\n",
        "    elif action == 3:  # left\n",
        "      new_state = (y, x - 1)\n",
        "    else:\n",
        "      raise ValueError(\n",
        "          'Invalid action: {} is not 0, 1, 2, or 3.'.format(action))\n",
        "\n",
        "    new_y, new_x = new_state\n",
        "    step_type = dm_env.StepType.MID\n",
        "    if self._layout[new_y, new_x] == -1:  # wall\n",
        "      reward = self._penalty_for_walls\n",
        "      discount = self._discount\n",
        "      new_state = (y, x)\n",
        "    elif self._layout[new_y, new_x] == 0:  # empty cell\n",
        "      reward = 0.\n",
        "      discount = self._discount\n",
        "    else:  # a goal\n",
        "      reward = self._layout[new_y, new_x]\n",
        "      discount = 0.\n",
        "      new_state = self._start_state\n",
        "      step_type = dm_env.StepType.LAST\n",
        "\n",
        "    self._state = new_state\n",
        "    self._num_episode_steps += 1\n",
        "    if (self._max_episode_length is not None and\n",
        "        self._num_episode_steps >= self._max_episode_length):\n",
        "      step_type = dm_env.StepType.LAST\n",
        "    return dm_env.TimeStep(\n",
        "        step_type=step_type,\n",
        "        reward=np.float32(reward),\n",
        "        discount=discount,\n",
        "        observation=self.get_obs())\n",
        "\n",
        "  def plot_grid(self, add_start=True):\n",
        "    plt.figure(figsize=(4, 4))\n",
        "    plt.imshow(self._layout <= -1, interpolation='nearest')\n",
        "    ax = plt.gca()\n",
        "    ax.grid(0)\n",
        "    plt.xticks([])\n",
        "    plt.yticks([])\n",
        "    # Add start/goal\n",
        "    if add_start:\n",
        "      plt.text(\n",
        "          self._start_state[1],\n",
        "          self._start_state[0],\n",
        "          r'$\\mathbf{S}$',\n",
        "          fontsize=16,\n",
        "          ha='center',\n",
        "          va='center')\n",
        "    plt.text(\n",
        "        self._goal_state[1],\n",
        "        self._goal_state[0],\n",
        "        r'$\\mathbf{G}$',\n",
        "        fontsize=16,\n",
        "        ha='center',\n",
        "        va='center')\n",
        "    h, w = self._layout.shape\n",
        "    for y in range(h - 1):\n",
        "      plt.plot([-0.5, w - 0.5], [y + 0.5, y + 0.5], '-k', lw=2)\n",
        "    for x in range(w - 1):\n",
        "      plt.plot([x + 0.5, x + 0.5], [-0.5, h - 0.5], '-k', lw=2)\n",
        "\n",
        "  def plot_state(self, return_rgb=False):\n",
        "    self.plot_grid(add_start=False)\n",
        "    # Add the agent location\n",
        "    plt.text(\n",
        "        self._state[1],\n",
        "        self._state[0],\n",
        "        u'😃',\n",
        "        fontname='symbola',\n",
        "        fontsize=18,\n",
        "        ha='center',\n",
        "        va='center',\n",
        "    )\n",
        "    if return_rgb:\n",
        "      fig = plt.gcf()\n",
        "      plt.axis('tight')\n",
        "      plt.subplots_adjust(0, 0, 1, 1, 0, 0)\n",
        "      fig.canvas.draw()\n",
        "      data = np.fromstring(fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')\n",
        "      w, h = fig.canvas.get_width_height()\n",
        "      data = data.reshape((h, w, 3))\n",
        "      plt.close(fig)\n",
        "      return data\n",
        "\n",
        "  def plot_policy(self, policy):\n",
        "    action_names = [\n",
        "        r'$\\uparrow$', r'$\\rightarrow$', r'$\\downarrow$', r'$\\leftarrow$'\n",
        "    ]\n",
        "    self.plot_grid()\n",
        "    plt.title('Policy Visualization')\n",
        "    h, w = self._layout.shape\n",
        "    for y in range(h):\n",
        "      for x in range(w):\n",
        "        # if ((y, x) != self._start_state) and ((y, x) != self._goal_state):\n",
        "        if (y, x) != self._goal_state:\n",
        "          action_name = action_names[policy[y, x]]\n",
        "          plt.text(x, y, action_name, ha='center', va='center')\n",
        "\n",
        "  def plot_greedy_policy(self, q):\n",
        "    greedy_actions = np.argmax(q, axis=2)\n",
        "    self.plot_policy(greedy_actions)\n",
        "\n",
        "\n",
        "def build_gridworld_task(task,\n",
        "                         discount=0.9,\n",
        "                         penalty_for_walls=-5,\n",
        "                         observation_type=ObservationType.STATE_INDEX,\n",
        "                         max_episode_length=200):\n",
        "  \"\"\"Construct a particular Gridworld layout with start/goal states.\n",
        "\n",
        "  Args:\n",
        "      task: string name of the task to use. One of {'simple', 'obstacle', \n",
        "        'random_goal'}.\n",
        "      discount: Discounting factor included in all Timesteps.\n",
        "      penalty_for_walls: Reward added when hitting a wall (should be negative).\n",
        "      observation_type: Enum observation type to use. One of:\n",
        "        * ObservationType.STATE_INDEX: int32 index of agent occupied tile.\n",
        "        * ObservationType.AGENT_ONEHOT: NxN float32 grid, with a 1 where the \n",
        "          agent is and 0 elsewhere.\n",
        "        * ObservationType.GRID: NxNx3 float32 grid of feature channels. \n",
        "          First channel contains walls (1 if wall, 0 otherwise), second the \n",
        "          agent position (1 if agent, 0 otherwise) and third goal position\n",
        "          (1 if goal, 0 otherwise)\n",
        "        * ObservationType.AGENT_GOAL_POS: float32 tuple with \n",
        "          (agent_y, agent_x, goal_y, goal_x).\n",
        "      max_episode_length: If set, will terminate an episode after this many \n",
        "        steps.\n",
        "  \"\"\"\n",
        "  tasks_specifications = {\n",
        "      'simple': {\n",
        "          'layout': [\n",
        "              [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, -1, -1, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, -1, -1, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, -1, -1, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1],\n",
        "          ],\n",
        "          'start_state': (2, 2),\n",
        "          'goal_state': (7, 2)\n",
        "      },\n",
        "      'obstacle': {\n",
        "          'layout': [\n",
        "              [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, -1, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, -1, 0, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, -1, -1, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, -1, -1, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1],\n",
        "          ],\n",
        "          'start_state': (2, 2),\n",
        "          'goal_state': (2, 8)\n",
        "      },\n",
        "      'random_goal': {\n",
        "          'layout': [\n",
        "              [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, -1, -1, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, -1, -1, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, -1, -1, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, 0, 0, 0, 0, 0, 0, 0, 0, -1],\n",
        "              [-1, -1, -1, -1, -1, -1, -1, -1, -1, -1],\n",
        "          ],\n",
        "          'start_state': (2, 2),\n",
        "          # 'randomize_goals': True\n",
        "      },\n",
        "  }\n",
        "  return GridWorld(\n",
        "      discount=discount,\n",
        "      penalty_for_walls=penalty_for_walls,\n",
        "      observation_type=observation_type,\n",
        "      max_episode_length=max_episode_length,\n",
        "      **tasks_specifications[task])\n",
        "\n",
        "\n",
        "def setup_environment(environment):\n",
        "  # Make sure the environment outputs single-precision floats.\n",
        "  environment = wrappers.SinglePrecisionWrapper(environment)\n",
        "\n",
        "  # Grab the spec of the environment.\n",
        "  environment_spec = specs.make_environment_spec(environment)\n",
        "\n",
        "  return environment, environment_spec"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZizdE9SQS-cN",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "We will use two distinct tabular GridWorlds:\n",
        "* `simple` where the goal is at the bottom left of the grid, little navigation required.\n",
        "* `obstacle` where the goal is behind an obstacle to avoid.\n",
        "\n",
        "You can visualize the grid worlds by running the cell below. \n",
        "\n",
        "Note that `S` indicates the start state and `G` indicates the goal. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7Xdnh3Odc63Q",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title Visualise gridworlds { form-width: \"30%\" }\n",
        "\n",
        "# Instantiate two tabular environments, a simple task, and one that involves\n",
        "# the avoidance of an obstacle.\n",
        "simple_grid = build_gridworld_task(\n",
        "    task='simple', observation_type=ObservationType.GRID)\n",
        "obstacle_grid = build_gridworld_task(\n",
        "    task='obstacle', observation_type=ObservationType.GRID)\n",
        "\n",
        "# Plot them.\n",
        "simple_grid.plot_grid()\n",
        "plt.title('Simple')\n",
        "\n",
        "obstacle_grid.plot_grid()\n",
        "plt.title('Obstacle');"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RTsiWgDSCL7C",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "In this environment, the agent has four possible  <font color='blue'>**Actions**</font>: `up`, `right`, `down`, and `left`.  <font color='green'>**Reward**</font> is `-5` for bumping into a wall, `+10` for reaching the goal, and `0` otherwise. The episode ends when the agent reaches the goal, and otherwise continues. **Discount** on continuing steps, is $\\gamma = 0.9$. \n",
        "\n",
        "Before we start building an agent to interact with this environment, let's first look at the types of objects the environment either returns (e.g. observations) or consumes (e.g. actions). The `environment_spec` will show you the form of the *observations*, *rewards* and *discounts* that the environment exposes and the form of the *actions* that can be taken.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rmKop4FECVV6",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "environment, environment_spec = setup_environment(simple_grid)\n",
        "\n",
        "print('actions:\\n', environment_spec.actions, '\\n')\n",
        "print('observations:\\n', environment_spec.observations, '\\n')\n",
        "print('rewards:\\n', environment_spec.rewards, '\\n')\n",
        "print('discounts:\\n', environment_spec.discounts, '\\n')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0VVTmep2UK6U",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "We first set the environment to its initial location by calling the `reset` method which returns the first observation. \n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rHden9m9FNPK",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "environment.reset()\n",
        "environment.plot_state()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pXb7u9epFWnX",
        "colab_type": "text"
      },
      "source": [
        "Now we want to take an action using the `step` method to interact with the environment which returns a `TimeStep` \n",
        "namedtuple with fields:\n",
        "\n",
        "```none\n",
        "step_type, reward, discount, observation\n",
        "``` \n",
        "\n",
        "We can then visualise the updated state of the grid. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LY1eopIWFe95",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "timestep = environment.step(1)\n",
        "environment.plot_state()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pSFDZPksEGpl",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Run loop  { form-width: \"30%\" }\n",
        "\n",
        "\n",
        "def run_loop(environment,\n",
        "             agent,\n",
        "             num_episodes=None,\n",
        "             num_steps=None,\n",
        "             logger_time_delta=1.,\n",
        "             label='training_loop',\n",
        "             log_loss=False,\n",
        "             ):\n",
        "  \"\"\"Perform the run loop.\n",
        "\n",
        "  We are following the Acme run loop.\n",
        "\n",
        "  Run the environment loop for `num_episodes` episodes. Each episode is itself\n",
        "  a loop which interacts first with the environment to get an observation and\n",
        "  then give that observation to the agent in order to retrieve an action. Upon\n",
        "  termination of an episode a new episode will be started. If the number of\n",
        "  episodes is not given then this will interact with the environment\n",
        "  infinitely.\n",
        "\n",
        "  Args:\n",
        "    environment: dm_env used to generate trajectories.\n",
        "    agent: acme.Actor for selecting actions in the run loop.\n",
        "    num_steps: number of episodes to run the loop for. If `None` (default), runs\n",
        "      without limit.\n",
        "    num_episodes: number of episodes to run the loop for. If `None` (default),\n",
        "      runs without limit.\n",
        "    logger_time_delta: time interval (in seconds) between consecutive logging\n",
        "      steps.\n",
        "    label: optional label used at logging steps.\n",
        "  \"\"\"\n",
        "  logger = loggers.TerminalLogger(label=label, time_delta=logger_time_delta)\n",
        "  iterator = range(num_episodes) if num_episodes else itertools.count()\n",
        "  all_returns = []\n",
        "  \n",
        "  num_total_steps = 0\n",
        "  for episode in iterator:\n",
        "    # Reset any counts and start the environment.\n",
        "    start_time = time.time()\n",
        "    episode_steps = 0\n",
        "    episode_return = 0\n",
        "    episode_loss = 0\n",
        "\n",
        "    timestep = environment.reset()\n",
        "    \n",
        "    # Make the first observation.\n",
        "    agent.observe_first(timestep)\n",
        "\n",
        "    # Run an episode.\n",
        "    while not timestep.last():\n",
        "      # Generate an action from the agent's policy and step the environment.\n",
        "      action = agent.select_action(timestep.observation)\n",
        "      timestep = environment.step(action)\n",
        "\n",
        "      # Have the agent observe the timestep and let the agent update itself.\n",
        "      agent.observe(action, next_timestep=timestep)\n",
        "      agent.update()\n",
        "\n",
        "      # Book-keeping.\n",
        "      episode_steps += 1\n",
        "      num_total_steps += 1\n",
        "      episode_return += timestep.reward\n",
        "\n",
        "      if log_loss:\n",
        "        episode_loss += agent.last_loss\n",
        "\n",
        "      if num_steps is not None and num_total_steps >= num_steps:\n",
        "        break\n",
        "\n",
        "    # Collect the results and combine with counts.\n",
        "    steps_per_second = episode_steps / (time.time() - start_time)\n",
        "    result = {\n",
        "        'episode': episode,\n",
        "        'episode_length': episode_steps,\n",
        "        'episode_return': episode_return,\n",
        "    }\n",
        "    if log_loss:\n",
        "      result['loss_avg'] = episode_loss/episode_steps\n",
        "\n",
        "    all_returns.append(episode_return)\n",
        "\n",
        "    # Log the given results.\n",
        "    logger.write(result)\n",
        "    \n",
        "    if num_steps is not None and num_total_steps >= num_steps:\n",
        "      break\n",
        "  return all_returns"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_gatpjQ8QA_H",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Evaluation loop { form-width: \"30%\" }\n",
        "\n",
        "def evaluate(environment, agent, evaluation_episodes):\n",
        "  frames = []\n",
        "\n",
        "  for episode in range(evaluation_episodes):\n",
        "    timestep = environment.reset()\n",
        "    episode_return = 0\n",
        "    steps = 0\n",
        "    while not timestep.last():\n",
        "      frames.append(environment.plot_state(return_rgb=True))\n",
        "\n",
        "      action = agent.select_action(timestep.observation)\n",
        "      timestep = environment.step(action)\n",
        "      steps += 1\n",
        "      episode_return += timestep.reward\n",
        "    print(\n",
        "        f'Episode {episode} ended with reward {episode_return} in {steps} steps'\n",
        "    )\n",
        "  return frames\n",
        "\n",
        "def display_video(frames, filename='temp.mp4', frame_repeat=1):\n",
        "  \"\"\"Save and display video.\"\"\"\n",
        "  # Write video\n",
        "  with imageio.get_writer(filename, fps=60) as video:\n",
        "    for frame in frames:\n",
        "      for _ in range(frame_repeat):\n",
        "        video.append_data(frame)\n",
        "  # Read video and display the video\n",
        "  video = open(filename, 'rb').read()\n",
        "  b64_video = base64.b64encode(video)\n",
        "  video_tag = ('<video  width=\"320\" height=\"240\" controls alt=\"test\" '\n",
        "               'src=\"data:video/mp4;base64,{0}\">').format(b64_video.decode())\n",
        "  return IPython.display.HTML(video_tag)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_0YgLdsi3kXw",
        "colab_type": "text"
      },
      "source": [
        "## Agent\n",
        "\n",
        "We will be implementing Tabular & Function Approximation agents. Tabular agents are purely in Python while for Function Approximation agents, we will use JAX."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4MAYbEvtJ1kT",
        "colab_type": "text"
      },
      "source": [
        "### Agent Implementation\n",
        "\n",
        "Each agent implements the following functions:\n",
        "\n",
        "\n",
        "\n",
        "> `__init__(self, number_of_actions, number_of_states, ...)`\n",
        "\n",
        "\n",
        "The constructor will provide the agent the number of actions and number of\n",
        "states.\n",
        "\n",
        "> `select_action(self, observation)`:\n",
        "\n",
        "This is the policy used by the actor to interact with the environment.\n",
        "\n",
        "> `observe_first(self, timestep)`:\n",
        "\n",
        "This function provides the agent with initial timestep in a given episode. Note\n",
        "that this is not the result of an action choice by the agent, hence it will only\n",
        "have `timestep.observation` set to a proper value.\n",
        "\n",
        "> `observe(self, action, next_timestep)`:\n",
        "\n",
        "This function provides the agent with the timestep that resulted from the given\n",
        "action choice. The timestep provides a `reward`, a `discount`, and an\n",
        "`observation`, all results of the previous action choice.\n",
        "\n",
        "Note: `timestep.step_type` will be either `MID` or `LAST` and should be used to\n",
        "determine whether this is the last observation in the episode.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7XD9bXC3UCHd",
        "colab_type": "text"
      },
      "source": [
        "### Random Agent\n",
        "\n",
        "We can just choose actions randomly to move around this environment."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0lU-ybzz4Ng7",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Build a random Agent  { form-width: \"30%\" }\n",
        "\n",
        "# Uniform random policy\n",
        "def random_policy(q):\n",
        "  return np.random.randint(4)\n",
        "\n",
        "# (Do not worry about the details here, we will explain the Actor class below)\n",
        "class RandomAgent(acme.Actor):\n",
        "  def select_action(self, observation):\n",
        "    return random_policy(None)    \n",
        "    \n",
        "  def observe_first(self, timestep):\n",
        "    \"\"\"The agent is being notified that environment was reset.\"\"\"\n",
        "    pass\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    \"\"\"The agent is being notified of an environment step.\"\"\"\n",
        "    pass\n",
        "\n",
        "  def update(self):    \n",
        "    \"\"\"Agent should update its parameters.\"\"\"\n",
        "    pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oxjzoRO03jGH",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Visualise agent's behaviour { form-width: \"30%\" }\n",
        "\n",
        "# This is how the random policy moves around\n",
        "frames = evaluate(environment, RandomAgent(), evaluation_episodes=1)\n",
        "display_video(frames, frame_repeat=5)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YPc0CrguF4GV",
        "colab_type": "text"
      },
      "source": [
        "# RL Lab - Part 1: Tabular Agents\n",
        "\n",
        "The first set of execises are based on the simpler case where the number of states is small enough for our agents to maintain a table of values for each individual state that it will ever encounter.\n",
        "\n",
        "In particular, we will consider the case where the GridWorld has a fixed layout, and the goal is always at the same location, hence the state is fully determined by the location of the agent. As such, the <font color='red'>observation</font> from the environment is changed to be an integer corresponding to each one of approximately 50 locations on the grid.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zL8J6nVc2zlq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "# Environment\n",
        "grid = build_gridworld_task(\n",
        "    task='simple',\n",
        "    observation_type=ObservationType.STATE_INDEX,\n",
        "    max_episode_length=200)\n",
        "environment, environment_spec = setup_environment(grid)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sAiZn23xICKh",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "## 1.0: Overview\n",
        "\n",
        "We will cover three basic RL tabular algorithms:\n",
        "- Policy iteration\n",
        "- SARSA Agent\n",
        "- Q-learning Agent\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JhHsnLcFID1u",
        "colab_type": "text"
      },
      "source": [
        "## 1.1: Policy iteration\n",
        "\n",
        "The first RL learning algorithm we will explore is **policy iteration**, which is repeating (1) Policy Evaluation and (2) Greedy Improvement until convergence.\n",
        "\n",
        "<center><img src=\"https://drive.google.com/uc?id=1lP2dFEXCBgYW744S3Lr3zMzVfOEYowdJ\" width=\"300\" /></center>\n",
        "\n",
        "For this exercise, we'll show you how to implement the \"first 2 arrows\", we will not repeat these steps to convergence yet.\n",
        "\n",
        "### 1. Policy Evaluation \n",
        "\n",
        "The purpose here is to evaluate a given policy $\\pi_e$:\n",
        "\n",
        "Compute the value function associated with following/employing this policy in a given MDP.\n",
        "\n",
        "$$ Q^{\\pi_e}(\\color{red}{s},\\color{blue}{a}) = \\mathbb{E}_{\\tau \\sim P^{\\pi_b}} \\left[ \\sum_t \\gamma^t \\color{green}{R_t}| s_0=\\color{red}s,a=\\color{blue}{a_0} \\right]$$\n",
        "\n",
        "where $\\tau = \\{\\color{red}{s_0}, \\color{blue}{a_0}, \\color{green}{r_0}, \\color{red}{s_1}, \\color{blue}{a_1}, \\color{green}{r_1}, \\cdots \\}$.\n",
        "\n",
        "\n",
        "Algorithm:\n",
        "\n",
        "**Initialize** $Q(\\color{red}{s}, \\color{blue}{a})$ for all $\\color{red}{s}$ ∈ $\\mathcal{\\color{red}S}$ and $\\color{blue}a$ ∈ $\\mathcal{\\color{blue}A}(\\color{red}s)$\n",
        "\n",
        "**Loop forever**:\n",
        "\n",
        "1. $\\color{red}{s} \\gets{}$current (nonterminal) state\n",
        " \n",
        "2. $\\color{blue}{a} \\gets{} \\text{behaviour_policy }\\pi_b(\\color{red}s)$\n",
        " \n",
        "3. Take action $\\color{blue}{a}$; observe resulting reward $\\color{green}{r}$, discount $\\gamma$, and state, $\\color{red}{s'}$\n",
        "\n",
        "4. Compute TD-error: $\\delta = \\color{green}R + \\gamma Q(\\color{red}{s'}, \\underbrace{\\pi_e(\\color{red}{s'}}_{\\color{blue}{a'}})) − Q(\\color{red}s, \\color{blue}a)$\n",
        "\n",
        "4. Update Q-value with a small $\\alpha$ step: $Q(\\color{red}s, \\color{blue}a) \\gets Q(\\color{red}s, \\color{blue}a) + \\alpha \\delta$\n",
        "\n",
        "### 2. Greedy Policy Improvement\n",
        "\n",
        "Once a good approximation to the Q-value of a policy is obtained, we can improve this policy by simply changing action selection towards those that are evaluated higher. \n",
        "\n",
        "$$ \\pi_{greedy} (\\color{blue}a|\\color{red}s) = \\arg\\max_\\color{blue}a Q^{\\pi_e}(\\color{red}s,\\color{blue}a) $$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cqrSos8dDPFb",
        "colab_type": "text"
      },
      "source": [
        "### Create a policy evaluation agent\n",
        "\n",
        "An ACME `Actor` is the part of our framework that directly interacts with an environment by generating actions. Here we borrow a figure from Acme to show how this interaction occurs:\n",
        "\n",
        "<center><img src=\"https://drive.google.com/uc?id=1T7FTpA9RgDYFkciDFZK4brNyURZN_ZGp\" width=\"500\" /></center>\n",
        "\n",
        "While you can always write your own actor, we also provide a number of useful premade versions. \n",
        "\n",
        "Tabular agents implement a function `q_values()` returning a matrix of Q values\n",
        "of shape: (`number_of_states`, `number_of_actions`)\n",
        "\n",
        "In this section, we will implement a `PolicyEvalAgent` as an ACME actor: given an `evaluation_policy` and a `behaviour_policy`, it will use the `behaviour_policy` to choose actions, and it will use the corresponding trajectory data to evaluate the `evaluation_policy` (i.e. compute the Q-values as if you were following the `evaluation_policy`)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iGtP3XRLF3qE",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Coding Task]** Policy Evaluation Agent { form-width: \"30%\" }\n",
        "class PolicyEvalAgent(acme.Actor):\n",
        "\n",
        "  def __init__(\n",
        "      self, \n",
        "      number_of_states, \n",
        "      number_of_actions,\n",
        "      evaluated_policy, \n",
        "      behaviour_policy=random_policy, \n",
        "      step_size=0.1):\n",
        "    self._state = None\n",
        "    self._number_of_states = number_of_states\n",
        "    self._number_of_actions = number_of_actions\n",
        "    self._step_size = step_size\n",
        "    self._behaviour_policy = behaviour_policy\n",
        "    self._evaluated_policy = evaluated_policy\n",
        "    \n",
        "    # ============ YOUR CODE HERE =============\n",
        "    # initialize your q-values (this is a table of state and action pairs\n",
        "    # Note: this can be random, but the code was tested w/ zero-initialization \n",
        "    pass\n",
        "    \n",
        "  @property\n",
        "  def q_values(self):\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    pass\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    pass\n",
        "    \n",
        "  def observe_first(self, timestep):\n",
        "    self._state = timestep.observation\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    s = self._state\n",
        "    a = action\n",
        "    r = next_timestep.reward\n",
        "    g = next_timestep.discount\n",
        "    next_s = next_timestep.observation\n",
        "\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    # Compute TD-Error.\n",
        "    # self._td_error =\n",
        "    pass\n",
        "\n",
        "  def update(self):    \n",
        "    # ============ YOUR CODE HERE =============\n",
        "    # Q-value table update.\n",
        "    pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BXMd87q0JVdr",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Solution]** Policy Evaluation Agent{ form-width: \"30%\" }\n",
        "class PolicyEvalAgent(acme.Actor):\n",
        "\n",
        "  def __init__(\n",
        "      self, number_of_states, number_of_actions, \n",
        "      evaluated_policy, \n",
        "      behaviour_policy=random_policy, \n",
        "      step_size=0.1):\n",
        "    self._state = None\n",
        "    self._number_of_states = number_of_states\n",
        "    self._number_of_actions = number_of_actions\n",
        "    self._step_size = step_size\n",
        "    self._behaviour_policy = behaviour_policy\n",
        "    self._evaluated_policy = evaluated_policy\n",
        "    \n",
        "    self._q = np.zeros((number_of_states, number_of_actions))\n",
        "    self._action = None\n",
        "    self._next_state = None\n",
        "\n",
        "  @property\n",
        "  def q_values(self):\n",
        "    return self._q\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    return self._behaviour_policy(self._q[observation])\n",
        "    \n",
        "  def observe_first(self, timestep):\n",
        "    self._state = timestep.observation\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    s = self._state\n",
        "    a = action\n",
        "    r = next_timestep.reward\n",
        "    g = next_timestep.discount\n",
        "    next_s = next_timestep.observation\n",
        "    \n",
        "    # Compute TD-Error.\n",
        "    self._action = a\n",
        "    self._next_state = next_s\n",
        "    \n",
        "    next_a = self._evaluated_policy(self._q[next_s])\n",
        "    self._td_error = r + g * self._q[next_s, next_a] - self._q[s, a]\n",
        "    \n",
        "    \n",
        "  def update(self):\n",
        "    # Q-value table update.\n",
        "    s = self._state\n",
        "    a = self._action\n",
        "    self._q[s, a] += self._step_size * self._td_error\n",
        "    self._state = self._next_state"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FA8FRfY-Dsth",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Helper functions for visualisation  { form-width: \"30%\" }\n",
        "\n",
        "map_from_action_to_subplot = lambda a: (2, 6, 8, 4)[a]\n",
        "map_from_action_to_name = lambda a: (\"up\", \"right\", \"down\", \"left\")[a]\n",
        "\n",
        "def plot_values(values, colormap='pink', vmin=-1, vmax=10):\n",
        "  plt.imshow(values, interpolation=\"nearest\", cmap=colormap, vmin=vmin, vmax=vmax)\n",
        "  plt.yticks([])\n",
        "  plt.xticks([])\n",
        "  plt.colorbar(ticks=[vmin, vmax])\n",
        "\n",
        "def plot_state_value(action_values, epsilon=0.1):\n",
        "  q = action_values\n",
        "  fig = plt.figure(figsize=(4, 4))\n",
        "  vmin = np.min(action_values)\n",
        "  vmax = np.max(action_values)\n",
        "  v = (1 - epsilon) * np.max(q, axis=-1) + epsilon * np.mean(q, axis=-1)\n",
        "  plot_values(v, colormap='summer', vmin=vmin, vmax=vmax)\n",
        "  plt.title(\"$v(s)$\")\n",
        "\n",
        "def plot_action_values(action_values, epsilon=0.1):\n",
        "  q = action_values\n",
        "  fig = plt.figure(figsize=(8, 8))\n",
        "  fig.subplots_adjust(wspace=0.3, hspace=0.3)\n",
        "  vmin = np.min(action_values)\n",
        "  vmax = np.max(action_values)\n",
        "  dif = vmax - vmin\n",
        "  for a in [0, 1, 2, 3]:\n",
        "    plt.subplot(3, 3, map_from_action_to_subplot(a))\n",
        "    \n",
        "    plot_values(q[..., a], vmin=vmin - 0.05*dif, vmax=vmax + 0.05*dif)\n",
        "    action_name = map_from_action_to_name(a)\n",
        "    plt.title(r\"$q(s, \\mathrm{\" + action_name + r\"})$\")\n",
        "    \n",
        "  plt.subplot(3, 3, 5)\n",
        "  v = (1 - epsilon) * np.max(q, axis=-1) + epsilon * np.mean(q, axis=-1)\n",
        "  plot_values(v, colormap='summer', vmin=vmin, vmax=vmax)\n",
        "  plt.title(\"$v(s)$\")\n",
        "      \n",
        "\n",
        "def smooth(x, window=10):\n",
        "  return x[:window*(len(x)//window)].reshape(len(x)//window, window).mean(axis=1)\n",
        "  \n",
        "def plot_stats(stats, window=10):\n",
        "  plt.figure(figsize=(16,4))\n",
        "  plt.subplot(121)\n",
        "  xline = range(0, len(stats.episode_lengths), window)\n",
        "  plt.plot(xline, smooth(stats.episode_lengths, window=window))\n",
        "  plt.ylabel('Episode Length')\n",
        "  plt.xlabel('Episode Count')\n",
        "  plt.subplot(122)\n",
        "  plt.plot(xline, smooth(stats.episode_rewards, window=window))\n",
        "  plt.ylabel('Episode Return')\n",
        "  plt.xlabel('Episode Count')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KsuM-9sEKeT1",
        "colab_type": "text"
      },
      "source": [
        "We will first see how this works on the `simple` GridWorld task.\n",
        "\n",
        "**Task 1**: Run the policy evaluation agent, evaluating the uniformly random policy on the `simple` task.\n",
        "\n",
        "Try different number of training steps, e.g. $\\texttt{num_steps} = 1e3, 1e5$. \n",
        "\n",
        "Visualise the resulting value functions $Q(\\color{red}s,\\color{blue}a)$.\n",
        "The plotting function is provided for you and it takes in a table of q-values."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YMumNsJIKhn_",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "num_steps = 1e3  #@param {type:\"number\"}\n",
        "\n",
        "# environment\n",
        "grid = build_gridworld_task(task='simple')\n",
        "environment, environment_spec = setup_environment(grid)\n",
        "\n",
        "# agent \n",
        "agent = PolicyEvalAgent(\n",
        "    number_of_states=environment_spec.observations.num_values, \n",
        "    number_of_actions=environment_spec.actions.num_values, \n",
        "    evaluated_policy=random_policy,\n",
        "    behaviour_policy=random_policy,\n",
        "    step_size=0.1)\n",
        "\n",
        "# run experiment and get the value functions from agent\n",
        "returns = run_loop(environment=environment, agent=agent, num_steps=int(num_steps))\n",
        "\n",
        "# get the q-values\n",
        "q = agent.q_values.reshape(grid._layout.shape + (4,))\n",
        "\n",
        "# visualize value functions\n",
        "print('AFTER {} STEPS ...'.format(num_steps))\n",
        "plot_action_values(q, epsilon=1.)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0EfsRbSUN0B1",
        "colab_type": "text"
      },
      "source": [
        "### Greedy Policy Improvement\n",
        "\n",
        "**Task 2**: Compute and Visualise the greedy policy based on the above evaluation, at the end of training.\n",
        "\n",
        "\n",
        "$$ \\pi_{greedy} (\\color{blue}a|\\color{red}s) = \\arg\\max_\\color{blue}a Q^{\\pi_e}(\\color{red}s,\\color{blue}a) $$\n",
        "\n",
        "**Q:** What do you observe? How does it compare to the behaviour policy we started from?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qQFPI_0c1YXo",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title **[Coding task]** Greedy policy\n",
        "def greedy(q_values):\n",
        "  pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "iTMx-QHU1f_j",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title **[Solution]** Greedy policy\n",
        "def greedy(q_values):\n",
        "  return np.argmax(q_values)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bCcnLnvLOBYZ",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title Visualize the policy on `simple` { form-width: \"30%\" }\n",
        "\n",
        "# Do here whatever works for you, but you should be able to see what the agent\n",
        "# would do at each step/state.\n",
        "\n",
        "pi = np.zeros(grid._layout_dims, dtype=np.int32)\n",
        "for i in range(grid._layout_dims[0]):\n",
        "  for j in range(grid._layout_dims[1]):\n",
        "    pi[i, j] = greedy(q[i, j])\n",
        "    \n",
        "grid.plot_policy(pi)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W3fDpu9eOUm0",
        "colab_type": "text"
      },
      "source": [
        "**Task 3**: Now try on the harder `obstacle` task and visualise the resulting value functions and the greedy policy on top of these values at the end of training.\n",
        "\n",
        "**Q:** What do you observe? \n",
        "- How does this policy compare with the optimal one?\n",
        "- Try running the training process longer -- what do you observe?"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Y4dbI0DLOeqC",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "num_steps = 1e5 #@param {type:\"number\"}\n",
        "\n",
        "# environment\n",
        "grid = build_gridworld_task(task='obstacle')\n",
        "environment, environment_spec = setup_environment(grid)\n",
        "\n",
        "# agent \n",
        "agent = PolicyEvalAgent(\n",
        "    number_of_states=environment_spec.observations.num_values, \n",
        "    number_of_actions=environment_spec.actions.num_values, \n",
        "    evaluated_policy=random_policy,\n",
        "    behaviour_policy=random_policy,\n",
        "    step_size=0.1)\n",
        "\n",
        "# run experiment and get the value functions from agent\n",
        "returns = run_loop(environment=environment, agent=agent, num_steps=int(num_steps))\n",
        "\n",
        "# get the q-values\n",
        "q = agent.q_values.reshape(grid._layout.shape + (4,))\n",
        "\n",
        "# visualize value functions\n",
        "print('AFTER {} STEPS ...'.format(num_steps))\n",
        "plot_action_values(q, epsilon=1.)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "D4KxWhW8PCD9",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title Visualise the greedy policy on `obstacle` { form-width: \"30%\" }\n",
        "grid.plot_greedy_policy(q)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EcrhrNnIr3kX",
        "colab_type": "text"
      },
      "source": [
        "## 1.2 On-policy control: SARSA Agent\n",
        "In this section, we are focusing on control RL algorithms, which perform the evaluation and improvement of the policy synchronously. That is, the policy that is being evaluated improves as the agent is using it to interact with the environent.\n",
        "\n",
        "\n",
        "The first algorithm we are going to be looking at is SARSA. This is an **on-policy algorithm** -- i.e: the data collection is done by leveraging the policy we're trying to optimize (and not just another fixed behaviour policy). \n",
        "\n",
        "As discussed during lectures, a greedy policy with respect to a given estimate of $Q^\\pi$ fails to explore the environment as needed; we will use instead an $\\epsilon$-greedy policy WRT $Q^\\pi$.\n",
        "\n",
        "### SARSA Algorithm\n",
        "\n",
        "\n",
        "**Initialize** $Q(\\color{red}{s}, \\color{blue}{a})$ for all $\\color{red}{s}$ ∈ $\\mathcal{\\color{red}S}$ and $\\color{blue}a$ ∈ $\\mathcal{\\color{blue}A}(\\color{red}s)$\n",
        "\n",
        "**Loop forever**:\n",
        "\n",
        "1. $\\color{red}s \\gets{}$current (nonterminal) state\n",
        " \n",
        "2. $\\color{blue}a \\gets{} \\text{epsilon_greedy}(Q(\\color{red}s, \\cdot))$\n",
        " \n",
        "3. Take action $\\color{blue}a$; observe resultant reward $\\color{green}r$, discount $\\gamma$, and state, $\\color{red}{s'}$\n",
        "\n",
        "4. $Q(\\color{red}s, \\color{blue}a) \\gets Q(\\color{red}s, \\color{blue}a) + \\alpha (\\color{green}r + \\gamma Q(\\color{red}{s'}, \\color{blue}{a'}) − Q(\\color{red}s, \\color{blue}a))$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xNfVHzosN2P0",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title **[Coding Task]** Epilson-greedy policy { form-width: \"30%\" }\n",
        "# Input(s): Q(s,:), epsilon\n",
        "# Output:   Sampled action based on epsilon-Greedy(Q(s,:))\n",
        "def epsilon_greedy(q_values, epsilon=0.1):\n",
        "  pass\n",
        "  #return the epsilon greedy action"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XWqlIWbwN7Mk",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title **[Solution]** Epilson-greedy policy { form-width: \"30%\" }\n",
        "def epsilon_greedy(q_values, epsilon=0.1):\n",
        "  if epsilon < np.random.random():\n",
        "    return np.argmax(q_values)\n",
        "  else:\n",
        "    return np.random.randint(np.array(q_values).shape[-1])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7bmAV4Kcr7Zz",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Coding Task]** SARSA Agent  { form-width: \"30%\" }\n",
        "class SarsaAgent(acme.Actor):\n",
        "\n",
        "  def __init__(\n",
        "      self, number_of_states, number_of_actions, epsilon, step_size=0.1): \n",
        "    self._q = np.zeros((number_of_states, number_of_actions))\n",
        "    self._step_size = step_size\n",
        "    self._epsilon = epsilon\n",
        "    self._state = None\n",
        "    self._action = None\n",
        "    self._next_state = None\n",
        "    \n",
        "  @property\n",
        "  def q_values(self):\n",
        "    return self._q\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    return epsilon_greedy(self._q[observation], self._epsilon)\n",
        "    \n",
        "  def observe_first(self, timestep):\n",
        "    self._state = timestep.observation\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    s = self._state\n",
        "    a = action\n",
        "    r = next_timestep.reward\n",
        "    g = next_timestep.discount\n",
        "    next_s = next_timestep.observation\n",
        "    \n",
        "    # ============ YOUR CODE HERE =============\n",
        "    # Online Q-value update\n",
        "    # self._td_error =\n",
        "    pass\n",
        "\n",
        "  def update(self):\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    # Q-value table update\n",
        "    pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JtlH1tU7sCEm",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Solution]** SARSA Agent { form-width: \"30%\" }\n",
        "class SarsaAgent(acme.Actor):\n",
        "\n",
        "  def __init__(\n",
        "      self, number_of_states, number_of_actions, epsilon, step_size=0.1): \n",
        "    self._q = np.zeros((number_of_states, number_of_actions))\n",
        "    self._number_of_states = number_of_states\n",
        "    self._number_of_actions = number_of_actions\n",
        "    self._step_size = step_size\n",
        "    self._epsilon = epsilon\n",
        "    self._state = None\n",
        "    self._action = None\n",
        "    self._next_state = None\n",
        "    \n",
        "  @property\n",
        "  def q_values(self):\n",
        "    return self._q\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    return epsilon_greedy(self._q[observation], self._epsilon)\n",
        "    \n",
        "  def observe_first(self, timestep):\n",
        "    self._state = timestep.observation\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    s = self._state\n",
        "    a = action\n",
        "    r = next_timestep.reward\n",
        "    g = next_timestep.discount\n",
        "    next_s = next_timestep.observation\n",
        "    next_a = epsilon_greedy(self._q[next_s], self._epsilon)\n",
        "    \n",
        "    # Online Q-value update\n",
        "    self._action = a\n",
        "    self._next_state = next_s\n",
        "    self._td_error = r + g * self._q[next_s, next_a] - self._q[s, a]\n",
        "\n",
        "  def update(self):\n",
        "    s = self._state\n",
        "    a = self._action\n",
        "    self._q[s, a] += self._step_size * self._td_error\n",
        "    self._state = self._next_state"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K8eBOcXZu1fM",
        "colab_type": "text"
      },
      "source": [
        "### **Task**: Run your SARSA agent on the `obstacle` environment"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "xKYEB2d2uGaa",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "num_steps = 1e5 #@param {type:\"number\"}\n",
        "\n",
        "# environment\n",
        "grid = build_gridworld_task(task='obstacle')\n",
        "environment, environment_spec = setup_environment(grid)\n",
        "\n",
        "# agent \n",
        "agent = SarsaAgent(\n",
        "    number_of_states=environment_spec.observations.num_values, \n",
        "    number_of_actions=environment_spec.actions.num_values, \n",
        "    epsilon=0.1,\n",
        "    step_size=0.1)\n",
        "\n",
        "# run experiment and get the value functions from agent\n",
        "returns = run_loop(environment=environment, agent=agent, num_steps=int(num_steps))\n",
        "\n",
        "# get the q-values\n",
        "q = agent.q_values.reshape(grid._layout.shape + (4,))\n",
        "\n",
        "# visualize value functions\n",
        "print('AFTER {} STEPS ...'.format(num_steps))\n",
        "plot_action_values(q, epsilon=1.)\n",
        "\n",
        "# visualise the greedy policy\n",
        "grid.plot_greedy_policy(q)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pFGX_zGcvb8D",
        "colab_type": "text"
      },
      "source": [
        "## 1.3 Off-policy control: Q-learning Agent\n",
        "\n",
        "Reminder: Q-learning is a very powerful and general algorithm, that enables control (figuring out the optimal policy/value function) both on and off-policy.\n",
        "\n",
        "**Initialize** $Q(\\color{red}{s}, \\color{blue}{a})$ for all $\\color{red}{s} \\in \\color{red}{\\mathcal{S}}$ and $\\color{blue}{a} \\in \\color{blue}{\\mathcal{A}}(\\color{red}{s})$\n",
        "\n",
        "**Loop forever**:\n",
        "\n",
        "1. $\\color{red}{s} \\gets{}$current (nonterminal) state\n",
        " \n",
        "2. $\\color{blue}{a} \\gets{} \\text{behaviour_policy}(\\color{red}{s})$\n",
        " \n",
        "3. Take action $\\color{blue}{a}$; observe resultant reward $\\color{green}{R}$, discount $\\gamma$, and state, $\\color{red}{s'}$\n",
        "\n",
        "4. $Q(\\color{red}{s}, \\color{blue}{a}) \\gets Q(\\color{red}{s}, \\color{blue}{a}) + \\alpha (\\color{green}{R} + \\gamma \\max_{\\color{blue}{a'}} Q(\\color{red}{s'}, \\color{blue}{a'}) − Q(\\color{red}{s}, \\color{blue}{a}))$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I6s820jAwoVA",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Coding Task]** Q-Learning Agent  { form-width: \"30%\" }\n",
        "class QLearningAgent(acme.Actor):\n",
        "\n",
        "  def __init__(\n",
        "      self, number_of_states, number_of_actions, behaviour_policy, step_size=0.1): \n",
        "    self._q = np.zeros((number_of_states, number_of_actions))\n",
        "    self._step_size = step_size\n",
        "    self._behaviour_policy = behaviour_policy\n",
        "    self._state = None\n",
        "    self._action = None\n",
        "    self._next_state = None\n",
        "    \n",
        "  @property\n",
        "  def q_values(self):\n",
        "    return self._q\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    return self._behaviour_policy(self._q[observation])\n",
        "    \n",
        "  def observe_first(self, timestep):\n",
        "    self._state = timestep.observation\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    s = self._state\n",
        "    a = action\n",
        "    r = next_timestep.reward\n",
        "    g = next_timestep.discount\n",
        "    next_s = next_timestep.observation\n",
        "    \n",
        "    # ============ YOUR CODE HERE =============\n",
        "    # Offline Q-value update\n",
        "    # self._td_error =\n",
        "    pass\n",
        "\n",
        "  def update(self):\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    pass"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ak1T5PNV8Pbk",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Solution]** Q-Learning Agent { form-width: \"30%\" }\n",
        "class QLearningAgent(acme.Actor):\n",
        "\n",
        "  def __init__(\n",
        "      self, number_of_states, number_of_actions, behaviour_policy, step_size=0.1): \n",
        "    self._q = np.zeros((number_of_states, number_of_actions))\n",
        "    self._step_size = step_size\n",
        "    self._behaviour_policy = behaviour_policy\n",
        "    self._state = None\n",
        "    self._action = None\n",
        "    self._next_state = None\n",
        "    \n",
        "  @property\n",
        "  def q_values(self):\n",
        "    return self._q\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    return self._behaviour_policy(self._q[observation])\n",
        "    \n",
        "  def observe_first(self, timestep):\n",
        "    self._state = timestep.observation\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    s = self._state\n",
        "    a = action\n",
        "    r = next_timestep.reward\n",
        "    g = next_timestep.discount\n",
        "    next_s = next_timestep.observation\n",
        "    \n",
        "    # Offline Q-value update\n",
        "    self._action = a\n",
        "    self._next_state = next_s\n",
        "    self._td_error = r + g * np.max(self._q[next_s]) - self._q[s, a]\n",
        "\n",
        "  def update(self):\n",
        "    s = self._state\n",
        "    a = self._action\n",
        "    self._q[s, a] += self._step_size * self._td_error\n",
        "    self._state = self._next_state"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2RqdV3rjwcAh",
        "colab_type": "text"
      },
      "source": [
        "### **Task 1**: Run your Q-learning agent on `obstacle`\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LL4PgT-jwi3-",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "epsilon = 1  #@param {type:\"number\"} \n",
        "num_steps = 1e5  #@param {type:\"number\"}\n",
        "\n",
        "# environment\n",
        "grid = build_gridworld_task(task='obstacle')\n",
        "environment, environment_spec = setup_environment(grid)\n",
        "\n",
        "# behavior policy\n",
        "behavior_policy = lambda qval: epsilon_greedy(qval, epsilon=epsilon)\n",
        "\n",
        "# agent\n",
        "agent = QLearningAgent(\n",
        "    number_of_states=environment_spec.observations.num_values,\n",
        "    number_of_actions=environment_spec.actions.num_values,\n",
        "    behaviour_policy=behavior_policy,\n",
        "    step_size=0.1)\n",
        "\n",
        "# run experiment and get the value functions from agent\n",
        "returns = run_loop(environment=environment, agent=agent, num_steps=int(num_steps))\n",
        "\n",
        "# get the q-values\n",
        "q = agent.q_values.reshape(grid._layout.shape + (4,))\n",
        "\n",
        "# visualize value functions\n",
        "print('AFTER {} STEPS ...'.format(num_steps))\n",
        "plot_action_values(q, epsilon=epsilon)\n",
        "\n",
        "# visualise the greedy policy\n",
        "grid.plot_greedy_policy(q)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cMk2ArG-weg_",
        "colab_type": "text"
      },
      "source": [
        "### **Task 2:** Experiment with different levels of 'greediness'\n",
        "* The default was $\\epsilon=1.$, what does this correspond to?\n",
        "* Try also $\\epsilon =0.1, 0.5$. What do you observe? Does the behaviour policy affect the training in any way?"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lqg1n48y81ei",
        "colab_type": "text"
      },
      "source": [
        "## 1.4 **[Homework]** Experience Replay\n",
        "\n",
        "Implement an agent that uses **Experience Replay** to learn action values, at each step:\n",
        "* select actions randomly\n",
        "* accumulate all observed transitions *(s, a, r, s')* in the environment in a *replay buffer*,\n",
        "* apply an online Q-learning \n",
        "* apply multiple Q-learning updates based on transitions sampled from the *replay buffer* (in addition to the online updates).\n",
        "\n",
        "\n",
        "**Initialize** $Q(\\color{red}s, \\color{blue}a)$ for all $\\color{red}{s} ∈ \\mathcal{\\color{red}S}$ and $\\color{blue}a ∈ \\mathcal{\\color{blue}A}(\\color{red}s)$\n",
        "\n",
        "**Loop forever**:\n",
        "\n",
        "1. $\\color{red}{s} \\gets{}$current (nonterminal) state\n",
        " \n",
        "2. $\\color{blue}{a} \\gets{} \\text{random_action}(\\color{red}{s})$\n",
        " \n",
        "3. Take action $\\color{blue}{a}$; observe resultant reward $\\color{green}{r}$, discount $\\gamma$, and state, $\\color{red}{s'}$\n",
        "\n",
        "4. $Q(\\color{red}{s}, \\color{blue}{a}) \\gets Q(\\color{red}{s}, \\color{blue}{a}) + \\alpha (\\color{green}{r} + \\gamma Q(\\color{red}{s'}, \\color{blue}{a'}) − Q(\\color{red}{s}, \\color{blue}{a}))$\n",
        "\n",
        "5. $\\text{ReplayBuffer.append_transition}(s, a, r, \\gamma, s')$\n",
        "\n",
        "6. Loop repeat n times:\n",
        "\n",
        "  1. $\\color{red}{s}, \\color{blue}{a}, \\color{green}{r}, \\gamma, \\color{red}{s'} \\gets \\text{ReplayBuffer}.\\text{sample_transition}()$\n",
        "  \n",
        "  4. $Q(\\color{red}{s}, \\color{blue}{a}) \\gets Q(\\color{red}{s}, \\color{blue}{a}) + \\alpha (\\color{green}{r} + \\gamma \\max_\\color{blue}{a'} Q(\\color{red}{s'}, \\color{blue}{a'}) − Q(\\color{red}{s}, \\color{blue}{a}))$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ietFnV739JwD",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Coding Task]** Q-learning AGENT with a simple replay buffer { form-width: \"30%\" }\n",
        "class ReplayQLearningAgent(acme.Actor):\n",
        "\n",
        "  def __init__(\n",
        "      self, number_of_states, number_of_actions, behaviour_policy, \n",
        "      num_offline_updates=0, step_size=0.1): \n",
        "    self._q = np.zeros((number_of_states, number_of_actions))\n",
        "    self._step_size = step_size\n",
        "    self._behaviour_policy = behaviour_policy\n",
        "    self._num_offline_updates = num_offline_updates\n",
        "    self._state = None\n",
        "    self._action = None\n",
        "    self._next_state = None\n",
        "    self._replay_buffer = []\n",
        "    \n",
        "  @property\n",
        "  def q_values(self):\n",
        "    return self._q\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    return self._behaviour_policy(self._q[observation])\n",
        "    \n",
        "  def observe_first(self, timestep):\n",
        "    self._state = timestep.observation\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    s = self._state\n",
        "    a = action\n",
        "    r = next_timestep.reward\n",
        "    g = next_timestep.discount\n",
        "    next_s = next_timestep.observation\n",
        "    \n",
        "    # Offline Q-value update\n",
        "    self._action = a\n",
        "    self._next_state = next_s\n",
        "    self._td_error = r + g * np.max(self._q[next_s]) - self._q[s, a]\n",
        "\n",
        "    if self._num_offline_updates > 0:\n",
        "      # ============ YOUR CODE HERE =============\n",
        "      # Update replay buffer.\n",
        "      pass\n",
        "    \n",
        "  def update(self):\n",
        "    s = self._state\n",
        "    a = self._action\n",
        "    self._q[s, a] += self._step_size * self._td_error\n",
        "    self._state = self._next_state\n",
        "\n",
        "    # Offline Q-value update\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "I6Lunsx1-kmf",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Solution]**  Q-learning AGENT with a simple replay buffer { form-width: \"30%\" }\n",
        "class ReplayQLearningAgent(acme.Actor):\n",
        "\n",
        "  def __init__(\n",
        "      self, number_of_states, number_of_actions, behaviour_policy, \n",
        "      num_offline_updates=0, step_size=0.1): \n",
        "    self._q = np.zeros((number_of_states, number_of_actions))\n",
        "    self._step_size = step_size\n",
        "    self._behaviour_policy = behaviour_policy\n",
        "    self._num_offline_updates = num_offline_updates\n",
        "    self._state = None\n",
        "    self._action = None\n",
        "    self._next_state = None\n",
        "    self._replay_buffer = []\n",
        "    \n",
        "  @property\n",
        "  def q_values(self):\n",
        "    return self._q\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    return self._behaviour_policy(self._q[observation])\n",
        "    \n",
        "  def observe_first(self, timestep):\n",
        "    self._state = timestep.observation\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    s = self._state\n",
        "    a = action\n",
        "    r = next_timestep.reward\n",
        "    g = next_timestep.discount\n",
        "    next_s = next_timestep.observation\n",
        "    \n",
        "    # Offline Q-value update\n",
        "    self._action = a\n",
        "    self._next_state = next_s\n",
        "    self._td_error = r + g * np.max(self._q[next_s]) - self._q[s, a]\n",
        "\n",
        "    if self._num_offline_updates > 0:\n",
        "      self._replay_buffer.append((s, a, r, g, next_s))\n",
        "\n",
        "  def update(self):\n",
        "    s = self._state\n",
        "    a = self._action\n",
        "    self._q[s, a] += self._step_size * self._td_error\n",
        "    self._state = self._next_state\n",
        "\n",
        "    # Offline Q-value update\n",
        "    if len(self._replay_buffer) > self._num_offline_updates:\n",
        "      for i in range(self._num_offline_updates):\n",
        "        idx = np.random.randint(0, len(self._replay_buffer))\n",
        "        s, a, r, g, next_s = self._replay_buffer[idx]\n",
        "        td_error = r + g * np.max(self._q[next_s]) - self._q[s, a]\n",
        "        self._q[s, a] += self._step_size * td_error"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k3J6CE2M_AdF",
        "colab_type": "text"
      },
      "source": [
        "### **Task**: Compare Q-learning with/without experience replay\n",
        "\n",
        "Use a small number of training steps (e.g. `num_steps = 1e3`) and vary `num_offline_updates` between `0` and `30`."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9yLCXKBH_F0j",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "num_offline_updates = 0#@param {type:\"integer\"}\n",
        "num_steps = 1e3\n",
        "\n",
        "grid = build_gridworld_task(task='obstacle')\n",
        "environment, environment_spec = setup_environment(grid)\n",
        "\n",
        "agent = ReplayQLearningAgent(\n",
        "    number_of_states=environment_spec.observations.num_values,\n",
        "    number_of_actions=environment_spec.actions.num_values,\n",
        "    behaviour_policy=random_policy,\n",
        "    num_offline_updates=num_offline_updates,\n",
        "    step_size=0.1)\n",
        "\n",
        "# run experiment and get the value functions from agent\n",
        "returns = run_loop(environment=environment, agent=agent, num_steps=int(num_steps))\n",
        "\n",
        "q = agent.q_values.reshape(grid._layout.shape + (4,))\n",
        "plot_action_values(q)\n",
        "\n",
        "grid.plot_greedy_policy(q)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rkn2ud_0Pn2o",
        "colab_type": "text"
      },
      "source": [
        "# RL Lab - Part 2: Function Approximation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yxqnvCLoe3KU",
        "colab_type": "text"
      },
      "source": [
        "<img src=\"https://drive.google.com/uc?id=1oqIQNM_tMPmP8l38C_3yp5uUego3S8kV\" width=\"500\" />\n",
        "\n",
        "So far we only considered look-up tables. In all previous cases every state and action pair $(\\color{red}{s}, \\color{blue}{a})$, had an entry in our Q table. Again, this is possible in this environment as the number of states is a equal to the number of cells in the grid. But this is not scalable to situations where, say, the goal location changes or the obstacles are in different locations at every episode (consider how big the table should be in this situation?).\n",
        "\n",
        "As example (not covered in this tutorial) is ATARI from pixels, where the number of possible frames an agent can see is exponential in the number of pixels on the screen.\n",
        "\n",
        "<center><img width=\"200\" alt=\"portfolio_view\" src=\"https://miro.medium.com/max/1760/1*XyIpmXXAjbXerDzmGQL1yA.gif\"></center>\n",
        "\n",
        "But what we **really** want is just being able to *compute* the Q-value, when fed with a particular $(\\color{red}{s}, \\color{blue}{a})$ pair. So if we had a way to get a function to do this work instead of keeping a big table, we'd get around this problem.\n",
        "\n",
        "To address this, we can use **Function Approximation** as a way to generalize Q-values over some representation of the very large state space, and **train** them to output the values they should. In this section, we will explore Q-Learning with function approximation, which, although theoretically proven to diverge for some degenerate MDPs, can yield impressive results in very large environments. In particular, we will look at [Neural Fitted Q (NFQ) Iteration](http://ml.informatik.uni-freiburg.de/former/_media/publications/rieecml05.pdf).\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OTAYVPnaJN0t",
        "colab_type": "text"
      },
      "source": [
        "## 2.1 NFQ agent"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-omtUOQCS8VI",
        "colab_type": "text"
      },
      "source": [
        "[Neural Fitted Q Iteration](http://ml.informatik.uni-freiburg.de/former/_media/publications/rieecml05.pdf) was one of the first papers to demonstrate how to leverage recent advances in Deep Learning to approximate the Q-value by a neural network $^1$.\n",
        "\n",
        "We represent $Q(\\color{red}s, \\color{blue}a)$ as a neural network $f()$. which given a vector $\\color{red}s$, will output a vector of Q-values for all possible actions $\\color{blue}a$.$^2$\n",
        "\n",
        "When introducing function approximations, and neural networks in particular, we need to have a loss to optimize. But looking back at the tabular setting above, you can see that we already have some notion of error: the **TD error**.\n",
        "\n",
        "By training our neural network to output values such that the *TD error is minimized*, we will also satisfy the Bellman Optimality Equation, which is a good sufficient condition to enforce, so that we may obtain an optimal policy.\n",
        "Thanks to automatic differentiation, we can just write the TD error as a loss (e.g. with a $L2$ loss, but others would work too), compute its gradient (which are now gradients with respect to individual parameters of the neural network) and slowly improve our Q-value approximation:\n",
        "\n",
        "$$Loss = \\mathbb{E}\\left[ \\left( \\color{green}{r} + \\gamma \\max_\\color{blue}{a'} Q(\\color{red}{s'}, \\color{blue}{a'}) − Q(\\color{red}{s}, \\color{blue}{a})  \\right)^2\\right]$$\n",
        "\n",
        "\n",
        "NFQ builds on Q-learning, but if one were to update the Q-values online directly, the training can be unstable and very slow.\n",
        "Instead, NFQ uses a Replay buffer, similar to what you just implemented above, to update the Q-value in a batched setting.\n",
        "\n",
        "When it was introduced, it also was entirely off-policy (i.e. one would use a random policy to collect data), and is prone to unstability when applied to more complex environments (e.g. when the input are pixels or the tasks are longer and complicated).\n",
        "But it is a good stepping stone to the more complex agents used today. Here, we will look at a slightly different and modernised implementation of NFQ.\n",
        "\n",
        "<br />\n",
        "\n",
        "---\n",
        "\n",
        "<sub>*$^1$ if you read the NFQ paper, they use a \"control\" notation, where there is a \"cost to minimize\", instead of \"rewards to maximize\", so don't be surprised if signs/max/min do not correspond.* </sub>\n",
        "\n",
        "<sub>*$^2$ we could feed it $\\color{blue}a$ as well and ask $f$ for a single scalar value, but given we have a fixed number of actions and we usually need to take an $argmax$ over them, it's easiest to just output them all in one pass.*</sub>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_NjO3wD-Sphk",
        "colab_type": "text"
      },
      "source": [
        "<center><img src=\"https://drive.google.com/uc?id=1ivTQBHWkYi_J9vWwXFd2sSWg5f2TB5T-\" width=\"400\" /></center> "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KULsnljicr9t",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Coding Task]** NFQ Agent  { form-width: \"30%\" }\n",
        "\n",
        "Transitions = collections.namedtuple('Transitions',\n",
        "                                     ['s_t', 'a_t', 'r_t', 'd_t', 's_tp1'])\n",
        "TrainingState = namedtuple('TrainingState', 'params, opt_state, step')\n",
        "\n",
        "\n",
        "class NeuralFittedQAgent(acme.Actor):\n",
        "\n",
        "  def __init__(self,\n",
        "               q_network,\n",
        "               observation_spec,\n",
        "               replay_capacity=100000,\n",
        "               epsilon=0.1,\n",
        "               batch_size=1,\n",
        "               learning_rate=3e-4):\n",
        "\n",
        "    self._observation_spec = observation_spec\n",
        "    self.epsilon = epsilon\n",
        "    self._batch_size = batch_size\n",
        "    self._replay_buffer = ReplayBuffer(replay_capacity)\n",
        "    self.last_loss = 0\n",
        "\n",
        "    # Setup Network and loss with Haiku\n",
        "    self._rng = hk.PRNGSequence(1)\n",
        "    self._q_network = hk.transform(q_network)\n",
        "\n",
        "    # Initialize network\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    # init_params = \n",
        "\n",
        "    # Setup optimizer\n",
        "    self._optimizer = optix.adam(learning_rate)\n",
        "    initial_optimizer_state = self._optimizer.init(initial_params)\n",
        "\n",
        "    self._state = TrainingState(\n",
        "        params=initial_params, opt_state=initial_optimizer_state, step=0)\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=(0,))\n",
        "  def _policy(self, params: hk.Params, rng_key: jnp.ndarray,\n",
        "              observation: jnp.ndarray, epsilon: float):\n",
        "    # You can use rlax.epsilon_greedy here\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    pass\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    return self._policy(self._state.params, next(self._rng), observation,\n",
        "                        self.epsilon)\n",
        "\n",
        "  def q_values(self, observation):\n",
        "    return jnp.squeeze(\n",
        "        self._q_network.apply(self._state.params, observation[None, ...]),\n",
        "        axis=0)\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=(0,))\n",
        "  def _loss(self, params: hk.Params, transitions: Transitions):\n",
        "\n",
        "    def _td_error(q_s, q_next_s, a, r, d):\n",
        "      \"\"\"TD error for a single transition.\"\"\"\n",
        "      # ============ YOUR CODE HERE =============\n",
        "      pass\n",
        "\n",
        "    batch_td_error = jax.vmap(_td_error)\n",
        "\n",
        "    # Compute batched Q-values [Batch, actions]\n",
        "    q_s = self._q_network.apply(params, transitions.s_t)\n",
        "    q_next_s = self._q_network.apply(params, transitions.s_tp1)\n",
        "    # Get batched td errors\n",
        "    td_errors = batch_td_error(q_s, q_next_s, transitions.a_t, transitions.r_t,\n",
        "                               transitions.d_t)\n",
        "    losses = 0.5 * td_errors**2.  # [Batch]\n",
        "\n",
        "    return jnp.mean(losses)\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=(0,))\n",
        "  def _train_step(self, state: TrainingState, transitions: Transitions):\n",
        "    # Do one learning step on the batch of transitions\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    # Use jax.value_and_grad to compute gradients and values from _loss, \n",
        "    # and optix.apply_updates to compute new parameters for the network.\n",
        "   \n",
        "    new_state = TrainingState(\n",
        "        params=new_params, opt_state=new_opt_state, step=state.step + 1)\n",
        "    return new_state, loss\n",
        "\n",
        "  def update(self):\n",
        "    if self._replay_buffer.is_ready(self._batch_size):\n",
        "      # Collect a minibatch of random transitions\n",
        "      transitions = Transitions(*self._replay_buffer.sample(self._batch_size))\n",
        "      # Compute loss and update parameters\n",
        "      self._state, self.last_loss = self._train_step(self._state, transitions)\n",
        "\n",
        "  def observe_first(self, timestep):\n",
        "    self._replay_buffer.push(timestep, None)\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    self._replay_buffer.push(next_timestep, action)\n",
        "\n",
        "\n",
        "class ReplayBuffer(object):\n",
        "  \"\"\"A simple Python replay buffer.\"\"\"\n",
        "\n",
        "  def __init__(self, capacity):\n",
        "    self._prev = None\n",
        "    self._action = None\n",
        "    self._latest = None\n",
        "    self.buffer = collections.deque(maxlen=capacity)\n",
        "\n",
        "  def push(self, timestep, action):\n",
        "    self._prev = self._latest\n",
        "    self._action = action\n",
        "    self._latest = timestep\n",
        "\n",
        "    if action is not None:\n",
        "      self.buffer.append(\n",
        "          (self._prev.observation, self._action, self._latest.reward,\n",
        "           self._latest.discount, self._latest.observation))\n",
        "\n",
        "  def sample(self, batch_size):\n",
        "    obs_tm1, a_tm1, r_t, discount_t, obs_t = zip(\n",
        "        *random.sample(self.buffer, batch_size))\n",
        "    return (jnp.stack(obs_tm1), jnp.asarray(a_tm1), jnp.asarray(r_t),\n",
        "            jnp.asarray(discount_t), jnp.stack(obs_t))\n",
        "\n",
        "  def is_ready(self, batch_size):\n",
        "    return batch_size <= len(self.buffer)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DWrRFI_qLmmt",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title **[Solution]** NFQ Agent  { form-width: \"30%\" }\n",
        "\n",
        "Transitions = collections.namedtuple('Transitions',\n",
        "                                     ['s_t', 'a_t', 'r_t', 'd_t', 's_tp1'])\n",
        "TrainingState = namedtuple('TrainingState', 'params, opt_state, step')\n",
        "\n",
        "\n",
        "class NeuralFittedQAgent(acme.Actor):\n",
        "\n",
        "  def __init__(self,\n",
        "               q_network,\n",
        "               observation_spec,\n",
        "               replay_capacity=100000,\n",
        "               epsilon=0.1,\n",
        "               batch_size=1,\n",
        "               learning_rate=3e-4):\n",
        "\n",
        "    self._observation_spec = observation_spec\n",
        "    self.epsilon = epsilon\n",
        "    self._batch_size = batch_size\n",
        "    self._replay_buffer = ReplayBuffer(replay_capacity)\n",
        "    self.last_loss = 0\n",
        "\n",
        "    # Setup Network and loss with Haiku\n",
        "    self._rng = hk.PRNGSequence(1)\n",
        "    self._q_network = hk.transform(q_network)\n",
        "\n",
        "    # Initialize network\n",
        "    dummy_observation = observation_spec.generate_value()\n",
        "    initial_params = self._q_network.init(\n",
        "        next(self._rng), dummy_observation[None, ...])\n",
        "\n",
        "    # Setup optimizer\n",
        "    self._optimizer = optix.adam(learning_rate)\n",
        "    initial_optimizer_state = self._optimizer.init(initial_params)\n",
        "\n",
        "    self._state = TrainingState(\n",
        "        params=initial_params, opt_state=initial_optimizer_state, step=0)\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=(0,))\n",
        "  def _policy(self, params: hk.Params, rng_key: jnp.ndarray,\n",
        "              observation: jnp.ndarray, epsilon: float):\n",
        "    q_values = self._q_network.apply(params, observation[None, ...])\n",
        "    actions = rlax.epsilon_greedy(epsilon).sample(rng_key, q_values)\n",
        "    return jnp.squeeze(actions, axis=0)\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    return self._policy(self._state.params, next(self._rng), observation,\n",
        "                        self.epsilon)\n",
        "\n",
        "  def q_values(self, observation):\n",
        "    return jnp.squeeze(\n",
        "        self._q_network.apply(self._state.params, observation[None, ...]),\n",
        "        axis=0)\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=(0,))\n",
        "  def _loss(self, params: hk.Params, transitions: Transitions):\n",
        "\n",
        "    def _td_error(q_s, q_next_s, a, r, d):\n",
        "      \"\"\"TD error for a single transition.\"\"\"\n",
        "      target_s = r + d * jnp.max(q_next_s)\n",
        "      td_error = jax.lax.stop_gradient(target_s) - q_s[a]\n",
        "      # Task: think of why we are not using td_error = target_s - q_s[a]? \n",
        "      return td_error\n",
        "\n",
        "    batch_td_error = jax.vmap(_td_error)\n",
        "\n",
        "    # Compute batched Q-values [Batch, actions]\n",
        "    q_s = self._q_network.apply(params, transitions.s_t)\n",
        "    q_next_s = self._q_network.apply(params, transitions.s_tp1)\n",
        "    # Get batched td errors\n",
        "    td_errors = batch_td_error(q_s, q_next_s, transitions.a_t, transitions.r_t,\n",
        "                               transitions.d_t)\n",
        "    losses = 0.5 * td_errors**2.  # [Batch]\n",
        "\n",
        "    return jnp.mean(losses)\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=(0,))\n",
        "  def _train_step(self, state: TrainingState, transitions: Transitions):\n",
        "    # Do one learning step on the batch of transitions\n",
        "    compute_loss_and_grad = jax.value_and_grad(self._loss)\n",
        "    loss, dloss_dparams = compute_loss_and_grad(state.params, transitions)\n",
        "    updates, new_opt_state = self._optimizer.update(dloss_dparams,\n",
        "                                                    state.opt_state)\n",
        "    new_params = optix.apply_updates(state.params, updates)\n",
        "\n",
        "    new_state = TrainingState(\n",
        "        params=new_params, opt_state=new_opt_state, step=state.step + 1)\n",
        "    return new_state, loss\n",
        "\n",
        "  def update(self):\n",
        "    if self._replay_buffer.is_ready(self._batch_size):\n",
        "      # Collect a minibatch of random transitions\n",
        "      transitions = Transitions(*self._replay_buffer.sample(self._batch_size))\n",
        "      # Compute loss and update parameters\n",
        "      self._state, self.last_loss = self._train_step(self._state, transitions)\n",
        "\n",
        "  def observe_first(self, timestep):\n",
        "    self._replay_buffer.push(timestep, None)\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    self._replay_buffer.push(next_timestep, action)\n",
        "\n",
        "\n",
        "class ReplayBuffer(object):\n",
        "  \"\"\"A simple Python replay buffer.\"\"\"\n",
        "\n",
        "  def __init__(self, capacity):\n",
        "    self._prev = None\n",
        "    self._action = None\n",
        "    self._latest = None\n",
        "    self.buffer = collections.deque(maxlen=capacity)\n",
        "\n",
        "  def push(self, timestep, action):\n",
        "    self._prev = self._latest\n",
        "    self._action = action\n",
        "    self._latest = timestep\n",
        "\n",
        "    if action is not None:\n",
        "      self.buffer.append(\n",
        "          (self._prev.observation, self._action, self._latest.reward,\n",
        "           self._latest.discount, self._latest.observation))\n",
        "\n",
        "  def sample(self, batch_size):\n",
        "    obs_tm1, a_tm1, r_t, discount_t, obs_t = zip(\n",
        "        *random.sample(self.buffer, batch_size))\n",
        "    return (jnp.stack(obs_tm1), jnp.asarray(a_tm1), jnp.asarray(r_t),\n",
        "            jnp.asarray(discount_t), jnp.stack(obs_t))\n",
        "\n",
        "  def is_ready(self, batch_size):\n",
        "    return batch_size <= len(self.buffer)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MQoI1y88Mfsz",
        "colab_type": "text"
      },
      "source": [
        "### **Task: Train a NFQ agent**\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "g7QmF3UGgYJa",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Training the NFQ Agent.  { form-width: \"30%\" }\n",
        "epsilon = 1. #@param {type:\"number\"}\n",
        "\n",
        "max_episode_length = 200\n",
        "\n",
        "# Environment\n",
        "grid = build_gridworld_task(\n",
        "    task='simple',\n",
        "    observation_type=ObservationType.AGENT_GOAL_POS,\n",
        "    max_episode_length=max_episode_length)\n",
        "environment, environment_spec = setup_environment(grid)\n",
        "\n",
        "# Define function approximation for the Q-values\n",
        "# i.e. Q_a(s) for a in num_actions.\n",
        "def q_network(observation: np.ndarray):\n",
        "  \"\"\"Outputs action values given an observation.\"\"\"\n",
        "  model = hk.Sequential([\n",
        "      hk.Flatten(),  # Flattens everything except the batch dimension\n",
        "      hk.nets.MLP([50, 50, environment_spec.actions.num_values])\n",
        "  ])\n",
        "  return model(observation)\n",
        "\n",
        "# Build the trainable Q-learning agent\n",
        "agent = NeuralFittedQAgent(\n",
        "    q_network,\n",
        "    environment_spec.observations,\n",
        "    epsilon=epsilon,\n",
        "    replay_capacity=100000,\n",
        "    batch_size=10,\n",
        "    learning_rate=1e-3)\n",
        "\n",
        "returns = run_loop(\n",
        "    environment=environment,\n",
        "    agent=agent,\n",
        "    num_episodes=300,\n",
        "    logger_time_delta=1.,\n",
        "    log_loss=True)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YWbMwjdgmxGe",
        "colab_type": "text"
      },
      "source": [
        "### Evaluate the policy it learned"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bZM2TNJ0PB6F",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Evaluating the agent.  { form-width: \"30%\" }\n",
        "\n",
        "# Change epsilon to be more greedy\n",
        "agent.epsilon = 0.05\n",
        "\n",
        "# Look at a few episodes\n",
        "frames = evaluate(environment, agent, evaluation_episodes=5)\n",
        "display_video(frames, frame_repeat=10)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vYmDVoZ4sDjJ",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Visualise the learned Q values\n",
        "\n",
        "# Evaluate the policy for every state, similar to tabular agents above.\n",
        "\n",
        "environment.reset()\n",
        "pi = np.zeros(grid._layout_dims, dtype=np.int32)\n",
        "q = np.zeros(grid._layout_dims + (4,))\n",
        "for y in range(grid._layout_dims[0]):\n",
        "  for x in range(grid._layout_dims[1]):\n",
        "    # Hack observation to see what the Q-network would output at that point.\n",
        "    environment.set_state(x, y)\n",
        "    obs = environment.get_obs()\n",
        "    q[y, x] = np.asarray(agent.q_values(obs))\n",
        "    pi[y, x] = np.asarray(agent.select_action(obs))\n",
        "    \n",
        "plot_action_values(q)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ek3MNCu0LGBE",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Compare the greedy policy with the behaviour policy { form-width: \"30%\" }\n",
        "\n",
        "grid.plot_greedy_policy(q)\n",
        "plt.title('Greedy policy using the learnt Q-values')\n",
        "\n",
        "grid.plot_policy(pi)\n",
        "_ = plt.title(\"Policy using the agent's behaviour policy\")\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Clv_QlpgoY1J",
        "colab_type": "text"
      },
      "source": [
        "# RL Lab - Part 3: Deep Reinforcement Learning"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gjR8zkBdjIrB",
        "colab_type": "text"
      },
      "source": [
        "\n",
        "<!-- <center><img src=\"https://drive.google.com/uc?id=1ivTQBHWkYi_J9vWwXFd2sSWg5f2TB5T-\" width=\"500\" /></center>  -->\n",
        "\n",
        "<center><img src=\"https://media.springernature.com/full/springer-static/image/art%3A10.1038%2Fnature14236/MediaObjects/41586_2015_Article_BFnature14236_Fig1_HTML.jpg\" width=\"500\" /></center> \n",
        "\n",
        "In this subsection, we will look at an advanced deep RL Agent based on the following publication, [Playing Atari with Deep Reinforcement Learning](https://deepmind.com/research/publications/playing-atari-deep-reinforcement-learning), which introduced the first deep learning model to successfully learn control policies directly from high-dimensional pixel inputs using RL.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "BukOfOsmtSQn"
      },
      "source": [
        " ## 3.1 Create an ACME DQN agent"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "NbHdPc-nxO2j",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Create the environment\n",
        "grid = build_gridworld_task(\n",
        "    task='simple', \n",
        "    observation_type=ObservationType.GRID,\n",
        "    max_episode_length=200)\n",
        "environment, environment_spec = setup_environment(grid)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "3Jcjk1w6oHVX",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Construct the agent and a training loop  { form-width: \"30%\" }\n",
        "\n",
        "# Build agent networks\n",
        "def network(x):\n",
        "  model = hk.Sequential([\n",
        "      hk.Conv2D(32, kernel_shape=[4,4], stride=[2,2], padding='VALID'),\n",
        "      jax.nn.relu,\n",
        "      hk.Conv2D(64, kernel_shape=[3,3], stride=[1,1], padding='VALID'),\n",
        "      jax.nn.relu,\n",
        "      hk.Flatten(),\n",
        "      hk.nets.MLP([50, 50, environment_spec.actions.num_values])\n",
        "  ])\n",
        "  return model(x)\n",
        "\n",
        "# Avoid logging from Acme\n",
        "class DummyLogger(object):\n",
        "\n",
        "  def write(self, data):\n",
        "    pass\n",
        "\n",
        "# Use library agent implementation.\n",
        "agent = dqn.DQN(\n",
        "    environment_spec=environment_spec,\n",
        "    network=network,\n",
        "    batch_size=10,\n",
        "    samples_per_insert=2,\n",
        "    epsilon=0.05,\n",
        "    min_replay_size=10,)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0dHDdPDr3QxI",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title Run a training loop  { form-width: \"30%\" }\n",
        "\n",
        "# Run a `num_episodes` training episodes.\n",
        "# Rerun this cell until the agent has learned the given task.\n",
        "returns = run_loop(environment=environment, agent=agent, num_episodes=100, num_steps=100000, logger_time_delta=0.2 )"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8ksVITeN5_Vq",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title Visualise the learned Q values { form-width: \"30%\" }\n",
        "\n",
        "# get agent parameters\n",
        "params = agent._learner.get_variables([])[0]\n",
        "\n",
        "# Evaluate the policy for every state, similar to tabular agents above.\n",
        "pi = np.zeros(grid._layout_dims, dtype=np.int32)\n",
        "q = np.zeros(grid._layout_dims + (4,))\n",
        "for y in range(grid._layout_dims[0]):\n",
        "  for x in range(grid._layout_dims[1]):\n",
        "    # Hack observation to see what the Q-network would output at that point.\n",
        "    environment.set_state(x, y)\n",
        "    obs = environment.get_obs()\n",
        "    q[y, x] = np.asarray(agent._learner._forward(params, np.expand_dims(obs, axis=0)))\n",
        "    pi[y, x] = np.asarray(agent.select_action(obs))\n",
        "    \n",
        "plot_action_values(q)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6PQaQej4LsU-",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Compare the greedy policy with the agent's policy { form-width: \"30%\" }\n",
        "\n",
        "grid.plot_greedy_policy(q)\n",
        "plt.title('Greedy policy using the learnt Q-values')\n",
        "\n",
        "grid.plot_policy(pi)\n",
        "plt.title(\"Policy using the agent's policy\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "acqPbd8zXH_K",
        "colab_type": "text"
      },
      "source": [
        "## 3.2 **[Advanced]** DQN Algorithm.\n",
        "\n",
        "The following coding exercise implements the loss function described in the DQN paper. This loss function is used by a learner class to compute gradients for the parameters $\\theta_i$ of the Q-network $Q( \\cdot; \\theta_i)$:\n",
        "\n",
        "```none\n",
        "loss(params: hk.Params, target_params: hk.Params, sample: reverb.ReplaySample)\n",
        "```\n",
        "which, at iteration `i` computes the DQN loss $L_i$ on the parameters $\\theta_i$, based on a the set of target parameters $\\theta_{i-1}$ and a given batch of sampled trajectories `sample`. As described in the manuscript, the loss function is defined as:\n",
        "\n",
        "$$L_i (\\theta_i) = \\mathbb{E}_{\\color{red}{s},\\color{blue}{a} \\sim \\rho(\\cdot)} \\left[ \\left( y_i - Q(\\color{red}{s},\\color{blue}{a} ;\\theta_i) \\right)^2\\right]$$\n",
        "\n",
        "where the target $y_i$ is computed using a bootstrap value computed from Q-value network with target parameters:\n",
        "\n",
        "$$ y_i = \\mathbb{E}_{\\color{red}{s'} \\sim \\mathcal{E}} \\left[ \\color{green}{r} + \\gamma \\max_{\\color{blue}{a'} \\in \\color{blue}{\\mathcal{A}}} Q(\\color{red}{s'}, \\color{blue}{a'} ; \\theta^{\\text{target}}_i) \\; | \\; \\color{red}{s}, \\color{blue}{a} \\right] $$\n",
        "\n",
        "The batch of data `sample` is prepackaged by the agent to match the sampling distributions $\\rho$ and $\\mathcal{E}$. To get the explicit data items, use the following:\n",
        "\n",
        "```none\n",
        "o_tm1, a_tm1, r_t, d_t, o_t = sample.data\n",
        "```\n",
        "\n",
        "The function is expected to return  \n",
        "* `mean_loss` is the mean of the above loss over the batched data,\n",
        "* (`keys`, `priorities`) will pair the `keys` corresponding to each batch item to the absolute TD-error used to compute the `mean_loss` above. The agent uses these to update priorities for samples in the replay buffer.\n",
        "\n",
        "\n",
        "**Note**. A full implementation of a DQN agent is outside the scope of this tutorial, but we encoruage you to explore the code (in a cell below) to understand where the learner fits with other the services used by the agent. Moreover, if you feel ambitious, we prepared a separate exercise where you are expected to implement the learner itself (see  *DQN Learner [Coding Task - Hard]*).\n",
        "\n",
        "\n",
        "\n",
        "**[Optional]**\n",
        "- use a Double-Q Learning Loss function instead of the original published loss (see [`rlax.double_q_learning`](https://github.com/deepmind/rlax/blob/870cba1ea8ad36725f4f3a790846298657b6fd4b/rlax/_src/value_learning.py#L233)) for more details.\n",
        "- for more stable optimization, use the Huber Loss instead of $L_2$, as prescribed in the manuscript (see [`rlax.huber_loss`](https://github.com/deepmind/rlax/blob/870cba1ea8ad36725f4f3a790846298657b6fd4b/rlax/_src/clipping.py#L31)).\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "liZKoQbaj4Wq",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title **[Coding Task - Easy]** DQN Loss function  { form-width: \"30%\" }\n",
        "\n",
        "TrainingState =  namedtuple('TrainingState', 'params, target_params, opt_state, step')\n",
        "LearnerOutputs =  namedtuple('LearnerOutputs', 'keys, priorities')\n",
        "\n",
        "class DQNLearner(acme.Learner):\n",
        "  \"\"\"DQN learner.\"\"\"\n",
        "\n",
        "  _state: TrainingState\n",
        "\n",
        "  def __init__(self,\n",
        "               network,\n",
        "               obs_spec,\n",
        "               discount,\n",
        "               importance_sampling_exponent,\n",
        "               target_update_period,\n",
        "               data_iterator,\n",
        "               optimizer,\n",
        "               rng,\n",
        "               replay_client,\n",
        "               max_abs_reward=1.,\n",
        "               huber_loss_parameter=1.,\n",
        "               ):\n",
        "    \"\"\"Initializes the learner.\"\"\"\n",
        "\n",
        "    # Transform network into a pure function.\n",
        "    network = hk.transform(network)\n",
        "\n",
        "    def loss(params: hk.Params, target_params: hk.Params,\n",
        "             sample: reverb.ReplaySample):\n",
        "      o_tm1, a_tm1, r_t, d_t, o_t = sample.data\n",
        "      keys, probs = sample.info[:2]\n",
        "      \n",
        "      # ============ YOUR CODE HERE =============\n",
        "      # return mean_loss, (keys, priorities)\n",
        "\n",
        "      pass\n",
        "      \n",
        "\n",
        "    def sgd_step(state, samples):\n",
        "      # Compute gradients on the given loss function and update the network\n",
        "      # using the optimizer provided at init time.\n",
        "      grad_fn = jax.grad(loss, has_aux=True)\n",
        "      gradients, (keys, priorities) = grad_fn(state.params, state.target_params,\n",
        "                                              samples)\n",
        "      updates, new_opt_state = optimizer.update(gradients, state.opt_state)\n",
        "      new_params = optix.apply_updates(state.params, updates)\n",
        "\n",
        "      # Update the internal state for the learner with (1) network parameters,\n",
        "      # (2) parameters of the target network, (3) the state of the optimizer,\n",
        "      # (4) Numbers of SGD steps performed by the agent.  \n",
        "      new_state = TrainingState(\n",
        "          params=new_params,\n",
        "          target_params=state.target_params,\n",
        "          opt_state=new_opt_state,\n",
        "          step=state.step + 1)\n",
        "\n",
        "      outputs = LearnerOutputs(keys=keys, priorities=priorities)\n",
        "      return new_state, outputs\n",
        "\n",
        "    # Internalise agent components (replay buffer, networks, optimizer).\n",
        "    self._replay_client = replay_client\n",
        "    self._iterator = data_iterator\n",
        "\n",
        "    # Since sampling is base on a priority experience replay, we need to pass\n",
        "    # the absolute td-loss values to the replay client to update priorities\n",
        "    # accordingly.\n",
        "    def update_priorities(outputs: LearnerOutputs):\n",
        "      for key, priority in zip(outputs.keys, outputs.priorities):\n",
        "        self._replay_client.mutate_priorities(\n",
        "            table='priority_table', \n",
        "            updates={key: priority})\n",
        "        \n",
        "    self._update_priorities = update_priorities\n",
        "\n",
        "    # Internalise the hyperparameters.\n",
        "    self._target_update_period = target_update_period\n",
        "\n",
        "    # Internalise logging/counting objects.\n",
        "    self._counter = counting.Counter()\n",
        "    self._logger = loggers.TerminalLogger('learner', time_delta=1.)\n",
        "\n",
        "    # Initialise parameters and optimiser state.\n",
        "    def initialization_fn(values):\n",
        "      values = tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), values)\n",
        "      # Add batch dim.\n",
        "      return tree_util.tree_map(lambda x: jnp.expand_dims(x, axis=0), values)\n",
        "\n",
        "    initial_params = network.init(next(rng), initialization_fn(obs_spec))\n",
        "    initial_target_params = network.init(next(rng), initialization_fn(obs_spec))\n",
        "    initial_opt_state = optimizer.init(initial_params)\n",
        "\n",
        "    self._state = TrainingState(\n",
        "        params=initial_params,\n",
        "        target_params=initial_target_params,\n",
        "        opt_state=initial_opt_state,\n",
        "        step=0)\n",
        "\n",
        "    self._forward = jax.jit(network.apply)\n",
        "    self._sgd_step = jax.jit(sgd_step)\n",
        "    \n",
        "  def step(self):\n",
        "    samples = next(self._iterator)\n",
        "    # Do a batch of SGD.\n",
        "    self._state, outputs = self._sgd_step(self._state, samples)\n",
        "\n",
        "    # Update our counts and record it.\n",
        "    result = self._counter.increment(steps=1)\n",
        "\n",
        "    # Periodically update target network parameters.\n",
        "    if self._state.step % self._target_update_period == 0:\n",
        "      self._state = self._state._replace(target_params=self._state.params)\n",
        "\n",
        "    # Update priorities in replay.\n",
        "    self._update_priorities(outputs)\n",
        "\n",
        "    # Write to logs.\n",
        "    self._logger.write(result)\n",
        "\n",
        "  def get_variables(self):\n",
        "    \"\"\"Network variables after a number of SGD steps.\"\"\"\n",
        "    return self._state.params"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jbULeKGwURlL",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title **[Coding Task - Hard]** DQN Learner  { form-width: \"30%\" }\n",
        "\n",
        "TrainingState =  namedtuple('TrainingState', 'params, target_params, opt_state, step')\n",
        "LearnerOutputs =  namedtuple('LearnerOutputs', 'keys, priorities')\n",
        "\n",
        "class DQNLearner(acme.Learner):\n",
        "  \"\"\"DQN learner.\"\"\"\n",
        "\n",
        "  _state: TrainingState\n",
        "\n",
        "  def __init__(self,\n",
        "               network,\n",
        "               obs_spec,\n",
        "               discount,\n",
        "               importance_sampling_exponent,\n",
        "               target_update_period,\n",
        "               data_iterator,\n",
        "               optimizer,\n",
        "               rng,\n",
        "               replay_client,\n",
        "               max_abs_reward=1.,\n",
        "               huber_loss_parameter=1.,\n",
        "               ):\n",
        "    \"\"\"Initializes the learner.\"\"\"\n",
        "\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    # Use provided params to initialize any jax functions used in the `step`\n",
        "    # function.\n",
        "    \n",
        "    # Internalise agent components (replay buffer, networks, optimizer).\n",
        "    self._replay_client = replay_client\n",
        "    self._iterator = data_iterator\n",
        "\n",
        "    # Since sampling is base on a priority experience replay, we need to pass\n",
        "    # the absolute td-loss values to the replay client to update priorities\n",
        "    # accordingly.\n",
        "    def update_priorities(outputs: LearnerOutputs):\n",
        "      for key, priority in zip(outputs.keys, outputs.priorities):\n",
        "        self._replay_client.mutate_priorities(\n",
        "            table='priority_table', \n",
        "            updates={key: priority})\n",
        "        \n",
        "    self._update_priorities = update_priorities\n",
        "\n",
        "    # Internalise the hyperparameters.\n",
        "    self._target_update_period = target_update_period\n",
        "\n",
        "    # Internalise logging/counting objects.\n",
        "    self._counter = counting.Counter()\n",
        "    self._logger = loggers.TerminalLogger('learner', time_delta=1.)\n",
        "\n",
        "    # Initialise parameters and optimiser state.\n",
        "\n",
        "    # Transform network into a pure function.\n",
        "    network = hk.transform(network)\n",
        "\n",
        "    def initialization_fn(values):\n",
        "      values = tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), values)\n",
        "      # Add batch dim.\n",
        "      return tree_util.tree_map(lambda x: jnp.expand_dims(x, axis=0), values)\n",
        "\n",
        "    initial_params = network.init(next(rng), initialization_fn(obs_spec))\n",
        "    initial_target_params = network.init(next(rng), initialization_fn(obs_spec))\n",
        "    initial_opt_state = optimizer.init(initial_params)\n",
        "\n",
        "    self._state = TrainingState(\n",
        "        params=initial_params,\n",
        "        target_params=initial_target_params,\n",
        "        opt_state=initial_opt_state,\n",
        "        step=0)\n",
        "\n",
        "  def step(self):\n",
        "    samples = next(self._iterator)\n",
        "    \n",
        "    # Do a batch of SGD and update self._state accordingly.\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    \n",
        "    # Update our counts and record it.\n",
        "    result = self._counter.increment(steps=1)\n",
        "\n",
        "    # Periodically update target network parameters.\n",
        "    # ============ YOUR CODE HERE =============\n",
        "    \n",
        "    # Update priorities in replay.\n",
        "    self._update_priorities(outputs)\n",
        "\n",
        "    # Write to logs.\n",
        "    self._logger.write(result)\n",
        "\n",
        "  def get_variables(self):\n",
        "    \"\"\"Network variables after a number of SGD steps.\"\"\"\n",
        "    return self._state.params"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "cellView": "form",
        "id": "eZuigF_bD0DP",
        "colab": {}
      },
      "source": [
        "# @title **[Solution]** DQN Learner  { form-width: \"30%\" }\n",
        "\n",
        "TrainingState =  namedtuple('TrainingState', 'params, target_params, opt_state, step')\n",
        "LearnerOutputs =  namedtuple('LearnerOutputs', 'keys, priorities')\n",
        "\n",
        "class DQNLearner(acme.Learner):\n",
        "  \"\"\"DQN learner.\"\"\"\n",
        "\n",
        "  _state: TrainingState\n",
        "\n",
        "  def __init__(self,\n",
        "               network,\n",
        "               obs_spec,\n",
        "               discount,\n",
        "               importance_sampling_exponent,\n",
        "               target_update_period,\n",
        "               data_iterator,\n",
        "               optimizer,\n",
        "               rng,\n",
        "               replay_client,\n",
        "               max_abs_reward=1.,\n",
        "               huber_loss_parameter=1.,\n",
        "               ):\n",
        "    \"\"\"Initializes the learner.\"\"\"\n",
        "\n",
        "    # Transform network into a pure function.\n",
        "    network = hk.transform(network)\n",
        "\n",
        "    def loss(params: hk.Params, target_params: hk.Params,\n",
        "             sample: reverb.ReplaySample):\n",
        "      o_tm1, a_tm1, r_t, d_t, o_t = sample.data\n",
        "      keys, probs = sample.info[:2]\n",
        "\n",
        "      # Forward pass.\n",
        "      q_tm1 = network.apply(params, o_tm1)\n",
        "      q_t_value = network.apply(target_params, o_t)\n",
        "      q_t_selector = network.apply(params, o_t)\n",
        "\n",
        "      # Cast and clip rewards.\n",
        "      d_t = (d_t * discount).astype(jnp.float32)\n",
        "      r_t = jnp.clip(r_t, -max_abs_reward, max_abs_reward).astype(jnp.float32)\n",
        "\n",
        "      # Compute double Q-learning n-step TD-error.\n",
        "      batch_error = jax.vmap(rlax.double_q_learning)\n",
        "      td_error = batch_error(q_tm1, a_tm1, r_t, d_t, q_t_value, q_t_selector)\n",
        "      batch_loss = rlax.huber_loss(td_error, huber_loss_parameter)\n",
        "\n",
        "      # Importance weighting.\n",
        "      importance_weights = (1. / probs).astype(jnp.float32)\n",
        "      importance_weights **= importance_sampling_exponent\n",
        "      importance_weights /= jnp.max(importance_weights)\n",
        "\n",
        "      # Reweight.\n",
        "      mean_loss = jnp.mean(importance_weights * batch_loss)  # []\n",
        "\n",
        "      priorities = jnp.abs(td_error).astype(jnp.float64)\n",
        "\n",
        "      return mean_loss, (keys, priorities)\n",
        "\n",
        "    def sgd_step(state, samples):\n",
        "      # Compute gradients on the given loss function and update the network\n",
        "      # using the optimizer provided at init time.\n",
        "      grad_fn = jax.grad(loss, has_aux=True)\n",
        "      gradients, (keys, priorities) = grad_fn(state.params, state.target_params,\n",
        "                                              samples)\n",
        "      updates, new_opt_state = optimizer.update(gradients, state.opt_state)\n",
        "      new_params = optix.apply_updates(state.params, updates)\n",
        "\n",
        "      # Update the internal state for the learner with (1) network parameters,\n",
        "      # (2) parameters of the target network, (3) the state of the optimizer,\n",
        "      # (4) Numbers of SGD steps performed by the agent.  \n",
        "      new_state = TrainingState(\n",
        "          params=new_params,\n",
        "          target_params=state.target_params,\n",
        "          opt_state=new_opt_state,\n",
        "          step=state.step + 1)\n",
        "\n",
        "      outputs = LearnerOutputs(keys=keys, priorities=priorities)\n",
        "\n",
        "      return new_state, outputs\n",
        "\n",
        "    # Internalise agent components (replay buffer, networks, optimizer).\n",
        "    self._replay_client = replay_client\n",
        "    self._iterator = data_iterator\n",
        "\n",
        "    # Since sampling is base on a priority experience replay, we need to pass\n",
        "    # the absolute td-loss values to the replay client to update priorities\n",
        "    # accordingly.\n",
        "    def update_priorities(outputs: LearnerOutputs):\n",
        "      for key, priority in zip(outputs.keys, outputs.priorities):\n",
        "        self._replay_client.mutate_priorities(\n",
        "            table='priority_table', \n",
        "            updates={key: priority})\n",
        "        \n",
        "    self._update_priorities = update_priorities\n",
        "\n",
        "    # Internalise the hyperparameters.\n",
        "    self._target_update_period = target_update_period\n",
        "\n",
        "    # Internalise logging/counting objects.\n",
        "    self._counter = counting.Counter()\n",
        "    self._logger = loggers.TerminalLogger('learner', time_delta=1.)\n",
        "\n",
        "    # Initialise parameters and optimiser state.\n",
        "    def initialization_fn(values):\n",
        "      values = tree_util.tree_map(lambda x: jnp.zeros(x.shape, x.dtype), values)\n",
        "      # Add batch dim.\n",
        "      return tree_util.tree_map(lambda x: jnp.expand_dims(x, axis=0), values)\n",
        "\n",
        "    initial_params = network.init(next(rng), initialization_fn(obs_spec))\n",
        "    initial_target_params = network.init(next(rng), initialization_fn(obs_spec))\n",
        "    initial_opt_state = optimizer.init(initial_params)\n",
        "\n",
        "    self._state = TrainingState(\n",
        "        params=initial_params,\n",
        "        target_params=initial_target_params,\n",
        "        opt_state=initial_opt_state,\n",
        "        step=0)\n",
        "\n",
        "    self._forward = jax.jit(network.apply)\n",
        "    self._sgd_step = jax.jit(sgd_step)\n",
        "    \n",
        "  def step(self):\n",
        "    samples = next(self._iterator)\n",
        "    # Do a batch of SGD.\n",
        "    self._state, outputs = self._sgd_step(self._state, samples)\n",
        "\n",
        "    # Update our counts and record it.\n",
        "    result = self._counter.increment(steps=1)\n",
        "\n",
        "    # Periodically update target network parameters.\n",
        "    if self._state.step % self._target_update_period == 0:\n",
        "      self._state = self._state._replace(target_params=self._state.params)\n",
        "\n",
        "    # Update priorities in replay.\n",
        "    self._update_priorities(outputs)\n",
        "\n",
        "    # Write to logs.\n",
        "    self._logger.write(result)\n",
        "\n",
        "  def get_variables(self):\n",
        "    \"\"\"Network variables after a number of SGD steps.\"\"\"\n",
        "    return self._state.params"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ywObWtqgaSXx",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title DQN Agent implementation (use for reference only) { form-width: \"30%\" }\n",
        "class DQN(acme.Actor):\n",
        "  def __init__(\n",
        "    self,\n",
        "    environment_spec,\n",
        "    network,\n",
        "    batch_size=256,\n",
        "    prefetch_size=4,\n",
        "    target_update_period=100,\n",
        "    samples_per_insert=32.0,\n",
        "    min_replay_size=1000,\n",
        "    max_replay_size=1000000,\n",
        "    importance_sampling_exponent=0.2,\n",
        "    priority_exponent=0.6,\n",
        "    n_step=5,\n",
        "    epsilon=0.,\n",
        "    learning_rate=1e-3,\n",
        "    discount=0.99,\n",
        "  ):\n",
        "    # Create a replay server to add data to. This is initialized as a\n",
        "    # table, and a Learner (defined separately) will be in charge of updating\n",
        "    # sample priorities based on the corresponding learner loss. \n",
        "    replay_table = reverb.Table(\n",
        "        name='priority_table',\n",
        "        sampler=reverb.selectors.Prioritized(priority_exponent),\n",
        "        remover=reverb.selectors.Fifo(),\n",
        "        max_size=max_replay_size,\n",
        "        rate_limiter=reverb.rate_limiters.MinSize(1))\n",
        "    self._server = reverb.Server([replay_table], port=None)\n",
        "    address = f'localhost:{self._server.port}'\n",
        "\n",
        "    # Use ACME reverb adder as a tool to add transition data into the replay\n",
        "    # buffer defined above.\n",
        "    self._adder = adders.NStepTransitionAdder(\n",
        "        client=reverb.Client(address),\n",
        "        n_step=n_step,\n",
        "        discount=discount)\n",
        "\n",
        "    # ACME datasets provides an interface to easily sample from a replay server.\n",
        "    dataset = datasets.make_reverb_dataset(\n",
        "        client=reverb.TFClient(address),\n",
        "        environment_spec=environment_spec,\n",
        "        batch_size=batch_size,\n",
        "        prefetch_size=prefetch_size,\n",
        "        transition_adder=True)\n",
        "    data_iterator = dataset.as_numpy_iterator()\n",
        "\n",
        "    # Create a learner that updates the parameters (and initializes them).\n",
        "    self._learner = DQNLearner(\n",
        "        network=network,\n",
        "        obs_spec=environment_spec.observations,\n",
        "        rng=hk.PRNGSequence(1),\n",
        "        optimizer=optix.adam(learning_rate),\n",
        "        discount=discount,\n",
        "        importance_sampling_exponent=importance_sampling_exponent,\n",
        "        target_update_period=target_update_period,\n",
        "        data_iterator=data_iterator,\n",
        "        replay_client=reverb.Client(address),\n",
        "    )\n",
        "    \n",
        "    # Create a feed forward actor that obtains its variables from the DQNLearner\n",
        "    # above.\n",
        "    def policy(params, key, observation):\n",
        "      action_values = hk.transform(network).apply(params, observation)\n",
        "      return rlax.epsilon_greedy(epsilon).sample(key, action_values)\n",
        "\n",
        "    self._policy = policy\n",
        "    self._rng = hk.PRNGSequence(1)\n",
        " \n",
        "    # We'll ignore the first min_observations when determining whether to take\n",
        "    # a step and we'll do so by making sure num_observations >= 0.\n",
        "    self._num_observations = -max(batch_size, min_replay_size)\n",
        "\n",
        "    observations_per_step = float(batch_size) / samples_per_insert\n",
        "    if observations_per_step >= 1.0:\n",
        "      self._observations_per_update = int(observations_per_step)\n",
        "      self._learning_steps_per_update = 1\n",
        "    else:\n",
        "      self._observations_per_update = 1\n",
        "      self._learning_steps_per_update = int(1.0 / observations_per_step)\n",
        "\n",
        "  def select_action(self, observation):\n",
        "    observation = tree_util.tree_map(lambda x: jnp.expand_dims(x, axis=0), \n",
        "                                     observation)\n",
        "    \n",
        "    key = next(self._rng)\n",
        "    params = self._learner.get_variables()\n",
        "    action = self._policy(params, key, observation)\n",
        "    action = tree_util.tree_map(lambda x: np.array(x).squeeze(axis=0), action)\n",
        "    return action \n",
        "\n",
        "  def observe_first(self, timestep):\n",
        "    self._adder.add_first(timestep)\n",
        "\n",
        "  def observe(self, action, next_timestep):\n",
        "    self._num_observations += 1\n",
        "    self._adder.add(action, next_timestep)\n",
        "\n",
        "  def update(self):\n",
        "    # Only allow updates after some minimum number of observations have been and\n",
        "    # then at some period given by observations_per_update.\n",
        "    if (self._num_observations >= 0 and\n",
        "        self._num_observations % self._observations_per_update == 0):\n",
        "      self._num_observations = 0\n",
        "\n",
        "      # Run a number of learner steps (usually gradient steps).\n",
        "      for _ in range(self._learning_steps_per_update):\n",
        "        self._learner.step()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0iBoBqLvcy14",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "# @title Run a training loop  { form-width: \"30%\" }\n",
        "\n",
        "# Run a `num_episodes` training episodes.\n",
        "# Rerun this cell until the agent has learned the given task.\n",
        "\n",
        "grid = build_gridworld_task(\n",
        "    task='simple', \n",
        "    observation_type=ObservationType.GRID, \n",
        "    max_episode_length=100,\n",
        ")\n",
        "environment, environment_spec = setup_environment(grid)\n",
        "\n",
        "agent = DQN(\n",
        "    environment_spec=environment_spec,\n",
        "    network=network,\n",
        "    batch_size=16,\n",
        "    samples_per_insert=2,\n",
        "    epsilon=0.1,\n",
        "    min_replay_size=100)\n",
        "\n",
        "returns = run_loop(environment=environment, agent=agent, num_episodes=200, \n",
        "         logger_time_delta=0.2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iJi7LDrn0eO4",
        "colab_type": "text"
      },
      "source": [
        "### DQN agent on the Gym Cartpole environment\n",
        "\n",
        "Here we show that you can apply what you learned to other environments such as Cartpole in [Gym](https://gym.openai.com/).\n",
        "\n",
        "\n",
        "<center><img src=\"https://user-images.githubusercontent.com/10624937/42135683-dde5c6f0-7d13-11e8-90b1-8770df3e40cf.gif\" height=\"250\" /></center>\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DIERzZVk0xIh",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Construct the agent and run the training loop { form-width: \"30%\" }\n",
        "\n",
        "#  Try different parameters to see how learning is affected.\n",
        "\n",
        "env = gym_wrapper.GymWrapper(gym.make('CartPole-v0'))\n",
        "env = wrappers.SinglePrecisionWrapper(env)\n",
        "\n",
        "environment, environment_spec = setup_environment(env)\n",
        "\n",
        "# Build agent networks\n",
        "def network(x):\n",
        "  model = hk.Sequential([\n",
        "      hk.Flatten(),\n",
        "      hk.nets.MLP([100, environment_spec.actions.num_values])\n",
        "  ])\n",
        "  return model(x)\n",
        "\n",
        "agent = dqn.DQN(\n",
        "    environment_spec=environment_spec,\n",
        "    network=network,\n",
        "    batch_size=64,\n",
        "    epsilon=0.01,\n",
        "    learning_rate=1e-4,\n",
        "    min_replay_size=100)\n",
        "\n",
        "returns = run_loop(environment=environment, agent=agent, num_episodes=300, \n",
        "         logger_time_delta=0.2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dDmLcICc98Z8",
        "colab_type": "code",
        "cellView": "form",
        "colab": {}
      },
      "source": [
        "#@title Visualise training curve { form-width: \"30%\" }\n",
        "\n",
        "# Compute rolling average over returns\n",
        "returns_avg = pd.Series(returns).rolling(10, center=True).mean()\n",
        "\n",
        "plt.figure(figsize=(8, 5))\n",
        "plt.plot(range(len(returns)), returns_avg)\n",
        "plt.xlabel('Episodes')\n",
        "plt.ylabel('Total reward');"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rzqrYxAtH11S",
        "colab_type": "text"
      },
      "source": [
        "# Want to learn more?\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "odBz1OO0JIXY",
        "colab_type": "text"
      },
      "source": [
        "This Colab was inspired by the [EEML 2019 RL practical](https://github.com/eemlcommunity/PracticalSessions2019/blob/master/rl/RL_Tutorial.ipynb) and the [Acme tutorial](https://github.com/deepmind/acme/blob/master/examples/tutorial.ipynb). \n",
        "\n",
        "Books and lecture notes\n",
        "*   [Reinforcement Learning: an Introduction by Sutton & Barto](http://incompleteideas.net/book/RLbook2018.pdf)\n",
        "* [Algorithms for Reinforcement Learning by Csaba Szepesvari](https://sites.ualberta.ca/~szepesva/papers/RLAlgsInMDPs.pdf)\n",
        "\n",
        "Lectures and course \n",
        "*   [RL Course by David Silver](https://www.youtube.com/playlist?list=PLzuuYNsE1EZAXYR4FJ75jcJseBmo4KQ9-)\n",
        "*   [Reinforcement Learning Course | UCL & DeepMind](https://www.youtube.com/playlist?list=PLqYmG7hTraZBKeNJ-JE_eyJHZ7XgBoAyb)\n",
        "*   [Emma Brunskill Stanford RL Course](https://www.youtube.com/playlist?list=PLoROMvodv4rOSOPzutgyCTapiGlY2Nd8u)\n",
        "*   [RL Course on Coursera by Martha White & Adam White](https://www.coursera.org/specializations/reinforcement-learning)\n",
        "\n",
        "More practical:\n",
        "* [Spinning Up in Deep RL by Josh Achiam](https://spinningup.openai.com/en/latest/)\n",
        "*   [Acme white paper](http://go/arxiv/2006.00979)\n",
        "\n",
        "\n",
        "\n"
      ]
    }
  ]
}