Actor Critic.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Actor Critic Algorithms\n",
    "In the DQN algorithm, we learned a Q-function by minimizing Bellman errors for an implicit policy that always took the action that maximized the Q-function. However, this scheme requires a discrete action spaces to allow for us to easily compute the optimal action at each state, unlike generic policy gradient algorithms that also worked with continuous action spaces.\n",
    "\n",
    "In this section, we will explore actor-critic algorithms which maintain an explicit policy (actor) like the policy gradient algorithms, learns a Q-function (critic) capturing the values of the _current policy_, and uses this learned Q-function to update the policy. Using a learned critic can provide much lower variance updates for the policy compared to using Monte-Carlo retun estimates, and also allows us to reuse our data by training the actor and critic on _off-policy_ data for more sample efficiency. We can thus take many more policy updates with an actor critic algorithm using our learned critic, instead of needing to wait and gather fresh samples every time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# As usual, a bit of setup\n",
    "import os\n",
    "import shutil\n",
    "import time\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "import deeprl.infrastructure.pytorch_util as ptu\n",
    "from deeprl.infrastructure.rl_trainer import RL_Trainer\n",
    "from deeprl.infrastructure.trainers import AC_Trainer\n",
    "\n",
    "from deeprl.policies.MLP_policy import MLPPolicyAC\n",
    "\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "def rel_error(x, y):\n",
    "    \"\"\" returns relative error \"\"\"\n",
    "    return np.max(np.abs(x - y) / (np.maximum(1e-8, np.abs(x) + np.abs(y))))\n",
    "\n",
    "def remove_folder(path):\n",
    "    # check if folder exists\n",
    "    if os.path.exists(path): \n",
    "        print(\"Clearing old results at {}\".format(path))\n",
    "        # remove if exists\n",
    "        shutil.rmtree(path)\n",
    "    else:\n",
    "        print(\"Folder {} does not exist yet. No old results to delete\".format(path))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ac_base_args_dict = dict(\n",
    "    env_name = 'Hopper-v2', #@param ['Ant-v2', 'Humanoid-v2', 'Walker2d-v2', 'HalfCheetah-v2', 'Hopper-v2']\n",
    "    exp_name = 'test_ac', #@param\n",
    "    save_params = False, #@param {type: \"boolean\"}\n",
    "    \n",
    "    ## PDF will tell you how to set ep_len\n",
    "    ## and discount for each environment\n",
    "    ep_len = 200, #@param {type: \"integer\"}\n",
    "    discount = 0.99, #@param {type: \"number\"}\n",
    "\n",
    "    # Training\n",
    "    num_agent_train_steps_per_iter = 1000, #@param {type: \"integer\"})\n",
    "    n_iter = 100, #@param {type: \"integer\"})\n",
    "\n",
    "    # batches & buffers\n",
    "    batch_size = 1000, #@param {type: \"integer\"})\n",
    "    eval_batch_size = 1000, #@param {type: \"integer\"}\n",
    "    train_batch_size = 256, #@param {type: \"integer\"}\n",
    "    max_replay_buffer_size = 1000000, #@param {type: \"integer\"}\n",
    "\n",
    "    #@markdown actor network\n",
    "    n_layers = 2, #@param {type: \"integer\"}\n",
    "    size = 256, #@param {type: \"integer\"}\n",
    "    entropy_weight=0, #@param {type: \"number\"}\n",
    "    learning_rate = 3e-4, #@param {type: \"number\"}\n",
    "    \n",
    "    # critic network\n",
    "    critic_n_layers = 2, #@param {type: \"integer\"}\n",
    "    critic_size = 256, #@param {type: \"integer\"}\n",
    "    target_update_rate = 5e-3,\n",
    "\n",
    "    #@markdown logging\n",
    "    video_log_freq = -1, #@param {type: \"integer\"}\n",
    "    scalar_log_freq = 1, #@param {type: \"integer\"}\n",
    "\n",
    "    #@markdown gpu & run-time settings\n",
    "    no_gpu = False, #@param {type: \"boolean\"}\n",
    "    which_gpu = 0, #@param {type: \"integer\"}\n",
    "    seed = 2, #@param {type: \"integer\"}\n",
    "    logdir = 'test',\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First fill out the target value calculation in the compute_target_value method of <code>critics/bootstrapped_continuous_critic.py</code>. Compared to the DQN critic, the key difference is that we are now estimating the value of the current policy, instead of the optimal policy as in DQN or Q-learning.\n",
    "\n",
    "To train our critic to evaluate the current policy $\\pi$, we simply sample actions from the current policy in our target value. For each sample $(s,a,s')$, our loss will be\n",
    "$$L(Q_{\\theta}(s, a), r(s,a) + \\gamma \\mathbb{E}_{a'\\sim \\pi(s')}[Q_{\\bar \\theta} (s', a')]),$$\n",
    "where $L$ is our loss function (for example squared error or the smooth L1 loss).\n",
    "In this assignment, we will simply sample a single action from the policy to estimate the target value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test bellman error for policy evaluation\n",
    "ac_dim = 3\n",
    "ob_dim = 11\n",
    "N = 5\n",
    "\n",
    "np.random.seed(0)\n",
    "obs = np.random.normal(size=(N, ob_dim))\n",
    "acts = np.random.choice(ac_dim, size=(N,))\n",
    "next_obs = np.random.normal(size=(N, ob_dim))\n",
    "rewards = np.random.normal(size=N)\n",
    "terminals = np.zeros(N)\n",
    "terminals[0] = 1\n",
    "\n",
    "ac_args = dict(ac_base_args_dict)\n",
    "\n",
    "env_str = 'Hopper'\n",
    "ac_args['env_name'] = '{}-v2'.format(env_str)\n",
    "ac_args['entropy_weight'] = 0.1\n",
    "actrainer = AC_Trainer(ac_args)\n",
    "critic = actrainer.rl_trainer.agent.critic\n",
    "\n",
    "class DummyDist:\n",
    "    def sample(self):\n",
    "        return ptu.from_numpy(1 + np.zeros(shape=(N, ac_dim)))\n",
    "\n",
    "def dummy_actor(next_obs):\n",
    "    return DummyDist()\n",
    "\n",
    "# assumes you call actor(next_obs) to get the distribution, then call distribution.sample()\n",
    "target_vals = critic.compute_target_value(ptu.from_numpy(next_obs), \n",
    "                                          ptu.from_numpy(rewards), \n",
    "                                          ptu.from_numpy(terminals), \n",
    "                                          dummy_actor)\n",
    "target_vals = ptu.to_numpy(target_vals)\n",
    "expected_targets = np.array([-0.9167948, -0.11123351, -0.36787638, -2.1131861,  -0.13868617])\n",
    "\n",
    "target_error = rel_error(target_vals, expected_targets)\n",
    "print(\"Target value error\", target_error, \"should be on the order of 1e-6 or lower\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this section, we will also update our target network parameters as an exponential moving average of the critic parameters, instead of simply copying the current parameters periodically as in DQN. Generally, either method for target networks tends to work with appropriately chosen update rates.\n",
    "\n",
    "Fill out the update_target_parameter_ema method in <code>critics/bootstrapped_continuous_critic.py</code>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test target network update\n",
    "ac_args = dict(ac_base_args_dict)\n",
    "\n",
    "env_str = 'Hopper'\n",
    "ac_args['env_name'] = '{}-v2'.format(env_str)\n",
    "ac_args['entropy_weight'] = 0.1\n",
    "actrainer = AC_Trainer(ac_args)\n",
    "critic = actrainer.rl_trainer.agent.critic\n",
    "\n",
    "critic.target_update_rate = 0.5\n",
    "\n",
    "# at initialization, target and critic networks are the same\n",
    "for p in critic.critic_network.parameters():\n",
    "    p.data += 1.\n",
    "    \n",
    "critic.update_target_network_ema()\n",
    "\n",
    "for p, target_p in zip(critic.critic_network.parameters(), critic.target_network.parameters()):\n",
    "    assert np.all(ptu.to_numpy((p-target_p)) == 0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will implement the actor update using the learned critic instead of Monte Carlo returns. \n",
    "\n",
    "To update our policy at a particular state $s$, our previous policy gradient (using the reward to go estimator) took a step on the objective (treating $Q^{\\pi}$ as a function that didn't depend on $\\pi$ and using the results of a single trajectory to estimate $Q^\\pi$)\n",
    "$$\\mathbb E_{a \\sim \\pi_{\\theta}(s)}[Q^{\\pi_\\theta}(s, a)],$$\n",
    "using the REINFORCE gradient estimator \n",
    "$$\\mathbb E_{a \\sim \\pi_{\\theta}(s)}[\\nabla_{\\theta} \\log \\pi_{\\theta}(a\\vert s) Q^{\\pi}(s, a)].$$\n",
    "This estimator only relied on the estimated value $Q^{\\pi_{\\theta}}(s,a)$, so was very general and could be applied with Monte Carlo estimates of $Q^{\\pi_\\theta}$.\n",
    "\n",
    "One way to estimate policy gradients with an actor critic algorithm would be to directly replace the Monte Carlo estimate of $Q$ with the learned critic $Q_{\\phi}$, and continue using the REINFORCE gradient estimator.\n",
    "However, we note that we can explicitly compute derivatives of our learned critic $Q(s, a)$ with respect to the action $a$, which can enable potentially better gradient estimates. \n",
    "\n",
    "In order to take advantage of this, we would also need to differentiate sampled actions $a$ with respect to our policy parameters, which we can through a technique known as the _reparameterization trick_ or the _pathwise_ estimator. \n",
    "The idea is that if our policy sampled actions according to $a \\sim \\mathcal N(\\mu_{\\theta}(s), \\sigma^2_{\\theta}(s))$, we can rewrite $a = f_{\\theta}(z)$, where $z \\sim \\mathcal N(0, 1)$, and $f(z) = \\mu_{\\theta}(s) + z \\cdot \\sigma_{\\theta}(s)$. Now all the randomness comes from sampling $z$, which doesn't depend on our policy, so we can now differentiate the sampled action $a$ with respect to our policy parameters $\\theta$ by simply differentiating through the function $f$ applied at the random noise $z$. \n",
    "\n",
    "Using the chain rule then allows to directly estimate gradients of \n",
    "$$\\mathbb{E}_{a \\sim \\pi_{\\theta}(s)}[Q_{\\phi}(s,a)]$$\n",
    "by drawing samples from $\\pi$ and differentiating $Q_{\\phi}(s,a)$ on the samples. \n",
    "\n",
    "Implement the actor update using this pathwise estimator in the update method of the MLPPolicyAC class in <code>policies/MLP_policy.py</code> (Hint: see the rsample function in for torch.distributions). Note that our implementation samples states uniformly from the entire replay buffer, not necessarily from the state distribution of the current policy. While this means we are no longer taking unbiased policy gradients (our estimates were already biased anyways due to using a learned critic), it works well in practice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compute actor update using the policy gradient. \n",
    "# For this test to pass, make sure you only call sample once per actor update to not throw off \n",
    "# the actor samples expected for the updates in the this test.\n",
    "torch.manual_seed(0)\n",
    "ac_dim = 2\n",
    "ob_dim = 3\n",
    "batch_size = 5\n",
    "\n",
    "np.random.seed(0)\n",
    "obs = np.random.normal(size=(N, ob_dim))\n",
    "\n",
    "policy = MLPPolicyAC(\n",
    "            ac_dim=ac_dim,\n",
    "            ob_dim=ob_dim,\n",
    "            n_layers=1,\n",
    "            size=2,\n",
    "            learning_rate=0.25,\n",
    "            entropy_weight=0.)\n",
    "\n",
    "def dummy_critic(obs, acts):\n",
    "    return torch.sum(acts + 1) + torch.sum(obs)\n",
    "\n",
    "initial_loss = policy.update(obs, dummy_critic)['Actor Training Loss']\n",
    "expected_initial_loss = -17.083496\n",
    "\n",
    "print(\"Initial loss error\", rel_error(expected_initial_loss, initial_loss), \"should be on the order of 1e-6 or less.\")\n",
    "for i in range(5):\n",
    "    loss = policy.update(obs, dummy_critic)['Actor Training Loss']\n",
    "    print(loss)\n",
    "\n",
    "expected_final_loss = -30.103575\n",
    "\n",
    "print(\"Final loss error\", rel_error(expected_final_loss, loss), \"should be on the order of 1e-6 or less.\")\n",
    "    \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we'll train our actor critic agent on the HalfCheetah task. You should see your policies generally get over 600 returns. \n",
    "\n",
    "We note that these actor critic algorithms, since they make use of off-policy updates, can be much more sample efficient than the basic policy gradient algorith we saw earlier. In our actor critic algorithms here, we only take 1000 new samples from the environment per iteration, while the policy gradient algorithms often needed many more samples per iteration to estimate the Monte Carlo returns (for example, we used 10000 in the Hopper experiments with policy gradient).\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ac_args = dict(ac_base_args_dict)\n",
    "\n",
    "env_str = 'HalfCheetah'\n",
    "ac_args['env_name'] = '{}-v2'.format(env_str)\n",
    "ac_args['n_iter'] = 50\n",
    "\n",
    "# Delete all previous logs\n",
    "remove_folder('logs/actor_critic/{}'.format(env_str))\n",
    "\n",
    "for seed in range(3):\n",
    "    print(\"Running actor critic experiment with seed\", seed)\n",
    "    ac_args['seed'] = seed\n",
    "    ac_args['logdir'] = 'logs/actor_critic/{}/seed{}'.format(env_str, seed)\n",
    "    actrainer = AC_Trainer(ac_args)\n",
    "    actrainer.run_training_loop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Visualize Actor Critic results on Halfheetah\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir logs/actor_critic/HalfCheetah"
   ]
  }
 ],
 "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}