Policy Gradients.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Policy Gradients\n",
    "The goal in policy gradient algorithms is to maximize the expected returns of a policy $\\pi_\\theta$ with parameters $\\theta$. Letting $\\tau=((s_0, a_0, r_0), \\ldots, (s_T, a_T, r_T) )$ denote a trajectory and $R(\\tau)$ the return of $\\tau$, this objective can be written as\n",
    "$$\\max_{\\theta} \\mathbb E_{\\tau \\sim \\pi_{\\theta}}[R(\\tau)].$$\n",
    "\n",
    "Using the REINFORCE trick, we can compute the policy gradient (the gradient of expected policy returns) as\n",
    "$$\\sum_{t=0}^T \\mathbb E_{s_t, a_t \\sim \\pi(\\tau)} \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t \\vert s_t) R(\\tau).$$\n",
    "\n",
    "We can then estimate this with a very simple scheme.\n",
    "We first sample a trajectory $\\tau = ((s_t, a_t, r_t))_{t=0}^\\infty$ from our current policy, compute the discounted return of the trajectory as $R$, then take a stochastic estimate of the policy gradient as \n",
    "\n",
    "$$\\sum_{t=0}^T \\mathbb \\nabla_{\\theta} \\log \\pi_{\\theta}(a_t \\vert s_t) R(\\tau).$$\n",
    "We can then repeat sample more trajectories to average the estimate over multiple samples.\n",
    "In practice, we will often use _discounted_ returns $\\tilde R(\\tau) = \\sum_{t=0}^T \\gamma^t r_t$ where $\\gamma$ is the discount factor and our policy gradient estimate will simply replace the undiscounted returns with $\\tilde R(\\tau)$.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#@title imports\n",
    "# As usual, a bit of setup\n",
    "import os\n",
    "import shutil\n",
    "import time\n",
    "import numpy as np\n",
    "import gym\n",
    "import torch\n",
    "\n",
    "import deeprl.infrastructure.pytorch_util as ptu\n",
    "\n",
    "from deeprl.infrastructure.rl_trainer import RL_Trainer\n",
    "from deeprl.infrastructure.trainers import PG_Trainer\n",
    "from deeprl.infrastructure.trainers import BC_Trainer\n",
    "\n",
    "from deeprl.agents.pg_agent import PGAgent\n",
    "from deeprl.policies.MLP_policy import MLPPolicyPG\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": [
    "pg_base_args_dict = dict(\n",
    "    env_name = 'Hopper-v2', #@param ['Ant-v2', 'Humanoid-v2', 'Walker2d-v2', 'HalfCheetah-v2', 'Hopper-v2']\n",
    "    exp_name = 'test_pg', #@param\n",
    "    save_params = False, #@param {type: \"boolean\"}\n",
    "    \n",
    "    ep_len = 200, #@param {type: \"integer\"}\n",
    "    discount = 0.95, #@param {type: \"number\"}\n",
    "\n",
    "    reward_to_go = True, #@param {type: \"boolean\"}\n",
    "    nn_baseline = False, #@param {type: \"boolean\"}\n",
    "    dont_standardize_advantages = True, #@param {type: \"boolean\"}\n",
    "\n",
    "    # Training\n",
    "    num_agent_train_steps_per_iter = 1, #@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 = 1000, #@param {type: \"integer\"}\n",
    "    max_replay_buffer_size = 1000000, #@param {type: \"integer\"}\n",
    "\n",
    "    #@markdown network\n",
    "    n_layers = 2, #@param {type: \"integer\"}\n",
    "    size = 64, #@param {type: \"integer\"}\n",
    "    learning_rate = 5e-3, #@param {type: \"number\"}\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": [
    "## Implementing policy gradients\n",
    "We will first compute a very naive policy gradient calculation by taking the whole discounted return of a trajectory. Fill out the method <code>_discounted_return</code> in <code>pg_agent.py</code>. Your error should be 1e-6 or lower."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Test return computation\n",
    "pg_args = dict(pg_base_args_dict)\n",
    "\n",
    "env_str = 'CartPole'\n",
    "pg_args['env_name'] = '{}-v0'.format(env_str)\n",
    "pgtrainer = PG_Trainer(pg_args)\n",
    "pgagent = pgtrainer.rl_trainer.agent\n",
    "\n",
    "T = 10\n",
    "np.random.seed(0)\n",
    "rewards = np.random.normal(size=T)\n",
    "discounted_returns = pgagent._discounted_return(rewards)\n",
    "\n",
    "expected_return = 6.49674307\n",
    "return_error = rel_error(discounted_returns, expected_return)\n",
    "print(\"Error in return estimate is\", return_error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we'll consider a return estimate with lower variance by taking the discounted reward-to-go at each timestep instead of the entire discounted return. More precisely, instead of taking $\\sum_{t'=0}^T \\gamma^{t'} r_{t'}$ as the return estimate for all timesteps $t$, we will instead use $\\sum_{t'=t}^T \\gamma^{t' - t} r_{t'}$ for the return estimate at timestep $t$. Fill out the method <code>_discounted_cumsum</code> in <code>pg_agent.py</code>.   Your error should be 1e-6 or lower."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Test reward to go computations\n",
    "pg_args = dict(pg_base_args_dict)\n",
    "\n",
    "env_str = 'CartPole'\n",
    "pg_args['env_name'] = '{}-v0'.format(env_str)\n",
    "pgtrainer = PG_Trainer(pg_args)\n",
    "pgagent = pgtrainer.rl_trainer.agent\n",
    "\n",
    "T = 10\n",
    "np.random.seed(0)\n",
    "rewards = np.random.normal(size=T)\n",
    "discounted_cumsum = pgagent._discounted_cumsum(rewards)\n",
    "expected_cumsum = np.array([6.49674307, 4.98177971, 4.82276053, 4.04633952, 1.90046981, 0.03464402,\n",
    " 1.06518095, 0.12115003, 0.28684973, 0.4105985])\n",
    "\n",
    "return_error = rel_error(discounted_cumsum, expected_cumsum)\n",
    "print(\"Error in return estimate is\", return_error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we'll use our return estimates to compute a policy gradient. Fill out the surrogate loss computation in the <code>update</code> method in MLPPolicyPG class in <code>policies/MLP_policy.py</code>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Test policy gradient (check gradients match what we expect)\n",
    "torch.manual_seed(0)\n",
    "ac_dim = 2\n",
    "ob_dim = 3\n",
    "batch_size = 5\n",
    "\n",
    "policy = MLPPolicyPG(\n",
    "            ac_dim=ac_dim,\n",
    "            ob_dim=ob_dim,\n",
    "            n_layers=1,\n",
    "            size=2,\n",
    "            learning_rate=0.25)\n",
    "\n",
    "np.random.seed(0)\n",
    "obs = np.random.normal(size=(batch_size, ob_dim))\n",
    "acts = np.random.normal(size=(batch_size, ac_dim))\n",
    "advs = 1000 * np.random.normal(size=(batch_size,))\n",
    "\n",
    "first_weight_before = np.array(ptu.to_numpy(next(policy.mean_net.parameters())))\n",
    "print(\"Weight before update\", first_weight_before)\n",
    "\n",
    "for i in range(5):\n",
    "    loss = policy.update(obs, acts, advs)['Training Loss']\n",
    "\n",
    "print(loss)\n",
    "expected_loss = -6142.9116\n",
    "loss_error = rel_error(loss, expected_loss)\n",
    "print(\"Loss Error\", loss_error, \"should be on the order of 1e-6 or lower\")\n",
    "\n",
    "first_weight_after = ptu.to_numpy(next(policy.mean_net.parameters()))\n",
    "print('Weight after update', first_weight_after)\n",
    "\n",
    "weight_change = first_weight_after - first_weight_before\n",
    "print(\"Change in weights\", weight_change)\n",
    "\n",
    "expected_change = np.array([[ 1.035012, 1.0455959, 0.11085394],\n",
    "                            [-1.1532364, -0.5915445, 0.557522]])\n",
    "updated_weight_error = rel_error(weight_change, expected_change)\n",
    "print(\"Weight Update Error\", updated_weight_error, \"should be on the order of 1e-6 or lower\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can compare the two return estimators on a simple environment and compare how well they do. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pg_args = dict(pg_base_args_dict)\n",
    "\n",
    "env_str = 'CartPole'\n",
    "pg_args['env_name'] = '{}-v0'.format(env_str)\n",
    "pg_args['reward_to_go'] = False\n",
    "pg_args['n_iter'] = 100\n",
    "\n",
    "# Delete all previous logs\n",
    "remove_folder('logs/policy_gradient/{}/full_returns/'.format(env_str))\n",
    "\n",
    "for seed in range(3):\n",
    "    print(\"Running policy gradient experiment with seed\", seed)\n",
    "    pg_args['seed'] = seed\n",
    "    pg_args['logdir'] = 'logs/policy_gradient/{}/full_returns/seed{}'.format(env_str, seed)\n",
    "    pgtrainer = PG_Trainer(pg_args)\n",
    "    pgtrainer.run_training_loop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pg_args = dict(pg_base_args_dict)\n",
    "\n",
    "env_str = 'CartPole'\n",
    "pg_args['env_name'] = '{}-v0'.format(env_str)\n",
    "pg_args['reward_to_go'] = True\n",
    "pg_args['n_iter'] = 100\n",
    "\n",
    "# Delete all previous logs\n",
    "remove_folder('logs/policy_gradient/{}/return_to_go/'.format(env_str))\n",
    "\n",
    "for seed in range(3):\n",
    "    print(\"Running policy gradient experiment with seed\", seed)\n",
    "    pg_args['seed'] = seed\n",
    "    pg_args['logdir'] = 'logs/policy_gradient/{}/return_to_go/seed{}'.format(env_str, seed)\n",
    "    pgtrainer = PG_Trainer(pg_args)\n",
    "    pgtrainer.run_training_loop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We should see the reward to go estimator outperforming the full returns estimator, with some runs reaching the maximum reward of 200. There will likely however be high variance between runs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Visualize Policy Gradient results on CartPole\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir logs/policy_gradient/CartPole"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can also compare our estimators on a more complex task, though you will probably see that they don't perform well (not getting much above 200 returns). Note that on this more complex task, we use a much larger batch size to reduce variance in the policy gradients."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pg_args = dict(pg_base_args_dict)\n",
    "\n",
    "env_str = 'Hopper'\n",
    "pg_args['env_name'] = '{}-v2'.format(env_str)\n",
    "pg_args['learning_rate'] = 0.01\n",
    "pg_args['reward_to_go'] = False\n",
    "pg_args['batch_size'] = 10000\n",
    "pg_args['train_batch_size'] = 10000\n",
    "pg_args['n_iter'] = 100\n",
    "\n",
    "# Delete all previous logs\n",
    "remove_folder('logs/policy_gradient/{}/full_returns/'.format(env_str))\n",
    "\n",
    "for seed in range(3):\n",
    "    print(\"Running policy gradient experiment with seed\", seed)\n",
    "    pg_args['seed'] = seed\n",
    "    pg_args['logdir'] = 'logs/policy_gradient/{}/full_returns/seed{}'.format(env_str, seed)\n",
    "    pgtrainer = PG_Trainer(pg_args)\n",
    "    pgtrainer.run_training_loop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pg_args = dict(pg_base_args_dict)\n",
    "\n",
    "env_str = 'Hopper'\n",
    "pg_args['env_name'] = '{}-v2'.format(env_str)\n",
    "pg_args['learning_rate'] = 0.01\n",
    "pg_args['reward_to_go'] = True\n",
    "pg_args['batch_size'] = 10000\n",
    "pg_args['train_batch_size'] = 10000\n",
    "pg_args['n_iter'] = 100\n",
    "\n",
    "# Delete all previous logs\n",
    "remove_folder('logs/policy_gradient/{}/return_to_go/'.format(env_str))\n",
    "\n",
    "for seed in range(3):\n",
    "    print(\"Running policy gradient experiment with seed\", seed)\n",
    "    pg_args['seed'] = seed\n",
    "    pg_args['logdir'] = 'logs/policy_gradient/{}/return_to_go/seed{}'.format(env_str, seed)\n",
    "    pgtrainer = PG_Trainer(pg_args)\n",
    "    pgtrainer.run_training_loop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Visualize Policy Gradient results on Hopper\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir logs/policy_gradient/Hopper"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Variance Reduction with a Value Function Baseline\n",
    "We can further reduce the policy gradient variance by including state-dependent baselines. In this section, we will train a value function network to predict the value of the policy at a state, then use the value function as a baseline by subtracting it from our reward-to-go estimate.\n",
    "\n",
    "Implement the value function baseline loss in the update method of the MLPPolicyPG class in <code>policies/MLP_policy.py</code>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test value function gradient\n",
    "torch.manual_seed(0)\n",
    "ac_dim = 2\n",
    "ob_dim = 3\n",
    "batch_size = 5\n",
    "\n",
    "policy = MLPPolicyPG(\n",
    "            ac_dim=ac_dim,\n",
    "            ob_dim=ob_dim,\n",
    "            n_layers=1,\n",
    "            size=2,\n",
    "            learning_rate=0.25,\n",
    "            nn_baseline=True)\n",
    "\n",
    "np.random.seed(0)\n",
    "obs = np.random.normal(size=(batch_size, ob_dim))\n",
    "acts = np.random.normal(size=(batch_size, ac_dim))\n",
    "advs = 1000 * np.random.normal(size=(batch_size,))\n",
    "qvals = advs\n",
    "\n",
    "first_weight_before = np.array(ptu.to_numpy(next(policy.baseline.parameters())))\n",
    "print(\"Weight before update\", first_weight_before)\n",
    "\n",
    "for i in range(5):\n",
    "    loss = policy.update(obs, acts, advs, qvals=qvals)['Baseline Loss']\n",
    "\n",
    "print(loss)\n",
    "expected_loss = 0.925361\n",
    "loss_error = rel_error(loss, expected_loss)\n",
    "print(\"Loss Error\", loss_error, \"should be on the order of 1e-6 or lower\")\n",
    "\n",
    "first_weight_after = ptu.to_numpy(next(policy.baseline.parameters()))\n",
    "print('Weight after update', first_weight_after)\n",
    "\n",
    "weight_change = first_weight_after - first_weight_before\n",
    "print(\"Change in weights\", weight_change)\n",
    "\n",
    "expected_change = np.array([[ 0.38988823,  0.70297027,  0.2609921 ],\n",
    "                            [-1.0340402,  -0.84166795,  0.7254925 ]])\n",
    "updated_weight_error = rel_error(weight_change, expected_change)\n",
    "print(\"Weight Update Error\", updated_weight_error, \"should be on the order of 1e-6 or lower\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the estimate_advantage function in <code>agents/pg_agent.py</code>, fill out the advantage estimate using the baseline and test your implementation below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Test return computation\n",
    "pg_args = dict(pg_base_args_dict)\n",
    "\n",
    "env_str = 'CartPole'\n",
    "pg_args['env_name'] = '{}-v0'.format(env_str)\n",
    "pg_args['nn_baseline'] = True\n",
    "pgtrainer = PG_Trainer(pg_args)\n",
    "pgagent = pgtrainer.rl_trainer.agent\n",
    "\n",
    "obs_dim = 4\n",
    "N = 10\n",
    "np.random.seed(0)\n",
    "obs = np.random.normal(size=(N, obs_dim))\n",
    "qs = np.random.normal(size=N)\n",
    "\n",
    "baseline_advantages = pgagent.estimate_advantage(obs, qs)\n",
    "expected_advantages = np.array([-0.44662586, -0.89629588, -1.14574752,  2.43957172, -0.06601728,\n",
    "       -0.00501807, -0.74720337,  1.27468092, -1.20184486,  0.25312274])\n",
    "\n",
    "advantage_error = rel_error(expected_advantages, baseline_advantages)\n",
    "print(\"Advantage error\", advantage_error, \"should be on the order of 1e-6 or lower\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train your policies!\n",
    "In this section, we will train our policies using the reward-to-go estimator and learning a value function baseline. On Hopper, you should see your methods get over 300 rewards consistently, and sometimes over 400. Returns will tend to oscillate during training. You should also see that using a value function baseline greatly improves performance over our earlier experiments without it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pg_args = dict(pg_base_args_dict)\n",
    "\n",
    "env_str = 'Hopper'\n",
    "pg_args['env_name'] = '{}-v2'.format(env_str)\n",
    "pg_args['learning_rate'] = 0.01\n",
    "pg_args['reward_to_go'] = True\n",
    "pg_args['nn_baseline'] = True\n",
    "pg_args['batch_size'] = 10000\n",
    "pg_args['train_batch_size'] = 10000\n",
    "pg_args['n_iter'] = 100\n",
    "\n",
    "# Delete all previous logs\n",
    "remove_folder('logs/policy_gradient/{}/with_baseline/'.format(env_str))\n",
    "\n",
    "for seed in range(3):\n",
    "    print(\"Running policy gradient experiment with seed\", seed)\n",
    "    pg_args['seed'] = seed\n",
    "    pg_args['logdir'] = 'logs/policy_gradient/{}/with_baseline/seed{}'.format(env_str, seed)\n",
    "    pgtrainer = PG_Trainer(pg_args)\n",
    "    pgtrainer.run_training_loop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Plot learning curves\n",
    "### Visualize Policy Gradient results on Hopper\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir logs/policy_gradient/Hopper"
   ]
  }
 ],
 "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": 5
}