C3_W3_A1_Assignment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Q-Learning - Lunar Lander\n",
    "\n",
    "In this assignment, you will train an agent to land a lunar lander safely on a landing pad on the surface of the moon.\n",
    "\n",
    "\n",
    "# Outline\n",
    "- [ 1 - Import Packages <img align=\"Right\" src=\"./images/lunar_lander.gif\" width = 60% >](#1)\n",
    "- [ 2 - Hyperparameters](#2)\n",
    "- [ 3 - The Lunar Lander Environment](#3)\n",
    "  - [ 3.1 Action Space](#3.1)\n",
    "  - [ 3.2 Observation Space](#3.2)\n",
    "  - [ 3.3 Rewards](#3.3)\n",
    "  - [ 3.4 Episode Termination](#3.4)\n",
    "- [ 4 - Load the Environment](#4)\n",
    "- [ 5 - Interacting with the Gym Environment](#5)\n",
    "    - [ 5.1 Exploring the Environment's Dynamics](#5.1)\n",
    "- [ 6 - Deep Q-Learning](#6)\n",
    "  - [ 6.1 Target Network](#6.1)\n",
    "    - [ Exercise 1](#ex01)\n",
    "  - [ 6.2 Experience Replay](#6.2)\n",
    "- [ 7 - Deep Q-Learning Algorithm with Experience Replay](#7)\n",
    "  - [ Exercise 2](#ex02)\n",
    "- [ 8 - Update the Network Weights](#8)\n",
    "- [ 9 - Train the Agent](#9)\n",
    "- [ 10 - See the Trained Agent In Action](#10)\n",
    "- [ 11 - Congratulations!](#11)\n",
    "- [ 12 - References](#12)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_**NOTE:** To prevent errors from the autograder, you are not allowed to edit or delete non-graded cells in this lab. Please also refrain from adding any new cells. \n",
    "**Once you have passed this assignment** and want to experiment with any of the non-graded code, you may follow the instructions at the bottom of this notebook._"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"1\"></a>\n",
    "## 1 - Import Packages\n",
    "\n",
    "We'll make use of the following packages:\n",
    "- `numpy` is a package for scientific computing in python.\n",
    "- `deque` will be our data structure for our memory buffer.\n",
    "- `namedtuple` will be used to store the experience tuples.\n",
    "- The `gym` toolkit is a collection of environments that can be used to test reinforcement learning algorithms. We should note that in this notebook we are using `gym` version `0.24.0`.\n",
    "- `PIL.Image` and `pyvirtualdisplay` are needed to render the Lunar Lander environment.\n",
    "- We will use several modules from the `tensorflow.keras` framework for building deep learning models.\n",
    "- `utils` is a module that contains helper functions for this assignment. You do not need to modify the code in this file.\n",
    "\n",
    "Run the cell below to import all the necessary packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "KYbOPKRtfQOr"
   },
   "outputs": [],
   "source": [
    "import time\n",
    "from collections import deque, namedtuple\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "import PIL.Image\n",
    "import tensorflow as tf\n",
    "import utils\n",
    "\n",
    "from pyvirtualdisplay import Display\n",
    "from tensorflow.keras import Sequential\n",
    "from tensorflow.keras.layers import Dense, Input\n",
    "from tensorflow.keras.losses import MSE\n",
    "from tensorflow.keras.optimizers import Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "# Set up a virtual display to render the Lunar Lander environment.\n",
    "Display(visible=0, size=(840, 480)).start();\n",
    "\n",
    "# Set the random seed for TensorFlow\n",
    "tf.random.set_seed(utils.SEED)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"2\"></a>\n",
    "## 2 - Hyperparameters\n",
    "\n",
    "Run the cell below to set the hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "MEMORY_SIZE = 100_000     # size of memory buffer\n",
    "GAMMA = 0.995             # discount factor\n",
    "ALPHA = 1e-3              # learning rate  \n",
    "NUM_STEPS_FOR_UPDATE = 4  # perform a learning update every C time steps"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"3\"></a>\n",
    "## 3 - The Lunar Lander Environment\n",
    "\n",
    "In this notebook we will be using [OpenAI's Gym Library](https://www.gymlibrary.dev/). The Gym library provides a wide variety of environments for reinforcement learning. To put it simply, an environment represents a problem or task to be solved. In this notebook, we will try to solve the Lunar Lander environment using reinforcement learning.\n",
    "\n",
    "The goal of the Lunar Lander environment is to land the lunar lander safely on the landing pad on the surface of the moon. The landing pad is designated by two flag poles and its center is at coordinates `(0,0)` but the lander is also allowed to land outside of the landing pad. The lander starts at the top center of the environment with a random initial force applied to its center of mass and has infinite fuel. The environment is considered solved if you get `200` points. \n",
    "\n",
    "<br>\n",
    "<br>\n",
    "<figure>\n",
    "  <img src = \"images/lunar_lander.gif\" width = 40%>\n",
    "      <figcaption style = \"text-align: center; font-style: italic\">Fig 1. Lunar Lander Environment.</figcaption>\n",
    "</figure>\n",
    "\n",
    "\n",
    "\n",
    "<a name=\"3.1\"></a>\n",
    "### 3.1 Action Space\n",
    "\n",
    "The agent has four discrete actions available:\n",
    "\n",
    "* Do nothing.\n",
    "* Fire right engine.\n",
    "* Fire main engine.\n",
    "* Fire left engine.\n",
    "\n",
    "Each action has a corresponding numerical value:\n",
    "\n",
    "```python\n",
    "Do nothing = 0\n",
    "Fire right engine = 1\n",
    "Fire main engine = 2\n",
    "Fire left engine = 3\n",
    "```\n",
    "\n",
    "<a name=\"3.2\"></a>\n",
    "### 3.2 Observation Space\n",
    "\n",
    "The agent's observation space consists of a state vector with 8 variables:\n",
    "\n",
    "* Its $(x,y)$ coordinates. The landing pad is always at coordinates $(0,0)$.\n",
    "* Its linear velocities $(\\dot x,\\dot y)$.\n",
    "* Its angle $\\theta$.\n",
    "* Its angular velocity $\\dot \\theta$.\n",
    "* Two booleans, $l$ and $r$, that represent whether each leg is in contact with the ground or not.\n",
    "\n",
    "<a name=\"3.3\"></a>\n",
    "### 3.3 Rewards\n",
    "\n",
    "After every step, a reward is granted. The total reward of an episode is the sum of the rewards for all the steps within that episode.\n",
    "\n",
    "For each step, the reward:\n",
    "- is increased/decreased the closer/further the lander is to the landing pad.\n",
    "- is increased/decreased the slower/faster the lander is moving.\n",
    "- is decreased the more the lander is tilted (angle not horizontal).\n",
    "- is increased by 10 points for each leg that is in contact with the ground.\n",
    "- is decreased by 0.03 points each frame a side engine is firing.\n",
    "- is decreased by 0.3 points each frame the main engine is firing.\n",
    "\n",
    "The episode receives an additional reward of -100 or +100 points for crashing or landing safely respectively.\n",
    "\n",
    "<a name=\"3.4\"></a>\n",
    "### 3.4 Episode Termination\n",
    "\n",
    "An episode ends (i.e the environment enters a terminal state) if:\n",
    "\n",
    "* The lunar lander crashes (i.e if the body of the lunar lander comes in contact with the surface of the moon).\n",
    "\n",
    "* The absolute value of the lander's $x$-coordinate is greater than 1 (i.e. it goes beyond the left or right border)\n",
    "\n",
    "You can check out the [Open AI Gym documentation](https://www.gymlibrary.dev/environments/box2d/lunar_lander/) for a full description of the environment. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"4\"></a>\n",
    "## 4 - Load the Environment\n",
    "\n",
    "We start by loading the `LunarLander-v2` environment from the `gym` library by using the `.make()` method. `LunarLander-v2` is the latest version of the Lunar Lander environment and you can read about its version history in the [Open AI Gym documentation](https://www.gymlibrary.dev/environments/box2d/lunar_lander/#version-history)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "ILVMYKewfR0n"
   },
   "outputs": [],
   "source": [
    "env = gym.make('LunarLander-v2')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once we load the environment we use the `.reset()` method to reset the environment to the initial state. The lander starts at the top center of the environment and we can render the first frame of the environment by using the `.render()` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "env.reset()\n",
    "PIL.Image.fromarray(env.render(mode='rgb_array'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to build our neural network later on we need to know the size of the state vector and the number of valid actions. We can get this information from our environment by using the `.observation_space.shape` and `action_space.n` methods, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "x3fdqdG4CUu2"
   },
   "outputs": [],
   "source": [
    "state_size = env.observation_space.shape\n",
    "num_actions = env.action_space.n\n",
    "\n",
    "print('State Shape:', state_size)\n",
    "print('Number of actions:', num_actions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"5\"></a>\n",
    "## 5 - Interacting with the Gym Environment\n",
    "\n",
    "The Gym library implements the standard “agent-environment loop” formalism:\n",
    "\n",
    "<br>\n",
    "<center>\n",
    "<video src = \"./videos/rl_formalism.m4v\" width=\"840\" height=\"480\" controls autoplay loop poster=\"./images/rl_formalism.png\"> </video>\n",
    "<figcaption style = \"text-align:center; font-style:italic\">Fig 2. Agent-environment Loop Formalism.</figcaption>\n",
    "</center>\n",
    "<br>\n",
    "\n",
    "In the standard “agent-environment loop” formalism, an agent interacts with the environment in discrete time steps $t=0,1,2,...$. At each time step $t$, the agent uses a policy $\\pi$ to select an action $A_t$ based on its observation of the environment's state $S_t$. The agent receives a numerical reward $R_t$ and on the next time step, moves to a new state $S_{t+1}$.\n",
    "\n",
    "<a name=\"5.1\"></a>\n",
    "### 5.1 Exploring the Environment's Dynamics\n",
    "\n",
    "In Open AI's Gym environments, we use the `.step()` method to run a single time step of the environment's dynamics. In the version of `gym` that we are using the `.step()` method accepts an action and returns four values:\n",
    "\n",
    "* `observation` (**object**): an environment-specific object representing your observation of the environment. In the Lunar Lander environment this corresponds to a numpy array containing the positions and velocities of the lander as described in section [3.2 Observation Space](#3.2).\n",
    "\n",
    "\n",
    "* `reward` (**float**): amount of reward returned as a result of taking the given action. In the Lunar Lander environment this corresponds to a float of type `numpy.float64` as described in section [3.3 Rewards](#3.3).\n",
    "\n",
    "\n",
    "* `done` (**boolean**): When done is `True`, it indicates the episode has terminated and it’s time to reset the environment. \n",
    "\n",
    "\n",
    "* `info` (**dictionary**): diagnostic information useful for debugging. We won't be using this variable in this notebook but it is shown here for completeness.\n",
    "\n",
    "To begin an episode, we need to reset the environment to an initial state. We do this by using the `.reset()` method. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "# Reset the environment and get the initial state.\n",
    "current_state = env.reset()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once the environment is reset, the agent can start taking actions in the environment by using the `.step()` method. Note that the agent can only take one action per time step. \n",
    "\n",
    "In the cell below you can select different actions and see how the returned values change depending on the action taken. Remember that in this environment the agent has four discrete actions available and we specify them in code by using their corresponding numerical value:\n",
    "\n",
    "```python\n",
    "Do nothing = 0\n",
    "Fire right engine = 1\n",
    "Fire main engine = 2\n",
    "Fire left engine = 3\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false
   },
   "outputs": [],
   "source": [
    "# Select an action\n",
    "action = 0\n",
    "\n",
    "# Run a single time step of the environment's dynamics with the given action.\n",
    "next_state, reward, done, _ = env.step(action)\n",
    "\n",
    "# Display table with values.\n",
    "utils.display_table(current_state, action, next_state, reward, done)\n",
    "\n",
    "# Replace the `current_state` with the state after the action is taken\n",
    "current_state = next_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In practice, when we train the agent we use a loop to allow the agent to take many consecutive actions during an episode."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"6\"></a>\n",
    "## 6 - Deep Q-Learning\n",
    "\n",
    "In cases where both the state and action space are discrete we can estimate the action-value function iteratively by using the Bellman equation:\n",
    "\n",
    "$$\n",
    "Q_{i+1}(s,a) = R + \\gamma \\max_{a'}Q_i(s',a')\n",
    "$$\n",
    "\n",
    "This iterative method converges to the optimal action-value function $Q^*(s,a)$ as $i\\to\\infty$. This means that the agent just needs to gradually explore the state-action space and keep updating the estimate of $Q(s,a)$ until it converges to the optimal action-value function $Q^*(s,a)$. However, in cases where the state space is continuous it becomes practically impossible to explore the entire state-action space. Consequently, this also makes it practically impossible to gradually estimate $Q(s,a)$ until it converges to $Q^*(s,a)$.\n",
    "\n",
    "In the Deep $Q$-Learning, we solve this problem by using a neural network to estimate the action-value function $Q(s,a)\\approx Q^*(s,a)$. We call this neural network a $Q$-Network and it can be trained by adjusting its weights at each iteration to minimize the mean-squared error in the Bellman equation.\n",
    "\n",
    "Unfortunately, using neural networks in reinforcement learning to estimate action-value functions has proven to be highly unstable. Luckily, there's a couple of techniques that can be employed to avoid instabilities. These techniques consist of using a ***Target Network*** and ***Experience Replay***. We will explore these two techniques in the following sections."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"6.1\"></a>\n",
    "### 6.1 Target Network\n",
    "\n",
    "We can train the $Q$-Network by adjusting it's weights at each iteration to minimize the mean-squared error in the Bellman equation, where the target values are given by:\n",
    "\n",
    "$$\n",
    "y = R + \\gamma \\max_{a'}Q(s',a';w)\n",
    "$$\n",
    "\n",
    "where $w$ are the weights of the $Q$-Network. This means that we are adjusting the weights $w$ at each iteration to minimize the following error:\n",
    "\n",
    "$$\n",
    "\\overbrace{\\underbrace{R + \\gamma \\max_{a'}Q(s',a'; w)}_{\\rm {y~target}} - Q(s,a;w)}^{\\rm {Error}}\n",
    "$$\n",
    "\n",
    "Notice that this forms a problem because the $y$ target is changing on every iteration. Having a constantly moving target can lead to oscillations and instabilities. To avoid this, we can create\n",
    "a separate neural network for generating the $y$ targets. We call this separate neural network the **target $\\hat Q$-Network** and it will have the same architecture as the original $Q$-Network. By using the target $\\hat Q$-Network, the above error becomes:\n",
    "\n",
    "$$\n",
    "\\overbrace{\\underbrace{R + \\gamma \\max_{a'}\\hat{Q}(s',a'; w^-)}_{\\rm {y~target}} - Q(s,a;w)}^{\\rm {Error}}\n",
    "$$\n",
    "\n",
    "where $w^-$ and $w$ are the weights of the target $\\hat Q$-Network and $Q$-Network, respectively.\n",
    "\n",
    "In practice, we will use the following algorithm: every $C$ time steps we will use the $\\hat Q$-Network to generate the $y$ targets and update the weights of the target $\\hat Q$-Network using the weights of the $Q$-Network. We will update the weights $w^-$ of the the target $\\hat Q$-Network using a **soft update**. This means that we will update the weights $w^-$ using the following rule:\n",
    " \n",
    "$$\n",
    "w^-\\leftarrow \\tau w + (1 - \\tau) w^-\n",
    "$$\n",
    "\n",
    "where $\\tau\\ll 1$. By using the soft update, we are ensuring that the target values, $y$, change slowly, which greatly improves the stability of our learning algorithm."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"ex01\"></a>\n",
    "### Exercise 1\n",
    "\n",
    "In this exercise you will create the $Q$ and target $\\hat Q$ networks and set the optimizer. Remember that the Deep $Q$-Network (DQN) is a neural network that approximates the action-value function $Q(s,a)\\approx Q^*(s,a)$. It does this by learning how to map states to $Q$ values.\n",
    "\n",
    "To solve the Lunar Lander environment, we are going to employ a DQN with the following architecture:\n",
    "\n",
    "* An `Input` layer that takes `state_size` as input.\n",
    "\n",
    "* A `Dense` layer with `64` units and a `relu` activation function.\n",
    "\n",
    "* A `Dense` layer with `64` units and a `relu` activation function.\n",
    "\n",
    "* A `Dense` layer with `num_actions` units and a `linear` activation function. This will be the output layer of our network.\n",
    "\n",
    "\n",
    "In the cell below you should create the $Q$-Network and the target $\\hat Q$-Network using the model architecture described above. Remember that both the $Q$-Network and the target $\\hat Q$-Network have the same architecture.\n",
    "\n",
    "Lastly, you should set `Adam` as the optimizer with a learning rate equal to `ALPHA`. Recall that `ALPHA` was defined in the [Hyperparameters](#2) section. We should note that for this exercise you should use the already imported packages:\n",
    "```python\n",
    "from tensorflow.keras.layers import Dense, Input\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false
   },
   "outputs": [],
   "source": [
    "# UNQ_C1\n",
    "# GRADED CELL\n",
    "\n",
    "# Create the Q-Network\n",
    "q_network = Sequential([\n",
    "    ### START CODE HERE ### \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    ### END CODE HERE ### \n",
    "    ])\n",
    "\n",
    "# Create the target Q^-Network\n",
    "target_q_network = Sequential([\n",
    "    ### START CODE HERE ### \n",
    "    \n",
    "    \n",
    "    \n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "    ])\n",
    "\n",
    "### START CODE HERE ### \n",
    "optimizer = None\n",
    "### END CODE HERE ###"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "# UNIT TEST\n",
    "from public_tests import *\n",
    "\n",
    "test_network(q_network)\n",
    "test_network(target_q_network)\n",
    "test_optimizer(optimizer, ALPHA) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "  <summary><font size=\"3\" color=\"darkgreen\"><b>Click for hints</b></font></summary>\n",
    "    \n",
    "```python\n",
    "# Create the Q-Network\n",
    "q_network = Sequential([\n",
    "    Input(shape=state_size),                      \n",
    "    Dense(units=64, activation='relu'),            \n",
    "    Dense(units=64, activation='relu'),            \n",
    "    Dense(units=num_actions, activation='linear'),\n",
    "    ])\n",
    "\n",
    "# Create the target Q^-Network\n",
    "target_q_network = Sequential([\n",
    "    Input(shape=state_size),                       \n",
    "    Dense(units=64, activation='relu'),            \n",
    "    Dense(units=64, activation='relu'),            \n",
    "    Dense(units=num_actions, activation='linear'), \n",
    "    ])\n",
    "\n",
    "optimizer = Adam(learning_rate=ALPHA)                                  \n",
    "``` "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"6.2\"></a>\n",
    "### 6.2 Experience Replay\n",
    "\n",
    "When an agent interacts with the environment, the states, actions, and rewards the agent experiences are sequential by nature. If the agent tries to learn from these consecutive experiences it can run into problems due to the strong correlations between them. To avoid this, we employ a technique known as **Experience Replay** to generate uncorrelated experiences for training our agent. Experience replay consists of storing the agent's experiences (i.e the states, actions, and rewards the agent receives) in a memory buffer and then sampling a random mini-batch of experiences from the buffer to do the learning. The experience tuples $(S_t, A_t, R_t, S_{t+1})$ will be added to the memory buffer at each time step as the agent interacts with the environment.\n",
    "\n",
    "For convenience, we will store the experiences as named tuples."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "# Store experiences as named tuples\n",
    "experience = namedtuple(\"Experience\", field_names=[\"state\", \"action\", \"reward\", \"next_state\", \"done\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By using experience replay we avoid problematic correlations, oscillations and instabilities. In addition, experience replay also allows the agent to potentially use the same experience in multiple weight updates, which increases data efficiency."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"7\"></a>\n",
    "## 7 - Deep Q-Learning Algorithm with Experience Replay\n",
    "\n",
    "Now that we know all the techniques that we are going to use, we can put them together to arrive at the Deep Q-Learning Algorithm With Experience Replay.\n",
    "<br>\n",
    "<br>\n",
    "<figure>\n",
    "  <img src = \"images/deep_q_algorithm.png\" width = 90% style = \"border: thin silver solid; padding: 0px\">\n",
    "      <figcaption style = \"text-align: center; font-style: italic\">Fig 3. Deep Q-Learning with Experience Replay.</figcaption>\n",
    "</figure>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"ex02\"></a>\n",
    "### Exercise 2\n",
    "\n",
    "In this exercise you will implement line ***12*** of the algorithm outlined in *Fig 3* above and you will also compute the loss between the $y$ targets and the $Q(s,a)$ values. In the cell below, complete the `compute_loss` function by setting the $y$ targets equal to:\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "    y_j =\n",
    "    \\begin{cases}\n",
    "      R_j & \\text{if episode terminates at step  } j+1\\\\\n",
    "      R_j + \\gamma \\max_{a'}\\hat{Q}(s_{j+1},a') & \\text{otherwise}\\\\\n",
    "    \\end{cases}       \n",
    "\\end{equation}\n",
    "$$\n",
    "\n",
    "Here are a couple of things to note:\n",
    "\n",
    "* The `compute_loss` function takes in a mini-batch of experience tuples. This mini-batch of experience tuples is unpacked to extract the `states`, `actions`, `rewards`, `next_states`, and `done_vals`. You should keep in mind that these variables are *TensorFlow Tensors* whose size will depend on the mini-batch size. For example, if the mini-batch size is `64` then both `rewards` and `done_vals` will be TensorFlow Tensors with `64` elements.\n",
    "\n",
    "\n",
    "* Using `if/else` statements to set the $y$ targets will not work when the variables are tensors with many elements. However, notice that you can use the `done_vals` to implement the above in a single line of code. To do this, recall that the `done` variable is a Boolean variable that takes the value `True` when an episode terminates at step $j+1$ and it is `False` otherwise. Taking into account that a Boolean value of `True` has the numerical value of `1` and a Boolean value of `False` has the numerical value of `0`, you can use the factor `(1 - done_vals)` to implement the above in a single line of code. Here's a hint: notice that `(1 - done_vals)` has a value of `0` when `done_vals` is `True` and a value of `1` when `done_vals` is `False`. \n",
    "\n",
    "Lastly, compute the loss by calculating the Mean-Squared Error (`MSE`) between the `y_targets` and the `q_values`. To calculate the mean-squared error you should use the already imported package `MSE`:\n",
    "```python\n",
    "from tensorflow.keras.losses import MSE\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false
   },
   "outputs": [],
   "source": [
    "# UNQ_C2\n",
    "# GRADED FUNCTION: calculate_loss\n",
    "\n",
    "def compute_loss(experiences, gamma, q_network, target_q_network):\n",
    "    \"\"\" \n",
    "    Calculates the loss.\n",
    "    \n",
    "    Args:\n",
    "      experiences: (tuple) tuple of [\"state\", \"action\", \"reward\", \"next_state\", \"done\"] namedtuples\n",
    "      gamma: (float) The discount factor.\n",
    "      q_network: (tf.keras.Sequential) Keras model for predicting the q_values\n",
    "      target_q_network: (tf.keras.Sequential) Keras model for predicting the targets\n",
    "          \n",
    "    Returns:\n",
    "      loss: (TensorFlow Tensor(shape=(0,), dtype=int32)) the Mean-Squared Error between\n",
    "            the y targets and the Q(s,a) values.\n",
    "    \"\"\"\n",
    "\n",
    "    # Unpack the mini-batch of experience tuples\n",
    "    states, actions, rewards, next_states, done_vals = experiences\n",
    "    \n",
    "    # Compute max Q^(s,a)\n",
    "    max_qsa = tf.reduce_max(target_q_network(next_states), axis=-1)\n",
    "    \n",
    "    # Set y = R if episode terminates, otherwise set y = R + γ max Q^(s,a).\n",
    "    ### START CODE HERE ### \n",
    "    y_targets = None\n",
    "    ### END CODE HERE ###\n",
    "    \n",
    "    # Get the q_values and reshape to match y_targets\n",
    "    q_values = q_network(states)\n",
    "    q_values = tf.gather_nd(q_values, tf.stack([tf.range(q_values.shape[0]),\n",
    "                                                tf.cast(actions, tf.int32)], axis=1))\n",
    "        \n",
    "    # Compute the loss\n",
    "    ### START CODE HERE ### \n",
    "    loss = None \n",
    "    ### END CODE HERE ### \n",
    "    \n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "# UNIT TEST    \n",
    "test_compute_loss(compute_loss)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "  <summary><font size=\"3\" color=\"darkgreen\"><b>Click for hints</b></font></summary>\n",
    "    \n",
    "```python\n",
    "def compute_loss(experiences, gamma, q_network, target_q_network):\n",
    "    \"\"\" \n",
    "    Calculates the loss.\n",
    "    \n",
    "    Args:\n",
    "      experiences: (tuple) tuple of [\"state\", \"action\", \"reward\", \"next_state\", \"done\"] namedtuples\n",
    "      gamma: (float) The discount factor.\n",
    "      q_network: (tf.keras.Sequential) Keras model for predicting the q_values\n",
    "      target_q_network: (tf.keras.Sequential) Keras model for predicting the targets\n",
    "          \n",
    "    Returns:\n",
    "      loss: (TensorFlow Tensor(shape=(0,), dtype=int32)) the Mean-Squared Error between\n",
    "            the y targets and the Q(s,a) values.\n",
    "    \"\"\"\n",
    "\n",
    "    \n",
    "    # Unpack the mini-batch of experience tuples\n",
    "    states, actions, rewards, next_states, done_vals = experiences\n",
    "    \n",
    "    # Compute max Q^(s,a)\n",
    "    max_qsa = tf.reduce_max(target_q_network(next_states), axis=-1)\n",
    "    \n",
    "    # Set y = R if episode terminates, otherwise set y = R + γ max Q^(s,a).\n",
    "    y_targets = rewards + (gamma * max_qsa * (1 - done_vals))\n",
    "    \n",
    "    # Get the q_values\n",
    "    q_values = q_network(states)\n",
    "    q_values = tf.gather_nd(q_values, tf.stack([tf.range(q_values.shape[0]),\n",
    "                                                tf.cast(actions, tf.int32)], axis=1))\n",
    "    \n",
    "    # Calculate the loss\n",
    "    loss = MSE(y_targets, q_values)\n",
    "    \n",
    "    return loss\n",
    "\n",
    "``` \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"8\"></a>\n",
    "## 8 - Update the Network Weights\n",
    "\n",
    "We will use the `agent_learn` function below to implement lines ***12 -14*** of the algorithm outlined in [Fig 3](#7). The `agent_learn` function will update the weights of the $Q$ and target $\\hat Q$ networks using a custom training loop. Because we are using a custom training loop we need to retrieve the gradients via a `tf.GradientTape` instance, and then call `optimizer.apply_gradients()` to update the weights of our $Q$-Network. Note that we are also using the `@tf.function` decorator to increase performance. Without this decorator our training will take twice as long. If you would like to know more about how to increase performance with `@tf.function` take a look at the [TensorFlow documentation](https://www.tensorflow.org/guide/function).\n",
    "\n",
    "The last line of this function updates the weights of the target $\\hat Q$-Network using a [soft update](#6.1). If you want to know how this is implemented in code we encourage you to take a look at the `utils.update_target_network` function in the `utils` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def agent_learn(experiences, gamma):\n",
    "    \"\"\"\n",
    "    Updates the weights of the Q networks.\n",
    "    \n",
    "    Args:\n",
    "      experiences: (tuple) tuple of [\"state\", \"action\", \"reward\", \"next_state\", \"done\"] namedtuples\n",
    "      gamma: (float) The discount factor.\n",
    "    \n",
    "    \"\"\"\n",
    "    \n",
    "    # Calculate the loss\n",
    "    with tf.GradientTape() as tape:\n",
    "        loss = compute_loss(experiences, gamma, q_network, target_q_network)\n",
    "\n",
    "    # Get the gradients of the loss with respect to the weights.\n",
    "    gradients = tape.gradient(loss, q_network.trainable_variables)\n",
    "    \n",
    "    # Update the weights of the q_network.\n",
    "    optimizer.apply_gradients(zip(gradients, q_network.trainable_variables))\n",
    "\n",
    "    # update the weights of target q_network\n",
    "    utils.update_target_network(q_network, target_q_network)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"9\"></a>\n",
    "## 9 - Train the Agent\n",
    "\n",
    "We are now ready to train our agent to solve the Lunar Lander environment. In the cell below we will implement the algorithm in [Fig 3](#7) line by line (please note that we have included the same algorithm below for easy reference. This will prevent you from scrolling up and down the notebook):\n",
    "\n",
    "* **Line 1**: We initialize the `memory_buffer` with a capacity of $N =$ `MEMORY_SIZE`. Notice that we are using a `deque` as the data structure for our `memory_buffer`.\n",
    "\n",
    "\n",
    "* **Line 2**: We skip this line since we already initialized the `q_network` in [Exercise 1](#ex01).\n",
    "\n",
    "\n",
    "* **Line 3**: We initialize the `target_q_network` by setting its weights to be equal to those of the `q_network`.\n",
    "\n",
    "\n",
    "* **Line 4**: We start the outer loop. Notice that we have set $M =$ `num_episodes = 2000`. This number is reasonable because the agent should be able to solve the Lunar Lander environment in less than `2000` episodes using this notebook's default parameters.\n",
    "\n",
    "\n",
    "* **Line 5**: We use the `.reset()` method to reset the environment to the initial state and get the initial state.\n",
    "\n",
    "\n",
    "* **Line 6**: We start the inner loop. Notice that we have set $T =$ `max_num_timesteps = 1000`. This means that the episode will automatically terminate if the episode hasn't terminated after `1000` time steps.\n",
    "\n",
    "\n",
    "* **Line 7**: The agent observes the current `state` and chooses an `action` using an $\\epsilon$-greedy policy. Our agent starts out using a value of $\\epsilon =$ `epsilon = 1` which yields an $\\epsilon$-greedy policy that is equivalent to the equiprobable random policy. This means that at the beginning of our training, the agent is just going to take random actions regardless of the observed `state`. As training progresses we will decrease the value of $\\epsilon$ slowly towards a minimum value using a given $\\epsilon$-decay rate. We want this minimum value to be close to zero because a value of $\\epsilon = 0$ will yield an $\\epsilon$-greedy policy that is equivalent to the greedy policy. This means that towards the end of training, the agent will lean towards selecting the `action` that it believes (based on its past experiences) will maximize $Q(s,a)$. We will set the minimum $\\epsilon$ value to be `0.01` and not exactly 0 because we always want to keep a little bit of exploration during training. If you want to know how this is implemented in code we encourage you to take a look at the `utils.get_action` function in the `utils` module.\n",
    "\n",
    "\n",
    "* **Line 8**: We use the `.step()` method to take the given `action` in the environment and get the `reward` and the `next_state`. \n",
    "\n",
    "\n",
    "* **Line 9**: We store the `experience(state, action, reward, next_state, done)` tuple in our `memory_buffer`. Notice that we also store the `done` variable so that we can keep track of when an episode terminates. This allowed us to set the $y$ targets in [Exercise 2](#ex02).\n",
    "\n",
    "\n",
    "* **Line 10**: We check if the conditions are met to perform a learning update. We do this by using our custom `utils.check_update_conditions` function. This function checks if $C =$ `NUM_STEPS_FOR_UPDATE = 4` time steps have occured and if our `memory_buffer` has enough experience tuples to fill a mini-batch. For example, if the mini-batch size is `64`, then our `memory_buffer` should have more than `64` experience tuples in order to pass the latter condition. If the conditions are met, then the `utils.check_update_conditions` function will return a value of `True`, otherwise it will return a value of `False`.\n",
    "\n",
    "\n",
    "* **Lines 11 - 14**: If the `update` variable is `True` then we perform a learning update. The learning update consists of sampling a random mini-batch of experience tuples from our `memory_buffer`, setting the $y$ targets, performing gradient descent, and updating the weights of the networks. We will use the `agent_learn` function we defined in [Section 8](#8) to perform the latter 3.\n",
    "\n",
    "\n",
    "* **Line 15**: At the end of each iteration of the inner loop we set `next_state` as our new `state` so that the loop can start again from this new state. In addition, we check if the episode has reached a terminal state (i.e we check if `done = True`). If a terminal state has been reached, then we break out of the inner loop.\n",
    "\n",
    "\n",
    "* **Line 16**: At the end of each iteration of the outer loop we update the value of $\\epsilon$, and check if the environment has been solved. We consider that the environment has been solved if the agent receives an average of `200` points in the last `100` episodes. If the environment has not been solved we continue the outer loop and start a new episode.\n",
    "\n",
    "Finally, we wanted to note that we have included some extra variables to keep track of the total number of points the agent received in each episode. This will help us determine if the agent has solved the environment and it will also allow us to see how our agent performed during training. We also use the `time` module to measure how long the training takes. \n",
    "\n",
    "<br>\n",
    "<br>\n",
    "<figure>\n",
    "  <img src = \"images/deep_q_algorithm.png\" width = 90% style = \"border: thin silver solid; padding: 0px\">\n",
    "      <figcaption style = \"text-align: center; font-style: italic\">Fig 4. Deep Q-Learning with Experience Replay.</figcaption>\n",
    "</figure>\n",
    "<br>\n",
    "\n",
    "**Note:** With this notebook's default parameters, the following cell takes between 10 to 15 minutes to run. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "start = time.time()\n",
    "\n",
    "num_episodes = 2000\n",
    "max_num_timesteps = 1000\n",
    "\n",
    "total_point_history = []\n",
    "\n",
    "num_p_av = 100    # number of total points to use for averaging\n",
    "epsilon = 1.0     # initial ε value for ε-greedy policy\n",
    "\n",
    "# Create a memory buffer D with capacity N\n",
    "memory_buffer = deque(maxlen=MEMORY_SIZE)\n",
    "\n",
    "# Set the target network weights equal to the Q-Network weights\n",
    "target_q_network.set_weights(q_network.get_weights())\n",
    "\n",
    "for i in range(num_episodes):\n",
    "    \n",
    "    # Reset the environment to the initial state and get the initial state\n",
    "    state = env.reset()\n",
    "    total_points = 0\n",
    "    \n",
    "    for t in range(max_num_timesteps):\n",
    "        \n",
    "        # From the current state S choose an action A using an ε-greedy policy\n",
    "        state_qn = np.expand_dims(state, axis=0)  # state needs to be the right shape for the q_network\n",
    "        q_values = q_network(state_qn)\n",
    "        action = utils.get_action(q_values, epsilon)\n",
    "        \n",
    "        # Take action A and receive reward R and the next state S'\n",
    "        next_state, reward, done, _ = env.step(action)\n",
    "        \n",
    "        # Store experience tuple (S,A,R,S') in the memory buffer.\n",
    "        # We store the done variable as well for convenience.\n",
    "        memory_buffer.append(experience(state, action, reward, next_state, done))\n",
    "        \n",
    "        # Only update the network every NUM_STEPS_FOR_UPDATE time steps.\n",
    "        update = utils.check_update_conditions(t, NUM_STEPS_FOR_UPDATE, memory_buffer)\n",
    "        \n",
    "        if update:\n",
    "            # Sample random mini-batch of experience tuples (S,A,R,S') from D\n",
    "            experiences = utils.get_experiences(memory_buffer)\n",
    "            \n",
    "            # Set the y targets, perform a gradient descent step,\n",
    "            # and update the network weights.\n",
    "            agent_learn(experiences, GAMMA)\n",
    "        \n",
    "        state = next_state.copy()\n",
    "        total_points += reward\n",
    "        \n",
    "        if done:\n",
    "            break\n",
    "            \n",
    "    total_point_history.append(total_points)\n",
    "    av_latest_points = np.mean(total_point_history[-num_p_av:])\n",
    "    \n",
    "    # Update the ε value\n",
    "    epsilon = utils.get_new_eps(epsilon)\n",
    "\n",
    "    print(f\"\\rEpisode {i+1} | Total point average of the last {num_p_av} episodes: {av_latest_points:.2f}\", end=\"\")\n",
    "\n",
    "    if (i+1) % num_p_av == 0:\n",
    "        print(f\"\\rEpisode {i+1} | Total point average of the last {num_p_av} episodes: {av_latest_points:.2f}\")\n",
    "\n",
    "    # We will consider that the environment is solved if we get an\n",
    "    # average of 200 points in the last 100 episodes.\n",
    "    if av_latest_points >= 200.0:\n",
    "        print(f\"\\n\\nEnvironment solved in {i+1} episodes!\")\n",
    "        q_network.save('lunar_lander_model.h5')\n",
    "        break\n",
    "        \n",
    "tot_time = time.time() - start\n",
    "\n",
    "print(f\"\\nTotal Runtime: {tot_time:.2f} s ({(tot_time/60):.2f} min)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can plot the total point history along with the moving average to see how our agent improved during training. If you want to know about the different plotting options available in the `utils.plot_history` function we encourage you to take a look at the `utils` module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "E_EUXxurfe8m",
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Plot the total point history along with the moving average\n",
    "utils.plot_history(total_point_history)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "c_xwgaX5MnYt"
   },
   "source": [
    "<a name=\"10\"></a>\n",
    "## 10 - See the Trained Agent In Action\n",
    "\n",
    "Now that we have trained our agent, we can see it in action. We will use the `utils.create_video` function to create a video of our agent interacting with the environment using the trained $Q$-Network. The `utils.create_video` function uses the `imageio` library to create the video. This library produces some warnings that can be distracting, so, to suppress these warnings we run the code below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false
   },
   "outputs": [],
   "source": [
    "# Suppress warnings from imageio\n",
    "import logging\n",
    "logging.getLogger().setLevel(logging.ERROR)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the cell below we create a video of our agent interacting with the Lunar Lander environment using the trained `q_network`. The video is saved to the `videos` folder with the given `filename`. We use the `utils.embed_mp4` function to embed the video in the Jupyter Notebook so that we can see it here directly without having to download it.\n",
    "\n",
    "We should note that since the lunar lander starts with a random initial force applied to its center of mass, every time you run the cell below you will see a different video. If the agent was trained properly, it should be able to land the lunar lander in the landing pad every time, regardless of the initial force applied to its center of mass."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "id": "3Ttb_zLeJKiG"
   },
   "outputs": [],
   "source": [
    "filename = \"./videos/lunar_lander.mp4\"\n",
    "\n",
    "utils.create_video(filename, env, q_network)\n",
    "utils.embed_mp4(filename)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"11\"></a>\n",
    "## 11 - Congratulations!\n",
    "\n",
    "You have successfully used Deep Q-Learning with Experience Replay to train an agent to land a lunar lander safely on a landing pad on the surface of the moon. Congratulations!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a name=\"12\"></a>\n",
    "## 12 - References\n",
    "\n",
    "If you would like to learn more about Deep Q-Learning, we recommend you check out the following papers.\n",
    "\n",
    "\n",
    "* Mnih, V., Kavukcuoglu, K., Silver, D. et al. Human-level control through deep reinforcement learning. Nature 518, 529–533 (2015).\n",
    "\n",
    "\n",
    "* Lillicrap, T. P., Hunt, J. J., Pritzel, A., et al. Continuous Control with Deep Reinforcement Learning. ICLR (2016).\n",
    "\n",
    "\n",
    "* Mnih, V., Kavukcuoglu, K., Silver, D. et al. Playing Atari with Deep Reinforcement Learning. arXiv e-prints.  arXiv:1312.5602 (2013)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<details>\n",
    "  <summary><font size=\"2\" color=\"darkgreen\"><b>Please click here if you want to experiment with any of the non-graded code.</b></font></summary>\n",
    "    <p><i><b>Important Note: Please only do this when you've already passed the assignment to avoid problems with the autograder.</b></i>\n",
    "    <ol>\n",
    "        <li> On the notebook’s menu, click “View” > “Cell Toolbar” > “Edit Metadata”</li>\n",
    "        <li> Hit the “Edit Metadata” button next to the code cell which you want to lock/unlock</li>\n",
    "        <li> Set the attribute value for “editable” to:\n",
    "            <ul>\n",
    "                <li> “true” if you want to unlock it </li>\n",
    "                <li> “false” if you want to lock it </li>\n",
    "            </ul>\n",
    "        </li>\n",
    "        <li> On the notebook’s menu, click “View” > “Cell Toolbar” > “None” </li>\n",
    "    </ol>\n",
    "    <p> Here's a short demo of how to do the steps above: \n",
    "        <br>\n",
    "        <img src=\"https://lh3.google.com/u/0/d/14Xy_Mb17CZVgzVAgq7NCjMVBvSae3xO1\" align=\"center\" alt=\"unlock_cells.gif\">\n",
    "</details>"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "TensorFlow - Lunar Lander.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}