Policy_Gradients_REINFORCE_learned_baseline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install box2d-py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# REINFORCE with a Learned Baseline\n",
    "\n",
    "In this tutorial we will train an agent to play `LunarLander-v2` using REINFORCE with a learned baseline.\n",
    "\n",
    "## Introduction\n",
    "\n",
    "Previously we implemented the REINFORCE algorithm and discussed how introducing a baseline reduces variance and improves performance. We noted that the value function $v$ would be an ideal baseline (if we knew it) and used the average return as a rough approximation. In this tutorial we will try to learn the value function so that we can use it as a baseline in REINFORCE. Like the policy, we will parametrise the value function using a simple neural network with parameters $\\omega$. Recall that the policy gradient theorem gives us\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) \\left(G_t - v_\\omega(s_t) \\right) \\right], \\, \\text{ where } \\,\n",
    "G_t = \\sum_{t'=t}^T \\gamma^{t' - t} r(s_{t'}, a_{t'}).\n",
    "\\end{align}\n",
    "Here we have used our parametrised value function $v_\\omega$ as the baseline. **But how do we learn $v_\\omega$?**. Well, we can use monte-carlo learning! Since we can calculate the discounted returns $G_t$ we can just minimize $\\frac{1}{2}\\sum_{t=0}^T|G_t - v_\\omega(s_t)|^2$.\n",
    "\n",
    "**REINFORCE with a Learned Baseline**:\n",
    "1. sample a trajectory $\\tau = (s_0, a_0, r_1, s_1, \\ldots, s_{T}, r_{T})$ using the policy $\\pi_\\theta$.\n",
    "2. compute the vector of returns $[G_0, G_1, \\ldots, G_T]$.\n",
    "3. compute the policy gradient $\\nabla_\\theta J(\\theta) \\approx \\sum_{t=0}^T \\nabla_\\theta \\log \\pi_\\theta(a_t | s_t) \\left(G_t - v_\\omega(s_t) \\right)$.\n",
    "4. compute the gradient of the value function loss $\\nabla_\\omega \\mathcal{L}(\\omega) = \\nabla_\\omega \\frac{1}{2}\\sum_{t=0}^T|G_t - v_\\omega(s_t)|^2$.\n",
    "5. update policy parameters $\\theta \\leftarrow \\theta + \\alpha \\nabla_\\theta J(\\theta)$.\n",
    "6. update the value function parameters $\\omega \\leftarrow \\omega - \\alpha \\nabla_\\omega \\mathcal{L}(\\omega)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gym\n",
    "import numpy as np\n",
    "from collections import deque\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "import random\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\n",
    "\n",
    "In `LunarLander-v2` our goal is to land a space shuttle on the moon. We have 4 actions: fire the left thruster, fire the right thruster, fire the main thruster, or do nothing. The state space consists of: the shuttle's position (x,y)-cooridinate, its velocity, the angle of tilt, the angular velocity, and 2 boolean flags indicating whether the left and right legs of the shuttle are in contact with the ground. For more details refer to the OpenAi/gym github wiki page ([link](https://github.com/openai/gym/wiki/Leaderboard#lunarlander-v2)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env = gym.make('LunarLander-v2')\n",
    "print('Environment:', 'LunarLander-v2')\n",
    "print('\\t','action space:', env.action_space)\n",
    "print('\\t','observation space:', env.observation_space)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Watching a Random Policy In Action\n",
    "\n",
    "Let's see how a random policy performs in this enviroment:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_1 = gym.make('LunarLander-v2')\n",
    "state = env_1.reset()\n",
    "for t in range(500):\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",
    "    if done:\n",
    "        state = env_1.reset()\n",
    "env_1.close()\n",
    "del env_1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's some awful flying! We'll try to do better with REINFORCE now.\n",
    "\n",
    "## The Policy and Value Function Networks\n",
    "\n",
    "First, let's define our policy and value function. \n",
    "\n",
    "For efficiency we can use a single network with two heads. For a given state, the first head will output a Categorical distribution over the actions while the second head will return the value of the state.\n",
    "\n",
    "![two headed network](https://i.imgur.com/Z4Fq3cO.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, s_size=8, h_size=128, a_size=4):\n",
    "        super(Net, self).__init__()\n",
    "        # IMPLEMENT-ME\n",
    "        # The first layer should be a shared linear layer with\n",
    "        # an input size of env.observation_space.n and an output size of 128\n",
    "        self.fc_shared = ...\n",
    "        # The policy head should be a linear layer with input size of 128 \n",
    "        # and an output size of env.action_space.n\n",
    "        self.fc_policy = ...\n",
    "        # The value function head should be a linear layer with input size of 128\n",
    "        # and an output size of 1\n",
    "        self.fc_value_function = ...\n",
    "        \n",
    "    def forward(self, x):\n",
    "        # IMPLEMENT-ME\n",
    "        # Define the forward pass\n",
    "        # apply a ReLU activation after the shared layer\n",
    "        x = F.relu(self.fc_shared(x))\n",
    "        # apply the policy head layer (without an activation).\n",
    "        logits = ...\n",
    "        # apply the value function head layer (without an activation)\n",
    "        value = ...\n",
    "        # define a Categorical distribution over the actions\n",
    "        dist = ...\n",
    "        return dist, value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How do we use the net?\n",
    "\n",
    "For a given state our policy returns a tuple consisting of a pytorch `Categorial` object and a pytorch `Tensor`. To recap, we can use `sample` to sample an action from the distribution and `log_prob` to find the log probability of a particular action."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Net().to(device)\n",
    "state = env.reset()\n",
    "state = torch.from_numpy(state).float().to(device)\n",
    "dist, value = net(state)\n",
    "action = dist.sample()\n",
    "print('Sampled action: ', action.item())\n",
    "print('Log probability of action: ', dist.log_prob(action).item())\n",
    "print('Estimated value of the state: ', value.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing the Return\n",
    "\n",
    "Given a sequence of returns compute the vector of discounted returns $[G_0, G_1, \\ldots, G_T]$. Note that we alse use the trick of 'normalizing' the returns i.e. we subtract the mean and divide by the standard deviation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_returns(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",
    "    returns -= returns.mean()\n",
    "    returns /= returns.std()\n",
    "    return returns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## REINFORCE with a Learned Baseline\n",
    "\n",
    "1. sample a trajectory $\\tau = (s_0, a_0, r_1, s_1, \\ldots, s_{T}, r_{T})$ using the policy $\\pi_\\theta$.\n",
    "2. compute the vector of returns $[G_0, G_1, \\ldots, G_T]$.\n",
    "3. compute the policy gradient $\\nabla_\\theta J(\\theta) \\approx \\sum_{t=0}^T \\nabla_\\theta \\log \\pi_\\theta(a_t | s_t) \\left(G_t - v_\\omega(s_t) \\right)$.\n",
    "4. compute the gradient of the value function loss $\\nabla_\\omega \\mathcal{L}(\\omega) = \\nabla_\\omega \\frac{1}{2}\\sum_{t=0}^T|G_t - v_\\omega(s_t)|^2$.\n",
    "5. update policy parameters $\\theta \\leftarrow \\theta + \\alpha \\nabla_\\theta J(\\theta)$.\n",
    "6. update the value function parameters $\\omega \\leftarrow \\omega - \\alpha \\nabla_\\omega \\mathcal{L}(\\omega)$.\n",
    "\n",
    "In practice we combine steps 3. and 4. by defining the composite loss:\n",
    "\\begin{align}\n",
    "\\texttt{loss} = \\sum_{t=0}^T \\log \\pi_\\theta(a_t | s_t) \\left(G_t - \\text{detach}(v_\\omega(s_t)) \\right) + \\frac{1}{2}\\sum_{t=0}^T|G_t - v_\\omega(s_t)|^2.\n",
    "\\end{align}\n",
    "This also allows us to perform the policy and value function updates in a single step."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define some hyperparameters\n",
    "gamma = 0.99\n",
    "lr = 0.02\n",
    "seed = 401\n",
    "number_episodes = 1250"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def reinforce_learned_baseline(seed):\n",
    "    env = gym.make('LunarLander-v2')\n",
    "    \n",
    "    # set random seeds (for reproducibility)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    np.random.seed(seed)\n",
    "    env.seed(seed)\n",
    "    random.seed(seed)\n",
    "\n",
    "    # instantiate the policy and optimizer\n",
    "    net = Net().to(device)\n",
    "    optimizer = optim.Adam(net.parameters(), lr=lr)\n",
    "\n",
    "    scores = []\n",
    "    scores_deque = deque(maxlen=50)\n",
    "    for episode in range(1, number_episodes+1):\n",
    "        ##################################################################\n",
    "        # 1. Collect trajectories using our policy and save the rewards, #\n",
    "        # log probability, and the estimated value of each state.        #                                                #\n",
    "        ##################################################################\n",
    "        log_probs = []\n",
    "        values = []\n",
    "        rewards = []\n",
    "\n",
    "        state = env.reset()\n",
    "        for t in range(1000):\n",
    "            # convert state to a torch Tensor\n",
    "            state = torch.from_numpy(state).float().to(device)\n",
    "            # get the distribution over actions and the estimated value of state\n",
    "            dist, value = net(state)\n",
    "\n",
    "            # sample an action from the distribution\n",
    "            action = dist.sample()\n",
    "            \n",
    "            # compute the log probability\n",
    "            log_prob = dist.log_prob(action)\n",
    "            \n",
    "            # take a step in the environment\n",
    "            state, reward, done, _ = env.step(action.item())\n",
    "\n",
    "            # save the reward, log probabily, and value \n",
    "            rewards.append(reward)\n",
    "            log_probs.append(log_prob.unsqueeze(0))\n",
    "            values.append(value)\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. Compute the vector of discounted returns                    #\n",
    "        ##################################################################\n",
    "        returns = compute_returns(rewards, gamma)\n",
    "        returns = torch.from_numpy(returns).float().to(device)\n",
    "\n",
    "        ##################################################################\n",
    "        # 3. and 4. Compute the loss for gradient descent                #\n",
    "        ##################################################################\n",
    "        values = torch.cat(values)\n",
    "        log_probs = torch.cat(log_probs)\n",
    "\n",
    "        # IMPLEMENT-ME\n",
    "        # compute the difference between the returns and the values\n",
    "        delta = ...\n",
    "\n",
    "        # IMPLEMENT-ME\n",
    "        # compute the policy loss term. multiply the log probabilities by delta and sum\n",
    "        # (remeber to call .detach() on delta since we do not want the gradient to propogate\n",
    "        # to the value function network here)\n",
    "        policy_loss = ...\n",
    "\n",
    "        # IMPLEMENT-ME\n",
    "        # compute the value function loss term\n",
    "        value_function_loss = ...\n",
    "\n",
    "        # IMPLEMENT-ME\n",
    "        # compute the composite loss\n",
    "        loss = ...\n",
    "\n",
    "        #################################################################\n",
    "        # 4. and 5. update the policy and value function parameters     #\n",
    "        #################################################################\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if episode % 50 == 0:\n",
    "            print('Episode {}\\tAverage Score: {:.2f}'.format(episode, np.mean(scores_deque)))\n",
    "\n",
    "    return net, scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net, scores = reinforce_learned_baseline(seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Watching Our Agent in Action\n",
    "\n",
    "Finally, let's see how our agent performs in the `LunarLander` environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env_1 = gym.make('LunarLander-v2')\n",
    "state = env_1.reset()\n",
    "for t in range(2000):\n",
    "    state = torch.from_numpy(state).float().to(device)\n",
    "    dist, value = net(state)\n",
    "    action = dist.sample().item()\n",
    "    env_1.render()\n",
    "    state, reward, done, _ = env_1.step(action)\n",
    "    if done:\n",
    "        state = env_1.reset()\n",
    "env_1.close()\n",
    "del env_1"
   ]
  }
 ],
 "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.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}