Prelude

The aim of this post is to share my understanding of some of the conceptual and theoretical background behind implementations of the Proximal Policy Optimisation (PPO) reinforcement learning (RL) algorithm. PPO is widely used due to its stability and sample efficiency - popular applications include beating the Dota 2 world champions and aligning language models. While the PPO paper provides quite a general and straightforward overview of the algorithm, modern implementations of PPO use several additional techniques to achieve state-of-the-art performance in complex environments ^[1]. You might discover this if you try to implement the algorithm solely based on the paper. I try and present a coherent narrative here around these additional techniques.

I'd recommend reading parts one, two, and three of SpinningUp if you're new to reinforcement learning. There's a few longer-form educational resources that I'd recommend if you'd like a broader understanding of the field ^[2], but this isn't comprehensive. You should be familiar with common concepts and terminology in RL ^[3]. For clarity, I'll try to spell out any jargon I use here.

Recap

Policy Gradient Methods

PPO is an on-policy reinforcement learning algorithm. It directly learns a stochastic policy function parameterised by $\theta$ representing the likelihood of action $a$ in state $s$, ${\pi }_{\theta }\left(a\right|s)$. Consider that we have some differentiable function, $J\left(\theta \right)$, which is a continuous performance measure of the policy ${\pi }_{\theta }$. In the simplest case, we have $J\left(\theta \right)=E_{\tau \sim {\pi }_{\theta }}\left[R\right(\tau \left)\right]$, which is known as the return ^[4] over a trajectory ^[5], $\tau$. PPO is a kind of policy gradient method ^[6] which directly optimizes the policy parameters $\theta$ against $J\left(\theta \right)$. The policy gradient theorem shows that:

$${\nabla }_{\theta }J\left(\theta \right)=E\left[\sum _{t=0}^{inf}{\nabla }_{\theta }ln{\pi }_{\theta }\right(a_t|s_t\left)R_t\right]$$

In other words, the gradient of our performance measure $J\left(\theta \right)$ with respect to our policy parameters $\theta$ points in the direction of maximising the return $R_t$. Crucially, this shows that we can estimate the true gradient using an expectation of the sample gradient - the core idea behind the REINFORCE ^[7] algorithm. This is great. This expression has the more general form which substitutes $R_t$ for some lower-variance estimator of the total expected reward, $\Phi$ ^[8] :

$$\begin{array}{ccc}{\nabla }_{\theta }J\left(\theta \right)& =E\left[\sum _{t=0}^{inf}{\nabla }_{\theta }ln{\pi }_{\theta }\right(a_t|s_t\left){\Phi }_t\right]& \left(1\right)\end{array}$$

Modern implementations of PPO make the choice of ${\Phi }_t=A^{\pi }\left(s_{t,}{a}_t\right)$, the advantage function. This function estimates the advantage of a particular action in a given state over the expected value of following the policy, i.e. how much better is taking this action in this state over all other actions? Briefly described here, the advantage function takes the form

$$A^{\pi }(s,a)=Q^{\pi }(s,a)-V^{\pi }\left(s\right)$$

where $V\left(s\right)$ is the state-value function, and $Q(s,a)$ is the state-action -value function, or Q-function ^[9]. I've found it easier to intuit the nuances of PPO by following the narrative around its motivations and predecessor. PPO iterates on the Trust Region Policy Optimization (TRPO) method which constrains the objective function with respect to the size of the policy update. The TRPO objective function is defined as ^[10][11] :

$$\begin{array}{cc}J\left(\theta \right)=& E\left[\frac{{\pi }_{\theta }(a_t,s_t)}{{\pi }_{{\theta }_{old}}(a_t,s_t)}{A}_t\right]\\ \text{subject to}& E\left[\text{KL}\right({\pi }_{{\theta }_{old}}\left|\right|{\pi }_{\theta }\left)\right]\le \delta \end{array}$$

Where KL is the Kullback-Liebler divergence (a measure of distance between two probability distributions), and the size of policy update is defined as the ratio between the new policy and the old policy:

$$r\left(\theta \right)=\frac{{\pi }_{\theta }(a_t,s_t)}{{\pi }_{{\theta }_{old}}(a_t,s_t)}$$

