Semi-gradient-with-NN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "c3a8b2eedba23933bc9763fd7a7f93f8",
     "grade": false,
     "grade_id": "cell-b37957442e1fbd8b",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "# Assignment 2 - Semi-gradient TD with a Neural Network\n",
    "\n",
    "Welcome to Course 3 Programming Assignment 2. In the previous assignment, you implemented semi-gradient TD with State Aggregation for solving a **policy evaluation task**. In this assignment, you will implement **semi-gradient TD with a simple Neural Network** and use it for the same policy evaluation problem. \n",
    "\n",
    "You will implement an agent to evaluate a fixed policy on the 500-State Randomwalk. As you may remember from the previous assignment, the 500-state Randomwalk includes 500 states. Each episode begins with the agent at the center and terminates when the agent goes far left beyond state 1 or far right beyond state 500. At each time step, the agent selects to move either left or right with equal probability. The environment determines how much the agent moves in the selected direction.\n",
    "\n",
    "**In this assignment, you will:**\n",
    "- Implement stochastic gradient descent method for state-value prediction.\n",
    "- Implement semi-gradient TD with a neural network as the function approximator and Adam algorithm.\n",
    "- Compare performance of semi-gradient TD with a neural network and semi-gradient TD with tile-coding.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "2e46d3cb20b9eaffc48288df9563d1ff",
     "grade": false,
     "grade_id": "cell-9b62cd317dfd157f",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "## Packages\n",
    "We import the following libraries that are required for this assignment:\n",
    "\n",
    "- [numpy](www.numpy.org) : Fundamental package for scientific computing with Python.\n",
    "- [matplotlib](http://matplotlib.org) : Library for plotting graphs in Python.\n",
    "- [RL-Glue](http://www.jmlr.org/papers/v10/tanner09a.html) : Library for reinforcement learning experiments.\n",
    "- [tqdm](https://tqdm.github.io/) : A package to display progress bar when running experiments.\n",
    "- BaseOptimizer : An abstract class that specifies the optimizer API for Agent.\n",
    "- plot_script : Custom script to plot results.\n",
    "- RandomWalkEnvironment : The Randomwalk environment script from Course 3 Assignment 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "361a4918a4d8258187005559ca15798b",
     "grade": false,
     "grade_id": "cell-b48b54531e30224f",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# Do not modify this cell!\n",
    "\n",
    "# Import necessary libraries\n",
    "# DO NOT IMPORT OTHER LIBRARIES - This will break the autograder.\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import os, shutil\n",
    "from tqdm import tqdm\n",
    "\n",
    "from rl_glue import RLGlue\n",
    "from environment import BaseEnvironment\n",
    "from agent import BaseAgent\n",
    "from optimizer import BaseOptimizer\n",
    "import plot_script\n",
    "from randomwalk_environment import RandomWalkEnvironment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "77b652f0dd2324ab9327d0b5c583edf2",
     "grade": false,
     "grade_id": "cell-de641a52a225de33",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "## Section 1: Create semi-gradient TD with a Neural Network\n",
    "In this section, you will implement an Agent that learns with semi-gradient TD with a neural network. You will use a neural network with one hidden layer. The input of the neural network is the one-hot encoding of the state number. We use the one-hot encoding of the state number instead of the state number itself because we do not want to build the prior knowledge that integer number inputs close to each other have similar values. The hidden layer contains 100 rectifier linear units (ReLUs) which pass their input if it is bigger than one and return 0 otherwise. ReLU gates are commonly used in neural networks due to their nice properties such as the sparsity of the activation and having non-vanishing gradients. The output of the neural network is the estimated state value. It is a linear function of the hidden units as is commonly the case when estimating the value of a continuous target using neural networks.\n",
    "\n",
    "The neural network looks like this:\n",
    "![](nn_structure.png)\n",
    "\n",
    "\n",
    "For a given input, $s$, value of $s$ is computed by:\n",
    "$$\n",
    "\\begin{align} \n",
    "\\psi &= sW^{[0]} + b^{[0]} \\\\\n",
    "x &= \\textit{max}(0, \\psi) \\\\\n",
    "v &= xW^{[1]} + b^{[1]}\n",
    "\\end{align} \n",
    "$$\n",
    "\n",
    "where $W^{[0]}$, $b^{[0]}$, $W^{[1]}$, $b^{[1]}$  are the parameters of the network and will be learned when training the agent."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "4ba60f0d7c91468d560247475385d931",
     "grade": false,
     "grade_id": "cell-5f7ab00aa29454a6",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "## 1-1: Implement helper methods\n",
    "\n",
    "Before implementing the agent, you first implement some helper functions which you will later use in agent's main methods. \n",
    "\n",
    "### Implement `get_value()`\n",
    "First, you will implement get_value() method which feeds an input $s$ into the neural network and returns the output of the network $v$ according to the equations above. To implement get_value(), take into account the following notes:\n",
    "\n",
    "- `get_value()` gets the one-hot encoded state number denoted by s as an input. \n",
    "- `get_value()` receives the weights of the neural network as input, denoted by weights and structured as an array of dictionaries. Each dictionary corresponds to weights from one layer of the neural network to the next. Each dictionary includes $W$ and $b$. The shape of the elements in weights are as follows:\n",
    "    - weights[0][\"W\"]: num_states $\\times$ num_hidden_units\n",
    "    - weights[0][\"b\"]: 1 $\\times$ num_hidden_units\n",
    "    - weights[1][\"W\"]: num_hidden_units $\\times$ 1\n",
    "    - weights[1][\"b\"]: 1 $\\times$ 1\n",
    "\n",
    "- The input of the neural network is a sparse vector. To make computation faster, we take advantage of input sparsity. To do so, we provided a helper method `my_matmul()`. **Make sure that you use `my_matmul()` for all matrix multiplications except for element-wise multiplications in this notebook.**\n",
    "- The max operator used for computing $x$ is element-wise. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "b4f8b88337142352f0977c64e53467ab",
     "grade": false,
     "grade_id": "cell-5a1a1bbc13b00f36",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "def my_matmul(x1, x2):\n",
    "    \"\"\"\n",
    "    Given matrices x1 and x2, return the multiplication of them\n",
    "    \"\"\"\n",
    "    \n",
    "    result = np.zeros((x1.shape[0], x2.shape[1]))\n",
    "    x1_non_zero_indices = x1.nonzero()\n",
    "    if x1.shape[0] == 1 and len(x1_non_zero_indices[1]) == 1:\n",
    "        result = x2[x1_non_zero_indices[1], :]\n",
    "    elif x1.shape[1] == 1 and len(x1_non_zero_indices[0]) == 1:\n",
    "        result[x1_non_zero_indices[0], :] = x2 * x1[x1_non_zero_indices[0], 0]\n",
    "    else:\n",
    "        result = np.matmul(x1, x2)\n",
    "    return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "checksum": "18603564b0ad43c6dd7b6075a9a79bfc",
     "grade": false,
     "grade_id": "cell-f6e1321ef3928d09",
     "locked": false,
     "schema_version": 1,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#GRADED FUNCTION: [get_value]\n",
    "\n",
    "def get_value(s, weights):\n",
    "    \"\"\"\n",
    "    Compute value of input s given the weights of a neural network\n",
    "    \"\"\"\n",
    "    ### Compute the ouput of the neural network, v, for input s (3 lines)\n",
    "    \n",
    "    ### START CODE HERE ###\n",
    "        \n",
    "    out = my_matmul( s , weights[0][\"W\"] ) + weights[0][\"b\"]\n",
    "    out = np.maximum( 0 , out )\n",
    "    v = my_matmul( out , weights[1][\"W\"] ) + weights[1][\"b\"]\n",
    "\n",
    "    ### END CODE HERE ###\n",
    "    return v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "feeb6580e587e701ca835133a3780004",
     "grade": false,
     "grade_id": "cell-d422d7296850b470",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Run the following code to test your implementation of the `get_value()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "ba83e4def1dfa6f3eacc1935bc531aee",
     "grade": true,
     "grade_id": "cell-e0dacb873e8f91f4",
     "locked": true,
     "points": 10,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "## Test Code for get_value() ## \n",
    "\n",
    "# Suppose num_states = 5, num_hidden_layer = 1, and num_hidden_units = 10 \n",
    "num_hidden_layer = 1\n",
    "s = np.array([[0, 0, 0, 1, 0]])\n",
    "\n",
    "weights_data = np.load(\"asserts/get_value_weights.npz\")\n",
    "weights = [dict() for i in range(num_hidden_layer+1)]\n",
    "weights[0][\"W\"] = weights_data[\"W0\"]\n",
    "weights[0][\"b\"] = weights_data[\"b0\"]\n",
    "weights[1][\"W\"] = weights_data[\"W1\"]\n",
    "weights[1][\"b\"] = weights_data[\"b1\"]\n",
    "\n",
    "estimated_value = get_value(s, weights)\n",
    "print (\"Estimated value: {}\".format(estimated_value))\n",
    "assert(np.allclose(estimated_value, np.array([[-0.21915705]])))\n",
    "\n",
    "print (\"Passed the assert!\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "ae6d398c5d6051d7a6510c6bb16cbad9",
     "grade": false,
     "grade_id": "cell-f34843c42f4bbc15",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "**Expected output**:\n",
    "\n",
    "    Estimated value: [[-0.21915705]]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "2833eef7385a16de825b9e20120186a5",
     "grade": false,
     "grade_id": "cell-c99e57755233d068",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "### Implement `get_gradient()`\n",
    "You will also implement `get_gradient()` method which computes the gradient of the value function for a given input, using backpropagation. You will later use this function to update the value function. \n",
    "\n",
    "As you know, we compute the value of a state $s$ according to: \n",
    "$$\n",
    "\\begin{align} \n",
    "\\psi &= sW^{[0]} + b^{[0]} \\\\\n",
    "x &= \\textit{max}(0, \\psi) \\\\\n",
    "v &= xW^{[1]} + b^{[1]}\n",
    "\\end{align} \n",
    "$$\n",
    "\n",
    "To update the weights of the neural network ($W^{[0]}$, $b^{[0]}$, $W^{[1]}$, $b^{[1]}$), we compute the gradient of $v$ with respect to the weights according to:\n",
    "\n",
    "$$\n",
    "\\begin{align} \n",
    "\\frac{\\partial v}{\\partial W^{[0]}} &= s^T(W^{[1]T} \\odot I_{x>0}) \\\\\n",
    "\\frac{\\partial v}{\\partial b^{[0]}} &= W^{[1]T} \\odot I_{x>0} \\\\\n",
    "\\frac{\\partial v}{\\partial W^{[1]}} &= x^T \\\\\n",
    "\\frac{\\partial v}{\\partial b^{[1]}} &= 1\n",
    "\\end{align}\n",
    "$$\n",
    "where $\\odot$ denotes element-wise matrix multiplication and $I_{x>0}$ is the gradient of the ReLU activation function which is an indicator whose $i$th element is 1 if $x[i]>0$ and 0 otherwise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "checksum": "576b7d14329dfdb7b8697e48f4096105",
     "grade": false,
     "grade_id": "cell-940eb29a426e021b",
     "locked": false,
     "schema_version": 1,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#GRADED FUNCTION: [get_gradient]\n",
    "\n",
    "def get_gradient(s, weights):\n",
    "    \"\"\"\n",
    "    Given inputs s and weights, return the gradient of v with respect to the weights\n",
    "    \"\"\"\n",
    "\n",
    "    ### Compute the gradient of the value function with respect to W0, b0, W1, b1 for input s (6~8 lines)\n",
    "    # grads[0][\"W\"] = ?\n",
    "    # grads[0][\"b\"] = ?\n",
    "    # grads[1][\"W\"] = ?\n",
    "    # grads[1][\"b\"] = ?\n",
    "    # Note that grads[0][\"W\"], grads[0][\"b\"], grads[1][\"W\"], and grads[1][\"b\"] should have the same shape as \n",
    "    # weights[0][\"W\"], weights[0][\"b\"], weights[1][\"W\"], and weights[1][\"b\"] respectively\n",
    "    # Note that to compute the gradients, you need to compute the activation of the hidden layer (x)\n",
    "\n",
    "    grads = [dict() for i in range(len(weights))]\n",
    "\n",
    "    ### START CODE HERE ###\n",
    "    \n",
    "    y0 = my_matmul(s,weights[0][\"W\"])+weights[0][\"b\"]\n",
    "    \n",
    "    relu_back = np.maximum( 0,y0 )\n",
    "    relu_back = [ np.squeeze( np.where( [relu_back>0] , 1.0, 0.0 ) ) ]\n",
    "\n",
    "    grads[0][\"W\"] = my_matmul( s.T , weights[1][\"W\"].T*relu_back )\n",
    "    grads[0][\"b\"] = weights[1][\"W\"].T*relu_back\n",
    "    grads[1][\"W\"] = np.maximum( 0,y0 ).T\n",
    "    grads[1][\"b\"] = 1\n",
    "    \n",
    "    ### END CODE HERE ###\n",
    "\n",
    "    return grads\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "d28c022a9fef03b20db192090abc644b",
     "grade": false,
     "grade_id": "cell-c6eb98da1041cedf",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Run the following code to test your implementation of the `get_gradient()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "cd8e33658e276ea4cafb6d9d0803df73",
     "grade": true,
     "grade_id": "cell-217e4e4ad7daf7d6",
     "locked": true,
     "points": 10,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "## Test Code for get_gradient() ## \n",
    "\n",
    "# Suppose num_states = 5, num_hidden_layer = 1, and num_hidden_units = 2 \n",
    "num_hidden_layer = 1\n",
    "s = np.array([[0, 0, 0, 1, 0]])\n",
    "\n",
    "weights_data = np.load(\"asserts/get_gradient_weights.npz\")\n",
    "weights = [dict() for i in range(num_hidden_layer+1)]\n",
    "weights[0][\"W\"] = weights_data[\"W0\"]\n",
    "weights[0][\"b\"] = weights_data[\"b0\"]\n",
    "weights[1][\"W\"] = weights_data[\"W1\"]\n",
    "weights[1][\"b\"] = weights_data[\"b1\"]\n",
    "\n",
    "grads = get_gradient(s, weights)\n",
    "\n",
    "grads_answer = np.load(\"asserts/get_gradient_grads.npz\")\n",
    "\n",
    "print(\"grads[0][\\\"W\\\"]\\n\", grads[0][\"W\"], \"\\n\")\n",
    "print(\"grads[0][\\\"b\\\"]\\n\", grads[0][\"b\"], \"\\n\")\n",
    "print(\"grads[1][\\\"W\\\"]\\n\", grads[1][\"W\"], \"\\n\")\n",
    "print(\"grads[1][\\\"b\\\"]\\n\", grads[1][\"b\"], \"\\n\")\n",
    "\n",
    "assert(np.allclose(grads[0][\"W\"], grads_answer[\"W0\"]))\n",
    "assert(np.allclose(grads[0][\"b\"], grads_answer[\"b0\"]))\n",
    "assert(np.allclose(grads[1][\"W\"], grads_answer[\"W1\"]))\n",
    "assert(np.allclose(grads[1][\"b\"], grads_answer[\"b1\"]))\n",
    "\n",
    "print(\"Passed the asserts!\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "7e660355c7662a6d5fde8a39a82ae258",
     "grade": false,
     "grade_id": "cell-d14e4ea8eb9e5c71",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "**Expected output**:\n",
    "\n",
    "    grads[0][\"W\"]\n",
    "     [[0.         0.        ]\n",
    "     [0.         0.        ]\n",
    "     [0.         0.        ]\n",
    "     [0.76103773 0.12167502]\n",
    "     [0.         0.        ]] \n",
    "\n",
    "    grads[0][\"b\"]\n",
    "     [[0.76103773 0.12167502]] \n",
    "\n",
    "    grads[1][\"W\"]\n",
    "     [[0.69198983]\n",
    "     [0.82403662]] \n",
    "\n",
    "    grads[1][\"b\"]\n",
    "     [[1.]] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "ff5c416e4be90cb754c2163fe04cc495",
     "grade": false,
     "grade_id": "cell-3976c881b8145ebd",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "### Implement stochastic gradient descent method for state-value prediction\n",
    "In this section, you will implement stochastic gradient descent (SGD) method for state_value prediction. Here is the basic SGD update for state-value prediction with TD:\n",
    "\n",
    "$$\\mathbf{w_{t+1}} = \\mathbf{w_{t}} + \\alpha \\delta_t \\nabla \\hat{v}(S_t,\\mathbf{w_{t}})$$\n",
    "\n",
    "At each time step, we update the weights in the direction  $g_t = \\delta_t \\nabla \\hat{v}(S_t,\\mathbf{w_t})$ using a fixed step-size $\\alpha$. $\\delta_t = R_{t+1} + \\gamma \\hat{v}(S_{t+1},\\mathbf{w_{t}}) - \\hat{v}(S_t,\\mathbf{w_t})$ is the TD-error. $\\nabla \\hat{v}(S_t,\\mathbf{w_{t}})$ is the gradient of the value function with respect to the weights.\n",
    "\n",
    "The following cell includes the SGD class. You will complete the `update_weight()` method of SGD assuming that the weights and update g are provided.\n",
    "\n",
    "**As you know, in this assignment, we structured the weights as an array of dictionaries. Note that the updates $g_t$, in the case of TD, is $\\delta_t \\nabla \\hat{v}(S_t,\\mathbf{w_t})$. As a result, $g_t$ has the same structure as $\\nabla \\hat{v}(S_t,\\mathbf{w_t})$ which is also an array of dictionaries.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "checksum": "1edadbb4d5dc91d41ffe1a25a20d14d6",
     "grade": false,
     "grade_id": "cell-ec1ae0c2c9d60a19",
     "locked": false,
     "schema_version": 1,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#GRADED FUNCTION: [SGD]\n",
    "\n",
    "class SGD(BaseOptimizer):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "    \n",
    "    def optimizer_init(self, optimizer_info):\n",
    "        \"\"\"Setup for the optimizer.\n",
    "\n",
    "        Set parameters needed to setup the stochastic gradient descent method.\n",
    "\n",
    "        Assume optimizer_info dict contains:\n",
    "        {\n",
    "            step_size: float\n",
    "        }\n",
    "        \"\"\"\n",
    "        self.step_size = optimizer_info.get(\"step_size\")\n",
    "    \n",
    "    def update_weights(self, weights, g):\n",
    "        \"\"\"\n",
    "        Given weights and update g, return updated weights\n",
    "        \"\"\"\n",
    "        for i in range(len(weights)):\n",
    "            for param in weights[i].keys():\n",
    "                \n",
    "                ### update weights (1 line)\n",
    "                # weights[i][param] = None\n",
    "                \n",
    "                ### START CODE HERE ###\n",
    "                \n",
    "                weights[i][param] = weights[i][param] + self.step_size*g[i][param]\n",
    "                \n",
    "                ### END CODE HERE ###\n",
    "                \n",
    "        return weights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "02c0dea41c93d0f3735de4850fb20778",
     "grade": false,
     "grade_id": "cell-06b56bd3ac6717af",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Run the following code to test your implementation of the `update_weights()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "3209c121f74aeb7c51a10676e3b3b5dd",
     "grade": true,
     "grade_id": "cell-dcf2cc6e7528b324",
     "locked": true,
     "points": 10,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# Do not modify this cell!\n",
    "\n",
    "## Test Code for update_weights() ## \n",
    "\n",
    "# Suppose num_states = 5, num_hidden_layer = 1, and num_hidden_units = 2 \n",
    "num_hidden_layer = 1\n",
    "\n",
    "weights_data = np.load(\"asserts/update_weights_weights.npz\")\n",
    "weights = [dict() for i in range(num_hidden_layer+1)]\n",
    "weights[0][\"W\"] = weights_data[\"W0\"]\n",
    "weights[0][\"b\"] = weights_data[\"b0\"]\n",
    "weights[1][\"W\"] = weights_data[\"W1\"]\n",
    "weights[1][\"b\"] = weights_data[\"b1\"]\n",
    "\n",
    "g_data = np.load(\"asserts/update_weights_g.npz\")\n",
    "g = [dict() for i in range(num_hidden_layer+1)]\n",
    "g[0][\"W\"] = g_data[\"W0\"]\n",
    "g[0][\"b\"] = g_data[\"b0\"]\n",
    "g[1][\"W\"] = g_data[\"W1\"]\n",
    "g[1][\"b\"] = g_data[\"b1\"]\n",
    "\n",
    "test_sgd = SGD()\n",
    "optimizer_info = {\"step_size\": 0.3}\n",
    "test_sgd.optimizer_init(optimizer_info)\n",
    "updated_weights = test_sgd.update_weights(weights, g)\n",
    "\n",
    "# updated weights asserts\n",
    "updated_weights_answer = np.load(\"asserts/update_weights_updated_weights.npz\")\n",
    "\n",
    "print(\"updated_weights[0][\\\"W\\\"]\\n\", updated_weights[0][\"W\"], \"\\n\")\n",
    "print(\"updated_weights[0][\\\"b\\\"]\\n\", updated_weights[0][\"b\"], \"\\n\")\n",
    "print(\"updated_weights[1][\\\"W\\\"]\\n\", updated_weights[1][\"W\"], \"\\n\")\n",
    "print(\"updated_weights[1][\\\"b\\\"]\\n\", updated_weights[1][\"b\"], \"\\n\")\n",
    "\n",
    "assert(np.allclose(updated_weights[0][\"W\"], updated_weights_answer[\"W0\"]))\n",
    "assert(np.allclose(updated_weights[0][\"b\"], updated_weights_answer[\"b0\"]))\n",
    "assert(np.allclose(updated_weights[1][\"W\"], updated_weights_answer[\"W1\"]))\n",
    "assert(np.allclose(updated_weights[1][\"b\"], updated_weights_answer[\"b1\"]))\n",
    "\n",
    "print(\"Passed the asserts!\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "5da852f49bbf6f5ae1aa046063066bf3",
     "grade": false,
     "grade_id": "cell-59a172bdd7aaf622",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "**Expected output**:\n",
    "\n",
    "    updated_weights[0][\"W\"]\n",
    "     [[ 1.17899492  0.53656321]\n",
    "     [ 0.58008221  1.47666572]\n",
    "     [ 1.01909411 -1.10248056]\n",
    "     [ 0.72490408  0.06828853]\n",
    "     [-0.20609725  0.69034095]] \n",
    "\n",
    "    updated_weights[0][\"b\"]\n",
    "     [[-0.18484533  0.92844539]] \n",
    "\n",
    "    updated_weights[1][\"W\"]\n",
    "     [[0.70488257]\n",
    "     [0.58150878]] \n",
    "\n",
    "    updated_weights[1][\"b\"]\n",
    "     [[0.88467086]] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "9f54dedb2d48647d810a9bdd8a1788da",
     "grade": false,
     "grade_id": "cell-29620a7c315e087d",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "### Adam Algorithm\n",
    "In this assignment, instead of using SGD for updating the weights, we use a more advanced algorithm called Adam. The Adam algorithm improves the SGD update with two concepts: adaptive vector step-sizes and momentum. It keeps estimates of the mean and second moment of the updates, denoted by $\\mathbf{m}$ and $\\mathbf{v}$ respectively:\n",
    "$$\\mathbf{m_t} = \\beta_m \\mathbf{m_{t-1}} + (1 - \\beta_m)g_t \\\\\n",
    "\\mathbf{v_t} = \\beta_v \\mathbf{v_{t-1}} + (1 - \\beta_v)g^2_t\n",
    "$$\n",
    "\n",
    "Given that $\\mathbf{m}$ and $\\mathbf{v}$ are initialized to zero, they are biased toward zero. To get unbiased estimates of the mean and second moment, Adam defines $\\mathbf{\\hat{m}}$ and $\\mathbf{\\hat{v}}$ as:\n",
    "$$ \\mathbf{\\hat{m_t}} = \\frac{\\mathbf{m_t}}{1 - \\beta_m^t} \\\\\n",
    "\\mathbf{\\hat{v_t}} = \\frac{\\mathbf{v_t}}{1 - \\beta_v^t}\n",
    "$$\n",
    "\n",
    "The weights are then updated as follows:\n",
    "$$ \\mathbf{w_t} = \\mathbf{w_{t-1}} + \\frac{\\alpha}{\\sqrt{\\mathbf{\\hat{v_t}}}+\\epsilon} \\mathbf{\\hat{m_t}}\n",
    "$$\n",
    "\n",
    "When implementing the agent you will use the Adam algorithm instead of SGD because it is more efficient. We have already provided you the implementation of the Adam algorithm in the cell below. You will use it when implementing your agent. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "bc01d14ea352fe85e325285f778a05ea",
     "grade": false,
     "grade_id": "cell-201cc0242fbec19f",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "class Adam(BaseOptimizer):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "    \n",
    "    def optimizer_init(self, optimizer_info):\n",
    "        \"\"\"Setup for the optimizer.\n",
    "\n",
    "        Set parameters needed to setup the Adam algorithm.\n",
    "\n",
    "        Assume optimizer_info dict contains:\n",
    "        {\n",
    "            num_states: integer,\n",
    "            num_hidden_layer: integer,\n",
    "            num_hidden_units: integer,\n",
    "            step_size: float, \n",
    "            self.beta_m: float\n",
    "            self.beta_v: float\n",
    "            self.epsilon: float\n",
    "        }\n",
    "        \"\"\"\n",
    "        \n",
    "        self.num_states = optimizer_info.get(\"num_states\")\n",
    "        self.num_hidden_layer = optimizer_info.get(\"num_hidden_layer\")\n",
    "        self.num_hidden_units = optimizer_info.get(\"num_hidden_units\")\n",
    "\n",
    "        # Specify Adam algorithm's hyper parameters\n",
    "        self.step_size = optimizer_info.get(\"step_size\")\n",
    "        self.beta_m = optimizer_info.get(\"beta_m\")\n",
    "        self.beta_v = optimizer_info.get(\"beta_v\")\n",
    "        self.epsilon = optimizer_info.get(\"epsilon\")\n",
    "\n",
    "        self.layer_size = np.array([self.num_states, self.num_hidden_units, 1])\n",
    "\n",
    "        # Initialize Adam algorithm's m and v\n",
    "        self.m = [dict() for i in range(self.num_hidden_layer+1)]\n",
    "        self.v = [dict() for i in range(self.num_hidden_layer+1)]\n",
    "\n",
    "        for i in range(self.num_hidden_layer+1):\n",
    "\n",
    "            # Initialize self.m[i][\"W\"], self.m[i][\"b\"], self.v[i][\"W\"], self.v[i][\"b\"] to zero\n",
    "            self.m[i][\"W\"] = np.zeros((self.layer_size[i], self.layer_size[i+1]))\n",
    "            self.m[i][\"b\"] = np.zeros((1, self.layer_size[i+1]))\n",
    "            self.v[i][\"W\"] = np.zeros((self.layer_size[i], self.layer_size[i+1]))\n",
    "            self.v[i][\"b\"] = np.zeros((1, self.layer_size[i+1]))\n",
    "\n",
    "        # Initialize beta_m_product and beta_v_product to be later used for computing m_hat and v_hat\n",
    "        self.beta_m_product = self.beta_m\n",
    "        self.beta_v_product = self.beta_v\n",
    "\n",
    "    def update_weights(self, weights, g):\n",
    "        \"\"\"\n",
    "        Given weights and update g, return updated weights\n",
    "        \"\"\"\n",
    "        \n",
    "        for i in range(len(weights)):\n",
    "            for param in weights[i].keys():\n",
    "\n",
    "                ### update self.m and self.v\n",
    "                self.m[i][param] = self.beta_m * self.m[i][param] + (1 - self.beta_m) * g[i][param]\n",
    "                self.v[i][param] = self.beta_v * self.v[i][param] + (1 - self.beta_v) * (g[i][param] * g[i][param])\n",
    "\n",
    "                ### compute m_hat and v_hat\n",
    "                m_hat = self.m[i][param] / (1 - self.beta_m_product)\n",
    "                v_hat = self.v[i][param] / (1 - self.beta_v_product)\n",
    "\n",
    "                ### update weights\n",
    "                weights[i][param] += self.step_size * m_hat / (np.sqrt(v_hat) + self.epsilon)\n",
    "                \n",
    "        ### update self.beta_m_product and self.beta_v_product\n",
    "        self.beta_m_product *= self.beta_m\n",
    "        self.beta_v_product *= self.beta_v\n",
    "        \n",
    "        return weights\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "6b51319a92f51a4bea52e1916eae6c61",
     "grade": false,
     "grade_id": "cell-93e3b1233b92a93f",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "## 1-2: Implement Agent Methods\n",
    "In this section, you will implement `agent_init()`, `agent_start()`, `agent_step()`, and `agent_end()`.\n",
    "\n",
    "In `agent_init()`, you will:\n",
    "   \n",
    "   - specify the neural network structure by filling self.layer_size with the size of the input layer, hidden layer, and output layer.\n",
    "   - initialize the network's parameters. We show the parameters as an array of dictionaries, self.weights, where each dictionary corresponds to weights from one layer to the next. Each dictionary includes $W$ and $b$. To initialize the parameters, you will use a normal distribution with mean 0 and standard deviation $\\sqrt{\\frac{2}{\\text{# input of each node}}}$. This initialization heuristic is commonly used when using ReLU gates and helps keep the output of a neuron from getting too big or too small. To initialize the network's parameters, use **self.rand_generator.normal()** which draws random samples from a normal distribution. The parameters of self.rand_generator.normal are mean of the distribution, standard deviation of the distribution, and output shape in the form of tuple of integers.\n",
    "\n",
    "\n",
    "In `agent_start()`, you will:\n",
    "   - specify self.last_state and self.last_action.\n",
    "   \n",
    "In `agent_step()` and `agent_end()`, you will:\n",
    "   - compute the TD error using $v(S_t)$ and $v(S_{t+1})$. To compute the value function for $S_t$ and $S_{t+1}$, you will get their one-hot encoding using `one_hot()` method that we provided below. You feed the one-hot encoded state number to the neural networks using `get_value()` method that you implemented above. Note that `one_hot()` method returns the one-hot encoding of a state as a numpy array of shape (1, num_states).\n",
    "   - retrieve the gradients using `get_gradient()` function that you implemented.\n",
    "   - use Adam_algorithm that we provided to update the neural network's parameters, self.weights.\n",
    "   - use `agent_policy()` method to select actions with. (only in `agent_step()`)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "383049ffcbcce7abc10d684ac0c1f3d5",
     "grade": false,
     "grade_id": "cell-b2e6a488bd1b6961",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "def one_hot(state, num_states):\n",
    "    \"\"\"\n",
    "    Given num_state and a state, return the one-hot encoding of the state\n",
    "    \"\"\"\n",
    "    # Create the one-hot encoding of state\n",
    "    # one_hot_vector is a numpy array of shape (1, num_states)\n",
    "    \n",
    "    one_hot_vector = np.zeros((1, num_states))\n",
    "    one_hot_vector[0, int((state - 1))] = 1\n",
    "    \n",
    "    return one_hot_vector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "nbgrader": {
     "checksum": "00a274397e3680c8f750785edc7a5f38",
     "grade": false,
     "grade_id": "cell-bcf71c78d138fd44",
     "locked": false,
     "schema_version": 1,
     "solution": true
    }
   },
   "outputs": [],
   "source": [
    "#GRADED FUNCTION: [Agent]\n",
    "\n",
    "class TDAgent(BaseAgent):\n",
    "    def __init__(self):\n",
    "        self.name = \"td_agent\"\n",
    "        pass\n",
    "\n",
    "    def agent_init(self, agent_info={}):\n",
    "        \"\"\"Setup for the agent called when the experiment first starts.\n",
    "\n",
    "        Set parameters needed to setup the semi-gradient TD with a Neural Network.\n",
    "\n",
    "        Assume agent_info dict contains:\n",
    "        {\n",
    "            num_states: integer,\n",
    "            num_hidden_layer: integer,\n",
    "            num_hidden_units: integer,\n",
    "            step_size: float, \n",
    "            discount_factor: float,\n",
    "            self.beta_m: float\n",
    "            self.beta_v: float\n",
    "            self.epsilon: float\n",
    "            seed: int\n",
    "        }\n",
    "        \"\"\"\n",
    "    \n",
    "        # Set random seed for weights initialization for each run\n",
    "        self.rand_generator = np.random.RandomState(agent_info.get(\"seed\")) \n",
    "        \n",
    "        # Set random seed for policy for each run\n",
    "        self.policy_rand_generator = np.random.RandomState(agent_info.get(\"seed\"))\n",
    "\n",
    "        # Set attributes according to agent_info\n",
    "        self.num_states = agent_info.get(\"num_states\")\n",
    "        self.num_hidden_layer = agent_info.get(\"num_hidden_layer\")\n",
    "        self.num_hidden_units = agent_info.get(\"num_hidden_units\")\n",
    "        self.discount_factor = agent_info.get(\"discount_factor\")\n",
    "\n",
    "        ### Define the neural network's structure (1 line)\n",
    "        # Specify self.layer_size which shows the number of nodes in each layer\n",
    "        # self.layer_size = np.array([None, None, None])\n",
    "        # Hint: Checkout the NN diagram at the beginning of the notebook\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "        \n",
    "        self.layer_size = np.array( [self.num_states,self.num_hidden_units,1] )\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        # Initialize the neural network's parameter (2 lines)\n",
    "        self.weights = [dict() for i in range(self.num_hidden_layer+1)]\n",
    "        for i in range(self.num_hidden_layer+1):\n",
    "\n",
    "            ### Initialize self.weights[i][\"W\"] and self.weights[i][\"b\"] using self.rand_generator.normal()\n",
    "            # Note that The parameters of self.rand_generator.normal are mean of the distribution, \n",
    "            # standard deviation of the distribution, and output shape in the form of tuple of integers.\n",
    "            # To specify output shape, use self.layer_size.\n",
    "\n",
    "            ### START CODE HERE ###\n",
    "            \n",
    "            self.weights[i][\"W\"] = self.rand_generator.normal( 0 , np.sqrt( 2/self.layer_size[i] ) , (self.layer_size[i],self.layer_size[i+1])  )\n",
    "            self.weights[i][\"b\"] = self.rand_generator.normal( 0 , np.sqrt( 2/self.layer_size[i] ) , (1,self.layer_size[i+1]) )\n",
    "            \n",
    "            ### END CODE HERE ###\n",
    "        \n",
    "        \n",
    "        # Specify the optimizer\n",
    "        self.optimizer = Adam()\n",
    "        optimizer_info = {\"num_states\": agent_info[\"num_states\"],\n",
    "                          \"num_hidden_layer\": agent_info[\"num_hidden_layer\"],\n",
    "                          \"num_hidden_units\": agent_info[\"num_hidden_units\"],\n",
    "                          \"step_size\": agent_info[\"step_size\"],\n",
    "                          \"beta_m\": agent_info[\"beta_m\"],\n",
    "                          \"beta_v\": agent_info[\"beta_v\"],\n",
    "                          \"epsilon\": agent_info[\"epsilon\"]}\n",
    "        self.optimizer.optimizer_init(optimizer_info)\n",
    "        \n",
    "        self.last_state = None\n",
    "        self.last_action = None\n",
    "\n",
    "    def agent_policy(self, state):\n",
    "\n",
    "        ### Set chosen_action as 0 or 1 with equal probability. \n",
    "        chosen_action = self.policy_rand_generator.choice([0,1])    \n",
    "        return chosen_action\n",
    "\n",
    "    def agent_start(self, state):\n",
    "        \"\"\"The first method called when the experiment starts, called after\n",
    "        the environment starts.\n",
    "        Args:\n",
    "            state (Numpy array): the state from the\n",
    "                environment's evn_start function.\n",
    "        Returns:\n",
    "            The first action the agent takes.\n",
    "        \"\"\"\n",
    "        ### select action given state (using self.agent_policy()), and save current state and action (2 lines)\n",
    "        # self.last_state = ?\n",
    "        # self.last_action = ?\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "        \n",
    "        self.last_state = state\n",
    "        self.last_action = self.agent_policy( state )\n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        return self.last_action\n",
    "\n",
    "    def agent_step(self, reward, state):\n",
    "        \"\"\"A step taken by the agent.\n",
    "        Args:\n",
    "            reward (float): the reward received for taking the last action taken\n",
    "            state (Numpy array): the state from the\n",
    "                environment's step based, where the agent ended up after the\n",
    "                last step\n",
    "        Returns:\n",
    "            The action the agent is taking.\n",
    "        \"\"\"\n",
    "        \n",
    "        ### Compute TD error (5 lines)\n",
    "        # delta = None\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "        \n",
    "        hot_state = one_hot(state,self.num_states)\n",
    "        last_hot_state = one_hot(self.last_state,self.num_states)\n",
    "        delta = reward + self.discount_factor*get_value(hot_state,self.weights) - get_value(last_hot_state,self.weights)\n",
    "        \n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        ### Retrieve gradients (1 line)\n",
    "        # grads = None\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "        \n",
    "        grads = get_gradient( last_hot_state, self.weights )\n",
    "        \n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        ### Compute g (1 line)\n",
    "        g = [dict() for i in range(self.num_hidden_layer+1)]\n",
    "        for i in range(self.num_hidden_layer+1):\n",
    "            for param in self.weights[i].keys():\n",
    "\n",
    "                # g[i][param] = None\n",
    "                ### START CODE HERE ###\n",
    "                g[i][param] = delta*grads[i][param]\n",
    "                ### END CODE HERE ###\n",
    "\n",
    "        ### update the weights using self.optimizer (1 line)\n",
    "        # self.weights = None\n",
    "        \n",
    "        ### START CODE HERE ###\n",
    "        \n",
    "        self.weights = self.optimizer.update_weights( self.weights , g )\n",
    "        \n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        ### update self.last_state and self.last_action (2 lines)\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "        \n",
    "        self.last_state = state\n",
    "        self.last_action = self.agent_policy( state )\n",
    "\n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        return self.last_action\n",
    "\n",
    "    def agent_end(self, reward):\n",
    "        \"\"\"Run when the agent terminates.\n",
    "        Args:\n",
    "            reward (float): the reward the agent received for entering the\n",
    "                terminal state.\n",
    "        \"\"\"\n",
    "\n",
    "        ### compute TD error (3 lines)\n",
    "        # delta = None\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "        \n",
    "        last_hot_state = one_hot(self.last_state,self.num_states)\n",
    "        delta = reward - get_value( last_hot_state ,self.weights)\n",
    "        \n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        ### Retrieve gradients (1 line)\n",
    "        # grads = None\n",
    "\n",
    "        ### START CODE HERE ###\n",
    "        \n",
    "        grads = get_gradient( last_hot_state, self.weights )\n",
    "        \n",
    "        ### END CODE HERE ###\n",
    "\n",
    "        ### Compute g (1 line)\n",
    "        g = [dict() for i in range(self.num_hidden_layer+1)]\n",
    "        for i in range(self.num_hidden_layer+1):\n",
    "            for param in self.weights[i].keys():\n",
    "\n",
    "                # g[i][param] = None\n",
    "                ### START CODE HERE ###\n",
    "                g[i][param] = delta*grads[i][param]\n",
    "                ### END CODE HERE ###\n",
    "\n",
    "        ### update the weights using self.optimizer (1 line)\n",
    "        # self.weights = None\n",
    "        \n",
    "        ### START CODE HERE ###\n",
    "        \n",
    "        self.weights = self.optimizer.update_weights( self.weights , g )\n",
    "        \n",
    "        ### END CODE HERE ###\n",
    "\n",
    "    def agent_message(self, message):\n",
    "        if message == 'get state value':\n",
    "            state_value = np.zeros(self.num_states)\n",
    "            for state in range(1, self.num_states + 1):\n",
    "                s = one_hot(state, self.num_states)\n",
    "                state_value[state - 1] = get_value(s, self.weights)\n",
    "            return state_value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "d1e29604e5409950b1b655aed4c68358",
     "grade": false,
     "grade_id": "cell-abac904ca1f57c77",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Run the following code to test your implementation of the `agent_init()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "ebabcf523ef713941511e9dffe085908",
     "grade": true,
     "grade_id": "cell-3061ed3fc479ee08",
     "locked": true,
     "points": 10,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "## Test Code for agent_init() ## \n",
    "\n",
    "agent_info = {\"num_states\": 5,\n",
    "              \"num_hidden_layer\": 1,\n",
    "              \"num_hidden_units\": 2,\n",
    "              \"step_size\": 0.25,\n",
    "              \"discount_factor\": 0.9,\n",
    "              \"beta_m\": 0.9,\n",
    "              \"beta_v\": 0.99,\n",
    "              \"epsilon\": 0.0001,\n",
    "              \"seed\": 0\n",
    "             }\n",
    "\n",
    "test_agent = TDAgent()\n",
    "test_agent.agent_init(agent_info)\n",
    "\n",
    "print(\"layer_size: {}\".format(test_agent.layer_size))\n",
    "assert(np.allclose(test_agent.layer_size, np.array([agent_info[\"num_states\"], \n",
    "                                                    agent_info[\"num_hidden_units\"], \n",
    "                                                    1])))\n",
    "\n",
    "print(\"weights[0][\\\"W\\\"] shape: {}\".format(test_agent.weights[0][\"W\"].shape))\n",
    "print(\"weights[0][\\\"b\\\"] shape: {}\".format(test_agent.weights[0][\"b\"].shape))\n",
    "print(\"weights[1][\\\"W\\\"] shape: {}\".format(test_agent.weights[1][\"W\"].shape))\n",
    "print(\"weights[1][\\\"b\\\"] shape: {}\".format(test_agent.weights[1][\"b\"].shape), \"\\n\")\n",
    "\n",
    "assert(test_agent.weights[0][\"W\"].shape == (agent_info[\"num_states\"], agent_info[\"num_hidden_units\"]))\n",
    "assert(test_agent.weights[0][\"b\"].shape == (1, agent_info[\"num_hidden_units\"]))\n",
    "assert(test_agent.weights[1][\"W\"].shape == (agent_info[\"num_hidden_units\"], 1))\n",
    "assert(test_agent.weights[1][\"b\"].shape == (1, 1))\n",
    "\n",
    "print(\"weights[0][\\\"W\\\"]\\n\", (test_agent.weights[0][\"W\"]), \"\\n\")\n",
    "print(\"weights[0][\\\"b\\\"]\\n\", (test_agent.weights[0][\"b\"]), \"\\n\")\n",
    "print(\"weights[1][\\\"W\\\"]\\n\", (test_agent.weights[1][\"W\"]), \"\\n\")\n",
    "print(\"weights[1][\\\"b\\\"]\\n\", (test_agent.weights[1][\"b\"]), \"\\n\")\n",
    "\n",
    "\n",
    "agent_weight_answer = np.load(\"asserts/agent_init_weights_1.npz\")\n",
    "assert(np.allclose(test_agent.weights[0][\"W\"], agent_weight_answer[\"W0\"]))\n",
    "assert(np.allclose(test_agent.weights[0][\"b\"], agent_weight_answer[\"b0\"]))\n",
    "assert(np.allclose(test_agent.weights[1][\"W\"], agent_weight_answer[\"W1\"]))\n",
    "assert(np.allclose(test_agent.weights[1][\"b\"], agent_weight_answer[\"b1\"]))\n",
    "\n",
    "print(\"Passed the asserts!\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "ca8f82862187a73c5e8dbeed3ea47d23",
     "grade": false,
     "grade_id": "cell-5cc06cee020ffd5e",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "**Expected output**:\n",
    "\n",
    "    layer_size: [5 2 1]\n",
    "    weights[0][\"W\"] shape: (5, 2)\n",
    "    weights[0][\"b\"] shape: (1, 2)\n",
    "    weights[1][\"W\"] shape: (2, 1)\n",
    "    weights[1][\"b\"] shape: (1, 1) \n",
    "\n",
    "    weights[0][\"W\"]\n",
    "     [[ 1.11568467  0.25308164]\n",
    "     [ 0.61900825  1.4172653 ]\n",
    "     [ 1.18114738 -0.6180848 ]\n",
    "     [ 0.60088868 -0.0957267 ]\n",
    "     [-0.06528133  0.25968529]] \n",
    "\n",
    "    weights[0][\"b\"]\n",
    "     [[0.09110115 0.91976332]] \n",
    "\n",
    "    weights[1][\"W\"]\n",
    "     [[0.76103773]\n",
    "     [0.12167502]] \n",
    "\n",
    "    weights[1][\"b\"]\n",
    "     [[0.44386323]]\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "4314e3d66b268acdd963c634c64794f1",
     "grade": false,
     "grade_id": "cell-3b7fd0c37d10e8e5",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Run the following code to test your implementation of the `agent_start()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "a34fd74bd64b42a07cbc0290a4c1c5a6",
     "grade": true,
     "grade_id": "cell-8e05d0347d000b8f",
     "locked": true,
     "points": 5,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# Do not modify this cell!\n",
    "\n",
    "## Test Code for agent_start() ## \n",
    "\n",
    "agent_info = {\"num_states\": 500,\n",
    "              \"num_hidden_layer\": 1,\n",
    "              \"num_hidden_units\": 100,\n",
    "              \"step_size\": 0.1,\n",
    "              \"discount_factor\": 1.0,\n",
    "              \"beta_m\": 0.9,\n",
    "              \"beta_v\": 0.99,\n",
    "              \"epsilon\": 0.0001,\n",
    "              \"seed\": 10\n",
    "             }\n",
    "\n",
    "# Suppose state = 250\n",
    "state = 250\n",
    "\n",
    "test_agent = TDAgent()\n",
    "test_agent.agent_init(agent_info)\n",
    "test_agent.agent_start(state)\n",
    "\n",
    "print(\"Agent state: {}\".format(test_agent.last_state))\n",
    "print(\"Agent selected action: {}\".format(test_agent.last_action))\n",
    "\n",
    "assert(test_agent.last_state == 250)\n",
    "assert(test_agent.last_action == 1)\n",
    "\n",
    "print(\"Passed the asserts!\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "f1dd9a1eb96e3c5873282e1fea477909",
     "grade": false,
     "grade_id": "cell-c88b673865769cf2",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "**Expected output**:\n",
    "\n",
    "    Agent state: 250\n",
    "    Agent selected action: 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "6a30c17eee9895c80e343ba8740a56d3",
     "grade": false,
     "grade_id": "cell-f59067a616f5df77",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Run the following code to test your implementation of the `agent_step()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "0091506a29350418401174f6e09922d7",
     "grade": true,
     "grade_id": "cell-4d3501419a656e08",
     "locked": true,
     "points": 10,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# Do not modify this cell!\n",
    "\n",
    "## Test Code for agent_step() ## \n",
    "\n",
    "agent_info = {\"num_states\": 5,\n",
    "              \"num_hidden_layer\": 1,\n",
    "              \"num_hidden_units\": 2,\n",
    "              \"step_size\": 0.1,\n",
    "              \"discount_factor\": 1.0,\n",
    "              \"beta_m\": 0.9,\n",
    "              \"beta_v\": 0.99,\n",
    "              \"epsilon\": 0.0001,\n",
    "              \"seed\": 0\n",
    "             }\n",
    "\n",
    "test_agent = TDAgent()\n",
    "test_agent.agent_init(agent_info)\n",
    "\n",
    "# load initial weights\n",
    "agent_initial_weight = np.load(\"asserts/agent_step_initial_weights.npz\")\n",
    "test_agent.weights[0][\"W\"] = agent_initial_weight[\"W0\"]\n",
    "test_agent.weights[0][\"b\"] = agent_initial_weight[\"b0\"]\n",
    "test_agent.weights[1][\"W\"] = agent_initial_weight[\"W1\"]\n",
    "test_agent.weights[1][\"b\"] = agent_initial_weight[\"b1\"]\n",
    "\n",
    "# load m and v for the optimizer\n",
    "m_data = np.load(\"asserts/agent_step_initial_m.npz\")\n",
    "test_agent.optimizer.m[0][\"W\"] = m_data[\"W0\"]\n",
    "test_agent.optimizer.m[0][\"b\"] = m_data[\"b0\"]\n",
    "test_agent.optimizer.m[1][\"W\"] = m_data[\"W1\"]\n",
    "test_agent.optimizer.m[1][\"b\"] = m_data[\"b1\"]\n",
    "\n",
    "v_data = np.load(\"asserts/agent_step_initial_v.npz\")\n",
    "test_agent.optimizer.v[0][\"W\"] = v_data[\"W0\"]\n",
    "test_agent.optimizer.v[0][\"b\"] = v_data[\"b0\"]\n",
    "test_agent.optimizer.v[1][\"W\"] = v_data[\"W1\"]\n",
    "test_agent.optimizer.v[1][\"b\"] = v_data[\"b1\"]\n",
    "\n",
    "# Assume the agent started at State 3\n",
    "start_state = 3\n",
    "test_agent.agent_start(start_state)\n",
    "\n",
    "# Assume the reward was 10.0 and the next state observed was State 1\n",
    "reward = 10.0\n",
    "next_state = 1\n",
    "test_agent.agent_step(reward, next_state)\n",
    "\n",
    "# updated weights asserts\n",
    "print(\"updated_weights[0][\\\"W\\\"]\\n\", test_agent.weights[0][\"W\"], \"\\n\")\n",
    "print(\"updated_weights[0][\\\"b\\\"]\\n\", test_agent.weights[0][\"b\"], \"\\n\")\n",
    "print(\"updated_weights[1][\\\"W\\\"]\\n\", test_agent.weights[1][\"W\"], \"\\n\")\n",
    "print(\"updated_weights[1][\\\"b\\\"]\\n\", test_agent.weights[1][\"b\"], \"\\n\")\n",
    "\n",
    "agent_updated_weight_answer = np.load(\"asserts/agent_step_updated_weights.npz\")\n",
    "assert(np.allclose(test_agent.weights[0][\"W\"], agent_updated_weight_answer[\"W0\"]))\n",
    "assert(np.allclose(test_agent.weights[0][\"b\"], agent_updated_weight_answer[\"b0\"]))\n",
    "assert(np.allclose(test_agent.weights[1][\"W\"], agent_updated_weight_answer[\"W1\"]))\n",
    "assert(np.allclose(test_agent.weights[1][\"b\"], agent_updated_weight_answer[\"b1\"]))\n",
    "\n",
    "# last_state and last_action assert\n",
    "print(\"Agent last state:\", test_agent.last_state)\n",
    "print(\"Agent last action:\", test_agent.last_action, \"\\n\")\n",
    "\n",
    "assert(test_agent.last_state == 1)\n",
    "assert(test_agent.last_action == 1)\n",
    "\n",
    "print (\"Passed the asserts!\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "bfc86292e1cc7edb9cfe97addd6aacff",
     "grade": false,
     "grade_id": "cell-7063c509f6f48ff0",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "**Expected output**:\n",
    "\n",
    "    updated_weights[0][\"W\"]\n",
    "     [[ 1.10893459  0.30763738]\n",
    "     [ 0.63690565  1.14778865]\n",
    "     [ 1.23397791 -0.48152743]\n",
    "     [ 0.72792093 -0.15829832]\n",
    "     [ 0.15021996  0.39822163]] \n",
    "\n",
    "    updated_weights[0][\"b\"]\n",
    "     [[0.29798822 0.96254535]] \n",
    "\n",
    "    updated_weights[1][\"W\"]\n",
    "     [[0.76628754]\n",
    "     [0.11486511]] \n",
    "\n",
    "    updated_weights[1][\"b\"]\n",
    "     [[0.58530057]] \n",
    "\n",
    "    Agent last state: 1\n",
    "    Agent last action: 1 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "5ba003cd2b9ee7270177c75652db5d07",
     "grade": false,
     "grade_id": "cell-d734c2a22fa11660",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Run the following code to test your implementation of the `agent_end()` function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "bd232cf2650f88080cdfc39a0d6732c1",
     "grade": true,
     "grade_id": "cell-a8031e265fdf211d",
     "locked": true,
     "points": 5,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# Do not modify this cell!\n",
    "\n",
    "## Test Code for agent_end() ## \n",
    "\n",
    "agent_info = {\"num_states\": 5,\n",
    "              \"num_hidden_layer\": 1,\n",
    "              \"num_hidden_units\": 2,\n",
    "              \"step_size\": 0.1,\n",
    "              \"discount_factor\": 1.0,\n",
    "              \"beta_m\": 0.9,\n",
    "              \"beta_v\": 0.99,\n",
    "              \"epsilon\": 0.0001,\n",
    "              \"seed\": 0\n",
    "             }\n",
    "\n",
    "test_agent = TDAgent()\n",
    "test_agent.agent_init(agent_info)\n",
    "\n",
    "# load initial weights\n",
    "agent_initial_weight = np.load(\"asserts/agent_end_initial_weights.npz\")\n",
    "test_agent.weights[0][\"W\"] = agent_initial_weight[\"W0\"]\n",
    "test_agent.weights[0][\"b\"] = agent_initial_weight[\"b0\"]\n",
    "test_agent.weights[1][\"W\"] = agent_initial_weight[\"W1\"]\n",
    "test_agent.weights[1][\"b\"] = agent_initial_weight[\"b1\"]\n",
    "\n",
    "# load m and v for the optimizer\n",
    "m_data = np.load(\"asserts/agent_step_initial_m.npz\")\n",
    "test_agent.optimizer.m[0][\"W\"] = m_data[\"W0\"]\n",
    "test_agent.optimizer.m[0][\"b\"] = m_data[\"b0\"]\n",
    "test_agent.optimizer.m[1][\"W\"] = m_data[\"W1\"]\n",
    "test_agent.optimizer.m[1][\"b\"] = m_data[\"b1\"]\n",
    "\n",
    "v_data = np.load(\"asserts/agent_step_initial_v.npz\")\n",
    "test_agent.optimizer.v[0][\"W\"] = v_data[\"W0\"]\n",
    "test_agent.optimizer.v[0][\"b\"] = v_data[\"b0\"]\n",
    "test_agent.optimizer.v[1][\"W\"] = v_data[\"W1\"]\n",
    "test_agent.optimizer.v[1][\"b\"] = v_data[\"b1\"]\n",
    "\n",
    "# Assume the agent started at State 4\n",
    "start_state = 4\n",
    "test_agent.agent_start(start_state)\n",
    "\n",
    "# Assume the reward was 10.0 and reached the terminal state\n",
    "reward = 10.0\n",
    "test_agent.agent_end(reward)\n",
    "\n",
    "# updated weights asserts\n",
    "print(\"updated_weights[0][\\\"W\\\"]\\n\", test_agent.weights[0][\"W\"], \"\\n\")\n",
    "print(\"updated_weights[0][\\\"b\\\"]\\n\", test_agent.weights[0][\"b\"], \"\\n\")\n",
    "print(\"updated_weights[1][\\\"W\\\"]\\n\", test_agent.weights[1][\"W\"], \"\\n\")\n",
    "print(\"updated_weights[1][\\\"b\\\"]\\n\", test_agent.weights[1][\"b\"], \"\\n\")\n",
    "\n",
    "agent_updated_weight_answer = np.load(\"asserts/agent_end_updated_weights.npz\")\n",
    "assert(np.allclose(test_agent.weights[0][\"W\"], agent_updated_weight_answer[\"W0\"]))\n",
    "assert(np.allclose(test_agent.weights[0][\"b\"], agent_updated_weight_answer[\"b0\"]))\n",
    "assert(np.allclose(test_agent.weights[1][\"W\"], agent_updated_weight_answer[\"W1\"]))\n",
    "assert(np.allclose(test_agent.weights[1][\"b\"], agent_updated_weight_answer[\"b1\"]))\n",
    "\n",
    "print (\"Passed the asserts!\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "dd572a2b2edac4a1a09af8451f1ecc80",
     "grade": false,
     "grade_id": "cell-cf7301495f685eaf",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "**Expected output:**\n",
    "\n",
    "    updated_weights[0][\"W\"]\n",
    "     [[ 1.10893459  0.30763738]\n",
    "     [ 0.63690565  1.14778865]\n",
    "     [ 1.17531054 -0.51043162]\n",
    "     [ 0.75062903 -0.13736817]\n",
    "     [ 0.15021996  0.39822163]] \n",
    "\n",
    "    updated_weights[0][\"b\"]\n",
    "     [[0.30846523 0.95937346]] \n",
    "\n",
    "    updated_weights[1][\"W\"]\n",
    "     [[0.68861703]\n",
    "     [0.15986364]] \n",
    "\n",
    "    updated_weights[1][\"b\"]\n",
    "     [[0.586074]] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "d4a06a16ec5c30a856e81579136e369b",
     "grade": false,
     "grade_id": "cell-f8442d4f392c0918",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "## Section 2 - Run Experiment\n",
    "\n",
    "Now that you implemented the agent, we can run the experiment. Similar to Course 3 Programming Assignment 1, we will plot the learned state value function and the learning curve of the TD agent. To plot the learning curve, we use Root Mean Squared Value Error (RMSVE). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "9fcc43fdb0c69e8c5ac96f8a13fc07ed",
     "grade": false,
     "grade_id": "cell-137829b5646a5222",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "## 2-1: Run Experiment for Semi-gradient TD with a Neural Network\n",
    "\n",
    "We have already provided you the experiment/plot code, so you can go ahead and run the two cells below.\n",
    "\n",
    "Note that running the cell below will take **approximately 12 minutes**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "05de82f9ea256714e4af3bb2a08e8870",
     "grade": false,
     "grade_id": "cell-2a279281a516eb2c",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "outputs": [],
   "source": [
    "# Do not modify this cell!\n",
    "\n",
    "true_state_val = np.load('data/true_V.npy')    \n",
    "state_distribution = np.load('data/state_distribution.npy')\n",
    "\n",
    "def calc_RMSVE(learned_state_val):\n",
    "    assert(len(true_state_val) == len(learned_state_val) == len(state_distribution))\n",
    "    MSVE = np.sum(np.multiply(state_distribution, np.square(true_state_val - learned_state_val)))\n",
    "    RMSVE = np.sqrt(MSVE)\n",
    "    return RMSVE\n",
    "\n",
    "# Define function to run experiment\n",
    "def run_experiment(environment, agent, environment_parameters, agent_parameters, experiment_parameters):\n",
    "    \n",
    "    rl_glue = RLGlue(environment, agent)\n",
    "        \n",
    "    # save rmsve at the end of each episode\n",
    "    agent_rmsve = np.zeros((experiment_parameters[\"num_runs\"], \n",
    "                            int(experiment_parameters[\"num_episodes\"]/experiment_parameters[\"episode_eval_frequency\"]) + 1))\n",
    "    \n",
    "    # save learned state value at the end of each run\n",
    "    agent_state_val = np.zeros((experiment_parameters[\"num_runs\"], \n",
    "                                environment_parameters[\"num_states\"]))\n",
    "\n",
    "    env_info = {\"num_states\": environment_parameters[\"num_states\"],\n",
    "                \"start_state\": environment_parameters[\"start_state\"],\n",
    "                \"left_terminal_state\": environment_parameters[\"left_terminal_state\"],\n",
    "                \"right_terminal_state\": environment_parameters[\"right_terminal_state\"]}\n",
    "\n",
    "    agent_info = {\"num_states\": environment_parameters[\"num_states\"],\n",
    "                  \"num_hidden_layer\": agent_parameters[\"num_hidden_layer\"],\n",
    "                  \"num_hidden_units\": agent_parameters[\"num_hidden_units\"],\n",
    "                  \"step_size\": agent_parameters[\"step_size\"],\n",
    "                  \"discount_factor\": environment_parameters[\"discount_factor\"],\n",
    "                  \"beta_m\": agent_parameters[\"beta_m\"],\n",
    "                  \"beta_v\": agent_parameters[\"beta_v\"],\n",
    "                  \"epsilon\": agent_parameters[\"epsilon\"]\n",
    "                 }\n",
    "    \n",
    "    print('Setting - Neural Network with 100 hidden units')\n",
    "    os.system('sleep 1')\n",
    "\n",
    "    # one agent setting\n",
    "    for run in tqdm(range(1, experiment_parameters[\"num_runs\"]+1)):\n",
    "        env_info[\"seed\"] = run\n",
    "        agent_info[\"seed\"] = run\n",
    "        rl_glue.rl_init(agent_info, env_info)\n",
    "        \n",
    "        # Compute initial RMSVE before training\n",
    "        current_V = rl_glue.rl_agent_message(\"get state value\")\n",
    "        agent_rmsve[run-1, 0] = calc_RMSVE(current_V)\n",
    "        \n",
    "        for episode in range(1, experiment_parameters[\"num_episodes\"]+1):\n",
    "            # run episode\n",
    "            rl_glue.rl_episode(0) # no step limit\n",
    "\n",
    "            if episode % experiment_parameters[\"episode_eval_frequency\"] == 0:\n",
    "                current_V = rl_glue.rl_agent_message(\"get state value\")\n",
    "                agent_rmsve[run-1, int(episode/experiment_parameters[\"episode_eval_frequency\"])] = calc_RMSVE(current_V)\n",
    "            elif episode == experiment_parameters[\"num_episodes\"]: # if last episode\n",
    "                current_V = rl_glue.rl_agent_message(\"get state value\")\n",
    "\n",
    "        agent_state_val[run-1, :] = current_V\n",
    "\n",
    "    save_name = \"{}\".format(rl_glue.agent.name).replace('.','')\n",
    "    \n",
    "    if not os.path.exists('results'):\n",
    "                os.makedirs('results')\n",
    "    \n",
    "    # save avg. state value\n",
    "    np.save(\"results/V_{}\".format(save_name), agent_state_val)\n",
    "\n",
    "    # save avg. rmsve\n",
    "    np.savez(\"results/RMSVE_{}\".format(save_name), rmsve = agent_rmsve,\n",
    "                                                   eval_freq = experiment_parameters[\"episode_eval_frequency\"],\n",
    "                                                   num_episodes = experiment_parameters[\"num_episodes\"])\n",
    "\n",
    "\n",
    "# Run Experiment\n",
    "\n",
    "# Experiment parameters\n",
    "experiment_parameters = {\n",
    "    \"num_runs\" : 20,\n",
    "    \"num_episodes\" : 1000,\n",
    "    \"episode_eval_frequency\" : 10 # evaluate every 10 episode\n",
    "}\n",
    "\n",
    "# Environment parameters\n",
    "environment_parameters = {\n",
    "    \"num_states\" : 500,\n",
    "    \"start_state\" : 250,\n",
    "    \"left_terminal_state\" : 0,\n",
    "    \"right_terminal_state\" : 501,\n",
    "    \"discount_factor\" : 1.0\n",
    "}\n",
    "\n",
    "# Agent parameters\n",
    "agent_parameters = {\n",
    "    \"num_hidden_layer\": 1,\n",
    "    \"num_hidden_units\": 100,\n",
    "    \"step_size\": 0.001,\n",
    "    \"beta_m\": 0.9,\n",
    "    \"beta_v\": 0.999,\n",
    "    \"epsilon\": 0.0001,\n",
    "}\n",
    "\n",
    "current_env = RandomWalkEnvironment\n",
    "current_agent = TDAgent\n",
    "\n",
    "# run experiment\n",
    "run_experiment(current_env, current_agent, environment_parameters, agent_parameters, experiment_parameters)\n",
    "\n",
    "# plot result\n",
    "plot_script.plot_result([\"td_agent\"])\n",
    "\n",
    "shutil.make_archive('results', 'zip', 'results')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "0db92bf291ef2988b3f6beb6d1941fad",
     "grade": false,
     "grade_id": "cell-bea11a6380a836e3",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "You plotted the learning curve for 1000 episodes. As you can see the RMSVE is still decreasing. Here we provide the pre-computed result for 5000 episodes and 20 runs so that you can see the performance of semi-gradient TD with a neural network after being trained for a long time.\n",
    "\n",
    "![](nn_5000_episodes.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "46d058a429acfa3122e14091d263670f",
     "grade": false,
     "grade_id": "cell-95a7ccf1f6e2737a",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "Does semi-gradient TD with a neural network find a good approximation within 5000 episodes? \n",
    "\n",
    "As you may remember from the previous assignment, semi-gradient TD with 10-state aggregation converged within 100 episodes. Why is TD with a neural network slower?\n",
    "\n",
    "Would it be faster if we decrease the number of hidden units? Or what about if we increase the number of hidden units?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "17e4631fea0f972c3204ed058931e0bc",
     "grade": false,
     "grade_id": "cell-f2aec56a76cc3b45",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "## 2-2: Compare Performance of Semi-gradient TD with a Neural Network and Semi-gradient TD with Tile-coding\n",
    "\n",
    "In this section, we compare the performance of semi-gradient TD with a Neural Network and semi-gradient TD with tile-coding. Tile-coding is a kind of coarse coding that uses multiple overlapping partitions of the state space to produce features. For tile-coding, we used 50 tilings each with 6 tiles. We set the step-size for semi-gradient TD with tile-coding to $\\frac{0.1}{\\text{# tilings}}$. See the figure below for the comparison between semi-gradient TD with tile-coding and semi-gradient TD with a neural network and Adam algorithm. This result is for 5000 episodes and 20 runs:\n",
    "![](nn_vs_tc.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "a499a96e2a5d29d0fe94e360ad2b9947",
     "grade": false,
     "grade_id": "cell-100dd59ca8cf87fd",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "How are the results?\n",
    "\n",
    "Semi-gradient TD with tile-coding is much faster than semi-gradient TD with a neural network. Why?\n",
    "\n",
    "Which method has a lower RMSVE at the end of 5000 episodes?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": false,
    "editable": false,
    "nbgrader": {
     "checksum": "f0df5ef7680decccecfb89bd404d0e19",
     "grade": false,
     "grade_id": "cell-cc391c0e03ced2bb",
     "locked": true,
     "schema_version": 1,
     "solution": false
    }
   },
   "source": [
    "### Wrapping up! \n",
    "\n",
    "You have successfully implemented Course 3 Programming Assignment 2.\n",
    "\n",
    "You have implemented **semi-gradient TD with a Neural Network and Adam algorithm** in 500-state Random Walk. \n",
    "\n",
    "You also compared semi-gradient TD with a neural network and semi-gradient TD with tile-coding. \n",
    "\n",
    "From the experiments and lectures, you should be more familiar with some of the strengths and weaknesses of using neural networks as the function approximator for an RL agent. On one hand, neural networks are powerful function approximators capable of representing a wide class of functions. They are also capable of producing features without exclusively relying on hand-crafted mechanisms. On the other hand, compared to a linear function approximator with tile-coding, neural networks can be less sample efficient. When implementing your own Reinforcement Learning agents, you may consider these strengths and weaknesses to choose the proper function approximator for your problems.\n",
    "\n",
    "---\n",
    "\n",
    "**Note**: Apart from using the 'Submit' button in the notebook, you have to submit an additional `zip` file containing the 'npy' files that were generated from running the experiment cells. To do so:\n",
    "\n",
    "1. Generate the zip file by running the experiment cells in the notebook. On the top of the notebook, navigate to 'File->Open' to open the directory view of this assignment. Select \"results.zip\" and click on \"Download\". Alternatively, you can download just the results folder and run \"zip -jr results.zip results/\" (The flag 'j' is required by the grader!).\n",
    "2. Go to the \"My submission\" tab and click on \"+ Create submission\".\n",
    "3. Click on \"C3M2 Data-file Grader\" and upload your `results.zip`.\n",
    "\n",
    "**These account for 40% of the marks, so don't forget to do so!**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "coursera": {
   "course_slug": "prediction-control-function-approximation",
   "graded_item_id": "ZJrJN",
   "launcher_item_id": "jSYQa"
  },
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}