S66_driving_DRL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "V03PX8SoEitW",
   "metadata": {},
   "source": [
    "# Deep Reinforcement Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3BC8gmeu2etX",
   "metadata": {
    "colab_type": "text",
    "tags": [
     "no-tex"
    ]
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/gtbook/robotics/blob/main/S66_driving_DRL.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "pVeijfbAiYRG",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install -q -U gtbook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "czpSAiabQQ8g",
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [],
   "source": [
    "import plotly.express as px\n",
    "\n",
    "try:\n",
    "    import google.colab\n",
    "except:\n",
    "    import plotly.io as pio\n",
    "    pio.renderers.default = \"png\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "puubBEzEpHDR",
   "metadata": {},
   "source": [
    "```{index} learning; deep reinforcement learning\n",
    "```\n",
    "\n",
    "> Deep Q-learning and policy gradient."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af80ae2b",
   "metadata": {
    "tags": [
     "no-pdf"
    ]
   },
   "source": [
    "<img src=\"Figures6/S66-Autonomous_Vehicle_with_LIDAR_and_cameras-03.jpg\" alt=\"Splash image with steampunk autonomous car\" width=\"40%\" align=center style=\"vertical-align:middle;margin:10px 0px\">"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f74d0dd",
   "metadata": {},
   "source": [
    "```{index} pair: deep reinforcement learning; DRL\n",
    "```\n",
    "Deep reinforcement learning (DRL) applies the power of deep learning to bring reinforcement learning to much more complex domains than what we were able to tackle with the Markov Decision Processes and RL concepts introduced in Chapter 3. The use of large, expressive neural networks has allowed researchers and practitioners alike to work with high bandwidth sensors such as video streams and LIDAR, and bring the promise of RL into real-world domains such as autonomous driving. This is still a field of active discovery and research, however, and we can give but a brief introduction here about what is a vast literature and problem space."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a8OPWJHrGMhk",
   "metadata": {},
   "source": [
    "## RL and Autonomous Driving"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fL6Mb9yj3sih",
   "metadata": {},
   "source": [
    "```{index} lateral control, longitudinal control, lane switching\n",
    "```\n",
    "A simple example in the autonomous driving domain is *lane switching*. Suppose we are driving along at 3-lane highway, and we can see some ways ahead, and -using the rear-view mirror- some ways behind us. We are driving at a speed that is comfortable to us, but other cars have different ideas about their optimal driving speed. Hence, sometimes we would like to change lanes, and we could learn a policy to do this for us. As discussed in Section 6.5, this is **lateral control**. A more sophisticated example would also allow us to adapt our speed to the traffic pattern, but by relying on a smart cruise control system we could safely ignore the **longitudinal control** problem."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "LeD4sdVqLHW_",
   "metadata": {},
   "source": [
    "To turn this into a reinforcement learning problem, we first need to define a state space ${\\cal X}$ and an action space ${\\cal A}$. There are a variety of ways to engineer this aspect of the problem. For example, we could somehow encode the longitudinal distance and lane index for each of the K closest cars, where K is a parameter, say 5 or 10. One problem is that the number of cars that are *actually* present is variable, which is difficult to deal with. Another approach is to make this into an image processing problem, by creating a finite element representation of the road before and behind us, and marking each cell as occupied or not. The latter is fairly compatible with automotive sensors such as LIDAR."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "TizVe3hJHrJK",
   "metadata": {},
   "source": [
    "In terms of actions, the easiest approach is to have a number of *discrete* choices to go `left`, `right`, or `stay` in the current lane. We could be more sophisticated about it and have both \"aggressive\" and \"slow\" versions of these in addition to a default version, akin to the motion primitives we previously discussed.\n",
    "\n",
    "Actually implementing this on an autonomous vehicle, or even sketching an implementation in a notebook with recorded or simulated data, is beyond what we can accomplish in a jupyter notebook. Hence, we will be content below to sketch three popular foundational methods from deep reinforcement learning, without actually implementing them here. At the end of this chapter we provide some references where you can delve into these topics more deeply."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0xi2Y6T5YAtY",
   "metadata": {},
   "source": [
    "## Deep Q-Learning\n",
    "\n",
    "```{index} deep reinforcement learning; DQN\n",
    "```\n",
    "> DQN is an early deep learning RL method akin to Q-learning.\n",
    "\n",
    "Recall from Section 3.6 that we can define a policy in terms of **Q-values**, sometimes also called state-action values, and that we can define the optimal policy as \n",
    "\\begin{equation}\n",
    "\\pi^*(x) = \\arg \\max_a Q^*(x,a)\n",
    "\\end{equation}\n",
    "where $Q^*(x,a)$ denote the Q-values for the *optimal* policy. In Q-learning, we start with some random Q-values and then iteratively improve an estimate $\\hat{Q}(x,a)$ for the optimal Q-values by alpha-blending between old and new estimates:\n",
    "\\begin{equation}\n",
    "\\hat{Q}(x,a) \\leftarrow (1-\\alpha) \\hat{Q}(x,a) + \\alpha~\\text{target}(x,a,x').\n",
    "\\end{equation}\n",
    "Above, the \"target value\"\n",
    "\\begin{equation}\n",
    "\\text{target}(x,a,x') \\doteq R(x,a,x') + \\gamma \\max_{a'} \\hat{Q}(x',a')\n",
    "\\end{equation}\n",
    "is a value that we think is an improvement on the previous value $\\hat{Q}(x,a)$. Indeed: $\\text{target}(x,a,x')$ uses the *current* estimate of the Q-values for future states, but improves on this by using the *known* rewards $R(x,a,x')$ for the current action $a$ in the current state $x$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "pCKixeLwsh2Z",
   "metadata": {},
   "source": [
    "```{index} pair: deep Q-network; DQN\n",
    "```\n",
    "In the **deep Q-network** or DQN method we use a *supervised learning* approach to Q-learning. We train a neural network, parameterized by $\\theta$, to approximate the optimal Q-values:\n",
    "\\begin{equation}\n",
    "Q^*(x,a) \\approx \\hat{Q}(x,a; \\theta)\n",
    "\\end{equation}\n",
    "It might be worthwhile at this point to re-visit Section 5.6, where we introduced neural networks and how to train them using stochastic gradient descent (SGD).\n",
    "\n",
    "```{index} execution phase, experience replay\n",
    "```\n",
    "In the context of RL, the DQN method uses two additional ideas that are crucial in making the training converge to something sensible in difficult problems. The first is splitting the training into *execution* and *experience replay* phases:\n",
    "\n",
    "- during the **execution phase**, the policy is executed (possibly with some degree of randomness) and the experiences $(x,a,r,x')$, with $r$ the reward, are stored in a dataset $D$;\n",
    "- during **experience replay**, we *randomly sample* from these experiences to create mini-batches of data, which are in turn used to perform SGD on the parameters $\\theta$.\n",
    "\n",
    "The second idea is to calculate the target values \n",
    "\\begin{equation}\n",
    "\\text{target}(x,a,x') \\doteq R(x,a,x') + \\gamma \\max_{a'} \\hat{Q}(x',a'; \\theta^{old})\n",
    "\\end{equation}\n",
    "with the parameters $\\theta^{old}$ from the previous epoch, to provide a more stable approximation target.\n",
    "The mini-batch loss we minimize using SGD is then\n",
    "\\begin{equation}\n",
    "\\mathcal{L}_{\\text{DQN}}(\\theta; D) \\doteq \\sum_{(x,a,r,x')\\in D} [\\text{target}(x,a,x') - Q(x,a; \\theta)]^2\n",
    "\\end{equation}\n",
    "\n",
    "```{index} off-policy RL\n",
    "```\n",
    "With this basic scheme, a team from DeepMind was able to achieve human or super-human performance on about 50 Atari 2600 games in 2015 {cite:p}`Mnih15nature_dqn`.\n",
    "DQN is a so-called **off-policy** method, in that each execution phase uses the best policy we computed so far, but we can still replay earlier experiences gathered with \"lesser\" policies. Nothing in the experience replay phase references the policy: every experience leads to a valid Q-value backup and a valid supervised learning signal."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "D6PHabNMU4OO",
   "metadata": {},
   "source": [
    "## Policy Optimization\n",
    "\n",
    "```{index} deep reinforcement learning; policy optimization\n",
    "```\n",
    "> Policy optimization takes a black box optimization approach to a deep policy.\n",
    "\n",
    "```{index} stochastic policy\n",
    "```\n",
    "Whereas the above gets at an optimal policy indirectly, via deep Q-learning, a different and very popular idea is to directly parameterize the policy using a neural network, with weights $\\theta$. It is common to make this a **stochastic policy**,\n",
    "\\begin{equation}\n",
    "\\pi(a|x; \\theta)\n",
    "\\end{equation}\n",
    "where $a \\in {\\cal A}$ is an action, $x \\in {\\cal X}$ is a state, and the policy outputs a *probability* for each action $a$ based on the state $x$. One of the reasons to prefer stochastic policies is that they are differentiable, as they output continuous values rather than discrete actions. This allows us to optimize for them via gradient descent, as we explore in the next section.\n",
    "\n",
    "```{index} cross-entropy\n",
    "```\n",
    "In Chapter 5 we used *supervised* learning to train neural networks, and we just applied this for learning Q-values in DQN. It is useful to consider how this might work for training a *policy*. Recall from Section 5.6 that we defined the empirical cross-entropy loss as\n",
    "\\begin{equation}\n",
    "\\mathcal{L}_{\\text{CE}}(\\theta; D) \\doteq - \\sum_{(x,y=c)\\in D} \\sum_c \\log p_c(x;\\theta)\n",
    "\\end{equation}\n",
    "where $D$ is a dataset of supervised learning examples. Above $y$ takes on a discrete class value $c$, but another very common way to write this is via a \"one-hot encoding\". Let $y_c$ be 1 if $y=c$ and 0 otherwise, then our loss becomes:\n",
    "\\begin{equation}\n",
    "\\mathcal{L}_{\\text{CE}}(\\theta; D) = -\\sum_{(x,y)\\in D} \\sum_c y_c \\log p_c(x;\\theta)\n",
    "\\end{equation}\n",
    "\n",
    "This formulation is equivalent, but now we are summing over all classes for each data point, with $y_c$ acting as a *weight*, either one or zero. When someone is so kind as to give us the optimal action $y_a$ (as a one-hot encoding) for every state $x$ in some dataset $D$, we can apply this same loss function to a stochastic policy, obtaining\n",
    "\\begin{equation}\n",
    "\\mathcal{L}_{\\text{CE}}(\\theta; D) = -\\sum_{(x,y)\\in D} \\sum_a y_a \\log \\pi(a| x; \\theta)\n",
    "\\end{equation}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "So1rSw4zS-C5",
   "metadata": {},
   "source": [
    "```{index} policy optimization, off-policy RL\n",
    "```\n",
    "In **policy optimization** we gather data by rolling out a set of trajectories $\\tau_i$. In supervised learning we have a dataset $D$ and labels $y_c$, but we have to proceed a bit differently in a reinforcement learning setting. In particular, for **on-policy** RL we gather data by executing our current best guess for the policy for some rollout length or horizon $H$, and we do this many different times, each time obtaining a *trajectory* $\\tau_i$.\n",
    "That still leaves the training signal: where does that come from? \n",
    "The key idea is to estimate how good a particular action is by estimating the state-action values $Q$ from the rollout rewards.\n",
    "In detail, we estimate the expected discounted reward starting at $x_{it}$, and taking action $a_{it}$, as\n",
    "\\begin{equation}\n",
    "\\hat{Q}(x_{it},a_{it}) \\doteq \\sum_{k=t}^H \\gamma^{k-t}R(x_{ik},a_{ik},x_{ik}').\n",
    "\\end{equation}\n",
    "Note in each rollout we can only sum until $k=H$, so Q-values earlier in the rollout will be estimated more accurately. Regardless, we can then use these estimated Q-values as an alternative to the \"one or zero\" weight above, obtaining the **surrogate loss** $\\mathcal{L}(\\theta)$ by averaging over the trajectories $\\tau_i$:\n",
    "\\begin{equation}\n",
    "\\mathcal{L}(\\theta) = - \\sum_i \\sum_{t=1}^H \\hat{Q}(x_{it},a_{it}) \\log \\pi(a_{it}|x_{it} \\theta).\n",
    "\\end{equation}\n",
    "\n",
    "As you can see, once again estimated Q-values are involved, and they are expected to converge to the optimal Q-values over time.\n",
    "This time, however, we parameterize the *policy* instead. The Q-values act as weights in the surrogate loss function that looks very much like the supervised cross-entropy loss, but now our supervision comes from *experience*."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ME-fRw_8CDg8",
   "metadata": {},
   "source": [
    "Putting this all together yields the basic policy optimization algorithm:\n",
    "\n",
    "- Initialize $\\theta$\n",
    "- Until convergence:\n",
    "    1. roll out a number of trajectories $\\tau_i$ using the current policy $\\pi(a;x,\\theta)$\n",
    "    2. try to change the parameters $\\theta$ as to decrease the surrogate loss function $\\mathcal{L}(\\theta)$\n",
    "    \n",
    "A simple, gradient-free approach for step 2 is simple hill-climbing, aka stochastic search:\n",
    "\n",
    " - perturb $\\theta$ to $\\theta'$\n",
    " - set $\\theta \\leftarrow \\theta'$ *iff* $\\mathcal{L}(\\theta') < \\mathcal{L}(\\theta)$\n",
    "\n",
    "The *perturbation* step could be as simple as adding some Gaussian noise to the weights of the network.\n",
    "More sophisticated \"black-box\" approaches such as genetic algorithms or evolution strategies can be applied to this problem instead, but they all share the property that they are \"gradient-free\", which seems to be a sub-optimal strategy. Hence, we next look at a gradient descent approach to policy optimization."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "-sLUpvmQ2sNd",
   "metadata": {},
   "source": [
    "## Policy Gradient Methods\n",
    "\n",
    "```{index} deep reinforcement learning; policy gradient methods\n",
    "```\n",
    "> Policy gradient methods are akin to policy iteration, with a neural flavor.\n",
    "\n",
    "```{index} softmax function, logit\n",
    "```\n",
    "In a nutshell, policy gradient methods calculate the *gradient* of the surrogate loss $\\mathcal{L}(\\theta)$ defined above with respect to the policy parameters $\\theta$:\n",
    "\\begin{equation}\n",
    "\\nabla_\\theta \\mathcal{L}(\\theta) \\leftarrow - \\sum_i \\sum_{t=1}^H \\hat{Q}(x_{it},a_{it}) \\nabla_\\theta \\log \\pi(a_{it}|x_{it}, \\theta),\n",
    "\\end{equation}\n",
    "where $\\nabla_\\theta \\log \\pi(a_{it}|x_{it}, \\theta)$ is the gradient of the logarithm of the stochastic policy. This is easily obtained via back-propagation using any neural network framework of choice. In the case that actions are discrete, as in our example above, a stochastic policy network typically ends with a *softmax* function. Recall that if the network outputs a vector of raw scores (the *logits*) $z \\in \\mathbb{R}^K$, the softmax is defined as\n",
    "\\begin{equation}\n",
    "\\mathrm{softmax}(z)_i = \\frac{e^{z_i}}{\\sum_{j=1}^K e^{z_j}}, \\quad i = 1,\\dots,K.\n",
    "\\end{equation}\n",
    "Thus, the logits are the raw outputs before applying the softmax, and $\\nabla_\\theta \\log \\pi(a_{it}|x_{it}, \\theta)$ is computed with respect to *these* values.\n",
    "We then use gradient descent to update the policy parameters:\n",
    "\\begin{equation}\n",
    "\\theta \\leftarrow \\theta - \\alpha \\nabla_\\theta \\mathcal{L}(\\theta)\n",
    "\\end{equation}\n",
    "where $\\alpha$ is a learning rate.\n",
    "\n",
    "The algorithm above, using the estimated Q-values, is almost identical to the REINFORCE method {cite:p}`Williams92ml_reinforce`. That algorithm further improves on performance by not using the raw Q-values but rather the difference between the Q-values and some baseline policy. This has the effect of reducing the variance in the estimated Q-values due to using only a finite amount of data.\n",
    "The REINFORCE algorithm was introduced in 1992 and hence pre-dates the deep-learning revolution by about 20 years. It should also be said that in DRL, the neural networks that are used are typically not very deep. Several modern methods, such as \"proximal policy optimization\" (PPO) {cite:p}`Schulman17_PPO` apply a number of techniques to improve this basic method even further and make it more sample-efficient. PPO is now one of the most often-used DRL methods."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "S66_driving_DRL.ipynb",
   "provenance": []
  },
  "interpreter": {
   "hash": "c6e4e9f98eb68ad3b7c296f83d20e6de614cb42e90992a65aa266555a3137d0d"
  },
  "kernelspec": {
   "display_name": "Python 3.9.6 64-bit ('base': conda)",
   "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.9.18"
  },
  "latex_metadata": {
   "affiliation": "Georgia Institute of Technology",
   "author": "Frank Dellaert and Seth Hutchinson",
   "title": "Introduction to Robotics"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}