Actor-Critic_exercise.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Actor-Critic Methods\n",
    "\n",
    "**Background**:\n",
    "In this practical we will train an agent using the Actor-Critic 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 Actor-Critic approach to directly training a parameterised policy and state-value function to maximise expected future rewards.\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": [
    "## Introduction\n",
    "\n",
    "Although the REINFORCE-with-baseline method learns both a policy and a state-value function, we do not consider it to be an actor–critic method because its state-value function is used only as a baseline, not as a critic. That is, it is not used for bootstrapping (updating the value estimate for a state from the estimated values of subsequent states), but only as a baseline for the state whose estimate is being updated. This is a useful distinction, for only through bootstrapping do we introduce bias and an asymptotic dependence on the quality of the function approximation. As we have seen, the bias introduced through bootstrapping and reliance on the state representation is often beneficial because it reduces variance and accelerates learning. REINFORCE with baseline is unbiased and will converge asymptotically to a local minimum, but like all Monte Carlo methods it tends to learn slowly (produce estimates of high variance) and to be inconvenient to implement online or for continuing problems. Temporal-difference methods we can eliminate these inconveniences, and through multi-step methods we can flexibly choose the degree of bootstrapping. In order to gain these advantages in the case of policy gradient methods we use actor–critic methods with a bootstrapping critic.\n",
    "\n",
    "![Actor-Critic Model](https://raw.githubusercontent.com/andrecianflone/rl_at_ammi/master/images/Policy-iteration-and-actor-critic-learning.png \"Actor-Critic Model\")\n",
    "\n",
    "\n",
    "**NOTE**: \n",
    "\n",
    "* Combine ideas from policy and value function methods\n",
    "* Actor improvement - Policy parameterised by $\\pi$\n",
    "* Critic evaluation - Value function parameterised by $\\omega$\n",
    "    - Either $V(s; \\omega)$ or $Q(s, a; \\omega)$\n",
    "    \n",
    "**Actor-Critic pseudocode**:\n",
    "\n",
    "Input: parameterised forms for $\\pi_{\\theta}(s|a)$ and $V_{\\omega}(s)$\n",
    "<br>\n",
    "Input: learning rates $\\alpha_{\\omega} > 0$ and $\\alpha_{\\theta} > 0$\n",
    "\n",
    "For each episode:<br>\n",
    "$\\quad$ Initialise $s$<br>\n",
    "$\\quad$ For each time step:<br>\n",
    "$\\quad \\quad$Choose $a \\sim \\pi_{\\theta}(s|a)$<br>\n",
    "$\\quad \\quad$Take $a$, observe $s′$, $𝑟$<br>\n",
    "$\\quad \\quad \\delta \\leftarrow r + \\gamma V_{\\omega}(s′) - V_{\\omega}(s)$<br>\n",
    "$\\quad \\quad \\omega \\leftarrow \\omega + \\alpha_{\\omega} \\delta \\nabla_{\\omega} V_{\\omega}(𝑠)$<br>\n",
    "$\\quad \\quad \\theta \\leftarrow \\theta + \\alpha_{\\theta} \\delta \\nabla_{\\theta} \\log\\pi_\\theta(a|s)$<br>\n",
    "$\\quad \\quad s \\leftarrow s'$<br>\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install box2d-py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function\n",
    "import gym\n",
    "import numpy as np\n",
    "from collections import deque\n",
    "from IPython.display import clear_output\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "import torch\n",
    "torch.manual_seed(0) # set random seed\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 \n",
    "\n",
    "# for auto-reloading external modules\n",
    "# (if you're curious, see http://stackoverflow.com/questions/1907993/autoreload-of-modules-in-ipython)\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "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",
    "\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from multiprocessing_env import SubprocVecEnv\n",
    "\n",
    "num_envs = 8\n",
    "env_name = \"CartPole-v1\"\n",
    "\n",
    "def make_env(seed, rank):\n",
    "    def _thunk():\n",
    "        env = gym.make(env_name)\n",
    "        env.seed(seed+rank)\n",
    "        return env\n",
    "\n",
    "    return _thunk\n",
    "\n",
    "env = gym.make(env_name)\n",
    "env.seed(111)\n",
    "print('action space:', env.action_space)\n",
    "print('observation space:', env.observation_space)"
   ]
  },
  {
   "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": [
    "state = env.reset()\n",
    "for t in range(150):\n",
    "    # sample a random action\n",
    "    action = env.action_space.sample()\n",
    "    env.render()\n",
    "    state, reward, done, _ = env.step(action)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Not very good! ... Now let's improve things using Actor-Critic.\n",
    "\n",
    "## The Actor-Critic Networks\n",
    "\n",
    "Simple enough. Refer to [torch.nn](https://pytorch.org/docs/stable/nn.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ActorCritic(nn.Module):\n",
    "    def __init__(self, num_inputs, num_outputs, hidden_size, std=0.0):\n",
    "        super(ActorCritic, self).__init__()\n",
    "        \n",
    "        self.critic = nn.Sequential(\n",
    "            nn.Linear(num_inputs, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, 1)\n",
    "        )\n",
    "        \n",
    "        self.actor = nn.Sequential(\n",
    "            nn.Linear(num_inputs, hidden_size),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden_size, num_outputs),\n",
    "            nn.Softmax(),\n",
    "        )\n",
    "        \n",
    "    def forward(self, x):\n",
    "        value = self.critic(x)\n",
    "        probs = self.actor(x)\n",
    "        dist  = Categorical(probs)\n",
    "        return dist, value"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_env(model, vis=False):\n",
    "    state = env.reset()\n",
    "    if vis: env.render()\n",
    "    done = False\n",
    "    total_reward = 0\n",
    "    while not done:\n",
    "        state = torch.FloatTensor(state).unsqueeze(0).to(device)\n",
    "        dist, _ = model(state)\n",
    "        next_state, reward, done, _ = env.step(dist.sample().cpu().numpy()[0])\n",
    "        state = next_state\n",
    "        if vis: env.render()\n",
    "        total_reward += reward\n",
    "    return total_reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot(frame_idx, rewards):\n",
    "    clear_output(True)\n",
    "    plt.figure(figsize=(20.0, 10.0))\n",
    "    plt.subplot(131)\n",
    "    plt.title('frame %s. reward: %s' % (frame_idx, rewards[-1]))\n",
    "    plt.plot(rewards)\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Selecting actions with our Actor network\n",
    "\n",
    "For a given state our networks returns a pytorch `Categorial` object along with the `value` 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": [
    "actorCritic = ActorCritic(env.observation_space.shape[0], env.action_space.n, 256).to(device)\n",
    "state = env.reset()\n",
    "state = torch.from_numpy(state).float().to(device)\n",
    "dist, value = actorCritic(state)\n",
    "action = dist.sample()\n",
    "print(\"Sampled action:\", action.item())\n",
    "print(\"Log probability of action:\", dist.log_prob(action).item())\n",
    "print(\"Value:\", value.item())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Computing the return\n",
    "Calculate the bootstrapped return $\\sum^{n}_{t=0} \\gamma^t r(s_t, a_t) +  \\gamma^{n+1} V_{\\omega}(s_{n+1})$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_returns(next_value, rewards, masks, gamma=0.99):\n",
    "    #IMPLEMENT-ME\n",
    "    # Compute the return and dont forget the value of the next state\n",
    "    R = ...\n",
    "    returns = []\n",
    "    for step in reversed(range(len(rewards))):\n",
    "        R = ...\n",
    "        returns.insert(0, R)\n",
    "    return returns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Actor-Critic\n",
    "\n",
    "Now its time to implement the algorithm\n",
    "\n",
    "**Actor-Critic pseudocode**:\n",
    "\n",
    "Input: parameterised forms for $\\pi_{\\theta}(s|a)$ and $V_{\\omega}(s)$\n",
    "<br>\n",
    "Input: learning rates $\\alpha_{\\omega} > 0$ and $\\alpha_{\\theta} > 0$\n",
    "\n",
    "\n",
    "Initialise $s$<br>\n",
    "For each time step:<br>\n",
    "$ \\quad$Choose $a \\sim \\pi_{\\theta}(s|a)$<br>\n",
    "$ \\quad$Take $a$, observe $s′$, $𝑟$<br>\n",
    "$ \\quad \\delta \\leftarrow r + \\gamma V_{\\omega}(s′) - V_{\\omega}(s)$<br>\n",
    "$ \\quad \\omega \\leftarrow \\omega + \\alpha_{\\omega} \\delta \\nabla_{\\omega} V_{\\omega}(𝑠)$<br>\n",
    "$ \\quad \\theta \\leftarrow \\theta + \\alpha_{\\theta} \\delta \\nabla_{\\theta} \\log\\pi_\\theta(a|s)$<br>\n",
    "$ \\quad s \\leftarrow s'$<br>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Hyperameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hidden_size = 256\n",
    "learning_rate = 3e-4\n",
    "num_step_td_update = 5\n",
    "max_steps = 100000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def actor_critic(seed, verbose=True):\n",
    "    \n",
    "    # set random seeds (for reproducibility)\n",
    "    torch.manual_seed(seed)\n",
    "    torch.cuda.manual_seed_all(seed)\n",
    "    envs = [make_env(seed, i) for i in range(num_envs)]\n",
    "    envs = SubprocVecEnv(envs)\n",
    "    \n",
    "    # instantiate the policy and optimiser\n",
    "    num_inputs  = envs.observation_space.shape[0]\n",
    "    num_outputs = envs.action_space.n\n",
    "    model = ActorCritic(num_inputs, num_outputs, hidden_size).to(device)\n",
    "    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n",
    "    \n",
    "    current_step_number = 0\n",
    "    test_rewards = []\n",
    "    state = envs.reset()\n",
    "    \n",
    "    while current_step_number < max_steps:\n",
    "        \n",
    "        log_probs = []\n",
    "        values    = []\n",
    "        rewards   = []\n",
    "        masks     = []\n",
    "        entropy = 0\n",
    "\n",
    "        for _ in range(num_step_td_update):\n",
    "            #IMPLEMENT-ME\n",
    "            # get the distribution over actions for state and the value of the state\n",
    "            state = ...\n",
    "            dist, value = ...\n",
    "            \n",
    "            #IMPLEMENT-ME\n",
    "            # sample an action from the distribution\n",
    "            action = ...\n",
    "            # take a step in the environment\n",
    "            next_state, reward, done, _ = envs.step(action.cpu().numpy())\n",
    "                \n",
    "            # compute the log probability\n",
    "            log_prob = dist.log_prob(action)\n",
    "            # compute the entropy\n",
    "            entropy += dist.entropy().mean()\n",
    "            \n",
    "            # save the log probability, value and reward \n",
    "            log_probs.append(log_prob)\n",
    "            values.append(value)\n",
    "            rewards.append(torch.FloatTensor(reward).unsqueeze(1).to(device))\n",
    "            masks.append(torch.FloatTensor(1 - done).unsqueeze(1).to(device))\n",
    "\n",
    "            state = next_state\n",
    "            current_step_number += 1\n",
    "            \n",
    "            if current_step_number % 1000 == 0:\n",
    "                test_rewards.append(np.mean([test_env(model) for _ in range(10)]))\n",
    "                plot(current_step_number, test_rewards)\n",
    "\n",
    "        next_state = torch.FloatTensor(next_state).to(device)\n",
    "        _, next_value = model(next_state)\n",
    "   \n",
    "        # calculate the discounted return of the episode\n",
    "        returns = compute_returns(next_value, rewards, masks)\n",
    "\n",
    "        log_probs = torch.cat(log_probs)\n",
    "        returns   = torch.cat(returns).detach()\n",
    "        values    = torch.cat(values)\n",
    "\n",
    "        #IMPLEMENT-ME\n",
    "        # Compute the advantage\n",
    "        advantage = ...\n",
    "        \n",
    "        #IMPLEMENT-ME\n",
    "        # Compute the actor's and critic's loss\n",
    "        actor_loss  = ...\n",
    "        critic_loss = ...\n",
    "\n",
    "        #IMPLEMENT-ME\n",
    "        # Compute the sum the actor and critic loss \n",
    "        loss = ... \n",
    "        loss -= 0.001 * entropy\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "env.seed(945)\n",
    "model = actor_critic(345)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Seeing our learned policy in action\n",
    "\n",
    "Let's watch our agent!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "state = env.reset()\n",
    "for t in range(2000):\n",
    "    dist, _ = model(torch.from_numpy(state).float().to(device))\n",
    "    action = dist.sample()\n",
    "    env.render()\n",
    "    state, reward, done, _ = env.step(action.item())\n",
    "    if done:\n",
    "        break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "- Reinforcement Learning: An Introduction [Book](http://incompleteideas.net/book/bookdraft2017nov5.pdf)\n",
    "- Policy Gradient [Algorithms](https://lilianweng.github.io/lil-log/2018/04/08/policy-gradient-algorithms.html#actor-critic)"
   ]
  }
 ],
 "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}