practicals/introduction_to_reinforcement_learning.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "m2s4kN_QPQVe"
      },
      "source": [
        "# Introduction to Reinforcement Learning\n",
        "\n",
        "\n",
        "<center>\n",
        "<img src=\"https://raw.githubusercontent.com/jcformanek/jcformanek.github.io/master/docs/assets/images/rl_in_space.png\" width=\"100%\" />\n",
        "</center>\n",
        "\n",
        "<a href=\"https://colab.research.google.com/github/deep-learning-indaba/indaba-pracs-2022/blob/main/practicals/introduction_to_reinforcement_learning.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>\n",
        "\n",
        "© Deep Learning Indaba 2022. Apache License 2.0.\n",
        "\n",
        "**Authors:**\n",
        "Claude Formanek, Kale-ab Tessera, Sicelukwanda Zwane, Sebastian Bodenstein\n",
        "\n",
        "**Reviewers:**\n",
        "Ruan van der Merwe, Avishkar Bhoopchand\n",
        "\n",
        "**Introduction:** \n",
        "\n",
        "In this tutorial, we will be learning about Reinforcement Learning, a type of Machine Learning where an agent learns to choose actions in an environment that lead to maximal reward in the long run. RL has seen tremendous success on a wide range of challenging problems such as learning to play complex video games like [Atari](https://www.deepmind.com/blog/agent57-outperforming-the-human-atari-benchmark), [StarCraft II](https://www.deepmind.com/blog/alphastar-mastering-the-real-time-strategy-game-starcraft-ii) and [Dota II](https://openai.com/five/). \n",
        "\n",
        "In this introductory tutorial we will solve the classic [CartPole](https://www.gymlibrary.ml/environments/classic_control/cart_pole/) environment, where an agent must learn to balance a pole on a cart, using several different RL approaches. Along the way you will be introduced to some of the most important concepts and terminology in RL.\n",
        "\n",
        "**Topics:** \n",
        "* Reinforcement Learning\n",
        "* Random Policy Search\n",
        "* Policy Gradient\n",
        "* Q-Learning\n",
        "\n",
        "**Level:** \n",
        "\n",
        "Beginner\n",
        "\n",
        "**Aims/Learning Objectives:**\n",
        "\n",
        "* Understand the basic theory behind RL.\n",
        "* Implement a simple random policy search algorithm.\n",
        "* Implement a simple policy gradient RL algorithm.\n",
        "* Implement a simple Q-learning algorithm.\n",
        "\n",
        "**Prerequisites:**\n",
        "\n",
        "* Some familiarity with [JAX](https://github.com/google/jax).\n",
        "* Neural network basics.\n",
        "\n",
        "**Outline:** \n",
        "\n",
        "* Section 1: Key Concepts in Reinforcement Learning\n",
        "* Section 2: Random Policy Search\n",
        "* Section 3: Policy Gradient\n",
        "* Section 4: Deep Q-Learning\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6EqhIg1odqg0"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4boGA9rYdt9l"
      },
      "outputs": [],
      "source": [
        "# @title Install required packages (run me) { display-mode: \"form\" }\n",
        "# @markdown This may take a minute or two to complete.\n",
        "%%capture\n",
        "!pip install jaxlib\n",
        "!pip install jax\n",
        "!pip install git+https://github.com/deepmind/dm-haiku\n",
        "!pip install gym==0.25\n",
        "!pip install gym[box2d]\n",
        "!pip install optax\n",
        "!pip install matplotlib\n",
        "!pip install chex"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gwbqggmcRjMy"
      },
      "outputs": [],
      "source": [
        "# @title Import required packages (run me) { display-mode: \"form\" }\n",
        "%%capture\n",
        "import copy\n",
        "from shutil import rmtree # deleting directories\n",
        "import random\n",
        "import collections # useful data structures\n",
        "import numpy as np\n",
        "import gym # reinforcement learning environments\n",
        "from gym.wrappers import RecordVideo\n",
        "import jax\n",
        "import jax.numpy as jnp # jax numpy\n",
        "import haiku as hk # jax neural network library\n",
        "import optax # jax optimizer library\n",
        "import matplotlib.pyplot as plt # graph plotting library\n",
        "from IPython.display import HTML\n",
        "from base64 import b64encode\n",
        "import chex\n",
        "\n",
        "# Hide warnings\n",
        "import warnings\n",
        "warnings.filterwarnings('ignore')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-ZUp8i37dFbU"
      },
      "source": [
        "## Section 1: Key Concepts in Reinforcement Learning\n",
        "\n",
        "Reinforcement Learning (RL) is a subfield of Machine Learning (ML). Unlike fields like supervised learning, where we give examples of expected behaviour to our models, RL focuses on *goal-orientated* learning from interactions, through trial-and-error. RL algorithms learn what to do (i.e. which optimal actions to take) in an environment to maximise some reward signal. In settings like a video game, the reward signal could be the score of the game, i.e., RL algorithms will try to maximise the score in the game by choosing the best actions.  \n",
        "\n",
        "<center>\n",
        "<img src=\"https://miro.medium.com/max/1400/1*Ews7HaMiSn2l8r70eeIszQ.png\" width=\"40%\" />\n",
        "</center>\n",
        "\n",
        "[*Image Source*](https://towardsdatascience.com/multi-agent-deep-reinforcement-learning-in-15-lines-of-code-using-pettingzoo-e0b963c0820b)\n",
        "\n",
        "More precisely, in RL we have an **agent** which perceives an **observation** $o_t$ of the current state $s_t$ of the **environment** and must choose an **action** $a_t$ to take. The environment then transitions to a new state $s_{t+1}$ in response to the agent's action and also gives the agent a scalar reward $r_t$ to indicate how good or bad the chosen action was given the environment's state. The goal in RL is for the agent to maximise the amount of reward it receives from the environment over time. The subscript $t$ is used to indicate the timestep number, i.e., $s_0$ is the state of the environment at the initial timestep, and $a_{99}$ is the agent's action at the $99th$ timestep. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ghgy69hFRjMz"
      },
      "source": [
        "### Environment - OpenAI Gym\n",
        "As mentioned above, an environment receives an action $a_t$ from the agent and returns reward $r_t$ and observation $o_t$.\n",
        "\n",
        "OpenAI has provided a Python package called **Gym** that includes implementations of popular environments and a simple interface for an RL agent to interact with. To use a supported [gym environment](https://www.gymlibrary.ml/), all you need to do is pass the name of the environment to the function `gym.make(<environment_name>)`.\n",
        "\n",
        "In this tutorial, we will be using a simple environment called **CartPole**. In CartPole the task is for the agent to learn to balance a pole for as long as possible by moving a cart *left* or *right*.\n",
        "\n",
        "<img src=\"https://miro.medium.com/max/600/1*v8KcdjfVGf39yvTpXDTCGQ.gif\" width=\"30%\" />"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WfxzajMYRjMz"
      },
      "outputs": [],
      "source": [
        "# Create the environment\n",
        "env_name = \"CartPole-v0\"\n",
        "env = gym.make(env_name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y_BbftaJj3zu"
      },
      "source": [
        "### States and Observations - $s_t$ and $o_t$\n",
        "\n",
        "In RL, an agent perceives an observation of the environment's state. In some settings, the observation may include all the information underlying the environment's state. Such an environment is called **fully observed**. In other settings, the agent may only receive partial information about the environment's state in its observation. Such an environment is called **partially observed**. \n",
        "\n",
        "For the rest of this tutorial, we will assume the environment is fully observed and so we will use state $s_t$ and observation $o_t$ interchangeably. In Gym we get the initial observation from the environment by calling the function `env.reset()`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HdS8nqOgRjM0"
      },
      "outputs": [],
      "source": [
        "# Reset the environment\n",
        "s_0 = env.reset()\n",
        "print(\"Initial State::\", s_0)\n",
        "\n",
        "# Get environment obs space\n",
        "obs_shape = env.observation_space.shape\n",
        "print(\"Environment Obs Space Shape:\", obs_shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pUNX6mbotABo"
      },
      "source": [
        "In CartPole, the state of the environment is represented by four numbers; *angular position of the pole, angular velocity of the pole, position of the cart, velocity of the cart*. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hL1Nkgy7nUfn"
      },
      "source": [
        "### Actions - $a_t$\n",
        "\n",
        "In RL, actions are usually either **discrete** or **continuous**. Continuous actions are given by a vector of real numbers. Discrete actions are given by an integer value. In environments where we can count out the finite set of actions we usually use discrete actions. \n",
        "\n",
        "In CartPole there are only two actions; *left* and *right*. As such, the actions can be represented by integers $0$ and $1$. In gym we can easily get the list of possible actions as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uOLZqU_LpIXh"
      },
      "outputs": [],
      "source": [
        "# Get action space - e.g. discrete or continuous\n",
        "print(f\"Environment action space: {env.action_space}\")\n",
        "\n",
        "# Get num actions\n",
        "num_actions = env.action_space.n\n",
        "print(f\"Number of actions: {num_actions}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tRsflxbDpoPm"
      },
      "source": [
        "### The Agent's Policy - $\\pi$\n",
        "\n",
        "In RL the agent chooses actions based on the observations it receives. We can think of the agent's action selection process as a function that takes an observation as input and returns an action as output. In RL we usually call this function the agent's **policy** and denote it $\\pi(s_t)=a_t$. In RL we usually parametrise our policy in some way and then try to learn the optimal parameters. A parametrised policy is usually denoted $\\pi_\\theta$, where $\\theta$ is the set of parameters.\n",
        "\n",
        "**Exercise 1:** As an exercise, lets implement a simple policy which takes a set of parameters and an observation  as input, and returns an action. Assume the observation is a vector of four numbers like the CartPole observation and that the action should be either $0$ or $1$. Assume also, the parameters are a vector of four real-numbers. Then the action should be computed as follows.\n",
        "\n",
        "\n",
        "1.   Compute the [vector dot product](https://www.mathsisfun.com/algebra/vectors-dot-product.html) between the observation and the parameters.\n",
        "2.   If the result is greater than zero, return action $1$.\n",
        "3. \t Else return action $0$.\n",
        "\n",
        "In this prac we will try to use JAX as much as possible. So, try to use JAX methods for this task. Below are some useful methods you can use. You will need to complete the code in the block below by replacing the `...` with the correct code.\n",
        "\n",
        "We have also provided a code cell to check your solution. Finally, we also provide the solution to the coding task which you can check after you have given the task a try.\n",
        "\n",
        "**Useful methods:** \n",
        "* Compute the vector dot product with `jax.numpy.dot` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.dot.html)).\n",
        "* When you try to conditionally assigne a value of $0$ or $1$ to the action based on the result of the dot product, you should use `jax.lax.select` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.select.html#jax.lax.select)). The method `jax.lax.select` takes three arguments as input. The first argument is a statement that will evaluate to either `True` or `False`. If the statment is `True` then `jax.lax.select` will return its second argument. If the statement is `False`, `jax.lax.select` will return its third argument.\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KSGCd7XB1z8k"
      },
      "outputs": [],
      "source": [
        "def linear_policy(params, obs):\n",
        "  \"\"\"A simple linear policy\n",
        "  \n",
        "  Args:\n",
        "    params: a vector of four real-numbers that give the parameters of the policy\n",
        "    obs: a vector of four real-numbers that give the agent's observation\n",
        "\n",
        "  Returns:\n",
        "    a discrete action given by a 0 or 1\n",
        "  \"\"\"\n",
        "  # YOUR CODE\n",
        "  dot_product_result = ...\n",
        "\n",
        "  action = jax.lax.select(\n",
        "      ..., # boolean statement goes here\n",
        "      ..., # result when the statement is True goes here\n",
        "      ..., # result when the statement is False goes here\n",
        "  )\n",
        "  # END YOUR CODE\n",
        "  return action"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9EnnBvceb1f3"
      },
      "outputs": [],
      "source": [
        "# @title Check exercise 1 (run me) { display-mode: \"form\" }\n",
        "\n",
        "def check_linear_policy(linear_policy):\n",
        "  fixed_obs = jnp.array([1,1,2,4])\n",
        "  \n",
        "  # check case1 - negative dot product.\n",
        "  # weights\n",
        "  params1 = jnp.array([1,1,1,1])\n",
        "  params2 = jnp.array([-1,-1,-1,-1])\n",
        "\n",
        "  hint1 = f\"Incorrect answer, your linear policy is incorrect. The action when \\\n",
        "  obs={fixed_obs} and params={params1} should be 1\"\n",
        "\n",
        "  hint2 = f\"Incorrect answer, your linear policy is incorrect. The action when \\\n",
        "  obs={fixed_obs} and params={params2} should be 0\"\n",
        "\n",
        "  hint = None\n",
        "  if linear_policy(params1, fixed_obs) != 1:\n",
        "    hint = hint1\n",
        "  elif linear_policy(params2, fixed_obs) != 0:\n",
        "    hint = hint2\n",
        "\n",
        "  if hint is not None:\n",
        "    print(hint)\n",
        "  else:\n",
        "    print(\"Your function is correct!\")\n",
        "\n",
        "try:    \n",
        "  check_linear_policy(linear_policy)\n",
        "except Exception as e:\n",
        "    print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OCzRmXbz2HfY"
      },
      "outputs": [],
      "source": [
        "# @title  Solution exercise 1 { display-mode: \"form\" }\n",
        "\n",
        "def linear_policy(params, obs):\n",
        "  \"\"\"A simple linear policy\n",
        "  \n",
        "  Args:\n",
        "    params: a vector of four real-numbers that give the parameters of the policy\n",
        "    obs: a vector of four real-numbers that give the agent's observation\n",
        "\n",
        "  Returns:\n",
        "    a discrete action given by a 0 or 1\n",
        "  \"\"\"\n",
        "  # YOUR CODE\n",
        "  dot_product_result = jax.numpy.dot(params, obs)\n",
        "\n",
        "  action = jax.lax.select(\n",
        "      dot_product_result > 0, # boolean statement goes here\n",
        "      1, # result when the statement is True goes here\n",
        "      0, # result when the statement is False goes here\n",
        "  )\n",
        "  # END YOUR CODE\n",
        "  return action "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PkuvT-jf6Ieh"
      },
      "source": [
        "### The Environment Transition Function - $P$\n",
        "\n",
        "Now that we have a policy we can pass actions from the agent to the environment. The environment will then transition to a new state in response to the agent's action. \n",
        "\n",
        "In RL we model this process by using a **state transition function** $P$ which takes the current state $s_t$ and an action $a_t$ as input and returns the next state $s_{t+1}$ as output:\n",
        "\n",
        "<center>\n",
        " $s_{t+1}=P(s_t, a_t)$\n",
        "</center> \n",
        "\n",
        "In gym, we can pass actions to the environment by calling the `env.step(<action>)` function. The function will then return four values:\n",
        "- the **next observation**\n",
        "- the **reward** for the action taken\n",
        "- a boolean flag to indicate if the game is **done** \n",
        "- some **extra** information.\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hh0j9-Tk7olb"
      },
      "outputs": [],
      "source": [
        "# Get the initial obs by resetting the env\n",
        "initial_obs = env.reset()\n",
        "\n",
        "# Randomly sample actions from env\n",
        "action = env.action_space.sample()\n",
        "\n",
        "# Step the environment\n",
        "next_obs, reward, done, info = env.step(action)\n",
        "\n",
        "print(\"Observation:\", initial_obs)\n",
        "print(\"Action:\", action)\n",
        "print(\"Next observation:\", next_obs)\n",
        "print(\"Reward:\", reward)\n",
        "print(\"Game is done:\", done)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bX9iZtu48UYn"
      },
      "source": [
        "### Episode Return - $R_t$\n",
        "\n",
        "In RL we usually break an agent's interactions with the environment up into **episodes**.The sum of all rewards collected during an episode is what we call the episode's **return** - $R_t$:\n",
        "\n",
        "<center>\n",
        "$R_t=\\sum_{t=0}^Tr_t$,\n",
        "</center>\n",
        "\n",
        "where $r_t$ is the reward at time $t$ and $T$ is the last timestep. The goal in RL is for the agent to chose actions which maximise this expected future return $R_t$.  \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y6KZA1Nq9p47"
      },
      "source": [
        "### Agent-environment Loop\n",
        "Now that we know what a policy is and we know how to step the environment, let's close the agent-environment loop.\n",
        "\n",
        "**Exercise 2:** Write a function that runs one episode of CartPole by sequentially choosing actions and stepping the environment. You should use the linear policy we defined earlier to chose actions. The function should keep track of the reward received and output the return at the end of the episode.\n",
        "\n",
        "In CartPole the agent receives a reward of `1` for every timestep the pole is still upright. If the pole falls over, the game is over and the agent receives no more reward. The game is also over after `200` timesteps, so the maximum reward the agent can collect is `200`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Buy0X7mi-gHP"
      },
      "outputs": [],
      "source": [
        "def run_episode(env):\n",
        "  episode_return = 0 # counter to keep track of rewards\n",
        "  done = False # initially set to False\n",
        "  params = jnp.array([1,-2,2,-1]) # fixed policy parameters\n",
        "\n",
        "  ## YOUR CODE\n",
        "  \n",
        "  obs = ... # TODO: get the initial obs from the env\n",
        "\n",
        "  while not done: # loop until episode is done\n",
        "\n",
        "    action = ... # TODO: compute action using linear policy\n",
        "    action = np.array(action) # We need to the convert the action from the policy to a np.array\n",
        "\n",
        "    obs, reward, done, info = ... # TODO: step the environment\n",
        "\n",
        "    episode_return = ... # TODO: add reward to episode return\n",
        "\n",
        "  return episode_return"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bA2Orj9PVbKO"
      },
      "outputs": [],
      "source": [
        "# @title Check exercise 2 (run me) { display-mode: \"form\" }\n",
        "\n",
        "try:\n",
        "  env.seed(42)\n",
        "  if run_episode(env) == 31:\n",
        "    print(\"Looks correct!\")\n",
        "  else:\n",
        "    print(\"Looks like your implementation might be wrong.\")\n",
        "except Exception as e:\n",
        "    print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h_vu0GImyBT4"
      },
      "outputs": [],
      "source": [
        "#@title Soluction Exercise 2 { display-mode: \"form\" }\n",
        "def run_episode(env):\n",
        "  episode_return = 0 # counter to keep track of rewards\n",
        "  done = False # initially set to False\n",
        "  params = jnp.array([1,-2,2,-1]) # fixed policy parameters\n",
        "\n",
        "  ## YOUR CODE\n",
        "  \n",
        "  obs = env.reset() # TODO: get the initial obs from the env\n",
        "\n",
        "  while not done: # loop until episode is done\n",
        "\n",
        "    # HINT: You might need to the convert the action from your policy to a np.array\n",
        "    action = linear_policy(params, obs) # TODO: compute action using linear policy\n",
        "    action = np.array(action) # We need to the convert the action from the policy to a np.array\n",
        "\n",
        "    obs, reward, done, info = env.step(action) # TODO: step the environment\n",
        "\n",
        "    episode_return = episode_return + reward # TODO: add reward to episode return\n",
        "\n",
        "  return episode_return\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fUGdzHxJnZGl"
      },
      "source": [
        "In CartPole, the environment is considered solved when the agent can reliably achieve an episode return of 500. As you can see, our current policy is far from optimal.\n",
        "\n",
        "One way we can find an optimal policy is by randomly trying out different policies until we find one that is optimal. This strategy is called Random Policy Search and can be supprisingly effective.\n",
        "\n",
        "Before we implement Random Policy Search, let's quickly cover the general RL training loop we will be using  to implement the algorithms in the rest of this turorial.\n",
        "\n",
        "### A General Purpose RL Training Loop\n",
        "We have implemented a general purpose RL training loop for you. The training loop takes several arguments as input but the three most important for you to understand are `agent_select_action_func`, `agent_learn_func` and the `agent_memory`. \n",
        "\n",
        "* The `agent_select_action_func` is a function we define and can pass to the training loop. The function takes an observation and set of `agent_params` as input and should return an action.\n",
        "* The `agent_learn_func` is another method we define and pass to the training loop. It should take the agent's parameters and some \"memories\" as input and then update and return the agents new parameters.\n",
        "* The `agent_memory` is a general purpose module we define that can store some relevant information about the agent's experiences in the environment that can be used in the `agent_learn_func`.\n",
        "\n",
        "\n",
        "Below is the training loop function we have implemented for you. You are welcome to go through the code and try to understand it but this is not required. As such, we have hidden the code by default, just make sure that you run the code cell before moving on because the training loop is used throughout this prac."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZWBwz3zMRjM0"
      },
      "outputs": [],
      "source": [
        "#@title Training loop (run me) { display-mode: \"form\" }\n",
        "\n",
        "# NamedTuple to store transitions\n",
        "Transition = collections.namedtuple(\"Transition\", [\"obs\", \"action\", \"reward\", \"next_obs\", \"done\"])\n",
        "\n",
        "# Training Loop\n",
        "def run_training_loop(env_name, agent_params, agent_select_action_func, \n",
        "    agent_actor_state=None, agent_learn_func=None, agent_learner_state=None, \n",
        "    agent_memory=None, num_episodes=1000, evaluator_period=100, \n",
        "    evaluation_episodes=8, learn_steps_per_episode=1, \n",
        "    train_every_timestep=False, video_subdir=\"\",):\n",
        "    \"\"\"\n",
        "    This function runs several episodes in an environment and periodically does \n",
        "    some agent learning and evaluation.\n",
        "    \n",
        "    Args:\n",
        "        env: a gym environment.\n",
        "        agent_params: an object to store parameters that the agent uses.\n",
        "        agent_select_func: a function that does action selection for the agent.\n",
        "        agent_actor_state (optional): an object that stores the internal state \n",
        "            of the agents action selection function.\n",
        "        agent_learn_func (optional): a function that does some learning for the \n",
        "            agent by updating the agent parameters.\n",
        "        agent_learn_state (optional): an object that stores the internal state \n",
        "            of the agent learn function.\n",
        "        agent_memory (optional): an object for storing an retrieving historical \n",
        "            experience.\n",
        "        num_episodes: how many episodes to run.\n",
        "        evaluator_period: how often to run evaluation.\n",
        "        evaluation_episodes: how many evaluation episodes to run.\n",
        "        train_every_timestep: whether to train every timestep rather than at the end \n",
        "            of the episode.\n",
        "        video_subdir: subdirectory to store epsiode recordings.\n",
        "\n",
        "    Returns:\n",
        "        episode_returns: list of all the episode returns.\n",
        "        evaluator_episode_returns: list of all the evaluator episode returns.\n",
        "    \"\"\"\n",
        "\n",
        "    # Setup Cartpole environment and recorder\n",
        "    env = gym.make(env_name, render_mode=\"rgb_array\") # training environment\n",
        "    eval_env = gym.make(env_name, render_mode=\"rgb_array\") # evaluation environment\n",
        "\n",
        "    # Video dir\n",
        "    video_dir = \"./video\"+\"/\"+video_subdir\n",
        "\n",
        "    # Clear video dir\n",
        "    try:\n",
        "      rmtree(video_dir)\n",
        "    except:\n",
        "      pass\n",
        "\n",
        "    # Wrap in recorder\n",
        "    env = RecordVideo(env, video_dir+\"/train\", episode_trigger=lambda x: (x % evaluator_period) == 0)\n",
        "    eval_env = RecordVideo(eval_env, video_dir+\"/eval\", episode_trigger=lambda x: (x % evaluation_episodes) == 0)\n",
        "\n",
        "    # JAX random number generator\n",
        "    rng = hk.PRNGSequence(jax.random.PRNGKey(0))\n",
        "    env.seed(0) # seed environment for reproducability\n",
        "    random.seed(0)\n",
        "\n",
        "    episode_returns = [] # List to store history of episode returns.\n",
        "    evaluator_episode_returns = [] # List to store history of evaluator returns.\n",
        "    timesteps = 0\n",
        "    for episode in range(num_episodes):\n",
        "\n",
        "        # Reset environment.\n",
        "        obs = env.reset()\n",
        "        episode_return = 0\n",
        "        done = False\n",
        "\n",
        "        while not done:\n",
        "\n",
        "            # Agent select action.\n",
        "            action, agent_actor_state = agent_select_action_func(\n",
        "                                            next(rng), \n",
        "                                            agent_params, \n",
        "                                            agent_actor_state, \n",
        "                                            np.array(obs)\n",
        "                                        )\n",
        "\n",
        "            # Step environment.\n",
        "            next_obs, reward, done, _ = env.step(int(action))\n",
        "\n",
        "            # Pack into transition.\n",
        "            transition = Transition(obs, action, reward, next_obs, done)\n",
        "\n",
        "            # Add transition to memory.\n",
        "            if agent_memory: # check if agent has memory\n",
        "              agent_memory.push(transition)\n",
        "\n",
        "            # Add reward to episode return.\n",
        "            episode_return += reward\n",
        "\n",
        "            # Set obs to next obs before next environment step. CRITICAL!!!\n",
        "            obs = next_obs\n",
        "\n",
        "            # Increment timestep counter\n",
        "            timesteps += 1\n",
        "\n",
        "            # Maybe learn every timestep\n",
        "            if train_every_timestep and (timesteps % 4 == 0) and agent_memory and agent_memory.is_ready(): # Make sure memory is ready\n",
        "                # First sample memory and then pass the result to the learn function\n",
        "                memory = agent_memory.sample()\n",
        "                agent_params, agent_learner_state = agent_learn_func(\n",
        "                                                        next(rng), \n",
        "                                                        agent_params, \n",
        "                                                        agent_learner_state, \n",
        "                                                        memory\n",
        "                                                    )\n",
        "\n",
        "        episode_returns.append(episode_return)\n",
        "\n",
        "        # At the end of every episode we do a learn step.\n",
        "        if agent_memory and agent_memory.is_ready(): # Make sure memory is ready\n",
        "\n",
        "            for _ in range(learn_steps_per_episode):\n",
        "                # First sample memory and then pass the result to the learn function\n",
        "                memory = agent_memory.sample()\n",
        "                agent_params, agent_learner_state = agent_learn_func(\n",
        "                                                        next(rng), \n",
        "                                                        agent_params, \n",
        "                                                        agent_learner_state, \n",
        "                                                        memory\n",
        "                                                    )\n",
        "\n",
        "        if (episode % evaluator_period) == 0: # Do evaluation\n",
        "\n",
        "            evaluator_episode_return = 0\n",
        "            for eval_episode in range(evaluation_episodes):\n",
        "                obs = eval_env.reset()\n",
        "                done = False\n",
        "                while not done:\n",
        "                    action, _ = agent_select_action_func(\n",
        "                                    next(rng), \n",
        "                                    agent_params, \n",
        "                                    agent_actor_state, \n",
        "                                    np.array(obs), \n",
        "                                    evaluation=True\n",
        "                                )\n",
        "\n",
        "                    obs, reward, done, _ = eval_env.step(int(action))\n",
        "\n",
        "                    evaluator_episode_return += reward\n",
        "\n",
        "            evaluator_episode_return /= evaluation_episodes\n",
        "\n",
        "            evaluator_episode_returns.append(evaluator_episode_return)\n",
        "\n",
        "            logs = [\n",
        "                    f\"Episode: {episode}\",\n",
        "                    f\"Episode Return: {episode_return}\",\n",
        "                    f\"Average Episode Return: {np.mean(episode_returns[-20:])}\",\n",
        "                    f\"Evaluator Episode Return: {evaluator_episode_return}\"\n",
        "            ]\n",
        "\n",
        "            print(*logs, sep=\"\\t\") # Print the logs\n",
        "\n",
        "    env.close()\n",
        "    eval_env.close()\n",
        "\n",
        "    return episode_returns, evaluator_episode_returns"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CTPlIq4oDBPY"
      },
      "source": [
        "## Section 2: Random Policy Search (RPS)\n",
        "In Section 1 we used a fixed set of parameters for our policy. That is to say we, didn't learn $\\pi$'s parameters $\\theta$, we simply kept them fixed ( `params = [1,-2,2,-1]`). \n",
        "\n",
        "We will now implment Random Policy Search (RPS), which is an algorithm that  randomly tries different policy parameters and keeps track of the best parameters found so far. We will say that policy parameters $\\theta_A$ are better than parameters $\\theta_B$ if the average episode return achieved over the last 20 episodes by the policy with parameters $\\theta_A$ is greater than that of the policy with parameters $\\theta_B$. \n",
        "\n",
        "\n",
        "To keep track of the \"current\" parameters as well as the \"best\" parameters, we will use a [NamedTuple](https://www.geeksforgeeks.org/namedtuple-in-python/)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1DcaC-PQRjM1"
      },
      "outputs": [],
      "source": [
        "# Parameter container for Random Policy Search\n",
        "RandomPolicySearchParams = collections.namedtuple(\"RandomPolicySearchParams\", [\"current\", \"best\"])\n",
        "\n",
        "# TEST: store two different sets of parameters\n",
        "current_params = np.ones(obs_shape) * -1\n",
        "best_params = np.zeros(obs_shape)\n",
        "rps_params = RandomPolicySearchParams(current_params, best_params)\n",
        "\n",
        "# How to access the best or current params.\n",
        "print(f\"Best params: {rps_params.best}\")\n",
        "print(f\"Current params: {rps_params.current}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "v91pDDzGex9a"
      },
      "source": [
        "Next we will implement the following:\n",
        "  - **RPS select action function** - define how we choose actions given a set of parameters.\n",
        "  - **RPS memory module** - define what experiences to store from the environment interactions.\n",
        "  - **RPS learn function** - define how we update and improve our policy parameters."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0tExceGeGNYH"
      },
      "source": [
        "### RPS choose action function\n",
        "Let's implement a function called `random_policy_search_choose_action` which we can pass to the training loop. The function needs to take in several arguments inorder for it to interface nicely with our generalised training loop but you will only need to use three of them - `params`, `obs` and `evaluation`.\n",
        "\n",
        "- `params` is an instance of `RandomPolicySearchParams` with \"current\" and \"best\".\n",
        "- `obs` is the latest observation from the environment.\n",
        "- `evaluation` is a boolean value that indicates if we should use the \"current\" or \"best\" parameters.  When `evaluation==True` we should use the \"best\" parameters, else we should use the \"current\" parameters.\n",
        "\n",
        "**Exercise 3:** Implement the `random_policy_search_choose_action` function as described above. You should make use of the `linear_policy` method we defined earlier. You will also want to use `jax.lax.select()` to conditionally return the \"best\" action or the \"current\" action. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1kmwT35JRjM1"
      },
      "outputs": [],
      "source": [
        "def random_policy_search_choose_action(\n",
        "    key, \n",
        "    params, \n",
        "    actor_state, \n",
        "    obs, \n",
        "    evaluation=False\n",
        "):\n",
        "  \"\"\"Random policy search select action method.\n",
        "\n",
        "  Args:\n",
        "    key: a random number (seed). Not used in this function.\n",
        "    params: the agent's parameters. In this case an instance of `RandomPolicySearchParams`\n",
        "    actor_state: some extra information about the actor. Not used in this function.\n",
        "    obs: the latest observation.\n",
        "    evaluation: a boolean indicating whether to use the best \"parameters\" or the \"current\" ones.\n",
        "\n",
        "  Returns:\n",
        "    The chosen action and the updated actor_state. In this function the actor_state is not updated.\n",
        "  \"\"\"\n",
        "\n",
        "  # YOUR CODE\n",
        "\n",
        "  best_action = ...\n",
        "\n",
        "  current_action = ...\n",
        "\n",
        "  action = jax.lax.select(\n",
        "      ... ,\n",
        "      ... ,\n",
        "      ... \n",
        "  )\n",
        "\n",
        "  # END YOUR CODE\n",
        "\n",
        "  return action, actor_state"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YE4xoAgyI3kG"
      },
      "outputs": [],
      "source": [
        "# @title Check exercise 3 (run me) {display-mode: \"form\"}\n",
        "\n",
        "def check_random_policy_search_choose_action(choose_action):\n",
        "  key = None # not used\n",
        "  actor_state = None # not used\n",
        "\n",
        "  # obs\n",
        "  obs = np.ones(obs_shape)\n",
        "\n",
        "  evaluation=False\n",
        "  current_params = np.ones(obs_shape) * -1\n",
        "  best_params = np.ones(obs_shape)\n",
        "  rps_params = RandomPolicySearchParams(current_params, best_params)\n",
        "  action, actor_state = choose_action(key,rps_params,actor_state,obs,evaluation)\n",
        "  if action != 0:\n",
        "    return False\n",
        "\n",
        "  evaluation=True\n",
        "  current_params = np.ones(obs_shape) * -1\n",
        "  best_params = np.ones(obs_shape)\n",
        "  rps_params = RandomPolicySearchParams(current_params, best_params)\n",
        "  action, actor_state = choose_action(key,rps_params,actor_state,obs,evaluation)\n",
        "  if action != 1:\n",
        "    return False\n",
        "\n",
        "  return True\n",
        "\n",
        "try:\n",
        "  if check_random_policy_search_choose_action(random_policy_search_choose_action):\n",
        "    print(\"Your function looks correct.\")\n",
        "  else:\n",
        "    print(\"Your function looks incorrect.\")\n",
        "except Exception as e:\n",
        "    print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UH_9qJ08IFg6"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 3 {display-mode: \"form\"}\n",
        "\n",
        "def random_policy_search_choose_action(\n",
        "    key, \n",
        "    params, \n",
        "    actor_state, \n",
        "    obs, \n",
        "    evaluation=False\n",
        "):\n",
        "\n",
        "  best_action = linear_policy(params.best, obs)\n",
        "\n",
        "  current_action = linear_policy(params.current, obs)\n",
        "\n",
        "  action = jax.lax.select(\n",
        "      evaluation ,\n",
        "      best_action ,\n",
        "      current_action \n",
        "  )\n",
        "\n",
        "  return action, actor_state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6zu51ep7Sh0M"
      },
      "source": [
        "### RPS agent memory\n",
        "\n",
        "For the Random Policy Search algorithm we will need to keep track of the average episode return for the last 50 episodes. Remember that we said the \"current\" parameters will be deemed to be the \"best\"parameters if the average episode return when using those parameters is greater than the previous best average episode return.\n",
        "\n",
        "We will use a general purpose memory interface which is fairy simple. The memory module should have three methods. The first is a function `memory.push(<transition>)` that adds some information about the latest environment transition to the memory. The second is a function `memory.is_ready()` to check if the memory is ready to do some learning. Finally a function `memory.sample()` should return the latest set of memories that can be passed to the `agent_learn_func`.\n",
        "\n",
        "#### Average Episode Return Memory\n",
        "We have built a simple agent memory module for you. It stores the `epsisode_returns` of the last 20 episodes. Read through our implementation below and see if you can understand it. The `memory.sample()` method returns the average episode return over the last 20 episodes."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YVkTBIK5RjM1"
      },
      "outputs": [],
      "source": [
        "class AverageEpisodeReturnBuffer:\n",
        "\n",
        "    def __init__(self, num_episodes_to_store=50):\n",
        "        \"\"\"\n",
        "        This class implements an agent memory that stores the average episode \n",
        "        return over the last 50 episodes.\n",
        "        \"\"\"\n",
        "        self.num_episodes_to_store = num_episodes_to_store\n",
        "        self.episode_return_buffer = []\n",
        "        self.current_episode_return = 0\n",
        "\n",
        "    def push(self, transition):\n",
        "        self.current_episode_return += transition.reward\n",
        "\n",
        "        if transition.done: # If the episode is done\n",
        "            # Add episode return to buffer\n",
        "            self.episode_return_buffer.append(self.current_episode_return)\n",
        "\n",
        "            # Reset episode return\n",
        "            self.current_episode_return = 0\n",
        "\n",
        "\n",
        "    def is_ready(self):\n",
        "        return len(self.episode_return_buffer) == self.num_episodes_to_store\n",
        "\n",
        "    def sample(self):\n",
        "        average_episode_return = np.mean(self.episode_return_buffer)\n",
        "\n",
        "        # Clear episode return buffer\n",
        "        self.episode_return_buffer = []\n",
        "\n",
        "        return average_episode_return"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0oXBsSa2KjWE"
      },
      "source": [
        "### RPS learn function\n",
        "Fnally, we need to implement the `random_policy_search_learn` function for our Random Policy Search algorithm. The learn function is quite simple. All we need to do is check if the current parameters are better than the best parameters. If they are better, then set the best parameters to be the current parameters and randomly generate a new set of current parameters. \n",
        "\n",
        "**Exercise 4:** Write a function to randomly generates new weights using JAX. The weights should be sampled from the interval `[-2,2]`.\n",
        "\n",
        "**Useful functions:** \n",
        "*   `jax.random.uniform` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.random.uniform.html#jax.random.uniform))\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8V2yFM2XMjGW"
      },
      "outputs": [],
      "source": [
        "def get_new_random_weights(random_key, old_weights, minval=-2.0, maxval=2.0):\n",
        "    new_weights_shape = old_weights.shape # you will need to use these values\n",
        "    new_weights_dtype = old_weights.dtype # you will need to use these values\n",
        "\n",
        "    # YOUR CODE\n",
        "\n",
        "    new_params = ...\n",
        "\n",
        "    # END YOUR CODE\n",
        "\n",
        "    return new_params"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T8ifHTKGQlAd"
      },
      "outputs": [],
      "source": [
        "# @title Check exercise 4 (run me) {display-mode: \"form\"}\n",
        "\n",
        "def check_get_new_random_weights(get_new_random_weights):\n",
        "  old_weights = np.ones(obs_shape, \"float32\")\n",
        "  random_key = jax.random.PRNGKey(42)\n",
        "\n",
        "  # Case 1\n",
        "  new_weights = get_new_random_weights(random_key, old_weights, minval=-2.0, maxval=2.0)\n",
        "  \n",
        "  if jnp.array_equal(new_weights, jnp.array([ 0.29657745,1.4265499, -1.7621555, -1.7505779 ])):\n",
        "    print(\"Function is correct!\")\n",
        "  else:\n",
        "    print(\"Something is wrong.\")\n",
        "\n",
        "try:\n",
        "  check_get_new_random_weights(get_new_random_weights)\n",
        "except Exception as e:\n",
        "    print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "T05R_AWHLx_I"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 4 {display-mode: \"form\"}\n",
        "\n",
        "def get_new_random_weights(random_key, old_weights,minval=-2.0,maxval=2.0):\n",
        "    new_weights_shape = old_weights.shape\n",
        "    new_weights_dtype = old_weights.dtype\n",
        "    # Sample new weights\n",
        "    new_weights = jax.random.uniform(random_key,new_weights_shape,new_weights_dtype,minval=minval,\n",
        "                      maxval=maxval)\n",
        "    return new_weights"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2Djxc9j-LzkM"
      },
      "source": [
        "Our learn function receives a memory, in the form of the average episode return, from the `AverageEpisodeReturnMemory` we implemented earlier. We can use this to compare the current parameters to the best parameters. But we will also need to keep track of the best average episode return for the learn function. For that, we can use the `learn_state` argument which is passed to the `agent_learn_func` in our training loop. As with the `RandomPolicySearchParams`, we will use a NamedTuple to store the `best_average_episode_return` in the `learn_state`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4cSH4uYmRjM2"
      },
      "outputs": [],
      "source": [
        "# A NamedTuple to store the best average episode return so far\n",
        "RandomPolicyLearnState = collections.namedtuple(\n",
        "  \"RandomPolicyLearnState\", \n",
        "  [\"best_average_episode_return\"]\n",
        ")\n",
        "\n",
        "# Test\n",
        "initial_learn_state = RandomPolicyLearnState(best_average_episode_return=-float(\"inf\"))\n",
        "print(\"Initial best average episode return:\", initial_learn_state.best_average_episode_return)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PxZXdP8bOu6b"
      },
      "source": [
        "Now we have everything we need to implement the `random_policy_search_learn` function.\n",
        "\n",
        "**Exercise 5:** Implement the `random_policy_search_learn` function. The function should check if the \"current\" parameters are better than the \"best\" parameters by comparing the `current_average_episode_return` to the `best_average_episode_return`. The function should also update the `learn_state`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Te_Q3qzmZcN_"
      },
      "outputs": [],
      "source": [
        "def random_policy_search_learn(key, params, learn_state, memory):\n",
        "    best_params = params.best \n",
        "    current_params = params.current\n",
        "\n",
        "    current_average_episode_return = memory # the memory contains the average episode return\n",
        "    best_average_episode_return = learn_state.best_average_episode_return\n",
        "\n",
        "\n",
        "    # YOUR CODE\n",
        "\n",
        "    best_params = jax.lax.select(\n",
        "        ... ,\n",
        "        ... ,\n",
        "        ...\n",
        "    )\n",
        "        \n",
        "    best_average_episode_return = jax.lax.select(\n",
        "        ... ,\n",
        "        ... ,\n",
        "        ...\n",
        "    )\n",
        "    \n",
        "    # END YOUR CODE\n",
        "\n",
        "    # Generate new random parameters\n",
        "    new_params = get_new_random_weights(key, current_params)\n",
        "\n",
        "    # Bundle weights in RandomPolicySearchParams NamedTuple\n",
        "    params = RandomPolicySearchParams(current=new_params, best=best_params)\n",
        "\n",
        "    return params, RandomPolicyLearnState(best_average_episode_return)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CXHV4k2iBLBF"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 5 {display-mode: \"form\"}\n",
        "\n",
        "params = RandomPolicySearchParams(np.ones(obs_shape, \"float32\"), np.ones(obs_shape, \"float32\") * -1)\n",
        "learn_state = RandomPolicyLearnState(10)\n",
        "memory = 11\n",
        "key = jax.random.PRNGKey(42)\n",
        "\n",
        "try:\n",
        "  new_params, new_learn_state = random_policy_search_learn(key, params, learn_state, memory)\n",
        "\n",
        "  if not jnp.array_equal(new_params.current, jnp.array([ 0.29657745,  1.4265499 , -1.7621555 , -1.7505779 ])):\n",
        "    print(\"Your function is incorrect.\")\n",
        "\n",
        "  elif not jnp.array_equal(new_params.best, jnp.array([1., 1., 1., 1.])):\n",
        "    print(\"Your function is incorrect.\")\n",
        "\n",
        "  elif new_learn_state.best_average_episode_return != 11:\n",
        "    print(\"Your function is incorrect.\")\n",
        "\n",
        "  else:\n",
        "    print(\"Your function looks correct.\")\n",
        "except Exception as e:\n",
        "    print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-ElX1w1iQTrE"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 5 {display-mode: \"form\"}\n",
        "\n",
        "def random_policy_search_learn(key, params, learn_state, memory):\n",
        "    best_params = params.best \n",
        "    current_params = params.current\n",
        "\n",
        "    current_average_episode_return = memory # the memory contains the average episode return\n",
        "    best_average_episode_return = learn_state.best_average_episode_return\n",
        "\n",
        "\n",
        "    # YOUR CODE\n",
        "\n",
        "    best_params = jax.lax.select(\n",
        "        current_average_episode_return > best_average_episode_return,\n",
        "        current_params,\n",
        "        best_params,\n",
        "    )\n",
        "        \n",
        "    best_average_episode_return = jax.lax.select(\n",
        "        current_average_episode_return > best_average_episode_return,\n",
        "        current_average_episode_return,\n",
        "        best_average_episode_return\n",
        "    )\n",
        "    \n",
        "    # END YOUR CODE\n",
        "\n",
        "    # Generate new random parameters\n",
        "    new_params = get_new_random_weights(key, current_params)\n",
        "\n",
        "    # Bundle weights in RandomPolicySearchParams NamedTuple\n",
        "    params = RandomPolicySearchParams(current=new_params, best=best_params)\n",
        "\n",
        "    # Update learn_state\n",
        "    learn_state = RandomPolicyLearnState(best_average_episode_return)\n",
        "\n",
        "    return params, learn_state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5Ol_AxnMBdgP"
      },
      "source": [
        "### RPS training loop\n",
        "We can now put everything together by passing the `memory` module, the `learn` function and `choose_action` function to the training loop. To help speed up our algorithm we will use `jax.jit` on the `learn` function and the `choose_action` function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qx57tf7vRjM3"
      },
      "outputs": [],
      "source": [
        "# JIT the learn and choose action functions\n",
        "random_policy_search_learn_jit = jax.jit(random_policy_search_learn)\n",
        "random_policy_search_choose_action_jit = jax.jit(random_policy_search_choose_action)\n",
        "\n",
        "# Initialise the parameters\n",
        "initial_weights = np.ones(obs_shape, \"float32\")\n",
        "initial_params = RandomPolicySearchParams(initial_weights, initial_weights)\n",
        "\n",
        "# Initialise the learn state\n",
        "initial_learn_state = RandomPolicyLearnState(best_average_episode_return=-float(\"inf\"))\n",
        "\n",
        "# Initialise memory\n",
        "memory = AverageEpisodeReturnBuffer(num_episodes_to_store=50)\n",
        "\n",
        "# Run the training loop\n",
        "print(\"Starting training. This may take up to 5 minutes to complete.\")\n",
        "chex.clear_trace_counter()\n",
        "episode_return, evaluator_episode_returns = run_training_loop(\n",
        "                                        env_name,\n",
        "                                        initial_params, \n",
        "                                        random_policy_search_choose_action_jit, \n",
        "                                        None, # no actor state\n",
        "                                        random_policy_search_learn_jit, \n",
        "                                        initial_learn_state, \n",
        "                                        memory, \n",
        "                                        num_episodes=1001,\n",
        "                                        video_subdir=\"rps\"\n",
        "                                    )\n",
        "\n",
        "# Plot graph of evaluator episode returns\n",
        "plt.plot(np.linspace(0, 1000, len(evaluator_episode_returns)), evaluator_episode_returns)\n",
        "plt.title(\"Random Policy Search\")\n",
        "plt.xlabel(\"Episodes\")\n",
        "plt.ylabel(\"Episode Return\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TG10FG6uS05A"
      },
      "source": [
        "Hopefully, you found a set of optimal parameters on CartPole (episode return reaches `200`). In the cell below you can watch some videos of the agent doing the task."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "KgayDJ7KWC4C"
      },
      "outputs": [],
      "source": [
        "#@title Visualise Policy {display-mode: \"form\"}\n",
        "#@markdown Choose an episode number that is a multiple of 100 and less than or equal to 2000, and **run** this cell.\n",
        "episode_number = 0 #@param {type:\"number\"}\n",
        "\n",
        "assert (episode_number % 100) == 0, \"Episode number must be a multiple of 100 since we only record every 100th episode.\"\n",
        "assert episode_number < 1001, \"Episode number must be less than or equal to 2000\"\n",
        "\n",
        "eval_episode_number = int(episode_number / 100 * 8)\n",
        "video_path = f\"./video/rps/eval/rl-video-episode-{eval_episode_number}.mp4\"\n",
        "\n",
        "mp4 = open(video_path,'rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "HTML(\"\"\"\n",
        "<video width=400 controls>\n",
        "      <source src=\"%s\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\" % data_url)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Xwf1IKKbrrVn"
      },
      "source": [
        "So, Random Policy Search did pretty well on this task. However, there is very little (if any) real *learning* going on here. Next, let's look at implementing a simple RL algorithm instead, that can use its experiences to guide our search for an optimal policy, rather than just randomly searching for it."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XEnSjZVESrxc"
      },
      "source": [
        "## Section 3: Policy Gradients (PG)\n",
        "As discussed, the goal in RL is to find a policy which maximise the expected cummulative reward (return) the agent receives from the environment. We can write the expected return of a policy as:\n",
        "\n",
        "$J(\\pi_\\theta)=\\mathrm{E}_{\\tau\\sim\\pi_\\theta}\\ [R(\\tau)]$,\n",
        "\n",
        "where $\\pi_\\theta$ is a policy parametrised by $\\theta$, $\\mathrm{E}$ means *expectation*, $\\tau$ is shorthand for \"*episode*\", $\\tau\\sim\\pi_\\theta$ is shorthand for \"*episodes sampled using the policy* $\\pi_\\theta$\", and $R(\\tau)$ is the return of episode $\\tau$.\n",
        "\n",
        "Then, the goal in RL is to find the parameters $\\theta$ that maximise the function $J(\\pi_\\theta)$. One way to find these parameters is to perform gradient ascent on $J(\\pi_\\theta)$ with respect to the parameters $\\theta$: \n",
        "\n",
        "$\\theta_{k+1}=\\theta_k + \\alpha \\nabla J(\\pi_\\theta)|_{\\theta_{k}}$,\n",
        "\n",
        "where $\\nabla J(\\pi_\\theta)|_{\\theta_{k}}$ is the gradient of the expected return with respect to the policy parameters $\\theta_k$ and $\\alpha$ is the step size. This quantity, $\\nabla J(\\pi_\\theta)$, is also called the **policy gradient** and is very important in RL. If we can comput the policy gradient, then we will have a means by which to directly optimise our policy.\n",
        "\n",
        "As it turns out, there is a way for us to compute the policy gradient and the mathematical derivation can be found [here](https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html). But for this tutorial we will ommit the derivation and just give you the result:\n",
        "\n",
        "\n",
        "$\\nabla_{\\theta} J(\\pi_{\\theta})=\\underset{\\tau \\sim \\pi_{\\theta}}{\\mathrm{E}}[\\sum_{t=0}^{T} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_{t} \\mid s_{t}) R(\\tau)]$\n",
        "\n",
        "Informaly, the policy gradient is equal to the gradient of the log of the probability of the action chosen, multiplied by the return of the episode in which the action was taken.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BTnTzgtSuy-y"
      },
      "source": [
        "### REINFORCE\n",
        "REINFORCE is a simple RL algorithm that uses the policy gradient to find the optimal policy by increasing the probability of choosing actions (reinforcing actions) that tend to lead to high return episodes.\n",
        "\n",
        "**Exercise 6:** Implement a function that takes the probability of an action and the return of the episode the action was taken in and computes the log of the probability, multiplied by the return. Make sure you use JAX.\n",
        "\n",
        "**Useful functions:**\n",
        "*   `jax.numpy.log`([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.log.html))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bJObUsoUrOyV"
      },
      "outputs": [],
      "source": [
        "def compute_weighted_log_prob(action_prob, episode_return):\n",
        "\n",
        "    # YOUR CODE\n",
        "\n",
        "    log_porb = ...\n",
        "\n",
        "    weighted_log_prob = ... \n",
        "\n",
        "    # END YOUR CODE\n",
        "\n",
        "    return weighted_log_prob"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eZPoTwa1Gbm1"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 6 {display-mode: \"form\"}\n",
        "\n",
        "try:\n",
        "  action_prob = 0.8\n",
        "  episode_return = 100\n",
        "  result = compute_weighted_log_prob(action_prob, episode_return)\n",
        "  if result != -22.314354:\n",
        "    print(\"Your implementation looks incorrect.\")\n",
        "  else:\n",
        "    print(\"Looks correct.\")\n",
        "except Exception as e:\n",
        "    print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x7dTAlCdrJkf"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 6 {display-mode: \"form\"}\n",
        "\n",
        "def compute_weighted_log_prob(action_prob, episode_return):\n",
        "\n",
        "    # YOUR CODE\n",
        "\n",
        "    log_prob = jax.numpy.log(action_prob)\n",
        "\n",
        "    weighted_log_prob = log_prob * episode_return\n",
        "\n",
        "    # END YOUR CODE\n",
        "\n",
        "    return weighted_log_prob\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NmgW9UJ3tIpl"
      },
      "source": [
        "### Rewards-to-go\n",
        "Performing gradient ascent on the gradient of the log of the action probability, weighted by the return of the episode will tend to push up the probability of actions that were in episodes with high return, regardless of *where* in the episode the action was taken. This does not really make much sense because an action near the end of an episode may be reinforced because lots of reward was collected earlier on in the episode, *before* the action was taken. RL agents should really only reinforce actions on the basis of their *consequences*. Rewards obtained before taking an action have no bearing on how good that action was: only rewards that come after. The cummulative rewards received after an action was taken is called the **rewards-to-go** and can be computed as:\n",
        "\n",
        "$\\hat{R}_i=\\sum_{t=i}^Tr_t$\n",
        "\n",
        "Compare this to the episode return:\n",
        "\n",
        "$R(\\tau)=\\sum_{t=0}^Tr_t$\n",
        "\n",
        "We can improve the reliability of the policy gradient by substituting the episode return with the rewards-to-go. The policy gradient with rewards-to-go is given by:\n",
        "\n",
        "$\\nabla_{\\theta} J(\\pi_{\\theta})=\\underset{\\tau \\sim \\pi_{\\theta}}{\\mathrm{E}}[\\sum_{t=0}^{T} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_{t} \\mid s_{t}) \\hat{R}_t]$\n",
        "\n",
        "**Exercise 7:** Implement a function that takes a list of all the rewards obtained in an episode and computes the rewards-to-go. Don't worry about using JAX in this function. You can use regular Python operations like `for-loops`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nV1Hww8E3dUJ"
      },
      "outputs": [],
      "source": [
        "def compute_rewards_to_go(rewards):\n",
        "    \"\"\"\n",
        "    This function should take a list of rewards as input and \n",
        "    compute the rewards-to-go for each timestep.\n",
        "\n",
        "    EXAMPLE: compute_rewards_to_go([1,2,3,4]) = [10, 9, 7, 4]\n",
        "    \n",
        "    Arguments:\n",
        "        rewards[t] is the reward at time step t.\n",
        "\n",
        "    Returns:\n",
        "        rewards_to_go[t] should be the reward-to-go at timestep t.\n",
        "    \"\"\"\n",
        "\n",
        "    rewards_to_go = []\n",
        "\n",
        "    # YOUR CODE\n",
        "\n",
        "\n",
        "    # END YOUR CODE\n",
        "\n",
        "    return rewards_to_go"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hLVaVRp28YGI"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 7 {display-mode: \"form\"}\n",
        "\n",
        "try: \n",
        "  result = compute_rewards_to_go([1,2,3,4])\n",
        "\n",
        "  if result != [10, 9, 7, 4]:\n",
        "    print(\"There is a problem with your implementation.\")\n",
        "  else:\n",
        "    print(\"Looks correct.\")\n",
        "except Exception as e:\n",
        "    print(\"An Error Occured: {}\".format(e))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "C2qy2F3xRjM3"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 7 {display-mode: \"form\"}\n",
        "\n",
        "def compute_rewards_to_go(rewards):\n",
        "    rewards_to_go = []\n",
        "    for i in range(len(rewards)):\n",
        "        r2g = 0\n",
        "        for j in range(i, len(rewards)):\n",
        "            r2g += rewards[j]\n",
        "        rewards_to_go.append(r2g)\n",
        "    return rewards_to_go"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IboxN9MS65i5"
      },
      "source": [
        "### REINFORCE memory\n",
        "Next we will need to make a new agent memory to store the rewards-to-go $\\hat{R}_t$ along with the observation $o_t$ and action $a_t$ at every timestep. Below we implemented such a memory module for you. The function `memory.sample()` will return a batch of last 500 memories. You are welcome to read through the code to try and understand it, but it is not required. Therefore, we hide the code by default."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xhS4V6auRjM3"
      },
      "outputs": [],
      "source": [
        "# @title Memory implementation (run me) {display-mode: \"form\"}\n",
        "\n",
        "# NamedTuple to store memory\n",
        "EpisodeRewardsToGoMemory = collections.namedtuple(\"EpisodeRewardsToGoMemory\", [\"obs\", \"action\", \"reward_to_go\"])\n",
        "\n",
        "class EpisodeRewardsToGoBuffer:\n",
        "\n",
        "    def __init__(self, num_transitions_to_store=512, batch_size=256):\n",
        "        self.batch_size = batch_size\n",
        "        self.memory_buffer = collections.deque(maxlen=num_transitions_to_store)\n",
        "        self.current_episode_transition_buffer = []\n",
        "\n",
        "    def push(self, transition):\n",
        "        self.current_episode_transition_buffer.append(transition)\n",
        "\n",
        "        if transition.done:\n",
        "\n",
        "            episode_rewards = []\n",
        "            for t in self.current_episode_transition_buffer:\n",
        "                episode_rewards.append(t.reward)\n",
        "\n",
        "            r2g = compute_rewards_to_go(episode_rewards)\n",
        "\n",
        "            for i, t in enumerate(self.current_episode_transition_buffer):\n",
        "                memory = EpisodeRewardsToGoMemory(t.obs, t.action, r2g[i])\n",
        "                self.memory_buffer.append(memory)\n",
        "\n",
        "            # Reset episode buffer\n",
        "            self.current_episode_transition_buffer = []\n",
        "\n",
        "\n",
        "    def is_ready(self):\n",
        "        return len(self.memory_buffer) >= self.batch_size\n",
        "\n",
        "    def sample(self):\n",
        "        random_memory_sample = random.sample(self.memory_buffer, self.batch_size)\n",
        "\n",
        "        obs_batch, action_batch, reward_to_go_batch = zip(*random_memory_sample)\n",
        "\n",
        "        return EpisodeRewardsToGoMemory(\n",
        "            np.stack(obs_batch).astype(\"float32\"), \n",
        "            np.asarray(action_batch).astype(\"int32\"), \n",
        "            np.asarray(reward_to_go_batch).astype(\"int32\")\n",
        "        )\n",
        "\n",
        "\n",
        "# Instantiate Memory\n",
        "REINFORCE_memory = EpisodeRewardsToGoBuffer(num_transitions_to_store=512, batch_size=256)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Idkav_aSYXvz"
      },
      "source": [
        "### Policy neural network\n",
        "Next, we will use a simple neural network to aproximate the policy. Our policy neural network will have an input layer that takes the observation as input and passes it through two hidden layers and then outputs one scalar value for each of the possible actions. So, in CartPole the output layer will have size `2`.\n",
        "\n",
        "[Haiku](https://github.com/deepmind/dm-haiku) is a library for implementing neural networks is JAX. Below we have implemented a simple function to make the policy network for you. \n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j2XO7VkORjM4"
      },
      "outputs": [],
      "source": [
        "def make_policy_network(num_actions: int, layers=[20, 20]) -> hk.Transformed:\n",
        "  \"\"\"Factory for a simple MLP network for the policy.\"\"\"\n",
        "\n",
        "  def policy_network(obs):\n",
        "    network = hk.Sequential(\n",
        "        [\n",
        "            hk.Flatten(),\n",
        "            hk.nets.MLP(layers + [num_actions])\n",
        "        ]\n",
        "    )\n",
        "    return network(obs)\n",
        "\n",
        "  return hk.without_apply_rng(hk.transform(policy_network))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1GR2y8FjaG-G"
      },
      "source": [
        "Haiku networks have two important functions you need to know about. The first is the `network.init(<random_key>, <input>)`, which returns a set of random initial parameters. The second method is the `network.apply(<params>, <input>)` which passes an input through the network using the set of parameters provided."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fJrn9o-Vatkw"
      },
      "outputs": [],
      "source": [
        "# Example\n",
        "POLICY_NETWORK = make_policy_network(num_actions=num_actions, layers=[20,20])\n",
        "random_key = jax.random.PRNGKey(42) # random key\n",
        "dummy_obs = np.ones(obs_shape, \"float32\")\n",
        "\n",
        "# Initialise parameters\n",
        "REINFORCE_params = POLICY_NETWORK.init(random_key, dummy_obs)\n",
        "print(\"Initial params:\", REINFORCE_params.keys())\n",
        "\n",
        "# Pass input through the network\n",
        "output = POLICY_NETWORK.apply(REINFORCE_params, dummy_obs)\n",
        "print(\"Policy network output:\", output)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MlouUBvoeunz"
      },
      "source": [
        "The outputs of our policy network are [logits](https://qr.ae/pv4YTe). To convert this into a probability distribution over actions we pass the logits to the [softmax](https://en.wikipedia.org/wiki/Softmax_function) function.\n",
        "\n",
        "### REINFORCE choose action function\n",
        "\n",
        "**Exercise 8:** Complete the function below which takes a vector of logits and randomly samples an action from a categorical distibution given by the logits. \n",
        "\n",
        "**Useful functions:**\n",
        "*   `jax.random.categorical` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.random.categorical.html))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m3Z8DxUmeOGJ"
      },
      "outputs": [],
      "source": [
        "def sample_action(random_key, logits):\n",
        "    \n",
        "  # YOUR CODE HERE\n",
        "  action = ...\n",
        "\n",
        "  # END YOUR code\n",
        "\n",
        "  return action"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5opHJMO0D_Ub"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 8 {display-mode: \"form\"}\n",
        "\n",
        "try:\n",
        "  random_key = jax.random.PRNGKey(42) # random key\n",
        "  action = sample_action(random_key, np.array([1,2], \"float32\"))\n",
        "  if action != 1:\n",
        "    print(\"Your function is incorrect.\")\n",
        "  else:\n",
        "    print(\"Seems correct.\")\n",
        "except Exception as e:\n",
        "    print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2huSgyukg1N4"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 8 {display-mode: \"form\"}\n",
        "\n",
        "def sample_action(random_key, logits):\n",
        "    \n",
        "    # YOUR CODE HERE\n",
        "\n",
        "    action = jax.random.categorical(random_key, logits)\n",
        "\n",
        "    # END YOUR code\n",
        "\n",
        "    return action\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jP5UH87VRjM4"
      },
      "source": [
        "Now we can implement the `REINFORCE_choose_action` function. We will pass the observation through the policy network to compute the logits and then pass the logits to the `sample_action` function to choose and action."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UJTzrDAZ0Ul5"
      },
      "outputs": [],
      "source": [
        "def REINFORCE_choose_action(key, params, actor_state, obs, evaluation=False):\n",
        "  obs = jnp.expand_dims(obs, axis=0) # add dummy batch dim before passing through network\n",
        "\n",
        "  # Pass obs through policy network to compute logits\n",
        "  logits = POLICY_NETWORK.apply(params, obs)\n",
        "  logits = logits[0] # remove batch dim\n",
        "\n",
        "  # Randomly sample action\n",
        "  sampled_action = sample_action(key, logits)\n",
        "  \n",
        "  return sampled_action, actor_state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NI26SLAb7iRo"
      },
      "source": [
        "Now that we have  implemented the `REINFORCE_choose_action` function, all we have left to do is to make a `REINFORCE_learn` function. The learn function should use the `weighted_log_prob` function we made earlier to compute the policy gradient loss and apply the gradient updates to our neural network."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "36ALCJESQJ8e"
      },
      "source": [
        "### Policy gradient loss\n",
        "\n",
        "**Exercise 9:** Complete the `policy_gradient_loss` function below. The function should compute the action probabilities by passing the `logits` through the softmax function. Then you should extract the probability of the given `action` (using array indexing) and compute the `weighted_log_prob` using the function we made earlier.\n",
        "\n",
        "**Useful methods:**\n",
        "*   `jax.nn.softmax` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.nn.softmax.html))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9sUKkqx0RjM4"
      },
      "outputs": [],
      "source": [
        "def policy_gradient_loss(action, logits, reward_to_go):\n",
        "\n",
        "  # YOUR CODE\n",
        "\n",
        "  all_action_probs = ... # convert logits into probs\n",
        "\n",
        "  action_prob = ... # using array indexing to get prob of action\n",
        "\n",
        "  weighted_log_prob = ...\n",
        "\n",
        "  # END YOUR CODE\n",
        "  \n",
        "  loss = - weighted_log_prob # negative because we want gradient `ascent`\n",
        "  \n",
        "  return loss"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7AMJvau1FsM5"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 9 {display-mode: \"form\"}\n",
        "\n",
        "try:\n",
        "  result = policy_gradient_loss(1, np.array([1,2], \"float32\"), 10)\n",
        "  if result != 3.1326165:\n",
        "    print(\"Your implementation looks wrong.\")\n",
        "  else:\n",
        "    print(\"Looks correct.\")\n",
        "except Exception as e:\n",
        "  print(\"An Error Occured: {}\".format(e))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1HAxSFU0qBVo"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 9 {display-mode: \"form\"}\n",
        "\n",
        "def policy_gradient_loss(action, logits, reward_to_go):\n",
        "\n",
        "  # YOUR CODE\n",
        "\n",
        "  all_action_probs = jax.nn.softmax(logits) # convert logits into probs\n",
        "\n",
        "  action_prob = all_action_probs[action]\n",
        "\n",
        "  weighted_log_prob = compute_weighted_log_prob(action_prob, reward_to_go)\n",
        "\n",
        "  # END YOUR CODE\n",
        "  \n",
        "  loss = - weighted_log_prob # negative because we want gradient `ascent`\n",
        "  \n",
        "  return loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gzuqx1jJrwVx"
      },
      "source": [
        "When we do a policy gradient update step we are going to want to do it using a batch of experience, rather than just a single experience like above. We can use JAX's [vmap](https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html#jax.vmap) function to easily make our `policy_gradient_loss` function work on a batch of experience."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3yq4naLURjM4"
      },
      "outputs": [],
      "source": [
        "def batched_policy_gradient_loss(params, obs_batch, action_batch, reward_to_go_batch):\n",
        "    # Get logits by passing observation through network\n",
        "    logits_batch = POLICY_NETWORK.apply(params, obs_batch)\n",
        "\n",
        "    policy_gradient_loss_batch = jax.vmap(policy_gradient_loss)(action_batch, logits_batch, reward_to_go_batch) # add batch\n",
        "\n",
        "    # Compute mean loss over batch\n",
        "    mean_policy_gradient_loss = jnp.mean(policy_gradient_loss_batch)\n",
        "\n",
        "    return mean_policy_gradient_loss\n",
        "\n",
        "# TEST\n",
        "obs_batch = np.ones((3, *obs_shape), \"float32\")\n",
        "actions_batch = np.array([1,0,0])\n",
        "rew2go_batch = np.array([2.3, 4.3, 2.1])\n",
        "\n",
        "loss = batched_policy_gradient_loss(REINFORCE_params, obs_batch, actions_batch, rew2go_batch)\n",
        "\n",
        "print(\"Policy gradient loss on batch:\", loss)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VDhTH3culwqo"
      },
      "source": [
        "### Network Optimiser\n",
        "\n",
        "To apply policy gradient updates to our neural network we will use a JAX library called [Optax](https://github.com/deepmind/optax). Optax has an implementation of the [Adam optimizer](https://www.geeksforgeeks.org/intuition-of-adam-optimizer/) which we can use."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pxXINlMHP5Ic"
      },
      "outputs": [],
      "source": [
        "REINFORCE_OPTIMIZER = optax.adam(1e-3)\n",
        "\n",
        "# Initialise the optimiser\n",
        "REINFORCE_optim_state = REINFORCE_OPTIMIZER.init(REINFORCE_params)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ViENrHOALbCw"
      },
      "source": [
        "Now we have everything we need tp make the `REINFORCE_learn` function. We will store the state of the optimiser in the `learn_state`. We will compute the gradient of the policy gradient loss by using `jax.grad` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.grad.html))."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CQr2Uz5ORjM5"
      },
      "outputs": [],
      "source": [
        "# A NamedTuple to store the state of the optimiser\n",
        "REINFORCELearnState = collections.namedtuple(\"LearnerState\", [\"optim_state\"])\n",
        "\n",
        "\n",
        "def REINFORCE_learn(key, params, learner_state, memory):\n",
        "    \n",
        "  # Get the policy gradient by using `jax.grad()` on `batched_policy_gradient_loss`\n",
        "  grad_loss = jax.grad(batched_policy_gradient_loss)(params, memory.obs, memory.action, memory.reward_to_go)\n",
        "\n",
        "  # Get param updates using gradient and optimizer\n",
        "  updates, new_optim_state = REINFORCE_OPTIMIZER.update(grad_loss, learner_state.optim_state)\n",
        "\n",
        "  # Apply updates to params\n",
        "  params = optax.apply_updates(params, updates)\n",
        "\n",
        "  return params, REINFORCELearnState(new_optim_state) # update learner state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5an3U2NhRKgG"
      },
      "source": [
        "### REINFORCE training loop\n",
        "Now we can train our REINFORCE agent by putting everything together using the training loop. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vioIcVGsRjM5"
      },
      "outputs": [],
      "source": [
        "# JIT the choose_action and learn functions for more speed\n",
        "REINFORCE_learn_jit = jax.jit(REINFORCE_learn)\n",
        "REINFORCE_choose_action_jit = jax.jit(REINFORCE_choose_action)\n",
        "\n",
        "# Initial learn state\n",
        "REINFORCE_learn_state = REINFORCELearnState(REINFORCE_optim_state)\n",
        "\n",
        "# Run training loop\n",
        "print(\"Starting training. This may take up to 10 minutes to complete.\")\n",
        "episode_returns, evaluator_returns = run_training_loop(\n",
        "                                        env_name,\n",
        "                                        REINFORCE_params,\n",
        "                                        REINFORCE_choose_action_jit, \n",
        "                                        None, # action state not used\n",
        "                                        REINFORCE_learn_jit, \n",
        "                                        REINFORCE_learn_state, \n",
        "                                        REINFORCE_memory,\n",
        "                                        num_episodes=10_001,\n",
        "                                        learn_steps_per_episode=2,\n",
        "                                        video_subdir=\"reinforce\"\n",
        "                                      )\n",
        "\n",
        "# Plot the episode returns\n",
        "plt.plot(episode_returns)\n",
        "plt.xlabel(\"Episode\")\n",
        "plt.ylabel(\"Episode Return\")\n",
        "plt.title(\"REINFORCE\")\n",
        "plt.show()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "caKL3ngNr_Yh"
      },
      "outputs": [],
      "source": [
        "#@title Visualise Policy {display-mode: \"form\"}\n",
        "#@markdown Choose an episode number that is a multiple of 100 and less than or equal to 1000, and **run this cell**.\n",
        "\n",
        "episode_number = 100 #@param {type:\"number\"}\n",
        "\n",
        "assert (episode_number % 100) == 0, \"Episode number must be a multiple of 100 since we only record every 100th episode.\"\n",
        "assert episode_number < 1001, \"Episode number must be less than or equal to 1000\"\n",
        "\n",
        "eval_episode_number = int(episode_number / 100 * 8)\n",
        "video_path = f\"./video/reinforce/eval/rl-video-episode-{eval_episode_number}.mp4\"\n",
        "\n",
        "mp4 = open(video_path,'rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "HTML(\"\"\"\n",
        "<video width=400 controls>\n",
        "      <source src=\"%s\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\" % data_url)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D_HHzFTOc-Qr"
      },
      "source": [
        "## Section 4: Q-Learning\n",
        "Another common aproach to finding an optimal policy in an environment using RL is via Q-learning. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lF1nOZrUSzhE"
      },
      "source": [
        "### State-Action Value function\n",
        "In Q-learning the agent learns a function that approximates the **value** of state-action pairs. By *value* we mean the return you expect to receive if you start in a particular state $s_t$, take a particular action $a_t$, and then act according to a particular policy $\\pi$ forever after. The state-action value function of policy $\\pi$ is given by\n",
        "\n",
        "$Q_\\pi(s,a)=\\mathrm{E}_{\\tau\\sim\\pi}\\left[R(\\tau) \\mid s_0=s,\\ a_t=a\\right]$.\n",
        "\n",
        "We say that the value function $Q_\\pi(s,a)$ is the **optimal** value function if the policy $\\pi$ is an optimal policy. We denote the optimal value function as follows:\n",
        "\n",
        "$Q_\\ast(s,a)=\\max \\limits_\\pi \\  \\mathrm{E}_{\\tau\\sim\\pi}\\left[R(\\tau) \\mid s_0=s,\\ a_0=a\\right]$\n",
        "\n",
        "There is an important relationship between the optimal action $a_\\ast$ in a state $s$ and the optimal state-action value function $Q_\\ast$. Namely, the optimal action $a_\\ast$ in state $s$ is equal to the action that maximises the optimal state-action value function. This relationship naturally induces an optimal policy:\n",
        "\n",
        "$\\pi_\\ast(s)=\\arg \\max \\limits_a\\ Q_\\ast(s, a)$"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "H2x2tqvZSihz"
      },
      "source": [
        "### Greedy action selection\n",
        "\n",
        "**Exercise 10:** Let's implement a function that, given a vector of Q-values, returns the action with the largest Q-value (i.e. the greedy action).\n",
        "\n",
        "**Useful methods:**\n",
        "*   `jax.numpy.argmax` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.argmax.html))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nn9P1YdzTIDU"
      },
      "outputs": [],
      "source": [
        "# Implement a function takes q-values as input and returns the greedy_action\n",
        "def select_greedy_action(q_values):\n",
        "\n",
        "  # YOUR CODE\n",
        "  action = ...\n",
        "  # END YOUR CODE\n",
        "\n",
        "  return action"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZBzLr_G7QKXR"
      },
      "outputs": [],
      "source": [
        "# @title Check exercise 10 (run me) {display-mode: \"form\"}\n",
        "\n",
        "try:\n",
        "  q_values = jnp.array([1,1,3,4])\n",
        "  action = select_greedy_action(q_values)\n",
        "\n",
        "  if action != 3:\n",
        "    print(\"Incorrect answer, your greedy action selector looks wrong\")\n",
        "  else:\n",
        "    print(\"Looks good.\")\n",
        "except Exception as e:\n",
        "  print(\"An Error Occured: {}\".format(e))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9GsPv35lRjM6"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 10 {display-mode: \"form\"}\n",
        "\n",
        "def select_greedy_action(q_values):\n",
        "  \n",
        "  # YOUR CODE\n",
        "  action = jnp.argmax(q_values)\n",
        "  # END YOUR CODE\n",
        "\n",
        "  return action"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PFQUlqZ4ZyLp"
      },
      "source": [
        "### Q-Network\n",
        "Unlike in the policy gradient approach from the previous section, in Q-learning and other value-based RL methods we don't need a parameterisation for the policy, rather we parametrise the Q-function using a neural network $Q_\\theta$. We obtain a policy from the Q-network by always choosing the action with the *greatest* value:\n",
        "\n",
        "$\\hat{\\pi}_\\theta(s)=\\arg \\max \\limits_a\\ Q_{\\theta}(s, a)$\n",
        "\n",
        "As we did previously, we shall use haiku to make a neural network to approximate this Q-function. The network will take an observation as input and then output a Q-value for each of the available actions. So in the case of CartPole, the output of the network will have size $2$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9wU1soJYZyLp"
      },
      "outputs": [],
      "source": [
        "def build_network(num_actions: int, layers=[20, 20]) -> hk.Transformed:\n",
        "  \"\"\"Factory for a simple MLP network for approximating Q-values.\"\"\"\n",
        "\n",
        "  def q_network(obs):\n",
        "    network = hk.Sequential(\n",
        "        [hk.Flatten(),\n",
        "         hk.nets.MLP(layers + [num_actions])])\n",
        "    return network(obs)\n",
        "\n",
        "  return hk.without_apply_rng(hk.transform(q_network))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sUwG-4qTS1zx"
      },
      "source": [
        "Let's initialise our Q-network and get the initial parameters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1uvq5n2cS9lu"
      },
      "outputs": [],
      "source": [
        "# Initialise Q-network\n",
        "Q_NETWORK = build_network(num_actions=num_actions, layers=[20, 20]) # two actions\n",
        "\n",
        "dummy_obs = jnp.zeros((1,*obs_shape), jnp.float32) # a dummy observation like the one in CartPole\n",
        "\n",
        "random_key = jax.random.PRNGKey(42) # random key\n",
        "Q_NETWORK_PARAMS = Q_NETWORK.init(random_key, dummy_obs) # Get initial params\n",
        "\n",
        "print(\"Q-Learning params:\", Q_NETWORK_PARAMS.keys())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iqCUeZfhEfyP"
      },
      "source": [
        "Before we implement the loss function required for training our Q-network. Let's first discuss the intuition behind it. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vbLig3uSZyLp"
      },
      "source": [
        "### The Bellman Equations\n",
        "The value function can be written recursively as:\n",
        "\n",
        "$Q_{\\pi}(s, a) =\\underset{s^{\\prime} \\sim P}{\\mathrm{E}}\\left[r(s, a)+ \\underset{a^{\\prime} \\sim \\pi}{\\mathrm{E}}\\left[Q_{\\pi}\\left(s^{\\prime}, a^{\\prime}\\right)\\right]\\right]$,\n",
        "\n",
        "where $s' \\sim P$ is shorthand for saying that the next state $s'$ is sampled from the environment’s transition function $P(s'\\mid s,a)$. Intuitively, this equation says that the value of the action $a$ you took in the state $s$ is equal to the reward $r$ you expect to get, plus the value you expect to get in the next state $s`$ you land in given that you will choose your next action $a`$ with the policy $\\pi$. The Bellman equation for the optimal value function is:\n",
        "\n",
        "$Q_{*}(s, a) =\\underset{s^{\\prime} \\sim P}{\\mathrm{E}}\\left[r(s, a)+\\ \\underset{a^{\\prime}}{\\max}\\ Q_{*}(s^{\\prime}, a^{\\prime})\\right]$\n",
        "\n",
        "Notice that instead of chosing your next action $a`$ with policy $\\pi$, we choose the action with the greatest Q-value.\n",
        "\n",
        "\n",
        "For a more in-depth discussion of the Bellman Equations, see the [OpenAI Spinning Up](https://spinningup.openai.com/en/latest/spinningup/rl_intro.html) website."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DsOJi5G8ZyLp"
      },
      "source": [
        "### The Bellman Backup\n",
        "To learn to approximate the optimal Q-value function, we can use the right-hand side of the Bellman equation as an update rule. In other words, suppose we have a Q-network $Q_\\theta$, approximated using parameters $\\theta$, then we can iteratively update the parameters such that\n",
        "\n",
        "$Q_\\theta(s,a)\\leftarrow r(s, a) + \\underset{a'}{\\max}\\ Q_\\theta(s', a')$.\n",
        "\n",
        "Intuitively, this says that the approximation of the Q-value of action $a$ in state $s$ should be updated such that it is closer to being equal to the reward received from the environment $r(s, a)$ plus the value of the best possible action in the next state $s'$. We can perform this optimisation by minimising the difference between the left and right-hand side, with respect to the parameters $\\theta$ using gradient descent. We can measure the difference between the two values using the [squared-error](https://en.wikipedia.org/wiki/Mean_squared_error#Loss_function).\n",
        "\n",
        "**Exercise 11:** Implement the squared-error function.\n",
        "\n",
        "**Useful functions**\n",
        "* `jax.numpy.square` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.square.html))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PTto__ohZyLp"
      },
      "outputs": [],
      "source": [
        "def compute_squared_error(pred, target):\n",
        "  # YOUR CODE\n",
        "  squared_error = ...\n",
        "  # END YOUR CODE\n",
        "\n",
        "  return squared_error"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wGhX8XTFVPVU"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 11 {display-mode: \"form\"}\n",
        "\n",
        "try:\n",
        "  result = compute_squared_error(1, 4)\n",
        "\n",
        "  if result != 9:\n",
        "    print(\"Your implementation looks wrong.\")\n",
        "  else:\n",
        "    print(\"Looks good.\")\n",
        "except Exception as e:\n",
        "  print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ms7pHyHGZyLp"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 11 {display-mode: \"form\"}\n",
        "def compute_squared_error(pred, target):\n",
        "\n",
        "  # YOUR CODE\n",
        "  squared_error = jax.numpy.square(pred - target)\n",
        "  # END YOUR CODE\n",
        "  return squared_error"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ycpZVkgdZyLp"
      },
      "source": [
        "**Exercise 12:** Implement a function that computes the **Bellman target** (right-hand side of the Bellman equation). If the episode is at the last timestep (i.e. done==1.0), then the Bellman target should be equal to the reward, with no extra value at the end.\n",
        "\n",
        "**Useful functions**\n",
        "* `jax.numpy.max` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.max.html))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "961_OllWZyLp"
      },
      "outputs": [],
      "source": [
        "# Bellman target\n",
        "def compute_bellman_target(reward, done, next_q_values):\n",
        "  \"\"\"A function to compute the bellman target.\n",
        "  \n",
        "  Args:\n",
        "      reward: a scalar reward.\n",
        "      done: a scalar of value either 1.0 or 0.0, indicating if the transition is a terminal one.\n",
        "      next_q_values: a vector of q_values for the next state. One for each action.\n",
        "  Returns:\n",
        "      A scalar equal to the bellman target.\n",
        "  \n",
        "  \"\"\"\n",
        "  # YOUR CODE\n",
        "  bellman_target = ...\n",
        "  # END YOUR CODE\n",
        "\n",
        "  return bellman_target"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "E5RAhegOWAkC"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 12 {display-mode: \"form\"}\n",
        "\n",
        "try:\n",
        "  # not done\n",
        "  result1 = compute_bellman_target(1, 0.0, np.array([3,2], \"float32\"))\n",
        "\n",
        "  # done\n",
        "  result2 = compute_bellman_target(1, 1.0, np.array([3,2], \"float32\"))\n",
        "\n",
        "  if result1 != 4 or result2 != 1:\n",
        "    print(\"Your implementation looks wrong.\")\n",
        "  else:\n",
        "    print(\"Looks good.\")\n",
        "except Exception as e:\n",
        "  print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wbAGg-_yZyLq"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 12 {display-mode: \"form\"}\n",
        "\n",
        "# Bellman target\n",
        "def compute_bellman_target(reward, done, next_q_values):\n",
        "    \"\"\"A function to compute the bellman target.\n",
        "    \n",
        "    Args:\n",
        "        reward: a scalar reward.\n",
        "        done: a scalar of value either 1.0 or 0.0, indicating if the transition is a terminal one.\n",
        "        next_q_values: a vector of q_values for the next state. One for each action.\n",
        "    Returns:\n",
        "        A scalar equal to the bellman target.\n",
        "    \n",
        "    \"\"\"\n",
        "    # YOUR CODE\n",
        "    bellman_target = reward + (1.0 - done) * jax.numpy.max(next_q_values)\n",
        "    # END YOUR CODE\n",
        "\n",
        "    return bellman_target\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8sIjrHSJZyLq"
      },
      "source": [
        "We can now combine these two functions to compute the loss for Q-learning. The Q-learning loss is equal to the squared difference between the predicted Q-value of an action and its corresponding Bellman target.\n",
        "\n",
        "**Exercise 13:** Implement the Q-learning loss."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LJY_kpFcZyLq"
      },
      "outputs": [],
      "source": [
        "def q_learning_loss(q_values, action, reward, done, next_q_values):\n",
        "    \"\"\"Implementation of the Q-learning loss.T\n",
        "    \n",
        "    Args:\n",
        "        q_values: a vector of Q-values, one for each action.\n",
        "        action: an integer, giving the action that was chosen. q_values[action] is the value of the chose action.\n",
        "        done: is a scalar that indicates if this is a terminal transition.\n",
        "        next_q_values: a vector of Q-values in the next state.\n",
        "    Returns:\n",
        "        The squared difference between the q_value of the chosen action and the bellman target.\n",
        "    \"\"\"\n",
        "    # YOUR CODE\n",
        "    chosen_action_q_value = ... # q_value of action, use array indexing\n",
        "    bellman_target = ...\n",
        "    squared_error = ...\n",
        "    # END YOUR CODE\n",
        "    \n",
        "    return squared_error"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1aZRn1qaWx2M"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 13 {display-mode: \"form\"}\n",
        "\n",
        "try:\n",
        "  result = q_learning_loss(np.array([3,2], \"float32\"), 1, 2, 0.0, np.array([3,2], \"float32\"))\n",
        "\n",
        "  if result != 9.0:\n",
        "    print(\"Your implementation looks wrong.\")\n",
        "  else:\n",
        "    print(\"Looks good.\")\n",
        "except Exception as e:\n",
        "  print(\"An Error Occured: {}\".format(e))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3_KQleTsZyLq"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 13 {display-mode: \"form\"}\n",
        "\n",
        "def q_learning_loss(q_values, action, reward, done, next_q_values):\n",
        "    \"\"\"Implementation of the Q-learning loss.T\n",
        "    \n",
        "    Args:\n",
        "        q_values: a vector of Q-values, one for each action.\n",
        "        action: an integer, giving the action that was chosen. q_values[action] is the value of the chose action.\n",
        "        done: is a scalar that indicates if this is a terminal transition.\n",
        "        next_q_values: a vector of Q-values in the next state.\n",
        "    Returns:\n",
        "        The squared difference between the q_value of the chosen action and the bellman target.\n",
        "    \"\"\"\n",
        "    # YOUR CODE\n",
        "    chosen_action_q_value = q_values[action]\n",
        "    bellman_target = compute_bellman_target(reward, done, next_q_values)\n",
        "    squared_error = compute_squared_error(chosen_action_q_value, bellman_target)\n",
        "    # END YOUR CODE\n",
        "    \n",
        "    return squared_error"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b4YnSUfJZyLq"
      },
      "source": [
        "### Target Q-network\n",
        "Notice that when we compute the bellman target we are using our Q-network $Q_\\theta$ to compute the value for the next state $s_t$. We are basically using our latest approximation of the Q-function to compute the target of our next approximation. Using an approximation to compute the target for your next approximation, is called bootstrapping. Unfortunately, if we naively bootstrap like this, it can make training a neural network very unstable. To mitigage this we can instead use a different set of parameters $\\hat{\\theta}$ to compute the values at state $s_{t+1}$. We will keep the parameters $\\hat{\\theta}$ fixed and only periodically update them to be equal to the latest online parameters $\\theta$ every couple of training steps *(say 100)*. This serves to keep the bellman targets fixed for a couple training steps to help reduce the instability due to bootstrapping. \n",
        "\n",
        "\n",
        "We will need to keep track of the latest (online) parameters, as well as the target networks parameters. Lets make a `NamedTuple` to store these two values. We will also need to keep track of the number of learner steps we have taken, so that we know when to update the target network. Lets store a `count` of the learn steps in the `learn_state`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DvZqUKmq6L7k"
      },
      "outputs": [],
      "source": [
        "# Store online and target parameters\n",
        "QLearnParams = collections.namedtuple(\"Params\", [\"online\", \"target\"])\n",
        "\n",
        "# Q-learn-state\n",
        "QLearnState = collections.namedtuple(\"LearnerState\", [\"count\", \"optim_state\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CJWH2_kNZsau"
      },
      "source": [
        "We will once again be using Optax to optimize our neural network in JAX. We store the state of the optimizer in the `learn_state` above. Lets now instantiate optimizer and add the initial Q-network parameters to a `QLearnParams` object."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HwqKTN6BaAXE"
      },
      "outputs": [],
      "source": [
        "# Initialise Q-network optimizer\n",
        "Q_LEARN_OPTIMIZER = optax.adam(3e-4) # learning rate\n",
        "\n",
        "Q_LEARN_OPTIM_STATE = Q_LEARN_OPTIMIZER.init(Q_NETWORK_PARAMS) # initial optim state\n",
        "\n",
        "# Create Learn State\n",
        "Q_LEARNING_LEARN_STATE = QLearnState(0, Q_LEARN_OPTIM_STATE) # count set to zero initially\n",
        "\n",
        "# Add initial Q-network weights to QLearnParams object\n",
        "Q_LEARNING_PARAMS = QLearnParams(online=Q_NETWORK_PARAMS, target=Q_NETWORK_PARAMS) # target equal to online"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xj89M8LgZlFe"
      },
      "source": [
        "Now we can implement a simple function that updates target networks parameters to equal the latest online parameters every 100 training steps."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SKrg_3rO6cCL"
      },
      "outputs": [],
      "source": [
        "def update_target_params(learn_state, online_weights, target_weights):\n",
        "  \"\"\"A function to update target params every 100 training steps\"\"\"\n",
        "\n",
        "  target = jax.lax.cond(\n",
        "      jax.numpy.mod(learn_state.count, 100) == 0,\n",
        "      lambda x, y: x,\n",
        "      lambda x, y: y,\n",
        "      online_weights, \n",
        "      target_weights\n",
        "  )\n",
        "\n",
        "  params = QLearnParams(online_weights, target)\n",
        "\n",
        "  return params"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KoiaYSo9ZyLq"
      },
      "source": [
        "### Q-learning loss\n",
        "We now have everything we need to implement the `q_learn` function which takes some batch of transitions and does a step of Q-learning to update the network paramters. But first we use `jax.vmap` to modify the `q_learning_loss` function so that it accepts batches of transitions. In addition, we will compute the Q-values by passing the observations through the `Q_NETWORK` and the target Q-values using the target parameters of the `Q_NETWORK`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hnrsGppWZyLq"
      },
      "outputs": [],
      "source": [
        "def batched_q_learning_loss(online_params, target_params, obs, actions, rewards, next_obs, dones):\n",
        "    q_values = Q_NETWORK.apply(online_params, obs) # use the online parameters\n",
        "    next_q_values = Q_NETWORK.apply(target_params, next_obs) # use the target parameters\n",
        "    squared_error = jax.vmap(q_learning_loss)(q_values, actions, rewards, dones, next_q_values) # vmap q_learning_loss\n",
        "    mean_squared_error = jnp.mean(squared_error) # mean squared error over batch\n",
        "    return mean_squared_error"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CU8vo9ZebnEa"
      },
      "source": [
        "Now we can create the `q_learn` function which computes the gradient of the `batched_q_learning_loss` and then uses an Optax optimizer to update the network weights and then finally (maybe) updates the target parameters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6BYoX2W_ZyLr"
      },
      "outputs": [],
      "source": [
        "def q_learn(rng, params, learner_state, memory):\n",
        "  # Compute gradients\n",
        "  grad_loss = jax.grad(batched_q_learning_loss)(params.online, params.target, memory.obs, \n",
        "                                          memory.action, memory.reward, \n",
        "                                          memory.next_obs, memory.done,\n",
        "                                          ) # jax.grad\n",
        "\n",
        "  # Get updates\n",
        "  updates, opt_state = Q_LEARN_OPTIMIZER.update(grad_loss, learner_state.optim_state)\n",
        "\n",
        "  # Apply them\n",
        "  new_weights = optax.apply_updates(params.online, updates)\n",
        "\n",
        "  # Maybe update target network\n",
        "  params = update_target_params(learner_state, new_weights, params.target)\n",
        "\n",
        "  # Increment learner step counter\n",
        "  learner_state = QLearnState(learner_state.count + 1, opt_state)\n",
        "\n",
        "  return params, learner_state"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VpZsKHssZyLq"
      },
      "source": [
        "### Replay Buffer\n",
        "For Q-learning we will need an agent memory that stores entire transitions: `obs`, `action`, `reward`, `next_obs`, `done`. When we retrieve transitions from the memory, they should be chosen randomly from all of the transitions collected so far. In RL we often call such a module a **replay buffer**. One benefit of using a replay buffer like this is that experiences can be re-used several times for training unlike in the policy gradient algorithm REINFORCE, where we discarded memories after using them for learning."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8tv5dUH6ZyLr"
      },
      "outputs": [],
      "source": [
        "class TransitionMemory(object):\n",
        "  \"\"\"A simple Python replay buffer.\"\"\"\n",
        "\n",
        "  def __init__(self, max_size=10_000, batch_size=256):\n",
        "    self.batch_size = batch_size\n",
        "    self.buffer = collections.deque(maxlen=max_size)\n",
        "\n",
        "  def push(self, transition):\n",
        "\n",
        "    # add transition to the replay buffer\n",
        "    self.buffer.append(\n",
        "        (transition.obs, transition.action, transition.reward, \n",
        "          transition.next_obs, transition.done)\n",
        "    )\n",
        "\n",
        "  \n",
        "  def is_ready(self):\n",
        "    return self.batch_size <= len(self.buffer)\n",
        "\n",
        "  def sample(self):\n",
        "    # Randomly sample a batch of transitions from the buffer\n",
        "    random_replay_sample = random.sample(self.buffer, self.batch_size)\n",
        "\n",
        "    # Batch the transitions together\n",
        "    obs_batch, action_batch, reward_batch, next_obs_batch, done_batch = zip(*random_replay_sample)\n",
        "\n",
        "    return Transition(\n",
        "        np.stack(obs_batch).astype(\"float32\"), \n",
        "        np.asarray(action_batch).astype(\"int32\"), \n",
        "        np.asarray(reward_batch).astype(\"float32\"), \n",
        "        np.stack(next_obs_batch).astype(\"float32\"), \n",
        "        np.asarray(done_batch).astype(\"float32\")\n",
        "    )\n",
        "\n",
        "# Instantiate the memory\n",
        "Q_LEARNING_MEMORY = TransitionMemory(max_size=50_000, batch_size=256)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XbHk03VVUHAV"
      },
      "source": [
        "### Random exploration\n",
        "We almost have everything we need for a functioning Q-learning agent. But one problem is that if we always choose the action with the highest Q-value then the agent's policy will be completly [deterministic](https://www.quora.com/What-is-the-intuitive-difference-between-a-stochastic-model-and-a-deterministic-model). This means the agent will always choose the same strategy. This can pose a problem because at the start of training, the Q-network will be very inaccurate (i.e. a bad aproximation of the true Q-function). As such, the agent will consistently choose suboptimal actions. Moreover, the agent will never deviate from its suboptimal strategy and will never discover new, potentially more rewarding  actions. As a result, the Q-network remains inaccurate. Ideally, the agent should try out many different strategies so that it can observe the outcomes (rewards) of its actions in different states and so improve its approximation of the Q-function.\n",
        "\n",
        "One easy way to ensure that the agent tries out many different actions is to let it periodically choose some random actions, instead of the greedy (best) action all the time.\n",
        "\n",
        "**Exercise 14:** Implement a function that, given the number of possible (discrete) actions, returns a random action.\n",
        "\n",
        "**Useful methods:**\n",
        "\n",
        "*  `jax.random.randint` ([docs](https://jax.readthedocs.io/en/latest/_autosummary/jax.random.randint.html))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uUKkpMLXUtko"
      },
      "outputs": [],
      "source": [
        "def select_random_action(key, num_actions):\n",
        "    \n",
        "    # YOUR CODE\n",
        "    action = ...\n",
        "    # END YOUR CODE\n",
        "\n",
        "    return action"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NO6va6S2Y40E"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 14 {display-mode: \"form\"}\n",
        "\n",
        "try:\n",
        "  random_key1 = random_key = jax.random.PRNGKey(6) # random key\n",
        "  random_key2 = random_key = jax.random.PRNGKey(1000) # random key\n",
        "  result1 = select_random_action(random_key1, 2)\n",
        "  result2 = select_random_action(random_key2, 2)\n",
        "\n",
        "  if result1 != 1 or result2 != 0:\n",
        "    print(\"Your implementation looks wrong.\")\n",
        "  else:\n",
        "    print(\"Looks good.\")\n",
        "except:\n",
        "  print(\"Your implementation looks wrong.\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lrVadhcwRjM6"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 14 {display-mode: \"form\"}\n",
        "\n",
        "def select_random_action(key, num_actions):\n",
        "    # YOUR CODE\n",
        "    action = jax.random.randint(\n",
        "        key, \n",
        "        shape=(), \n",
        "        minval=0, \n",
        "        maxval=num_actions\n",
        "    )\n",
        "    # END YOUR CODE\n",
        "\n",
        "    return action"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8-kKDFT6XU6y"
      },
      "source": [
        "### $\\varepsilon$-greedy action selection\n",
        "At the start of training, when the accuracy of the Q-network is low, it is worthwhile for the agent to mostly take random actions so that it can learn about how goo/bad actions are. However, as the accuracy of the Q-network improves, the agent should start taking fewer random actions and instead start choosing the greedy actions with respect to the Q-values. Choosing the best actions given the current Q-network is referred to as **exploitation.** In RL we often call the ratio of random to greedy actions **epsilon** $\\varepsilon$. Epsilon is usually a decimal value in the interval $[0,1]$, where for example $\\varepsilon=0.4$ means that the agent chooses a random action 40% of the time and the greedy action 60% of the time. It is common in RL to linearly decrease the value of epsilon over time so that the agent becomes increasingly greedy as the accuracy of its Q-network improves through learning.\n",
        "\n",
        "\n",
        "**Exercise 15:** Implement a function that takes the number of timesteps as input and returns the current epsilon value."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_qejnbCocurG"
      },
      "outputs": [],
      "source": [
        "EPSILON_DECAY_TIMESTEPS = 3000 # decay epsilon over 3000 timesteps\n",
        "EPSILON_MIN = 0.1 # 10% exploration"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5ujSbssCZyLs"
      },
      "outputs": [],
      "source": [
        "def get_epsilon(num_timesteps):\n",
        "  # YOUR CODE\n",
        "  epsilon = ... # decay epsilon\n",
        "\n",
        "  epsilon = jax.lax.select(\n",
        "      epsilon < EPSILON_MIN,\n",
        "      ..., # if less than min then set to min\n",
        "      ... # else don't change epsilon\n",
        "  )\n",
        "  # END YOUR CODE\n",
        "\n",
        "  return epsilon"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hvu4zA64aUou"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 15 {display-mode: \"form\"}\n",
        "def check_get_epsilon(get_epsilon):\n",
        "  try:\n",
        "    result1 = get_epsilon(10)\n",
        "    result2 = get_epsilon(5_010)\n",
        "\n",
        "    if result1 != 0.99666667 or result2 != 0.1:\n",
        "      print(\"Your function looks wrong.\")\n",
        "    else:\n",
        "      print(\"Your function looks correct.\")\n",
        "  except:\n",
        "    print(\"Your function looks wrong.\")\n",
        "\n",
        "check_get_epsilon(get_epsilon)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QTpezAeUYFR7"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 15 {display-mode: \"form\"}\n",
        "\n",
        "def get_epsilon(num_timesteps):\n",
        "\n",
        "  # YOUR CODE\n",
        "  epsilon = 1.0 - num_timesteps / EPSILON_DECAY_TIMESTEPS\n",
        "\n",
        "  epsilon = jax.lax.select(\n",
        "      epsilon < EPSILON_MIN,\n",
        "      EPSILON_MIN,\n",
        "      epsilon\n",
        "  )\n",
        "  # END YOUR CODE\n",
        "\n",
        "  return epsilon\n",
        "\n",
        "# CHECK\n",
        "check_get_epsilon(get_epsilon)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t56oo58TVQ_s"
      },
      "source": [
        "**Exercise 16:** Now lets put these functions together to do epsilon-greedy action selection."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NlQx8K4vKUXj"
      },
      "outputs": [],
      "source": [
        "def select_epsilon_greedy_action(key, q_values, num_timesteps):  \n",
        "    num_actions = len(q_values) # number of available actions\n",
        "\n",
        "    # YOUR CODE HERE\n",
        "    epsilon = ... # get epsilon value\n",
        "\n",
        "    should_explore = ... # hint: a boolean expression to check if some random number is less than epsilon\n",
        "\n",
        "    action = jax.lax.select(\n",
        "        should_explore,\n",
        "        ..., # if should explore\n",
        "        ... # if should be greedy\n",
        "    )\n",
        "    # END YOUR CODE\n",
        "\n",
        "    return action"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mkVU2a5e7P6b"
      },
      "outputs": [],
      "source": [
        "#@title Check exercise 16 {display-mode: \"form\"}\n",
        "\n",
        "try:\n",
        "  rng = hk.PRNGSequence(jax.random.PRNGKey(42))\n",
        "  dummy_q_values = jnp.array([0,1], jnp.float32)\n",
        "  num_timesteps = 5010 # very greedy\n",
        "  actions1 = []\n",
        "  for i in range(10):\n",
        "      actions1.append(int(select_epsilon_greedy_action(next(rng), dummy_q_values, num_timesteps)))\n",
        "\n",
        "  num_timesteps = 0 # completly random\n",
        "  actions2 = []\n",
        "  for i in range(10):\n",
        "      actions2.append(int(select_epsilon_greedy_action(next(rng), dummy_q_values, num_timesteps)))\n",
        "\n",
        "  if actions1 != [1, 1, 0, 1, 1, 0, 1, 1, 1, 1] or actions2 != [0, 0, 0, 1, 1, 1, 1, 0, 0, 0]:\n",
        "    print(\"Looks like something might be incorrect!\")\n",
        "  else:\n",
        "    print(\"Looks correct!\")\n",
        "except:\n",
        "  print(\"Looks like something might be incorrect!\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bXDRQhORRjM6"
      },
      "outputs": [],
      "source": [
        "#@title Solution exercise 16 {display-mode: \"form\"}\n",
        "\n",
        "# Now make a function that takes an epsilon-greedy action\n",
        "\n",
        "def select_epsilon_greedy_action(key, q_values, num_timesteps):\n",
        "    num_actions = len(q_values) # number of available actions\n",
        "\n",
        "    # YOUR CODE\n",
        "    epsilon = get_epsilon(num_timesteps)\n",
        "\n",
        "    should_explore = jax.random.uniform(key, (1,))[0] < epsilon\n",
        "\n",
        "    num_actions = len(q_values)\n",
        "\n",
        "    action = jax.lax.select(\n",
        "        should_explore,\n",
        "        select_random_action(key, num_actions), \n",
        "        select_greedy_action(q_values)\n",
        "    )\n",
        "    # END YOUR CODE\n",
        "\n",
        "    return action"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y5W23MnobN9x"
      },
      "source": [
        "### Q-learning select action\n",
        "\n",
        "We now have everything we need to make the `q_learning_select_action` function. We will use the `actor_state` to store a counter which keeps track of the current number of timesteps. We can use the counter to decrement our `epsilon` value."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "81TysLc0RjM6"
      },
      "outputs": [],
      "source": [
        "# Actor state stores the current number of timesteps\n",
        "QActorState = collections.namedtuple(\"ActorState\", [\"count\"])\n",
        "\n",
        "def q_learning_select_action(key, params, actor_state, obs, evaluation=False):\n",
        "    obs = jnp.expand_dims(obs, axis=0) # add dummy batch dim\n",
        "    q_values = Q_NETWORK.apply(params.online, obs)[0] # remove batch dim\n",
        "\n",
        "    action = select_epsilon_greedy_action(key, q_values, actor_state.count)\n",
        "    greedy_action = select_greedy_action(q_values)\n",
        "\n",
        "    action = jax.lax.select(\n",
        "        evaluation,\n",
        "        greedy_action,\n",
        "        action\n",
        "    )\n",
        "\n",
        "    next_actor_state = QActorState(actor_state.count + 1) # increment timestep counter\n",
        "\n",
        "    return action, next_actor_state\n",
        "\n",
        "Q_LEARNING_ACTOR_STATE = QActorState(0) # counter set to zero"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z884-1oNRGEr"
      },
      "source": [
        "### Training\n",
        "We can now put everything together using the agent-environment loop. But first,lets jit the select action function and the learn function for some extra speed."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xbdHDbd1RjM8"
      },
      "outputs": [],
      "source": [
        "# Jit functions\n",
        "q_learning_select_action_jit = jax.jit(q_learning_select_action)\n",
        "q_learn_jit = jax.jit(q_learn)\n",
        "\n",
        "# Run environment loop\n",
        "print(\"Starting training. This may take up to 8 minutes to complete.\")\n",
        "episode_returns, evaluator_returns = run_training_loop(\n",
        "                                        env_name,\n",
        "                                        Q_LEARNING_PARAMS, \n",
        "                                        q_learning_select_action_jit, \n",
        "                                        Q_LEARNING_ACTOR_STATE,\n",
        "                                        q_learn_jit,\n",
        "                                        Q_LEARNING_LEARN_STATE,\n",
        "                                        Q_LEARNING_MEMORY,\n",
        "                                        num_episodes=1001,\n",
        "                                        train_every_timestep=True, # do learning after every timestep\n",
        "                                        video_subdir=\"q_learning\"\n",
        "                                    )\n",
        "\n",
        "plt.plot(episode_returns)\n",
        "plt.xlabel(\"Episodes\")\n",
        "plt.ylabel(\"Episode Return\")\n",
        "plt.title(\"Deep Q-Learning\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2k0-41wbpFDE"
      },
      "source": [
        "At this stage, the approximated Q-function hopefully converged to a decent policy for balancing the pole in the CartPole problem."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "t4v0ZbB4FQUs"
      },
      "outputs": [],
      "source": [
        "#@title Visualise Policy\n",
        "#@markdown Choose an episode number that is a multiple of 100 and less than or equal to 1000, and **run this cell**.\n",
        "\n",
        "episode_number = 0 #@param {type:\"number\"}\n",
        "\n",
        "assert (episode_number % 100) == 0, \"Episode number must be a multiple of 100 since we only record every 100th episode.\"\n",
        "assert episode_number < 1001, \"Episode number must be less than or equal to 1000\"\n",
        "\n",
        "eval_episode_number = int(episode_number / 100 * 32)\n",
        "video_path = f\"./video/q_learning/eval/rl-video-episode-{eval_episode_number}.mp4\"\n",
        "\n",
        "mp4 = open(video_path,'rb').read()\n",
        "data_url = \"data:video/mp4;base64,\" + b64encode(mp4).decode()\n",
        "HTML(\"\"\"\n",
        "<video width=400 controls>\n",
        "      <source src=\"%s\" type=\"video/mp4\">\n",
        "</video>\n",
        "\"\"\" % data_url)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CIIteLdFwG4B"
      },
      "source": [
        "This section attempts to summarise [Playing Atari with Deep Reinforcement Learning](https://arxiv.org/abs/1312.5602), the research paper where Deep-Q Learning was first introduced. To understand the concepts covered in this section better, we recommend you give the original paper a read."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fV3YG7QOZD-B"
      },
      "source": [
        "## Conclusion\n",
        "**Summary:**\n",
        "\n",
        "In this practical we learnt the basics of reinforcement learning (RL).\n",
        "\n",
        "In the first section we learnt some basic concepts such as environment observations, action selection strategies, rewards, and episodes. We learnt about rewards and that the goal in RL is to learn a policy which maximises some notion of cummulative reward that the agent receives from the environment (return). \n",
        "\n",
        "In the second section we searched for an optimal policy in CartPole using an algorithm called RandomSearch. Basically, we tried out different policies until we happened to find one that worked well. This method did not yield consistent results and success required immense luck.\n",
        "\n",
        "In the third section we learnt about policy gradients and how we can use gradient ascent to adjust the parameters in our agents policy in the direction which maximises the expected cummulative reward (return).\n",
        "\n",
        "Finally, in the fourth section we learnt about the state-action value function and how it is related to an optimal policy. We implemented an algorithm called Q-learning to learn the optimal state-action value function in CartPole. We learnt about the importance of using a target network and epsilon-greedy exploration.\n",
        "\n",
        "**Next Steps:** \n",
        "\n",
        "Now that you have successfully solved CartPole with two different RL algorithms, REINFORCE and Deep Q-Learning, we now encourage you to use what you have learnt to try and solve some more challenging environments. OpenAI Gym is a great place to find RL environments. [LunarLander](https://www.gymlibrary.ml/environments/box2d/lunar_lander/) is a great next step. You can go to the start of this notebook ande replace the environment with LunarLander by replacing `env = gym.make(\"CartPole-v1\")` with env = `gym.make(\"LunarLander-v2\")`. Note, you will need to increase the number of training episodes in order to learn a good policy in LunarLander because it is a significantly more challenging environment than CartPole.\n",
        "\n",
        "<center>\n",
        "<img src=\"https://miro.medium.com/max/1194/1*Dj2fkRjrMA0w9E-PuyETdg.gif\" width=\"60%\" />\n",
        "</center>\n",
        "\n",
        "In addition, there are many RL algorithms out there that make significant improvements to REINFORCE and Deep Q-Learning. See these resources:\n",
        "* [REINFORCE with baseline](https://spinningup.openai.com/en/latest/spinningup/rl_intro3.html#baselines-in-policy-gradients)\n",
        "* [Double Deep Q-Network](https://arxiv.org/pdf/1509.06461.pdf)\n",
        "* [Proximal Policy Optimisation (PPO)](https://arxiv.org/pdf/1707.06347.pdf)\n",
        "\n",
        "If you are looking for a more indepth online course for RL you can checkout these courses:\n",
        "* [Reinforcement Learning Foundations on LinkedIn Learning](https://www.linkedin.com/learning/reinforcement-learning-foundations) (made by one of our very own tutors, Khaulat Abdulhakeem)\n",
        "* [An introduction to Reinforcement Learning on FreeCodeCamp](https://www.freecodecamp.org/news/an-introduction-to-reinforcement-learning-4339519de419/)\n",
        "* [Reinforcement Learning Specialization on Coursera](https://www.coursera.org/specializations/reinforcement-learning)\n",
        "\n",
        "Finally, the most infuential textbook on RL is available for free online:\n",
        "* [Reinforcement Learning: An Introduction](http://incompleteideas.net/book/the-book-2nd.html) by Richard S. Sutton and Andrew G. Barto\n",
        "\n",
        "**Appendix:** \n",
        "\n",
        "N/a\n",
        "\n",
        "**References:** \n",
        "\n",
        "* [OpenAI Spinning Up](https://spinningup.openai.com/en/latest/)\n",
        "* [Deep Q-Network]()\n",
        "\n",
        "For other practicals from the Deep Learning Indaba, please visit [here](https://github.com/deep-learning-indaba/indaba-pracs-2022)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "o1ndpYE50BpG"
      },
      "source": [
        "## Feedback\n",
        "\n",
        "Please provide feedback that we can use to improve our practicals in the future."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "OIZvkhfRz9Jz"
      },
      "outputs": [],
      "source": [
        "#@title Generate Feedback Form. (Run Cell)\n",
        "from IPython.display import HTML\n",
        "HTML(\n",
        "\"\"\"\n",
        "<iframe \n",
        "\tsrc=\"https://forms.gle/bvLLPX74LMGrFefo9\",\n",
        "  width=\"80%\" \n",
        "\theight=\"1200px\" >\n",
        "\tLoading...\n",
        "</iframe>\n",
        "\"\"\"\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oglV4kHMWnIN"
      },
      "source": [
        "<img src=\"https://baobab.deeplearningindaba.com/static/media/indaba-logo-dark.d5a6196d.png\" width=\"50%\" />"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Copy of intro_to_rl.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3.8.13 ('dli-pracs')",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.8.13"
    },
    "vscode": {
      "interpreter": {
        "hash": "0f35138c66f99be0fd7a1210bb4aa94a6fbaa5f29d51ad43aec0e0ea0ff050f0"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}