Policy_Gradients_REINFORCE.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Policy-Gradients with the REINFORCE algorithm\n",
    "\n",
    "**Overview**:\n",
    "In this practical we will train an agent using the REINFORCE algorithm to learn to balance a pole in the OpenAI gym [Cartpole environment](https://gym.openai.com/envs/CartPole-v1).\n",
    "\n",
    "**Learning objectives**:\n",
    "* Understand the policy-gradient approach to directly training a parameterised policy to maximise expected future rewards.\n",
    "* Understand how the policy-gradient theorem allows us to improve the policy using on-policy.\n",
    "\n",
    "**What is expected of you**:\n",
    " * Go through the explanation, keeping the above learning objectives in mind.\n",
    " * Fill in the missing code (\"#IMPLEMENT-ME\") and train a model to solve the Cartpole-v1 environment in OpenAI gym (you solve it when reward=500)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# A Simple Policy-Gradient Cartpole Example\n",
    "\n",
    "## Introduction\n",
    "\n",
    "We have seen in your course that there are many different approaches to training RL agents. In this practical we will take a look at REINFORCE (Monte-Carlo) - a simple policy-based method. REINFORCE (and policy-based methods in general) directly optimise a parametrised policy in order to maximise future rewards.\n",
    "\n",
    "We will try to learn a policy $\\pi_\\theta(a | s)$ which outputs a distribution over the possible actions $a$, given the current state $s$ of the environment. The goal is find a set of parameters $\\theta$ to maximise the expected discounted return:\n",
    "\\begin{align}\n",
    "J(\\theta) = \\mathbb{E}_{\\tau \\sim p_\\theta} \\left[\\sum_{t=0}^T \\gamma^t r(s_t, a_t)\\right],\n",
    "\\end{align}\n",
    "where $\\tau$ is a trajectory sampled from $p_\\theta$. The **policy-gradient** theorem gives us the derivative of this objective function:\n",
    "\\begin{align}\n",
    "    \\nabla_\\theta J(\\theta) = \\mathbb{E}_{\\tau \\sim p_\\theta} \\left[\\left(\\sum_{t=0}^{T} \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t)\\right) \\left(\\sum_{t=0}^T \\gamma^t r(s_t, a_t) \\right) \\right].\n",
    "\\end{align}\n",
    "\n",
    "**NOTE**: \n",
    "* We have a policy $\\pi_\\theta(a|s)$ which tells the agent which action $a$ to take, given the state $s$, and it is parameterised in terms of parameters $\\theta$.\n",
    "* Our goal is to maximise $J(\\theta)$ by **choosing actions from this policy** that lead to high future rewards.\n",
    "* We'll use gradient-based optimisation to update the policy parameters $\\theta$. We therefore want the gradient of our objective w.r.t. (with respect to) the policy parameters.\n",
    "* We use the policy-gradient theorem to find a expression for the gradient. This is an expectation over trajectories from our policy and the environment.\n",
    "* Since we can now sample trajectories $(s_0, a_0, r_1, s_1, a_1, r_2, \\ldots)$ using our policy $\\pi_\\theta$, we can approximate this gradient using **[Monte-Carlo](https://en.wikipedia.org/wiki/Monte_Carlo_integration)** methods.\n",
    "\n",
    "This algorithm is called **Monte-Carlo REINFORCE**, and is one type of policy-gradient algorithm. Let's use this to solve the Cartpole environment!\n",
    "\n",
    "**Monte-Carlo REINFORCE**:\n",
    "\n",
    "for each episode:\n",
    "1. sample a trajectory $\\tau$ using the policy $\\pi_\\theta$.\n",
    "2. compute $\\nabla_\\theta J(\\theta) \\approx \\left(\\sum_{t=0}^{T} \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t)\\right) \\left(\\sum_{t=0}^T \\gamma^t r(s_t, a_t) \\right)$.\n",
    "3. update policy parameters $\\theta \\leftarrow \\theta + \\alpha \\nabla_\\theta J(\\theta)$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install Pytorch\n",
    "Website: https://pytorch.org/\n",
    "\n",
    "Command without cuda (only cpu)\n",
    "\n",
    "    pip install torch==1.3.1+cpu torchvision==0.4.2+cpu -f https://download.pytorch.org/whl/torch_stable.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import various packages\n",
    "from collections import deque\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import gym\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.distributions import Categorical"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use gpu if available\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print('device:', device)\n",
    "\n",
    "# configure matplotlib\n",
    "%matplotlib inline\n",
    "plt.rcParams['figure.figsize'] = (15.0, 10.0) # set default size of plots\n",
    "plt.rcParams['image.interpolation'] = 'nearest'\n",
    "plt.rcParams['image.cmap'] = 'gray'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The Environment - Cart Pole\n",
    "\n",
    "Cartpole is a standard benchmark in reinforcement learning and is a good sandbox for trying things out. The goal is to balance a pendulum on top of a moving cart. We have 2 actions - either push the cart to the left or push to the right. The state space consists of the cart's position and velocity and the pendulum's angle and angular velocity. For more details refer to the OpenAi/gym github wiki page ([link](https://github.com/openai/gym/wiki/CartPole-v0)) \n",
    "\n",
    "Let's create the environment and take a look at the state and action spaces."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('CartPole-v1')\n",
    "print('Environment:', 'CartPole-v1')\n",
    "print('\\t','action space:', env.action_space)\n",
    "print('\\t','observation space:', env.observation_space)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here we can see that there are 2 discrete actions and a **continuous** state space.\n",
    "\n",
    "**Observation**\n",
    "Type: Box(4)  \n",
    "*(The 'Box' type is continuous, with a value between Min and Max)*\n",
    "\n",
    "Num | Observation | Min | Max\n",
    "---|---|---|---\n",
    "0 | Cart Position | -2.4 | 2.4\n",
    "1 | Cart Velocity | -Inf | Inf\n",
    "2 | Pole Angle | ~ -41.8&deg; | ~ 41.8&deg;\n",
    "3 | Pole Velocity At Tip | -Inf | Inf\n",
    "\n",
    "**Actions**\n",
    "Type: Discrete(2)\n",
    "\n",
    "Num | Action\n",
    "--- | ---\n",
    "0 | Push cart to the left\n",
    "1 | Push cart to the right"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Taking a few steps\n",
    "\n",
    "To get a better feel for the environment, we will use a random policy to genrate a short trajectory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SUB = str.maketrans(\"t0123456789+\", \"ₜ₀₁₂₃₄₅₆₇₈₉₊\")\n",
    "\n",
    "state = env.reset()\n",
    "for i in range (5):\n",
    "    # sample random action\n",
    "    action = env.action_space.sample()\n",
    "    # take action in the environment\n",
    "    new_state, reward, done, _ = env.step(action)\n",
    "    \n",
    "    print('Step t=', i+1, ': (', '[st], at , rt, [st+1]'.translate(SUB),')')\n",
    "    print('(', state, ',', action, ',', reward, ',', new_state, ')')\n",
    "    print(\"\\n \\n\")\n",
    "    \n",
    "    state = new_state"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Watching a random policy agent play\n",
    "\n",
    "Let's also see how a random policy performs in this enviroment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_1 = gym.make('CartPole-v1')\n",
    "state = env_1.reset()\n",
    "for t in range(200):\n",
    "    # sample a random action\n",
    "    action = env_1.action_space.sample()\n",
    "    env_1.render()\n",
    "    state, reward, done, _ = env_1.step(action)\n",
    "env_1.close()\n",
    "del env_1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not very good! The pole only stayed up for a few time steps... \n",
    "\n",
    "Now let's improve this using REINFORCE.\n",
    "\n",
    "## The Policy\n",
    "\n",
    "We begin by parameterising the policy $\\pi_\\theta(a | s)$ as a simple neural network which takes the state (a vector of 4 elements provided by `gym`) as input, and produces a Categorical distribution over the possible actions as output. Simple enough. Refer to [torch.nn](https://pytorch.org/docs/stable/nn.html#module)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Policy(nn.Module):\n",
    "    def __init__(self, s_size=4, h_size=16, a_size=2):\n",
    "        super(Policy, self).__init__()\n",
    "        # IMPLEMENT-ME \n",
    "        # Define neural network layers. Refer to nn.Linear (https://pytorch.org/docs/stable/nn.html#torch.nn.Linear)\n",
    "        # We are going to use a neural network with one hidden layer of size 16.\n",
    "        # The first layer should have an input size of env.observation_space.shape and an output size of 16\n",
    "        self.fc1 = ...\n",
    "        # The second layer should have an input size of 16 and an output size of env.action_space.n\n",
    "        self.fc2 = ...\n",
    "\n",
    "    def forward(self, x):\n",
    "        # IMPLEMENT-ME \n",
    "        # Implement the forward pass\n",
    "        # apply a ReLU activation after the first linear layer\n",
    "        x = ...\n",
    "        # apply the second linear layer (without an activation).\n",
    "        # the outputs of the second layer will act as the log probabilities for the Categorial distribution.\n",
    "        x = ...\n",
    "        return Categorical(logits=x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Selecting actions with our policy\n",
    "\n",
    "For a given state our policy returns a pytorch `Categorial` object. We can sample from this distribution by calling it's `sample` method and we can find the log probability of an action using `log_prob`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = Policy().to(device) # .to(device) sends object to the device(GPU) if present  \n",
    "state = env.reset()\n",
    "# convert state (a numpy array) to a torch tensor\n",
    "state = torch.from_numpy(state).float().to(device)\n",
    "dist = policy(state) # distribution = ...\n",
    "action = dist.sample()\n",
    "\n",
    "print(\"Sampled action: \", action.item())\n",
    "print(\"Log probability of action: \", dist.log_prob(action).item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing the return\n",
    "\n",
    "Given a sequence of rewards $(r(s_0, a_0), \\ldots, r(s_T, a_T))$ we want to calculate the return,  $G = \\sum_{t=0}^T \\gamma^t r(s_t, a_t)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_returns(rewards, gamma):\n",
    "    # IMPLEMENT-ME\n",
    "    # compute the return using the above equation\n",
    "    return returns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## REINFORCE\n",
    "\n",
    "Now its time to implement the algorithm\n",
    "\n",
    "**Monte-Carlo REINFORCE**:\n",
    "\n",
    "for each episode:\n",
    "1. sample a trajectory $\\tau$ using the policy $\\pi_\\theta$.\n",
    "2. compute $\\nabla_\\theta J(\\theta) \\approx \\left(\\sum_{t=0}^{T} \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t)\\right) \\left(\\sum_{t=0}^T \\gamma^t r(s_t, a_t) \\right)$.\n",
    "3. update policy parameters $\\theta \\leftarrow \\theta + \\alpha \\nabla_\\theta J(\\theta)$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparameters \n",
    "learning_rate = 1e-2\n",
    "number_episodes = 1500\n",
    "max_episode_length = 1000\n",
    "gamma = 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reinforce(env, seed, verbose=True):\n",
    "    # set random seeds (for reproducibility)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    env.seed(seed)\n",
    "\n",
    "    # instantiate the policy and optimizer\n",
    "    policy = Policy().to(device)\n",
    "    optimizer = optim.Adam(policy.parameters(), lr=learning_rate)\n",
    "\n",
    "    scores = []\n",
    "    scores_deque = deque(maxlen=100) # ? \n",
    "    for episode in range(1, number_episodes+1):\n",
    "        #################################################################\n",
    "        # 1. Collect trajectories using our policy and save the rewards #\n",
    "        # and the log probability of each action taken.                 #\n",
    "        #################################################################\n",
    "        log_probs = []\n",
    "        rewards = []\n",
    "        state = env.reset()\n",
    "        for t in range(max_episode_length):\n",
    "            # IMPLEMENT-ME: get the distribution over actions for state\n",
    "            dist = ...\n",
    "\n",
    "            # IMPLEMENT-ME: sample an action from the distribution\n",
    "            action = ...\n",
    "\n",
    "            # IMPLEMENT-ME: compute the log probability\n",
    "            log_prob = ...\n",
    "\n",
    "            # IMPLEMENT-ME: take a step in the environment\n",
    "            state, reward, done, _ = ...\n",
    "\n",
    "            # save the reward and log probability\n",
    "            rewards.append(reward)\n",
    "            log_probs.append(log_prob.unsqueeze(0))\n",
    "\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "        # for reporting save the score\n",
    "        scores.append(sum(rewards))\n",
    "        scores_deque.append(sum(rewards))\n",
    "\n",
    "        #################################################################\n",
    "        # 2. evaluate the policy gradient                               #\n",
    "        #################################################################\n",
    "        \n",
    "        # IMPLEMENT-ME: calculate the discounted return of the trajectory\n",
    "        returns = ...\n",
    "\n",
    "        log_probs = torch.cat(log_probs)\n",
    "        # IMPLEMENT-ME: multiply the log probabilities by the returns and sum (see the policy-gradient theorem)\n",
    "        # Remember to multiply the result by -1 because we want to maximise the returns\n",
    "        policy_loss = ...\n",
    "\n",
    "        #################################################################\n",
    "        # 3. update the policy parameters (gradient descent)            #\n",
    "        #################################################################\n",
    "        optimizer.zero_grad()\n",
    "        policy_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # report the score to check that we're making progress\n",
    "        if episode % 50 == 0 and verbose:\n",
    "            print('Episode {}\\tAverage Score: {:.2f}'.format(episode, np.mean(scores_deque)))\n",
    "\n",
    "        if np.mean(scores_deque) >= 495.0 and verbose:\n",
    "            print('Environment solved in {:d} episodes!\\tAverage Score: {:.2f}'.format(episode, np.mean(scores_deque)))\n",
    "            break\n",
    "            \n",
    "    return policy, scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('CartPole-v1')\n",
    "policy, scores = reinforce(env, 42)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Seeing our learned policy in action\n",
    "\n",
    "Let's watch our learned policy balance the pole!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_1 = gym.make('CartPole-v1')\n",
    "state = env_1.reset()\n",
    "for t in range(2000):\n",
    "    dist = policy(torch.from_numpy(state).float().to(device))\n",
    "    action = dist.sample()\n",
    "    env_1.render()\n",
    "    state, reward, done, _ = env_1.step(action.item())\n",
    "    if done:\n",
    "        break\n",
    "env_1.close()\n",
    "del env_1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting the results\n",
    "\n",
    "Finally, let's plot the learning curve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def moving_average(a, n) :\n",
    "    ret = np.cumsum(a, dtype=float)\n",
    "    ret[n:] = ret[n:] - ret[:-n]\n",
    "    return ret / n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "x = np.arange(1, len(scores)+1)\n",
    "ax.plot(x, scores, label='Score')\n",
    "m_average = moving_average(scores, 50)\n",
    "ax.plot(x, m_average, label='Moving Average (w=100)', linestyle='--')\n",
    "plt.legend()\n",
    "plt.ylabel('Score')\n",
    "plt.xlabel('Episode #')\n",
    "plt.title('REINFORCE learning curve - CartPole-v1')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is what your graph should look like.\n",
    "![REINFORCE learning curve](https://i.imgur.com/GjjL4C1.png \"REINFORCE learning curve\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that at the end of training our policy consistantly (more or less) recieves returns of 500."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Investigating the variance of REINFORCE\n",
    "\n",
    "We noted in class that REINFORCE is a high variance algorithm. We can investigate the variance by running multiple trials and averaging the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('CartPole-v1')\n",
    "np.random.seed(53)\n",
    "seeds = np.random.randint(1000, size=5)\n",
    "all_scores = []\n",
    "for seed in seeds:\n",
    "    print(\"started training with seed: \", seed)\n",
    "    _, scores = reinforce(env, int(seed), verbose=False)\n",
    "    print(\"completed training with seed: \", seed)\n",
    "    all_scores.append(scores)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smoothed_scores = [moving_average(s, 50) for s in all_scores]\n",
    "smoothed_scores = np.array(smoothed_scores)\n",
    "mean = smoothed_scores.mean(axis=0)\n",
    "std = smoothed_scores.std(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "x = np.arange(1, len(mean)+1)\n",
    "ax.plot(x, mean, '-', color='blue')\n",
    "ax.fill_between(x, mean - std, mean + std, color='blue', alpha=0.2)\n",
    "plt.ylabel('Score')\n",
    "plt.xlabel('Episode #')\n",
    "plt.title('REINFORCE averaged over 5 seeds')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is what your graph should look like.\n",
    "![REINFORCE averaged over 5 seeds](https://i.imgur.com/5jSdIyV.png \"REINFORCE averaged over 5 seeds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reducing the variance of REINFORCE \n",
    "\n",
    "In class we saw a couple of tricks to reduce the variance of REINFORCE and improve its performance. Firstly, future actions should not change past decision. Present actions only impact the future. Therefore, we can change our objective function to reflect this:\n",
    "\\begin{align}\n",
    "    \\nabla_\\theta J(\\theta) = \\mathbb{E}_{\\tau \\sim p_\\theta} \\left[\\sum_{t=0}^{T} \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t) \\sum_{t'= t}^T \\gamma^{t'- t} r(s_{t'}, a_{t'})\\right].\n",
    "\\end{align}\n",
    "\n",
    "We can also reduce variance by subtracing a state dependent baseline to get:\n",
    "\\begin{align}\n",
    "    \\nabla_\\theta J(\\theta) = \\mathbb{E}_{\\tau \\sim p_\\theta} \\left[\\sum_{t=0}^{T} \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t) \\sum_{t'= t}^T \\left( \\gamma^{t'- t} r(s_{t'}, a_{t'}) - b(s_{t'}) \\right)\\right].\n",
    "\\end{align}\n",
    "\n",
    "For our baseline we'll use the average of the returns over the trajectory. As a final trick we normalise the returns by dividing by the standard deviation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_returns_baseline(rewards, gamma):\n",
    "    r = 0\n",
    "    returns = []\n",
    "    for step in reversed(range(len(rewards))):\n",
    "        r = rewards[step] + gamma * r\n",
    "        returns.insert(0, r)\n",
    "    returns = np.array(returns)\n",
    "    # IMPLEMENT-ME: normalize the returns by subtracting the mean and dividing by the standard deviation\n",
    "    ...\n",
    "    ...\n",
    "    ...\n",
    "    returns = ...\n",
    "    return returns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reinforce_baseline(env, seed, verbose=True):\n",
    "    # set random seeds (for reproducibility)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    env.seed(seed)\n",
    "\n",
    "    # instantiate the policy and optimizer\n",
    "    policy = Policy().to(device)\n",
    "    optimizer = optim.Adam(policy.parameters(), lr=learning_rate)\n",
    "\n",
    "    scores = []\n",
    "    scores_deque = deque(maxlen=100)\n",
    "    for episode in range(1, number_episodes+1):\n",
    "        #################################################################\n",
    "        # 1. Collect trajectories using our policy and save the rewards #\n",
    "        # and the log probability of each action taken.                 #\n",
    "        #################################################################\n",
    "        log_probs = []\n",
    "        rewards = []\n",
    "        state = env.reset()\n",
    "        for t in range(max_episode_length):\n",
    "            # IMPLEMENT-ME: get the distribution over actions for state\n",
    "            dist = ...\n",
    "\n",
    "            # IMPLEMENT-ME: sample an action from the distribution\n",
    "            action = ...\n",
    "\n",
    "            # IMPLEMENT-ME: compute the log probability\n",
    "            log_prob = ...\n",
    "\n",
    "            # IMPLEMENT-ME: take a step in the environment\n",
    "            state, reward, done, _ = ...\n",
    "\n",
    "            # save the reward and log probability\n",
    "            rewards.append(reward)\n",
    "            log_probs.append(log_prob.unsqueeze(0))\n",
    "\n",
    "            if done:\n",
    "                break\n",
    "\n",
    "        # for reporting save the score\n",
    "        scores.append(sum(rewards))\n",
    "        scores_deque.append(sum(rewards))\n",
    "\n",
    "        #################################################################\n",
    "        # 2. evaluate the policy gradient (with variance reduction)     #\n",
    "        #################################################################\n",
    "        \n",
    "        # calculate the discounted return of the trajectory\n",
    "        returns = compute_returns_baseline(rewards, gamma)\n",
    "        returns = torch.from_numpy(returns).float().to(device)\n",
    "\n",
    "        log_probs = torch.cat(log_probs)\n",
    "        # IMPLEMENT-ME: multiply the log probabilities by the returns and sum (see the policy-gradient theorem)\n",
    "        # Remember to multiply the result by -1 because we want to maximise the returns\n",
    "        policy_loss = ...\n",
    "\n",
    "        #################################################################\n",
    "        # 3. update the policy parameters (gradient descent)            #\n",
    "        #################################################################\n",
    "        optimizer.zero_grad()\n",
    "        policy_loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        # report the score to check that we're making progress\n",
    "        if episode % 50 == 0 and verbose:\n",
    "            print('Episode {}\\tAverage Score: {:.2f}'.format(episode, np.mean(scores_deque)))\n",
    "\n",
    "        if np.mean(scores_deque) >= 495.0 and verbose:\n",
    "            print('Environment solved in {:d} episodes!\\tAverage Score: {:.2f}'.format(episode, np.mean(scores_deque)))\n",
    "            break\n",
    "            \n",
    "    return policy, scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see if these changes give us any improvement."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('CartPole-v1')\n",
    "np.random.seed(53)\n",
    "seeds = np.random.randint(1000, size=5)\n",
    "all_scores_baseline = []\n",
    "for seed in seeds:\n",
    "    print(\"started training with seed: \", seed)\n",
    "    _, scores = reinforce_baseline(env ,int(seed), verbose=False)\n",
    "    print(\"completed training with seed: \", seed)\n",
    "    all_scores_baseline.append(scores)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Comparing the methods\n",
    "\n",
    "Finally we'll compare the performance of the two methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "smoothed_scores_baseline = [moving_average(s, 50) for s in all_scores_baseline]\n",
    "smoothed_scores_baseline = np.array(smoothed_scores_baseline)\n",
    "mean_baseline = smoothed_scores_baseline.mean(axis=0)\n",
    "std_baseline = smoothed_scores_baseline.std(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "x = np.arange(1, len(mean_baseline)+1)\n",
    "ax.plot(x, mean_baseline, '-', color='green', label='Variance reduced REINFORCE')\n",
    "ax.plot(x, mean, '-', color='blue', label='REINFORCE')\n",
    "ax.fill_between(x, mean_baseline - std_baseline, mean_baseline + std_baseline, color='green', alpha=0.2)\n",
    "ax.fill_between(x, mean - std, mean + std, color='blue', alpha=0.2)\n",
    "plt.ylabel('Score')\n",
    "plt.xlabel('Episode #')\n",
    "plt.legend()\n",
    "plt.title('Comparison of REINFORCE and Variance reduced REINFORCE (averaged over 5 seeds)')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is what your graph should look like.\n",
    "![Comparison of REINFORCE and Variance reduced REINFORCE (averaged over 5 seeds)](https://i.imgur.com/BaorVPg.png \"Comparison of REINFORCE and Variance reduced REINFORCE (averaged over 5 seeds)\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}