Policy gradient methods optimise policies through (ideally small) iterative gradient updates to parameters $\theta$. The old policy, ${\pi }_{{\theta }_{old}}(a_t,s_t)$, is the one used to generate the current trajectory, and the new policy, ${\pi }_{\theta }(a_t,s_t)$ is the policy currently being optimised ^[12]. If the advantage is positive, then the new policy becomes greedier relative to the old policy ^[13], and we have $r\left(\theta \right)>1$ - the inverse applies for negative advantage, where we have $r\left(\theta \right)<1$. The core principle of TRPO (and PPO) is to prevent large parameter updates which occur due to variance in environment and training dynamics, thus helping to stabilise optimisation. Updating using the ratio between old and new policies in this way allows for selective reinforcement or penalisation of actions whilst grounding updates relative to the original, stable policy ^[14].

PPO

PPO modifies the TRPO objective function by constraining the objective function:

$$J\left(\theta \right)=E\left[min\right(\frac{{\pi }_{\theta }(a_t,s_t)}{{\pi }_{{\theta }_{old}}(a_t,s_t)}{A}_t,clip(\frac{{\pi }_{\theta }(a_t,s_t)}{{\pi }_{{\theta }_{old}}(a_t,s_{t)}},1-ϵ,1+ϵ)A_t]$$

To break this somewhat dense equation down, we can substitute our earlier expression $r\left(\theta \right)$ in:

$$J\left(\theta \right)=E\left[\text{min}\right(r\left(\theta \right)A_t,\text{clip}\left(r\right(\theta ),1-ϵ,1+ϵ)A_t\left)\right]$$

The clip term constrains the policy ratio, $r$, to lie between $1-ϵ$ and $1+ϵ$ ^[15]. In this manner, the objective function disincentives greedy policy updates in the direction of improving the objective. When the new policy places lower density on actions compared to the previous policy, i.e. may be more conservative, the advantage update is smaller. When the new policy places higher density on actions compared to the previous policy i.e. is greedier, the advantage update is also smaller. This is why PPO is considered to place a lower pessimistic bound on policy updates. The policy is only updated by $1+ϵ$ or $1-ϵ$, depending on whether the advantage function is positive or negative, respectively.

So far, we've introduced the concept of policy gradients, objective functions in RL, and the core concepts behind PPO. Reinforcement learning algorithms place significant emphasis in reducing variance during optimisation ^[16]. This becomes apparent in estimating the advantage function which relies on rewards obtained during a trajectory. Practical implementations of policy gradient algorithms take additional steps to trade variance for bias here by also estimating an on-policy state-value function $V_{\varphi }^{\pi }\left(s_t\right)$, which is the expected return an agent receives from starting in state $s_t$, and following policy $\pi$ thereafter. Jointly learning a value function and policy function in this way is known as the Actor-Critic framework ^[17].

Value Learning and Actor-Critics

Value-function learning involves approximating the future rewards of following a policy from a current state. The value function is learned alongside the policy, and the simplest method is to minimise a mean-squared-error objective against the discounted return ^[18], $R\left(\tau \right)=\sum _{t=0}^{inf}{\gamma }^t{r}_t$:

$${\varphi }^{\ast }={\text{arg min}}_{\varphi }{E}_{s_{t,R_t\sim }\pi }\left[\right(V_{\varphi }\left(s_t\right)-{\gamma }^{t}r\left(s_t\right){)}^2]$$

We can now use this to learn a lower-variance estimator of the advantage function ^[19] :

$$\begin{array}{cc}{\hat{A}}^{\pi }(s,a)& =E_{s_{t+1}}\left[Q^{\pi }\right(s_t,a_t)-V_{\varphi }^{\pi }(s_t\left)\right]\\ & =E_{s_{t+1}}[r_t+\gamma V_{\varphi }^{\pi }(s_{t+1})-V_{\varphi }^{\pi }(s\left)\right]\end{array}$$

We end up with an estimator of the advantage function that solely relies on samples rewards and our learned approximate value function. In fact, our expression shows that the advantage function can be estimated with the temporal-difference ^[20] residual error, ${\delta }_t$, of the value function:

$$\begin{array}{cc}{V}_{\varphi }\left(s_t\right)& \leftarrow V_{\varphi }(s_{t)}+\alpha (r_t+\gamma V_{\varphi }\left(s_{t+1}\right)-V_{\varphi }\left(s_t\right))\\ V_{\varphi }\left(s_t\right)& \leftarrow V_{\varphi }(s_{t)}+\alpha {\delta }_t\\ \text{where}{\delta }_t& =r_t+\gamma V_{\varphi }\left(s_{t+1}\right)-V_{\varphi }\left(s_t\right)\end{array}$$

Generalised Advantage Estimation (GAE)

There's one more thing we can do to reduce variance. The current advantage estimator estimates reward by samples collected from a single trajectory. However, there is a huge amount of variance in the possible trajectories that may evolve from a given state. These trajectories may look similar in the short-term (except for policies early in the optimisation process which are far more random), but longer-term rewards can vary significantly. We could consider a lower-variance estimate of the advantage by only sampling rewards near to the present, and replacing longer-term reward estimates with our lower-variance, biased value function estimates. This is the central idea behind n-step returns^[21]. If we choose $n$ correctly, i.e. before trajectories start to diverge significantly, we can obtain better samples of near-term rewards.

Consider the term ${\delta }_t$ in our estimation of the advantage function. We take the initial reward observed from the environment, $r_t$, then bootstrap future estimated discounted rewards, and subtract a baseline estimated value function for the state ^[22]. For a single time-step, this can be denoted as:

$${\widehat{A_t}}^{\left(1\right)}=r_t+\gamma V_{\varphi }^{\pi }\left(s_{t+1}\right)-V_{\varphi }^{\pi }\left(s\right)={\delta }_t^{V_{\varphi }}$$

What if we sample rewards from multiple timesteps, and then estimate the future discounted rewards from then on? Let's denote ${\widehat{A_t}}^{\left(k\right)}$ as follows:

$$\begin{array}{cccc}{\widehat{A_t}}^{\left(1\right)}& =r_t+\gamma V_{\varphi }^{\pi }\left(s_{t+1}\right)-V_{\varphi }^{\pi }\left(s\right)& & ={\delta }_t^{V_{\varphi }}\\ {\widehat{A_t}}^{\left(2\right)}& =r_t+\gamma r_{t+1}+{\gamma }^2{V}_{\varphi }^{\pi }\left(s_{t+2}\right)-V_{\varphi }^{\pi }\left(s\right)& & ={\delta }_t^{V_{\varphi }}+\gamma {\delta }_{t+1}^{V_{\varphi }}\\ {\widehat{A_t}}^{\left(k\right)}& =r_t+\gamma r_{t+1}+...+r_{t+k-1}+{\gamma }^k{V}_{\varphi }^{\pi }\left(s_{t+k}\right)-V_{\varphi }^{\pi }\left(s\right)& & ={\delta }_t^{V_{\varphi }}+\gamma {\delta }_{t+1}^{V_{\varphi }}+{\gamma }^2{\delta }_{t+1}^{V_{\varphi }}\\ {\widehat{A_t}}^{\left(\infty \right)}& =r_t+\gamma r_{t+1}++{\gamma }^2{r}_{t+2}+...-V_{\varphi }^{\pi }\left(s\right)& & =\sum _{l=0}^{k-1}{\gamma }^l{\delta }_{t+1}^{V_{\varphi }}\end{array}$$

$n=1$ and $n=inf$ now become extreme cases in $n$-step return estimation. Observe that for $k=\infty$ we recover an unbiased, high-variance expectation of the infinite-horizon discounted return, minus our baseline estimate value function. GAE ^[23] introduces a discount factor, $\lambda \in [0,1]$, to take an exponentially weighted average over every $k$-th step estimator. Using notation from the paper ^[24] we can derive a generalized advantage estimator for cases where $0<\lambda <1$:

$$\begin{array}{cc}{\widehat{A_t}}^{\gamma \lambda }& =(1-\lambda )({\widehat{A_t}}^{\left(1\right)}+\lambda {\widehat{A_t}}^{\left(2\right)}+{\lambda }^2{\widehat{A_t}}^{\left(3\right)}+...)\\ & =(1-\lambda )({\delta }_t^{V_{\varphi }}+\lambda ({\delta }_t^{V_{\varphi }}+\gamma {\delta }_{t+1}^{V_{\varphi }})+{\lambda }^2({\delta }_t^{V_{\varphi }}+\gamma {\delta }_{t+1}^{V_{\varphi }}+{\gamma }^2{\delta }_{t+2}^{V_{\varphi }})+...)\\ & =(1-\lambda )\left({\delta }_t^{V_{\varphi }}\right(1+\lambda +{\lambda }^2+...)+\gamma {\delta }_{t+1}^{V_{\varphi }}(\lambda +{\lambda }^2+{\lambda }^3+...)+...)\\ & =(1-\lambda )({\delta }_t^{V_{\varphi }}\left(\frac{1}{1-\lambda }\right)+\gamma {\delta }_t^{V_{\varphi }}\left(\frac{\lambda }{1-\lambda }\right)+{\gamma }^2{\delta }_t^{V_{\varphi }}(\frac{{\lambda }^2}{1-\lambda }\left)\right)\\ & =\sum _{l=1}^{\infty }(\gamma \lambda {)}^l{\delta }_{t+1}^{V_{\varphi }}\end{array}$$

Great. As you may have noticed, there's two special cases for this expression - $\lambda =0$, and $\lambda =1$:

$$\begin{array}{cc}{\widehat{A_t}}^{\gamma \ast 0}& ={\widehat{A_t}}^{\left(1\right)}=r_t+\gamma V_{\varphi }^{\pi }\left(s_{t+1}\right)-V_{\varphi }^{\pi }\left(s\right)={\delta }_t\\ {\widehat{A_t}}^{\gamma \ast 1}& =\sum _{l=1}^{\infty }{\gamma }^l{\delta }_{t+1}^{V_{\varphi }}\\ & =(r_t+\gamma V_{\varphi }(s_{t+1})-V_{\varphi }(s_t\left)\right)\\ & +\gamma (r_{t+1}+\gamma V_{\varphi }(s_{t+2})-V_{\varphi }(s_{t+1}\left)\right)\\ & +{\gamma }^2(r_{t+2}+\gamma V_{\varphi }(s_{t+3})-V_{\varphi }(s_{t+2}\left)\right)+...\\ & =r_{t+}\gamma V_{\varphi }\left(s_{t+1}\right)-V_{\varphi }\left(s_t\right)\\ & +\gamma r_{t+1}+{\gamma }^2{V}_{\varphi }\left(s_{t+2}\right)-\gamma V_{\varphi }\left(s_{t+1}\right)\\ & +{\gamma }^2{r}_{t+2}+{\gamma }^3{V}_{\varphi }\left(s_{t+3}\right)-{\gamma }^2{V}_{\varphi }\left(s_{t+2}\right)+...\\ & =\sum _{l=1}^{\infty }{\gamma }^l{r}_{t+1}-V_{\varphi }\left(s_t\right)\end{array}$$

We see that $\lambda =0$ obtains the original biased low-variance actor-critic advantage estimator. For $\lambda =1$, we obtain a low-bias, high-variance advantage estimator, which is simply the discounted return minus our baseline estimated value function. For $\lambda \in (0,1)$, we obtain an advantage estimator which allows control of the bias-variance tradeoff. The authors note that the two parameters, $\gamma$ and $\lambda$, control variance in different ways. Lower values of $\gamma$ discount future rewards, which will always result in a biased estimate of the return, since there's an implicit prior that future rewards are less valuable. The authors note that $\gamma$ controls the strength of this prior regardless of how accurate our value function is. On the other hand, since $n$-step returns are unbiased estimators of the advantage function, lower values of $\lambda$ reduce variance when the value function is accurate. In other words, when $V^{\varphi }\left(s_t\right)=V\left(s_t\right)$ and for $0<\lambda <1$ , we obtain an unbiased estimator of the advantage function.

Pseudocode

Tying everything together, we can show the general process of updating our policy and value function using PPO for a single trajectory of fixed length $N$:<br><br>

Given policy estimator ${\pi }_{\theta }$, value function estimator $V^{\varphi }$, $\gamma$, $\lambda$, $N$ time-steps.

For $t=1,2,...,N:$

Run policy ${\pi }_{\theta }$ in environment and collect rewards, observations, actions, and value function estimates $r_t,s_t,a_r,v_t$ where $v_t=V^{\varphi }\left(s_t\right)$

Compute ${\delta }_t^{V^{\varphi }}=r_t+\gamma v_{t+1}-v_t$ for all $t$.

Compute generalized advantage estimate, $\widehat{A_t}=\sum _{l=0}^N(\gamma \lambda {)}^t{\delta }_t$ for all t.

Sample ${\pi }_{\theta }(a_t,s_t)$ $log$-probabilities from stored actions and states.

Optimise $\theta$ using $J\left(\theta \right)$, the PPO objective ^[25].

Optimise $\varphi$ using $L=\left(V^{\varphi }\right(s_t)-{\gamma }^t{r}_t{)}^2$ using stored $r_t$ and $s_t$ ^[26].

... repeat!

Feedback and corrections are very welcomed and appreciated. I'll follow this post with implementation details soon. For now, I'd recommend the excellent resource by Shengyi Huang on reproducing PPO implementations from popular RL libraries. If you're able to implement the policy update from the PPO paper, hopefully there's enough detail here that you're able to reproduce other implementations of PPO. Thank you so much for reading.