Skip to main content
$$ \newcommand{\R}{\mathbb{R}} \newcommand{\N}{\mathbb{N}} \newcommand{\Z}{\mathbb{Z}} $$

Lecture 18: policy gradients

Goal: synthesize optimal policies using function approximation.

References

  • Amari, S.-I. Natural Gradient Works Efficiently in Learning. Neural Comput. 10, 251–276 (1998).
  • Kakade, S. A natural policy gradient. Adv. Neural Inf. Process. Syst. 14, 1531–1538 (2001).
  • Fazel, M., Ge, R., Kakade, S. & Mesbahi, M. Global Convergence of Policy Gradient Methods for the Linear Quadratic Regulator. in Proceedings of International Conference on Machine Learning (ICML) vol. 80 1467–1476 (PMLR, 10–15 Jul 2018).

Problem formulation

Consider the optimization problem \(\min_{u} c(u)\)

This could have came from on MDP in which case \[ c(u)=E\left[\sum_{t=0}^{T} \mathcal{L}_{t}\left(x_{t}, u_{t}\right) \mid u_{t} \sim P_{u}(x), x_{t+1} \sim P_{x}\left(x_{t}, u_{t}\right)\right] \] or a trajectory optimization problem in which case \[ c(u)=\sum_{t=0}^{T} \mathcal{L}_{t}\left(x_{t}, u_{t}\right)\ \text{s.t.}\ x_{t}^{+}= F\left(x_{t}, u_{t}\right). \]

But it doesn’t really matter! The methods belaw are agnostic.

policy gradients

Consider policy gradient: \(u^{+}=u - \gamma D c^{\top}(u)\), and recall that this descent algorithm can be derived from an auxiliary optimization problem.

\[\min \langle D c(u), v\rangle \ \ \text{s.t.} \ \ \|v\|_{2} \leq \|D c(u)\|_{2}. \]

Since \(\langle x, y\rangle=\|x\|_{2}\|y\|_{2} \cos \theta\) where \(\theta\) is the angle between \(x \& y\), the solution is \(v^{*}=-Dc(u)^{\top} \in \mathbb{R}^{m}\)

What if we consider a different optimization problem? One idea is to use a local quadratic model for the cost landscape specified by the (non-zero) gradient and Hessian

\[ \min_{v} c(x) = D{c}(x) \cdot v + \frac{1}{2} v^{\top} D^{2} c(x) v. \]

The derivative of this cost with respect to \(v\) is \(Dc(x) + v^\top D^2 c(x)\), which is stationary if and only if \(v=-\left[D^{2} c(x)\right]^{-1} D_{c}(x)^{\top}\). This choice of \(v\) is termed the natural policy gradient.

Note: \(D^2 c(x)\) must be positive-definite for the optimization problem to be well-posed (i.e. to define a unique finite solution).

Let’s explore what makes this gradient “natural”. Consider a purely quadratic cost \(c(u)=\frac{1}{2}\left(u-u^{*}\right)^{\top} H\left(u-u^{*}\right)\) so that \(D c(u)=\left(u-u^{*}\right)^{\top} H\) and \(D^{2} c=H\). Let’s see how the policy gradient and natural policy gradient perform under a linear change of coordinates \(w = L u\) where \(L\) is an invertible matrix.

PG: \(u^{+}=u-\gamma H\left(u-u^{*}\right) \implies w^{+}=w-\gamma L H\left(L^{-1} w-u^{*}\right)\) which are different equations and will have different convergence rates depending on the choice of \(L\).

NPG: \(u^{+}=u-\gamma H^{-1} H\left(u-u^{*}\right) \implies w^{+}=w-\gamma L H^{-1} H\left(L^{-1} w-u^{*}\right)\) which both simplify to \(v - \gamma (v - u^*)\) regardless of the choice of \(L\).

Another term for the natural policy gradient is Newton-Raphson. These are second-order methods because they use second derivatives of the cost. They can be costly to compute but converge in fewer iterations, so there is a space-time tradeoff.

example: LQR

Let’s consider the Linear-Quadratic Regulator (LQR) problem where \(x^{+}=A x+B u,\ x_{0} \sim P,\ u_{t}=-K x_{t}\) where \(P\) is a probability distribution and the cost is \[ c(K)=E\left[\sum_{t=0}^{\infty} x_{t}^{\top} Q x_{t}+u_{t}^{\top} R u_{t}\right] \]

Proposition: \(D c(K)=2\left(\left(R+B^{\top} P_{K} B\right) K-B^{\top} P_{K} A\right) \Sigma_{K}\) where \[ P_{K}=Q+K^{\top} R K+(A-B K)^{\top} P_{K} (A-B K) \] is a Riccatti equation and \(\Sigma_{K}=E\left[\sum_{t=0}^{\infty} x_{t} x_{t}^{\top}\right]\) is the state covariance.

Proof: \[\begin{aligned} c\left(K ; x_{0}\right)=E\left[x_{0}^{\top} P_{K} x_{0}\right] & =E\left[x_{0}^{\top}\left(Q+K^{\top} R K\right) x_{0}+x_{0}^{\top}(A-B K)^{\top} P_{K}(A-B K) x_{0}\right] \\ & =E\left[x_{0}^{\top}\left(Q+K^{\top} R K\right) x_{0}+c\left(K ;(A-B K) x_{0}\right)\right] \end{aligned}\]

Recalling that \(D_{M} x^{\top} f(M) x=D f(M) x x^{\top}\) and \(D_{M} M^{\top} N M=N M\) we have

\[\begin{aligned} D_{K} c\left(K; x_{0}\right) & = E\left[2 R K x_{0} x_{0}^{\top} - 2 B^{\top} P_{K}(A-B K) x_{0} x_{0}^{\top}+D_{K} c(K ; (A-B K) x_{0}) \right] \\ & = E\left[2 \left( (R + B^\top P_K B) K - B^{\top} P_{K} A \right) x_{0} x_{0}^{\top}+ D_{K} c(K ; (A-B K) x_{0}) \right] \\ \end{aligned}\] And by induction with \(x_t = (A - B K)^t x_0\) \[\begin{aligned} D_{K} c\left(K; x_{t}\right) & = E\left[2 \left( (R + B^\top P_K B) K - B^{\top} P_{K} A \right) x_{t} x_{t}^{\top}+ D_{K} c(K ; x_{t+1}) \right] \\ \end{aligned}\]

and finally \[ D_{K} c\left(K; x_0\right) = E\left[2 \left( (R + B^\top P_K B) K - B^{\top} P_{K} A \right) \sum_{t = 0}^\infty x_{t} x_{t}^{\top} \right]. \]

This result gives us the following possible updates:

Policy gradient: \(K^{+}=K-\gamma D c(K).\)

Natural policy gradient: \(K^{+}=K-\gamma D c(K) \Sigma_{K}^{-1}.\)