DQN.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Deep Q Networks\n",
    "Previously, we trained policies using policy gradient algorithms, which directly estimated the gradient of the returns for the policy and performed stochastic gradient ascent. In this section, we will now implement deep Q-learning, which does not explicitly optimize a policy, but simply infers the policy from a learned Q function that is trained via dynamic programming.\n",
    "\n",
    "We will assume discrete action spaces for this notebook as to enable us to easily select actions that maximize the Q-function at given states."
   ]
  },
  {
   "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 torch\n",
    "import numpy as np\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 DQN_Trainer\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": [
    "dqn_base_args_dict = dict(\n",
    "    env_name = 'LunarLander-v3', #@param \n",
    "    exp_name = 'test_dqn', #@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.95, #@param {type: \"number\"}\n",
    "\n",
    "    # Training\n",
    "    num_agent_train_steps_per_iter = 1, #@param {type: \"integer\"})\n",
    "    num_critic_updates_per_agent_update = 1, #@param {type: \"integer\"}\n",
    "  \n",
    "    #@markdown Q-learning parameters\n",
    "    double_q = False, #@param {type: \"boolean\"}\n",
    "\n",
    "    # batches & buffers\n",
    "    batch_size = 32, #@param {type: \"integer\"})\n",
    "    batch_size_initial=1000,\n",
    "\n",
    "    #@markdown logging\n",
    "    video_log_freq = -1, #@param {type: \"integer\"}\n",
    "    scalar_log_freq = 1000, #@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": [
    "## DQN updates\n",
    "Recall in Q-learning, we attempt to solve the optimal state-action values (which we refer to as Q-values $Q(s,a)$), by finding solutions to the Bellman equation given by\n",
    "$$Q(s,a) = r(s,a) + \\gamma \\mathbb{E}_{s' \\sim p(s'\\vert s,a)}[\\max_{a'}Q(s', a')].$$\n",
    "\n",
    "Regular tabular Q-learning would take sample transitions $(s, a, r, s')$ and perform updates according to\n",
    "$$Q(s,a) \\leftarrow Q(s,a) + \\alpha (r(s,a) + \\gamma \\max_{a'} Q(s', a') - Q(s,a)),$$\n",
    "where $\\alpha$ is a stepsize parameter.\n",
    "\n",
    "This can be interpreted as updating $Q(s,a)$ by taking one gradient step on a squared Bellman error objective\n",
    "$$(r(s, a) +\\gamma \\max_{a'} \\tilde Q(s', a') - Q(s,a))^2,$$\n",
    "where $\\tilde Q$ is a copy of $Q$, but is not differentiated when taking the gradient step.\n",
    "\n",
    "Adapting this update to the setting where we use a neural network with parameters $\\theta$ to approximate $Q(s,a)$, we then train $\\theta$ with the loss function \n",
    "$$\\min_{\\theta} \\mathbb{E}_{s, a, s' \\sim D} [L(Q_{\\theta}(s,a), r(s,a) + \\gamma \\max_{a'} Q_{\\tilde \\theta}(s', a'))]$$\n",
    "where $D$ is our replay buffer containing past transitions we've experienced, $L$ is some loss function capturing how far the predicted Q-values are from the target values, and $\\tilde \\theta$ are the target Q function parameters, which are usually a delayed copy of $\\theta$ for stability reasons.\n",
    "\n",
    "We note our previous policy gradient algorithms were _on-policy_ algoritms, which meant they updated the policy using only the data collected from the most recent policy, and discard all the data after using it just once. In contrast, DQN uses _off-policy_ updates by sampling data from all past interactions, allowing for data reuse over time."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fill out the missing components for the basic Q-learning update in <code>critics/dqn_critic.py</code> (not including the double_q section)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Test DQN updates\n",
    "dqn_args = dict(dqn_base_args_dict)\n",
    "\n",
    "env_str = 'LunarLander'\n",
    "dqn_args['env_name'] = '{}-v3'.format(env_str)\n",
    "dqn_args['double_q'] = False\n",
    "dqntrainer = DQN_Trainer(dqn_args)\n",
    "dqnagent = dqntrainer.rl_trainer.agent\n",
    "critic = dqnagent.critic\n",
    "\n",
    "ob_dim = critic.ob_dim\n",
    "ac_dim = 6\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",
    "first_weight_before = np.array(ptu.to_numpy(next(critic.q_net.parameters())))\n",
    "print(\"Weight before update (first row)\", first_weight_before[0])\n",
    "\n",
    "\n",
    "loss = critic.update(obs, acts, next_obs, rewards, terminals)['Training Loss']\n",
    "expected_loss = 0.9408444\n",
    "loss_error = rel_error(loss, expected_loss)\n",
    "print(\"Initial loss\", loss)\n",
    "print(\"Initial Loss Error\", loss_error, \"should be on the order of 1e-6 or lower\")\n",
    "\n",
    "for i in range(4):\n",
    "    loss = critic.update(obs, acts, next_obs, rewards, terminals)['Training Loss']\n",
    "    print(loss)\n",
    "\n",
    "expected_loss = 0.7889254\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",
    "\n",
    "first_weight_after = np.array(ptu.to_numpy(next(critic.q_net.parameters())))\n",
    "print(\"Weight after update (first row)\", first_weight_after.shape)\n",
    "# Test DQN gradient\n",
    "print(first_weight_after[0])\n",
    "weight_change_partial = first_weight_after[0] - first_weight_before[0]\n",
    "expected_weight_change = np.array([-0.00491365, -0.00500049, -0.00499149, -0.00491229, -0.00490125,  0.00489534,\n",
    " -0.00282785, -0.00171614,  0.00485604])\n",
    "\n",
    "\n",
    "updated_weight_error = rel_error(weight_change_partial, expected_weight_change)\n",
    "print(\"Weight Update Error\", updated_weight_error, \"should be on the order of 1e-6 or lower\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Implement the missing components in the get_action method of <code>policies/argmax_policy.py</code> and the step_env method in <code>agents/dqn_agent.py</code> to allow our agent to interact with the environment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Test argmax policy\n",
    "dqn_args = dict(dqn_base_args_dict)\n",
    "\n",
    "env_str = 'LunarLander'\n",
    "dqn_args['env_name'] = '{}-v3'.format(env_str)\n",
    "dqn_args['double_q'] = False\n",
    "dqntrainer = DQN_Trainer(dqn_args)\n",
    "dqnagent = dqntrainer.rl_trainer.agent\n",
    "actor = dqnagent.actor\n",
    "\n",
    "ob_dim = critic.ob_dim\n",
    "ac_dim = 6\n",
    "N = 5\n",
    "\n",
    "np.random.seed(0)\n",
    "obs = np.random.normal(size=(N, ob_dim))\n",
    "\n",
    "actions = actor.get_action(obs)\n",
    "correct_actions = np.array([1, 0, 1, 0, 1])\n",
    "\n",
    "assert np.all(correct_actions == actions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now test our DQN implementation on the LunarLander environment. These experiments can take a while to run (over 10 minutes per seed) on CPU, so start early."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dqn_args = dict(dqn_base_args_dict)\n",
    "\n",
    "env_str = 'LunarLander'\n",
    "dqn_args['env_name'] = '{}-v3'.format(env_str)\n",
    "dqn_args['double_q'] = False\n",
    "\n",
    "# Delete all previous logs\n",
    "remove_folder('logs/dqn/{}/vanilla_dqn'.format(env_str))\n",
    "\n",
    "for seed in range(3):\n",
    "    print(\"Running DQN experiment with seed\", seed)\n",
    "    dqn_args['seed'] = seed\n",
    "    dqn_args['logdir'] = 'logs/dqn/{}/vanilla_dqn/seed{}'.format(env_str, seed)\n",
    "    dqntrainer = DQN_Trainer(dqn_args)\n",
    "    dqntrainer.run_training_loop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Visualize vanilla DQN results on Lunar Lander\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir logs/dqn/LunarLander/vanilla_dqn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Double DQN\n",
    "One potential issue with learning our Q functions with bootstrapping is _maximization bias_, where the learned Q-values tend to overestimate the actual expected future returns. The main idea is that when there is estimation error in the next state's Q-values, even if the values were correct on average, picking the action with the maximum Q-value would tend to select one where the value is overestimated. This overoptimistic value would then also get propagated via the Bellman backups to other states and actions, and can potentially slow down learning.\n",
    "\n",
    "Double DQN (https://arxiv.org/abs/1509.06461) proposes a simple solution to alleviate this _maximization bias_. Instead of taking the next action that maximizes the target network's Q-value, it selects the action to maximize the _current_ Q function at the next state, and then takes the target network's estimate of that action's value. \n",
    "\n",
    "Implement the double DQN target value in the update method in <code>critics/dqn_critic.py</code>."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#### Test DQN target value with double Q\n",
    "dqn_args = dict(dqn_base_args_dict)\n",
    "\n",
    "env_str = 'LunarLander'\n",
    "dqn_args['env_name'] = '{}-v3'.format(env_str)\n",
    "dqn_args['double_q'] = True\n",
    "dqntrainer = DQN_Trainer(dqn_args)\n",
    "dqnagent = dqntrainer.rl_trainer.agent\n",
    "critic = dqnagent.critic\n",
    "\n",
    "ob_dim = critic.ob_dim\n",
    "ac_dim = 6\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",
    "first_weight_before = np.array(ptu.to_numpy(next(critic.q_net.parameters())))\n",
    "print(\"Weight before update (first row)\", first_weight_before[0])\n",
    "\n",
    "\n",
    "loss = critic.update(obs, acts, next_obs, rewards, terminals)['Training Loss']\n",
    "expected_loss = 0.93894196\n",
    "loss_error = rel_error(loss, expected_loss)\n",
    "print(\"Initial loss\", loss)\n",
    "print(\"Initial Loss Error\", loss_error, \"should be on the order of 1e-6 or lower\")\n",
    "\n",
    "for i in range(4):\n",
    "    loss = critic.update(obs, acts, next_obs, rewards, terminals)['Training Loss']\n",
    "    print(loss)\n",
    "\n",
    "expected_loss = 0.7871182\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",
    "\n",
    "first_weight_after = np.array(ptu.to_numpy(next(critic.q_net.parameters())))\n",
    "print(\"Weight after update (first row)\", first_weight_after.shape)\n",
    "# Test DQN gradient\n",
    "print(first_weight_after[0])\n",
    "weight_change_partial = first_weight_after[0] - first_weight_before[0]\n",
    "print(weight_change_partial)\n",
    "expected_weight_change = np.array([-0.0049137, -0.00500057, -0.00499138, -0.00491226, -0.00490116,  0.00489506,\n",
    " -0.00284088, -0.00171939,  0.00485736])\n",
    "\n",
    "\n",
    "updated_weight_error = rel_error(weight_change_partial, expected_weight_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 now also run some experiments on LunarLander with Double DQN. You may be able to see that double DQN performs slightly better and more stably, but as there is very high variance, dont' worry if you do not."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run with double DQN\n",
    "dqn_args = dict(dqn_base_args_dict)\n",
    "\n",
    "env_str = 'LunarLander'\n",
    "dqn_args['env_name'] = '{}-v3'.format(env_str)\n",
    "dqn_args['double_q'] = True\n",
    "\n",
    "# Delete all previous logs\n",
    "remove_folder('logs/dqn/{}/double_dqn'.format(env_str))\n",
    "\n",
    "for seed in range(3):\n",
    "    print(\"Running DQN experiment with seed\", seed)\n",
    "    dqn_args['seed'] = seed\n",
    "    dqn_args['logdir'] = 'logs/dqn/{}/double_dqn/seed{}'.format(env_str, seed)\n",
    "    dqntrainer = DQN_Trainer(dqn_args)\n",
    "    dqntrainer.run_training_loop()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### Visualize all DQN results on Lunar Lander\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir logs/dqn/LunarLander/"
   ]
  }
 ],
 "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
}