RL5 - Policy Gradient methods.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Class 5: Policy Gradient methods.**\n",
    "\n",
    "1. [Learning parametric policies](#policies)\n",
    "2. [Policy gradient methods](#PG)\n",
    "3. [Objective function](#objective)\n",
    "4. [The Policy Gradient theorem](#theorem)\n",
    "5. [REINFORCE: Monte Carlo Policy Gradient](#reinforce)\n",
    "6. [Actor-Critic methods](#ac)\n",
    "7. [The link with Policy Iteration](#pi)\n",
    "8. [Deep Actor-Critic algorithms and references](#deep)\n",
    "9. [Practice](#practice)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id=\"policies\"></a>Learning parametric policies\n",
    "\n",
    "**Bottomline question:**<br>\n",
    "The previous classes have focussed on *action-value methods*; they aimed at estimating $Q^*$ in order to deduce $\\pi^*$. Could we directly optimize $\\pi$?\n",
    "\n",
    "Suppose we have a policy $\\pi_\\theta$ parameterized by a vector $\\theta$. Our goal is to find the parameter $\\theta^*$ corresponding to $\\pi^*$.\n",
    "\n",
    "Remarks:\n",
    "- $\\pi_\\theta$ might not be able to represent $\\pi^*$. We will take a shortcut and call $\\pi^*$ the best policy among the $\\pi_\\theta$ ones.\n",
    "- for discrete state and action space, the tabular policy representation is a special case of policy parameterization.\n",
    "- policy parameterization is a (possibly useful) way of introducing prior knowledge on the set of the desired policies.\n",
    "- the optimal deterministic policies might not belong to the policy subspace of $\\pi_\\theta$, thus it makes sense to consider stochastic policies for $\\pi_\\theta$.\n",
    "- for problems with significant policy approximation, the best approximate policy (among $\\pi_\\theta$ ones) may very well be stochastic.\n",
    "- it makes even more sense to consider stochastic policies that it opens the family of environments that we can tackle, like partially observable MDPs or multi-player games.\n",
    "\n",
    "For stochastic policies, we shall write $\\pi_\\theta(a|s)$.\n",
    "\n",
    "In the remainder of the class, we will assume that $\\pi_\\theta$ is differentiable with respect to $\\theta$.\n",
    "\n",
    "To directly optimize $\\theta$ we need a criterion $J(\\theta)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id=\"PG\"></a>Policy gradient methods\n",
    "\n",
    "Suppose now we define some performance metric $J(\\pi_\\theta) = J(\\theta)$. If $J$ is differentiable and a stochastic estimate $\\nabla_\\theta J(\\theta)$ of the gradient is available, then we can define the gradient ascent update procedure:\n",
    "$$\\theta \\leftarrow \\theta + \\alpha \\nabla_\\theta J(\\theta).$$\n",
    "\n",
    "We will call **policy gradient methods** all methods that follow such a procedure (whether or not they also learn a value function or not).\n",
    "\n",
    "Remarks: \n",
    "- Note that $J$ is a more general criterion that might differ from $Q$ in the definition above (even though it seems reasonable to assume both should be related). For example, $J$ could be defined as the value of a starting state (or a distribution of starting states) in episodic cases, or as the undiscounted reward over a certain horizon, or as the average reward.\n",
    "- Why is it interesting to look at policy gradient methods? Because for continuous actions there is no maximization step ($\\max_a Q(s,a)$) during evaluation but only a call to $\\pi_\\theta(s)$ (or a draw from $\\pi_\\theta(a|s)$). This makes Policy Gradient a method of choice for continuous actions domains (especially common in Robotics).\n",
    "- When do policy gradient approaches outperform value-based ones? It's hard to give a precise criterion; it really depends on the problem. One thing that comes into play is how easy it is to approximate the optimal policy or the optimal value function. If one is simpler than the other (by \"simpler\", we mean \"it is easier to find a parameterization whose spanned function space almost includes the function to approximate\"), then it is a good heuristic to try to approximate it. But this criterion might itself be hard to assess.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id=\"objective\"></a>Objective function\n",
    "\n",
    "Let's consider a starting state $x$, distributed according to $p_0(x)$, and let's define:\n",
    "$$J(\\theta) = \\mathbb{E}_x \\left[ V^{\\pi_\\theta}(x) \\right].$$\n",
    "\n",
    "Although it is not strictly necessary for the following sections, let's play a bit with this formulation. We will adopt the following notation:\n",
    "- We drop the $\\theta$ in $\\pi_\\theta$ (for the sake of readability).\n",
    "- $p^\\pi(s \\rightarrow s', k)$ is the probability of transitioning from $s$ to $s'$ in $k$ time steps, while applying policy $\\pi$.\n",
    "\n",
    "One can reformulate the definition of the value of state $x$ under policy $\\pi$:\n",
    "\\begin{align}\n",
    "V^\\pi(x) & = \\mathbb{E} \\left[ \\sum_t \\gamma^t r(s_t,a_t) | \\pi, x \\right],\\\\\n",
    " & = \\sum_t \\gamma^t \\mathbb{E}\\left( r(s_t,a_t) | \\pi, x \\right),\\\\\n",
    " & = \\sum_t \\gamma^t \\int_S \\int_A r(s,a) p^\\pi(x \\rightarrow s, t) \\pi(a|s) da ds,\\\\\n",
    " & = \\int_S \\sum_t \\gamma^t p^\\pi(x \\rightarrow s, t) \\int_A r(s,a) \\pi(a|s) da ds.\n",
    "\\end{align}\n",
    "\n",
    "Let's write $\\rho^\\pi(s)$ the improper discounted state distribution, given $\\pi$ and the starting state distribution $p_0(x)$:\n",
    "$$\\rho^\\pi(s) = \\int_S \\sum_t \\gamma^t p_0(x) p^\\pi(x \\rightarrow s, t) dx.$$\n",
    "\n",
    "Then we have:\n",
    "$$J(\\theta) = \\int_S \\rho^\\pi(s) \\int_A r(s,a) \\pi(a|s) da ds.$$\n",
    "\n",
    "In other terms:\n",
    "$$J(\\theta) = \\mathbb{E}_{\\substack{s\\sim\\rho^\\pi \\\\ a\\sim \\pi}} \\left[ r(s,a) \\right].$$"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id=\"theorem\"></a>The Policy Gradient theorem\n",
    "\n",
    "The crucial problem of computing $\\nabla_\\theta J(\\theta)$ lies in the fact that when $\\theta$ changes, both $\\pi$ and $\\rho^\\pi$ change in turn. So there seems to be no straighforward way of evaluating this gradient. One could fall back on a *finite differences* approach to estimating this gradient, but this would require trying out a series of increments $\\Delta \\theta$ which quickly becomes impractical (because the increment size is hard to tune, especially in stochastic systems, and also because of the sample inefficiency of the approach).\n",
    "\n",
    "Remark:\n",
    "- Let's not discard finite difference methods too quickly. They have their merits and showed great successes through methods such as PEGASUS (Ng and Jordan, 2000).\n",
    "\n",
    "The key result of this class is that one can express the gradient of $J(\\theta)$ as directly proportional to the value of $Q^\\pi$ and the gradient of $\\pi$:\n",
    "$$\\nabla_\\theta J(\\theta) \\propto \\mathbb{E}_{\\substack{s\\sim\\rho^\\pi \\\\ a\\sim \\pi}} \\left[ Q^\\pi(s,a) \\nabla_\\theta \\log\\pi(a|s)\\right]$$\n",
    "\n",
    "The proof of this result is simple but a bit tedious. We can however give the general intuition. Let's consider trajectories $\\tau = (s_0,a_0,r_0,...)$ drawn according to $\\pi$ from the starting state. Each of these trajectories has an overall payoff of $G(\\tau) = \\sum_t \\gamma^t r_t$ and is drawn with probability density $p_\\theta(\\tau)$. Then the objective function can be written:\n",
    "\\begin{align}\n",
    "J(\\theta) &= \\mathbb{E}_\\tau \\left[ G(\\tau) \\right]\\\\\n",
    " &= \\int G(\\tau) p_\\theta(\\tau) d\\tau\n",
    "\\end{align}\n",
    "\n",
    "So the objective function's gradient is:\n",
    "\\begin{align}\n",
    "\\nabla_\\theta J(\\theta) &= \\nabla_\\theta \\int G(\\tau) p_\\theta(\\tau) d\\tau,\\\\\n",
    " &= \\int G(\\tau) \\nabla_\\theta p_\\theta(\\tau) d\\tau,\\\\\n",
    " &= \\int G(\\tau) p_\\theta(\\tau) \\frac{\\nabla_\\theta p_\\theta(\\tau)}{p_\\theta(\\tau)} d\\tau,\\\\\n",
    " &= \\mathbb{E}_\\tau \\left[ G(\\tau) \\nabla \\log p_\\theta(\\tau) \\right].\n",
    "\\end{align}\n",
    "\n",
    "The transition from this expression on trajectories to an expression on states and actions provides the result."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id=\"reinforce\"></a>REINFORCE: Monte Carlo Policy Gradient\n",
    "\n",
    "Let's apply directly the Policy Gradient theorem. To compute the gradient, we can run the policy within the environment, this will provide us with states distributed according to $\\rho^\\pi$ and actions distributed according to $\\pi$. The full trajectory of states-actions-rewards provides a Monte Carlo estimate $G_t$ of $Q^\\pi(s_t,a_t)$ from any state $s_t$ traversed by the trajectory. In turn, this allows to compute $Q^\\pi(s_t,a_t) \\nabla_\\theta \\pi(a_t|s_t)$ for any of these states. The sum over all states provides the gradient estimate.\n",
    "\n",
    "This algorithm, introduced by Williams (1992) is called REINFORCE. It requires a finite-length trajectory and its pseudo-code goes as follows.\n",
    "1. Initialize policy parameter $\\theta$\n",
    "2. Generate a trajectory by playing $\\pi$: $s_0,a_0,r_0,...s_{T}$\n",
    "3. For $t\\in [1, 2, … , T]$:\n",
    "    1. Estimate return $G_t$\n",
    "    1. Update policy parameter: $\\theta \\leftarrow \\theta + \\alpha \\gamma^t G_t \\nabla_\\theta \\log \\pi(a_t|s_t)$\n",
    "    \n",
    "When one takes a look at the $G_t \\nabla_\\theta \\log \\pi(a_t|s_t) = G_t \\frac{\\nabla_\\theta \\pi(a_t|s_t)}{\\pi(a_t|s_t)}$ product, its interpretation is intuitive. $\\nabla_\\theta \\pi(a_t|s_t)$ is a vector in parameter space that points in the direction of greatest increase of $\\pi(a_t|s_t)$. The update will encourage taking a step in this direction if the action provided high return (through $G_t$), but will discourage moving in this direction if the action is already picked frequently (through $\\pi(a|s)$) so that other actions have a chance also.\n",
    "\n",
    "The gradient estimate of the policy gradient often has a high variance. A common practice consists in substracting an action-independent *baseline* $b(s)$ from the estimate of $Q^\\pi(s,a)$, yielding:\n",
    "$$\\nabla_\\theta J(\\theta) \\propto \\mathbb{E}_{\\substack{s\\sim\\rho^\\pi \\\\ a\\sim \\pi}} \\left[ \\left( Q^\\pi(s,a) -b(s) \\right) \\nabla_\\theta \\log\\pi(a|s)\\right]$$\n",
    "\n",
    "It is rather easy to remark that this baseline does not affect the gradient estimate since it's expected value is zero. However it can contribute to strongly decrease the estimate's variance. One common choice for such a baseline is the policy's value function $V^\\pi$. This turns REINFORCE into an *advantage* estimation algorithm where the advantage is the function defined as:\n",
    "$$A^\\pi(s,a) = Q^\\pi(s,a) - V^\\pi(s)$$\n",
    "\n",
    "But estimating this advantage function requires estimating $V^\\pi$, which in turn requires to maintain a (possibly parameterized) function $V$ on top of the current (parameterized) policy $\\pi$.\n",
    "\n",
    "We will leave this topic aside for now and will directly look at the generalization of the Policy Gradient theorem to Actor-Critic methods."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id=\"ac\"></a>Actor-Critic methods\n",
    "\n",
    "Suppose now that we don't want a Monte Carlo estimate of $Q^\\pi(s,a)$ in the Policy Gradient theorem, and are rather willing to store a function approximator for $Q^\\pi(s,a)$. This leads us to store both a policy and a value function. The value function *criticizes* the policy's selected actions, hence the names of *critic* and *actor*.\n",
    "\n",
    "Remark that the temporal difference at each time step $\\delta = r + \\gamma V^\\pi(s') - V^\\pi(s)$ is an estimate of the advantage $A^\\pi(s,a)$. Using this remark, a simple one-step Actor-Critic method based on TD(0) and a value function $V_w$ goes as follows:\n",
    "1. In $s$, draw $a \\sim \\pi$\n",
    "2. Observe $r, s'$\n",
    "3. Compute $\\delta = r + \\gamma V_w(s') - V_w(s)$\n",
    "4. Update critic's parameters (TD(0) step) $w \\leftarrow w + \\alpha \\delta \\nabla_w V_w(s)$\n",
    "5. Update actor's parameters (policy gradient theorem) $\\theta \\leftarrow \\theta + \\alpha \\delta \\nabla_\\theta \\log \\pi(a|s)$\n",
    "6. $s\\leftarrow s'$ and repeat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id=\"pi\"></a>The link with Policy Iteration\n",
    "\n",
    "Let's take a step back and reconsider the Actor-Critic architecture we have just introduced.\n",
    "\n",
    "Basically, on the one hand, we have a critic $V_w$ or $Q_w$ that aims at estimating the $V^\\pi$ or $Q^\\pi$ value function of policy $\\pi$. And on the other hand, we have an actor whose policy $\\pi_\\theta$ is incrementally improved so as to maximize $J(\\theta)$. This should sound familiar.\n",
    "\n",
    "In part, this is familiar because this resembles a lot the SARSA update. Let's take a minute to spot the differences.\n",
    "\n",
    "But more generally, this actually belongs to the class of approximate Policy Iteration algorithms. Let's recall the Policy Iteration procedure:\n",
    "1. Solve $Q=T^\\pi Q$\n",
    "2. Solve $\\pi = Greedy(Q)$\n",
    "3. Repeat\n",
    "\n",
    "And let's now allow for an approximate resolution of these steps, via gradient descent:\n",
    "1. Approximately solve $\\min \\|Q - T^\\pi Q \\|$ via gradient descent\n",
    "2. Approximately solve $\\pi = Greedy(Q)$ using the policy gradient theorem\n",
    "3. Repeat\n",
    "\n",
    "After each collected sample, the update of the Actor-Critic algorithm above performs exactly one gradient step on the critic and one gradient step for the actor.\n",
    "\n",
    "This perspective allows to define a much broader family of Actor-Critic methods that perform various number of gradient steps on the critic or the actor, use n-step returns, introduce a sequence of $Q_{i+1} = T^\\pi Q_i$ functions, soften the policy gradient steps, etc. in order to make the overall Actor-Critic algorithm more efficient and robust."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id=\"deep\"></a>Deep Actor-Critic algorithms and references\n",
    "\n",
    "At this stage, it is really tempting to throw a neural network at our critic and our actor, start collecting samples in a replay buffer, and try to design Deep Policy Gradient methods. Let's review some of the key ones from the litterature.\n",
    "\n",
    "- [Asynchronous Advantage Actor-Critic (A3C) (2016)](https://arxiv.org/abs/1602.01783). Builds a unique network that approximates $V$ and $\\pi$, replaces the replay buffer with an army of asynchronous actors that provide independent samples for the gradient computations. It is the direct adaptation of the Actor-Critic algorithm above. Its little brother A2C discards the asynchronous aspect, while keeping the good overall performance.\n",
    "- [Trust Region Policy Optimization (TRPO) (2015)](https://arxiv.org/abs/1502.05477). Imposes small policy gradient steps by introducing a \"maximum KL divergence between successive policies\" constraint in the actor's update.\n",
    "- [Proximal Policy Optimization (PPO) (2017)](https://arxiv.org/abs/1707.06347). Same philosophy as TRPO but simpler and more efficient. Instead of a KL divergence constraints, it imposes a gradient clipping for the policy gradient.\n",
    "- [Scalable trust-region method for deep reinforcement learning using Kronecker-factored approximation (ACKTR) (2017)](https://arxiv.org/abs/1708.05144). Uses Kronecker factorization to make TRPO's update more efficient.\n",
    "- [Sample Efficient Actor-Critic with Experience Replay (ACER) (2017)](https://arxiv.org/abs/1611.01224). Several improvements upon TRPO and A2C.\n",
    "- [Soft Actor Critic Algorithms (SAC) (2019)](https://arxiv.org/abs/1812.05905). Introduces an entropy regularization term in the objective function.\n",
    "- [Modified Actor Critic algoritms (MoPPO) (2019)](https://arxiv.org/abs/1907.01298). Casts the critic update in a modified policy iteration scheme by building the sequence of $Q_{i+1} = T^\\pi Q_i$ functions, applies this to PPO.\n",
    "\n",
    "One thing that was not covered in this class is the [Deterministic Policy Gradient theorem](http://proceedings.mlr.press/v32/silver14.html). which allows to perform policy gradient steps on deterministic policies (with many benefits). This family of algorithms spanned their own deep counterparts. Notably:\n",
    "- [Deep Deterministic Policy Gradients (DDPG) (2015)](https://arxiv.org/abs/1509.02971). Implements the DPG theorem on deep neural networks, with a replay buffer.\n",
    "- [Twin Delayed Deep Deterministic Policy Gradients (TD3) (2018)](https://arxiv.org/abs/1802.09477). Introduces three improvements over DDPG, namely a double critic update, two separate networks for the critic and target policy smoothing.\n",
    "- [Distributed Distributional Deep Deterministic Policy Gradients (D4PG) (2018)](https://arxiv.org/abs/1804.08617). Improves on DDPG with parallel actors, a distributional value function estimator, and batch normalization.\n",
    "\n",
    "One can remark that the (stochastic) Policy Gradient update is an *on-policy* update: it requires the samples to have been drawn from the current policy's stationary distribution. This was generalized to off-policy updates in the [Off-Policy Actor Critic](https://arxiv.org/abs/1205.4839) paper (2012) and has been used in most of the algorithms above.\n",
    "\n",
    "Lilian Weng keeps a nice overview (and zoo) of Actor-Critic methods [on her blog](https://lilianweng.github.io/lil-log/2018/04/08/policy-gradient-algorithms.html).\n",
    "\n",
    "Finally, to grasp an overview of Policy search methods (even beyond the scope of Policy Gradients) a good read is the [Policy Search in Continuous Action Domains: an Overview](https://arxiv.org/abs/1803.04706) paper (2019)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# <a id=\"practice\"></a>Practice\n",
    "\n",
    "Implement an Advantage Actor-Critic algorithm using a single network representing both the policy and the value function. Test your implementation on your favorite Atari game, or on a new gym environment (like the [box2d ones](https://gym.openai.com/envs/#box2d)).\n",
    "Draw inspiration from the [A3C](https://arxiv.org/abs/1602.01783) paper!"
   ]
  }
 ],
 "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}