We are reaching the end of our series, but we still have a few steps to do in our reinforcement learning (RL) journey before we can fully flesh out proximal policy optimization (PPO).
If you’ve been following along and comparing these posts to the Sutton & Barto textbook (or other RL literature), you might have noticed that we’ve purposely skipped a lot of algorithms (even popular ones) in the RL world, such as off-policy Monte Carlo control, expected SARSA, Dyna-Q, true online TD(λ), and so on. In RL, much like the rest of the computer science research world, new algorithms are proposed all the time. In most cases, these new algorithms offer a slight improvement or fixes to a problem on an older algorithm. My goal for this series was to provide concrete stepping stones from the absolute basics of RL to PPO, which meant I had to ignore all the algorithms that branched off that main path.
PPO is a synthesis of many of the ideas we’ve been building on throughout the series. It is an actor-critic method (like the one we saw in the previous post), but it includes a number of techniques and tricks to help improve sample efficiency, maintain training stability, and reduce hyperparameter sensitivity. It builds on the lessons learned from other actor-critic algorithms, like trust region policy optimization (TRPO). One of the most important tricks is the use of a surrogate objective function, which allows actor-critic methods to use batched rollout data practically (i.e. without having to perform a gradient ascent step for every timestep in the rollout data).
The rest of this post will focus on defining a surrogate objective function and showing how a ratio of the policies makes multiple updates possible from a single rollout.
Rollout Batches
We left the previous post with the one-step actor-critic algorithm that updates both actor and critic parameters after each timestep. While this opens up the possibility of continuing tasks (i.e. without having to wait for the episode to end), it proves to be an inefficient use of samples.
PPO collects a bunch of T timesteps (i.e. a rollout). It uses the states, actions, and rewards from these timesteps as a batch of data. Rather than update the actor and critic parameters once per timestep or once per batch, it splits the batch into minibatches and performs an update using each minibatch. Each time through the full batch is known as an epoch, so batch_size / minibatch_size updates are performed for each epoch. PPO performs updates across K epochs.
Note that the policy gradient we derived in post 12 assumed the trajectories were sampled from the current policy. When we make our first minibatch update (via a gradient ascent step), θ changes, which means the data in our batch was generated by a policy that no longer exists. By the second gradient step update, we’re using samples from a different policy. The rest of PPO exists to handle this scenario. It uses a correction factor to account for the change in action probabilities and a mechanism that keeps us from trusting that correction further than it deserves.
The Surrogate Objective Function
In post 12, we defined J(θ) to be the expected return under policy π, given a set of parameters θ.
\(J(\boldsymbol{\theta}) = \sum_{\tau} p(\tau; \boldsymbol{\theta}) \cdot G(\tau) = \mathbb{E}_{s_0 \sim p(s_0)} \left[ v_{\pi_{\boldsymbol{\theta}}}(s_0) \right]\)
This is a type of objective function, which is the value we attempt to maximize (if using gradient ascent, e.g. total discounted return in RL) or minimize (if using gradient descent, e.g. a loss function for supervised learning). In other words, we want to find values for the variables of differentiation (e.g. the parameters θ or ϕ) that either maximize or minimize the objective function.
In most cases, we need to differentiate the objective function in order to find the gradient. This gives us a direction to climb up or down the “hill” (i.e. how to change the variables of differentiation) of the objective function toward increasing or decreasing objective values. For example, we saw the update rule in the REINFORCE algorithm with the gradient of the objective function already set for us:
θ ← θ + α · γt · Gt · ∇θ log π(At|St;θ) Recall that we spent posts 12 and 13 mathematically proving that we can determine the gradient of the expected return under a given policy by sampling (rather than having to compute it analytically):
\(\nabla J(\boldsymbol{\theta}) = \mathbb{E}_{\tau \sim \pi_\theta}\left[ \sum_{t} \gamma^t \cdot \nabla_\theta \log \pi(a_t|s_t;\boldsymbol{\theta}) \cdot G_t \right]\)
In post 14, we used temporal difference (TD) error as an advantage function with the estimated state value of the next state as the baseline:
\(\hat{A}_t = \delta_t = R_{t+1} + \gamma \cdot V(S_{t+1}; \boldsymbol{\phi}) – V(S_t; \boldsymbol{\phi})\)
This provides lower variance than the full return-to-go Gt (at the cost of some bias), and, when used in combination with a critic neural network (NN), gives us the following update rules (with φ as the parameters for the critic and θ as the parameters for the actor):
δ ← Rₜ₊₁ + γ · V(Sₜ₊₁; φ) − V(Sₜ; φ) (V(Sₜ₊₁; φ) = 0 if Sₜ₊₁ is terminal)
φ ← φ + αφ · δ · ∇φ V(Sₜ; φ)
θ ← θ + αθ · I · δ · ∇θ log π(Aₜ|Sₜ; θ)In most implementations, the gradient is found numerically for us using an autograd library, like the one found in PyTorch. The above updates work fine when used with a single sample (i.e. the “single step” in single-step actor-critic). If you tried to do this on a batch of 2048 samples, you’d have to perform 2048 full backward passes through the network, which is computationally inefficient.
In our case, J(θ) cannot be computed at all without a model of the environment, and we spent all of post 12 showing how to use the expected value to get the gradient, which could be used in the updates for REINFORCE and one-step actor-critic.
In order to work with batches (and not have to perform a backward pass for each sample), we introduce a surrogate objective function: a substitute function used in place of the original objective function (usually when the original objective function is discontinuous, non-differentiable, or impractical to work with). The idea is that the gradient of the surrogate objective function and the gradient of the original objective function should be the same at a given set of learned parameters (e.g. φ or θ) so that performing a gradient ascent or descent step produces the same result.
We’ll use L(θ) to denote our surrogate objective function, which we’ll use instead of J(θ) (the total expected return that we want to maximize). We need to choose a function L(θ) such that:
\(\nabla_\theta L(\boldsymbol{\theta}) \approx \nabla_\theta J(\boldsymbol{\theta})\)
We use the approximate symbol (≈) here for two reasons. First, L(θ) is computed from a batch average rather than a true expectation, which is the same Monte Carlo sampling approximation we covered in post 7. Second, Ât is itself an estimate (denoted by the “hat” symbol modifier), which we’ll use in our particular surrogate function.
Many batched policy gradient methods use the following surrogate function. It allows us to approximate ∇θ J(θ) (by using ∇θ L(θ)) even if J(θ) ≠ L(θ).
\(L(\boldsymbol{\theta}) = \frac{1}{N}\sum\limits_{t} \hat{A}_t \cdot \log \pi_\theta(a_t|s_t)\)
Note the notation switch here: πθ(aₜ|aₜ) is the same as π(aₜ|aₜ; θ). Many modern deep RL papers (including the PPO paper) use the former notation. I’ll switch to this θ subscript version going forward so that it lines up with the notation used in the papers.
Ât is computed from the batch before any parameter updates happen, so it’s a fixed number as far as θ is concerned (with respect to differentiation). When we differentiate L(θ), the Ât term passes straight through as a constant multiplier.
\(\nabla_\theta L(\boldsymbol{\theta}) = \frac{1}{N}\sum\limits_{t} \hat{A}_t \cdot \nabla_\theta \log \pi_\theta(a_t|s_t)\)
Notice that because Ât is constant with respect to θ, differentiating the surrogate reproduces the same gradient estimator we used in the previous post. Also, the two approximations mentioned earlier (the batch average 1/N · Σt standing in for a true expectation and Ât standing in for the true advantage A(s,a)) are properties of that estimator, not of the surrogate. Switching to a surrogate objective adds no bias or variance of its own.
When using an autograd tool, we can calculate a single scalar value, L(θ), from the whole batch of samples, which the tool can use to compute the gradient. Note that actually computing the value of L(θ) does not give us any useful information. All we care about is that the gradients of the surrogate and original objective functions are the same (or close enough) so that when we perform a gradient ascent step, we’re moving the NN parameters (θ) toward maximizing the real objective, J(θ).
We can use this surrogate gradient in place of the original gradient in our actor update rule from the one-step actor-critic method in the previous post. Note that we maintain the original ∇θ log πθ(Aₜ|Sₜ), but we can now work with a minibatch of samples.
\(\boldsymbol{\theta} \leftarrow \boldsymbol{\theta} + \alpha_\theta \cdot \frac{1}{N}\sum\limits_{t} \hat{A}_t \cdot \nabla_\theta \log \pi_\theta(a_t|s_t)\)
You might have noticed that the discount factor γ (or the accumulated version I) has mysteriously disappeared. While it would be mathematically rigorous to keep it, most practical implementations of actor-critic methods drop this term to help the update move more quickly toward maximizing returns. Note that γ is usually kept in other parts of the algorithm, however.
Rather than compute a separate update for each timestep, we can now effectively compute one average gradient for the entire batch and perform a (more efficient) single update. We would then repeat this process K · batch_size / minibatch_size times.
There’s a catch, though: L(θ) breaks after the first update. L(θ) reproduces the correct gradient (i.e. ∇θ L(θ) ≈ ∇θ J(θ)) when the data came from the original policy (πθold). At the start of the second update, θ has changed, which means computing ∇θ L(θ) on the second update (and later) is no longer valid: you would compute probabilities from a policy that did not generate the rollout data. The fix is to introduce a correction factor.
Surrogate Correction Factor
To create a correction factor, we’ll need to rely on the concept of importance sampling, which is a way to estimate the expected value of some function (e.g. E[f(x)]) by sampling from a different distribution than what originally generated f(x). In other words, we want the expected value of f(x) when x is drawn from distribution p, but our samples were drawn from a different distribution q. We can introduce a new probability distribution, q(x), and use it to estimate the original E[f(x)] generated by p(x).
Here is a quick proof to show how that works mathematically:
\( \begin{align*} \mathbb{E}_{x \sim p}[f(x)] &= \sum\limits_{x} p(x) \cdot f(x) \\ &= \sum\limits_{x} q(x) \cdot \frac{p(x)}{q(x)} \cdot f(x) \\ &= \mathbb{E}_{x \sim q}\left[ \frac{p(x)}{q(x)} \cdot f(x) \right] \end{align*}\)
We start with the definition of expected value (sum over values of f(x), weighted by the individual probabilities). Inside the summation, we multiply each term by 1, as given by q(x) / q(x). We can apply the definition of expected value here, but instead of calculating the expectation over the original p(x) distribution, we’re doing it over the new q(x) distribution. So now, we can estimate the original expected value by sampling under q(x) and doing a little math.
There is one important caveat to keep in mind: q(x) > 0 wherever p(x) > 0. You cannot correct for outcomes that have no chance of being sampled in the first place.
In our case, p(x) is the current policy that we are updating, πθ, and q(x) is the original policy used to collect the rollout data, πθold.
We can use importance sampling to create a new surrogate function that takes both the old and new policies into account:
\(\begin{align*} L^{CPI}(\boldsymbol{\theta}) &= \frac{1}{N}\sum\limits_{t} \hat{A}_t \cdot \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)} \\ &= \hat{\mathbb{E}}_t \left[ \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)} \cdot \hat{A}_t \right] \end{align*}\)
Note: this ratio-based surrogate comes from the paper Approximately optimal approximate reinforcement learning. The CPI superscript means “conservative policy iteration.”
You’ll often see this written as Êt […]. This is the empirical average over the timesteps in the batch, which is the same thing as 1/N · Σt […]. Additionally, the new and old policy ratio is often given the term rt(θ). As a result, we can write the surrogate objective in the shortened form:
\(L^{CPI}(\boldsymbol{\theta}) = \hat{\mathbb{E}}_t \left[ r_t(\boldsymbol{\theta}) \cdot \hat{A}_t \right]\)
Something important to note here: on the first gradient step of each batch, πθ is the policy that collected the rollout data, so for this first step only, πθ = πθold. As a result, the ratio is 1, and it works the same as the general surrogate function we saw earlier (without the ratio). To demonstrate this, here is the gradient of the CPI surrogate objective:
\(\begin{align*} \nabla_\theta L^{CPI}(\boldsymbol{\theta}) &= \nabla_\theta \left[ \frac{1}{N}\sum\limits_{t} \hat{A}_t \cdot \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)} \right] \\ &= \frac{1}{N}\sum\limits_{t} \nabla_\theta \left[ \hat{A}_t \cdot \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)} \right] \end{align*}\)
With the assumption that πθ = πθold on the first gradient step, we get:
\(\begin{align*} \nabla_\theta L^{CPI}(\boldsymbol{\theta}) \Big|_{\boldsymbol{\theta} = \boldsymbol{\theta}_{old}} &= \frac{1}{N}\sum\limits_{t} \frac{\hat{A}_t}{\pi_{\theta_{old}}(a_t|s_t)} \cdot \nabla_\theta \pi_\theta(a_t|s_t) \\ &= \frac{1}{N}\sum\limits_{t} \hat{A}_t \cdot \frac{\nabla_\theta \pi_\theta(a_t|s_t)}{\pi_\theta(a_t|s_t)} \\ &= \frac{1}{N}\sum\limits_{t} \hat{A}_t \cdot \nabla_\theta \log \pi_\theta(a_t|s_t) \end{align*}\)
The final line should look familiar: it’s the basic surrogate gradient we derived earlier. So, as long as πθ = πθold (i.e. before the first gradient update), we showed that the gradient of the CPI surrogate is the same as the gradient of the basic surrogate:
\(\nabla_\theta L^{CPI}(\boldsymbol{\theta}) \Big|_{\boldsymbol{\theta} = \boldsymbol{\theta}_{old}} = \nabla_\theta L(\boldsymbol{\theta})\)
At that first step, the CPI surrogate gives us the same batched gradient we started with. However, once the policy moves away from the original policy (πθold) on subsequent updates, each sample’s contribution gets reweighted by how much more or less likely the current policy (πθ) is to have taken that action. In other words, the ratio rt(θ) allows us to update the current policy by using a rollout from an old policy.
Keep in mind that it does not correct for the fact that a different policy would have visited different states. As a result, the surrogate LCPI(θ) works as a good approximation for the real objective when the current policy πθ is reasonably close to the old policy πθold. So, how do we know what “too far” of a policy update looks like? PPO addresses this limitation by introducing a “clip” that prevents the ratio from moving the parameters too far in a single update.
Implementation of the Correction Factor
How do you actually calculate πθ(at | st) / πθold(at | st)? In code, this is relatively straightforward. We don’t need to divide one full probability distribution over another. We just need scalar values, taken at each timestep, and we perform basic division.
During the rollout, the agent takes action at from state st at timestep t. Just keep a record of log πθold(at | st) (the log-probability of the action actually taken) for that timestep along with the action and state. For our Gaussian policy (continuous action probability distribution we saw in post 12), it is the log-density of the sampled action under the network’s mean and standard deviation outputs.
At update time, run a forward pass using the current policy on the same state (st) to get the probability distribution (i.e. mean and standard deviation). Compute the log-density using the state and action actually taken in the rollout for that timestep to get log πθ(at | st). From there, we can then numerically compute the following for that timestep t:
\(r_t(\boldsymbol{\theta}) = \exp\left( \log \pi_\theta(a_t|s_t) – \log \pi_{\theta_{old}}(a_t|s_t) \right)\)
Probabilities can be extremely small, which creates an underflow risk during division. We subtract the logs and then exponentiate to avoid this risk (i.e. it’s more numerically stable).
Repeat this for all timesteps in the rollout for a given batch. Note that log πθold(at | st) is recorded once during a rollout while log πθ(at | st) is recomputed on each forward pass.
Conclusion
Surrogate objectives are common in RL and other parts of machine learning. However, they can require some time to fully grasp, which is why we spent an entire blog post addressing them. We also built up the idea of using a specific LCPI(θ) surrogate function that relies on a new vs. old policy ratio, which allows us to utilize samples collected under one policy to update the parameters in a new policy. As a result, we can now collect rollouts in batches and perform multiple updates to the actor and critic using a single batch.
In the next post, we’ll dive into the other mathematical machinery that makes PPO work.
If you have any questions or feedback, please let me know in the comments below!