RL for LLMs, Part 1: RL Foundations

In this post, we will review details about traditional RL and then over the next few posts, we will connect it to the LLM setting.

Think of the simplest scenario, kind of examples that RL is introduced in schools, like the balancing of a CartPole. The objective is to balance the pole for as long as possible -- that's the reward. The state is pole's position, velocity, angle, and angular velocity. The action is force applied to the base of the pole in order to give left or right velocity in order to balance the pole. And finally, the transition dynamics is given by the physics of the pole. This simple scenario is a Markov Decision Process (MDP). Generally speaking, a MDP is a tuple of $(\mathcal{S}, \mathcal{A}, P, R, \gamma)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $P$ is the transition dynamics, $R$ is the reward function, and $\gamma$ is the discount factor. A policy is a mapping from states to distribution over actions, $\pi: \mathcal{S} \to \mathcal{A} \mid \mathcal{S}$, so if you are in a given state, the policy will tell you which action to take in order to maximize the reward, and that is the goal of any RL setup -- and I will bold it so that we know where we start in RL : The goal of any RL setup is to find the policy that maximizes the sum of expected rewards over time.

Formally, the objective is $J(\pi) = E_{\pi} [\sum_{t=0}^{T} \gamma^t R_t]$ where $\gamma$ is the discount factor and $T$ is the horizon.

One can solve for the best policy or learn the algorithm for the best policy using this framework.

In this series of posts, our goal is to understand how traditional RL led to current state-of-the-art RL for LLMs, so this will be more of an in-depth review towards a specific direction rather than a comprehensive review of everything RL.

Goal: The goal is to select actions that maximize the expected sum of reward. But that's the final outcome, we need to find a way in order to achieve it, that happens by choosing actions that lead to high rewards. So let's define a trajectory as a sequence of states, actions, and rewards, $\tau = (s_0, a_0, r_0, s_1, a_1, r_1, \ldots, s_T, a_T, r_T, s_{end})$, where $s_{end}$ is the terminal state.

Say we get some reward at the end of the trajectory, $R$, now the question is how do we assign credit to the actions taken in this trajectory? There are two possible ways this happens (a): the environment gives us reward for each step in which case we compute something termed as reward-to-go and use tha to assign credit to each action, or (b) when we only get terminal reward. In this case again there are two ways to address this: (i.) each action in a trajectory receives the same credit or (ii.) we estimate the reward-to-go using different techniques and use that to assign credit to each action individually.

Now lets define reward-to-go, $G_t = \sum_{i=t}^{T} \gamma^{i-t} R_{i}$, this is the sum of rewards (also called realized return) from the current step to the end of trajectory. Rewards-to-go are a more accurate representation of reward as it removes the rewards that occured before the action. An action cannot cause past rewards. However, its application is uninformative in scenario when we don't have access to step-wise rewards and $\gamma = 1$ (no discounting), which is the common situation when doing RL for LLMs. In this case, at each step what we want to know is "when I took action 'a' in this step, what is the expected return averaging over all future states and actions" -- this naturally leads us to defining another term, action-value function:

$$Q^{\pi}(s, a) = E_{\pi}[G_t \mid S_{t} = s, A_{t} = a]$$

We can use $Q$ to improve a policy, by choosing actions that maximize the expected return, lets define that as

$$J(\theta) = \sum_{a} \pi_{\theta}(a \mid s)Q^{\pi}(s, a)$$

$$\nabla_{\theta} J(\theta) = \sum_{a} Q^{\pi}(s, a) \nabla_{\theta} \pi_{\theta}(a \mid s)$$

$$\nabla_{\theta} \pi_{\theta}(a \mid s) = \pi_{\theta}(a \mid s) \nabla_{\theta} \log \ \pi_{\theta} (a \mid s)$$

$$\nabla_{\theta} J(\theta) = \sum_{a} Q^{\pi}(s, a) \pi_{\theta}(a \mid s) \nabla_{\theta} \log \ \pi_{\theta} (a \mid s)$$

$$\nabla_{\theta} J(\theta) = E_{a} [Q^{\pi}(s, a) \nabla_{\theta} \log \ \pi_{\theta} (a \mid s) ]$$

This says sample an action 'a', estimate its Q value, and then update the policy in the direction of $Q^{\pi}(s, a) \nabla_{\theta} \log \ \pi_{\theta} (a \mid s)$, this updates the parameters as follows:

$$\theta \leftarrow \theta + \alpha * Q^{\pi}(s, a) \nabla_{\theta} \log \ \pi_{\theta} (a \mid s)$$

increasing the $\log \ \pi_{\theta} (a \mid s)$ increases the probability of that action and a larger Q-value gives the action a larger update. In case of LLM, $s_{t}$ is the current prefix, $a_{t}$ is the sampled token and $Q^{\pi}(s_{t}, a_{t})$ is the expected completion reward after selecting $a_{t}$, this the policy gradient update reinforces tokens that lead to greater expected reward. Note that the \log-score trick which gives us $\nabla_{\theta} \log \ \pi_{\theta} (a \mid s)$ is very easily computable, that is just the cross entropy loss of the sampled token which is done in even normal SFT. In SFT we try to minimize the $L = -\log \pi_{\theta} (a | s)$, which gives the gradient $- \nabla_{\theta} \log \ \pi_{\theta} (a \mid s)$, in RL, for a sampled action $a_{t}$, you instead optimize using $L = - Q^{\pi}(s, a) \log \pi_{\theta} (a | s)$, so its gradient is $- Q^{\pi}(s, a) \nabla_{\theta} \log \ \pi_{\theta} (a \mid s)$, so just the cross-entropy loss weighted by the $Q$ value at the sampled token.

For an entire trajectory, this is

$$\nabla_{\theta} J(\theta) = E_{\tau \sim \pi_{\theta}} [\sum_{t=0}^{T} Q^{\pi}(s_{t}, a_{t}) \nabla_{\theta} \log \ \pi_{\theta} (a_{t} \mid s_{t})]$$

So we have now landed at the first algorithmic approach, policy gradient theorem, which can be used in to optimize actions to increase reward.

Now let's try to further optimize this approach. One of the issues that plaques the vanilla policy gradient method is the noise. So the above equation of the action-value estimation or even the gradient update happens over an extreme large search-space and in practice we sample N trajectories and compute the gradients and updated over these sampled trajectories and the sampling process can determine the optimization landscape to a large extent.

Consider this example to understand this issue: imagine an environment when your rewards are between 0 and 1, so everything is positive and therefore each trajectory is pushed up by policy gradient update, except that good ones are pushed more than the bad ones (determined by the differences in the magnitude in their absolute reward). Now lets add a constant to all rewards, say shift = +100, so now all rewards are between 100 and 101. The true signal from the sampled data remains the same, but the variance in the optimizer updates becomes way larger, each trajectory is pushed up by 100 * log-prob gradient and the differences between good and bad trajectories becomes harder to distinguish, there a common mode in both the positive and negative trajectories which drowns the learning signal, i.e., high variance. Here's another way to see this issue: consider two actions $a_1$ and $a_2$ whose $Q(s, a_1)$ = 0 and $Q(s, a_2)$ = 1. In this situation it is clear that the optimizer updates should push up the value of $a_2$, however the gradient signal is the same is we shift these values, so now we get $Q(s, a_1)$ = 50 and $Q(s, a_2)$ = 51 and the optimizer will push up $a_1$ by 50 times its log-prob gradient when infact it should not be encouraged. To deal with this issue, we want to remove the common mode from all trajectories -- this preserves the same training signal while decreasing the variance.

A widely accepted mode that can be subtracted is the value function which is defined as the expected return from a particular state (irrespective of what action is taken), i.e.,

$$V^{\pi}(s) = E[G_{t} | S = s_{t} ]$$

This baseline-centered action-value is termed as advantage:

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

Computing Advantage

As we established above, advantage is the desired signal we want to optimize, however the environment rarely gives us advantage directly, it usually gives us rewards and that too mostly terminal. So how can we get the advantage? Reconsider the definition of advantage

$$A^{\pi}(s, a) \leftarrow Q^{\pi}(s, a) - V^{\pi}(s)$$

However this way of learning the value function has a few issues:

1. this approach returns an unbiased estimate of the return, but can have a high variance.
2. this approach has greater variance with longer episodes.
3. this approach cannot be used to update the value function until the trajectory finishes.

Therefore another approach was proposed to learn the value function called as TD-learning. In TD-learning instead of waiting till the end of the trajectory to estimate the return we use the step reward and the predicted value of the next state, i.e., $\hat{G_t} \leftarrow r_t + \gamma V_{\phi}(s_{t+1})$. And therefore the advantage become $A^{\pi}(s_t, a_t) = r_t + \gamma V_{\phi}(s_{t+1}) - V_{\phi}(s_t)$. For terminal-only rewards, $r_t$ = 0, so the $A^{\pi}(s_t, a_t) = \gamma V_{\phi}(s_{t+1}) - V_{\phi}(s_t)$ (at the actual terminal state, the terminal reward replaces the next-state prediction). Here's how TD-learning works:

1. First step is the target which is formed by $y_t^{TD} \leftarrow r_t + \gamma \ \text{stopgrad}(V_{\phi}(s_{t+1}))$
2. Update the value function to match the target, i.e., $\mathcal{L}_v(\phi) = \frac{1}{2}(V_{\phi}(s_t) - y_t^{TD})^2$
3. Defint TD error as $\delta_t = r_t + \gamma V_{\phi}(s_{t+1}) - V_{\phi}(s_t)$. The critic's update is $\phi \leftarrow \phi + \alpha * \delta_t * \nabla_{\phi}V_{\phi}(s_t)$
4. TD also gives us the advantage, consider $E[\delta_t \mid S = s_t, A = a_t] = E[r_t + \gamma V(s_{t+1}) - V(s_t) \mid S = s_t, A = a_t] = E [ r_t + \gamma V(s_{t+1}) \mid s_t, a_t] - V^{\pi}(s_t) = Q^{\pi}(s_t, a_t) - V^{\pi}(s_t) = A^{\pi}(s_t, a_t)$. Therefore the advantage is just the TD error.

The above described procedure is termed as one-step TD learning. However, this has issues of the opposite nature to Monte-Carlo estimate, i.e., it has quite high reliance on the critic which means high bias and it can be bad when the critic is inaccurate which they are in the initial part of the training. A middle-ground is reached with the n-step TD learning. In this the target is $y_t^n = \sum_{l=0}^{n-1} \gamma^{l} r_{t+l} + \gamma^n V_{\phi}(s_{t+n})$. In this case the advantage $\hat{A}_t^n = y_t^n - V_{\phi}(s_t)$. So this leads to the trade-off, small n leads to higher critic bias, lower variance and larger n leads to lower bias and higher variance.

Generalized advantage estimation (GAE) was proposed to smoothly interpolate between these two extremes. We know $\hat{A}_t^1 = y_t^1 - V_{\phi}(s_t), \ \hat{A}_t^2 = y_t^2 - V_{\phi}(s_t), \ldots$. GAE combines with using weights controlled by $\lambda$ (geometric weighted average), i.e., $\hat{A}_t^{GAE} = A_t^{(1)} + (\lambda) A_{t}^{(2)} + (\lambda)^2 A_{t}^{(3)} + \dots$ = $\hat{A}_t^{GAE} = \sum_{n=1}^{T-t} \lambda^{n-1} A_t^{(n)}$ (multiply the RHS by (1 - $\lambda$) so that the weights sum to 1). There is another way of representing the GAE which comes form $\hat{A}_t^{(1)} = \delta_t, \ \hat{A}_t^{(2)} = \delta_t + \gamma \delta_{t+1}, \ \hat{A}_t^{(3)} = \delta_t + \gamma \delta_{t+1} + \gamma^2 \delta_{t+2}.$ When taking geometric weighted average, we get $\hat{A}_t^{GAE} = (1 - \lambda)(A_t^{(1)} + \lambda A_t^{(2)} + \lambda^2 A_t^{(3)}) = (1 - \lambda) \delta_t + (\gamma \lambda) \delta_{t+1} + (\gamma \lambda)^2 \delta_{t+2} = \sum_{l=0}^{T-t} (\gamma \lambda)^l \delta_{t+l}$. Therefore

$$\hat{A}_t^{GAE} = (1 - \lambda)\sum_{n=1}^{\infty} \lambda^{n-1} \hat{A}_t^{(n)} = \sum_{l=0}^{\infty} (\gamma \lambda)^l \ \delta_{t+l}$$

To learn the critic with GAE, we use the same regression based approach as earlier, i.e., $\mathcal{L}_V(\phi) = \frac{1}{2}(V_{\phi}(s_t) - \text{stopgrad}(y_t^{GAE}))^2$, where $y_t^{GAE} = \hat{A}_t^{GAE} + V_t^{old}$.

Overall, this setup of learning is called the actor-critic setup, the actor $\pi_{\theta}$ is parametrized by $\theta$ choses the action and the critic $V_{\phi}$ parametrized by $\phi$ predicts the expected return.

Actor-Critic Based Optimization

The overall loop using an actor-critic based optimization has these steps:

1. Sample trajectories $\{\tau_{1}, \tau_{2}, \tau_{3}, \dots\}$ using the actor (policy) $\pi_{\theta_{old}}$
2. Observe rewards
3. Predict the value at each state $V_{old}(s_t)$
4. Compute TD errors $\delta_t = r_t + \gamma V_{old}(s_{t+1}) - V_{old}(s_t)$
5. Compute advantage (prefer GAE) $\hat{A}_t^{GAE} = (1 - \lambda)\sum_{n=1}^{\infty} \lambda^{n-1}\hat{A}_t^{(n)}$.
6. Construct critic target $y_{i, t} = \hat{A}_{i, t} + V_{old}(s_{i, t})$
7. Update the actor, $\theta_{new} \leftarrow \theta_{old} + \alpha \frac{1}{M} [\sum_{i,t} \text{stopgrad}(\hat{A}_{i,t}^{GAE}) \nabla_{\theta} \log(\pi(a_t \mid s_t))]$, where M is all the sampled state-action paris across all trajectories.
8. Update the critic, whose $\mathcal{L_V}(\phi) = \frac{1}{2M} \sum_{i,t}(V_{\phi}(s_{i, t}) - \text{stopgrad}(y_{i, t}))$ and update it using $\phi_{new} = \phi_{old} - \alpha \nabla_{\phi}\ \mathcal{L_V}(\phi)$

Now once we have updated the policy, the sampled trajectories $\{\tau_{1}, \tau_{2}, \tau_{3}, \dots\}$ were no longer generated by the new policy $\pi_{\theta_{new}}$ -- which means $\tau \sim \pi_{\theta_{old}}$ but we now want to optimize $\pi_{\theta_{new}}$. The policy gradient algorithm we have described above is on-policy, i.e., we need trajectories from the policy we want to optimize in order to optimize it. However, collecting trajectories for just one step optimization and then discarding them is not efficient. If possible, we would like to resue the rollouts we generated for several optimization steps. And this is reasonable because one optimization step usually shouldn't move the policy far enough that the sampled trajectories are completely unviable under its probabilities.

However usage of trajectories from earlier policy to update current policy brings up a few problems:

1. off-policyness: the mathematics we have established above for policy gradient theorem works under the assumption of on-policy sampling. Using off-policy samples does not hold up the same equations and the updates becomes not theoretically justified.
2. the number of updates that can be performed: Since the data is from old policy, we cannot use it to update the current policy many times, that has to be limited in order for the updates to be meaningful.

This leads to the introduction of a new concept, importance sampling. The importance sampling ratio computes how much less or more likely is the sampled action under the new policy compared to the old policy.

In general suppose we want to compute the expectation under a distribution $p$, $E_{x \ \sim \ p}[f(x)] = \sum_x p(x)f(x)$. But now our samples are not from distribution $p$, but another distribution $q$, then we rewrite this as $E_{x \ \sim \ p}[f(x)] = \sum_x q(x) \frac{p(x)}{q(x)} f(x) = E_{x \ \sim \ q} [\frac{p(x)}{q(x)}f(x)]$. The ratio corrects for the fact that the samples are not from $p$, but from $q$. When applied to RL, the ratio is equivalently written as $\rho_t(\theta) = \frac{\pi_{\theta}(a_t \mid s_t)}{\pi_{old}(a_t \mid s_t)}$.

Now the question is how do we use this concept of importance sampling?

We first compute the IS ratio using the formula described above and then take this into account when computing the loss. Recall the loss when the samples are on-policy: $\mathcal{L}(\theta) = -\frac{1}{M} \sum_{i, t}(\text{sg}(\hat{A}_{i, t}) \ \log \pi_{\theta}(a_{i,t} \mid s_{i, t}))$. Under IS this loss becomes $\mathcal{L}(\theta) = -\frac{1}{M} \sum_{i, t} \frac{\pi_{\theta}(a_{i, t} \mid s_{i, t})}{\pi_{old}(a_{i,t} \mid s_{i,t})} \text{sg}(\hat{A}_{i, t})$, so just multiply the advantage by the IS ratio.

So recall that the updates of RL looks similar to the SFT cross-entropy loss updates except that each token as a weight given by advantage of that token, now that advantage itself is weighted by the IS ratio of that token.

It is important to note that the above loss formulation only corrects the per-token ratio given a fixed prefix, it does not take into account the probability of the prefix itself. Correcting that would require the cumulative product of IS ratios at all prior token positions. So this formulation of the IS only considers the question "how likely the next token is given that we have this prefix", it does not consider "how likely the new policy is to reach the prefix in the first place". Why do we not do the full correction? This is because the cumulative products would become unstable. Even with just one per-token IS ratio being too small or too large the cumulative product can become exponentially large or small -- and may overflow or underflow numerically -- as trajectory length grows. And therefore we accept some bias from using the old policy's prefix distribution to trade-off some variance.

How long can we use previous sampled trajectories for?

So importance sampling allows us to reuse previous rollouts for optimizing current policy (under certain conditions) which helps us with the efficiency problem that we pointed out earlier.

However, we cannot use the stale data for infinitely long, eventually the updates push the current policy sufficiently far away from the trajectory sampling policy which makes the updates mathematically unjustified and the stored prefixes become less representative of the prefixes the new policy could generate.

And there is another issue which is the if the advantage of an action-state pair if greater than 0, and if we use the same prefixes, the update steps will keep encouraging this action-state pair until its probability reached 1 (unless the updates are constrained). During training we usually compute the advantage once from the rollout (using $\pi_{old}$) and treat it fixed (we'll discuss this). Let's define $p_{old} = \pi_{old}(a_{i,t} \mid s_{i, t})$ and $p_{\theta} = \pi_{\theta}(a_{i,t} \mid s_{i, t})$. The loss for the actor is $L_t(\theta) = -\rho_t(\theta) \ \text{sg}(\hat{A}_{i,t}) = -\frac{p_{\theta}}{p_{old}} \ \text{sg}(\hat{A}_{i,t})$.

Lets see what happens to the importance sampling ratio of an action-state pair whose $\hat{A}_{i,t} > 0$. So as actor is updated and its loss goes down (since advantage is fixed and positive), for the loss to go down the IS ratio $\rho_t(\theta)$ must increase. Therefore, everytime we optimize this same stored action-state pair, the objective continues to increase the probability of this action-state pair. Similarly if the $\hat{A}_{i,t} < 0$, then the objective continues to drive down the probability of this action-state pair.

So the number of times we can use the previously sampled trajectories must remain limited because of these reasons:

1. With updates, the new policy becomes sufficiently different that the stored prefixes become less representative of the sequences the new policy could generate.
2. The stored (i.e., fixed) advantage estimates continue to push the probability to 1 or 0 unless contrained.
3. The stored advantage estimates become stale under the new policy.

This raises a question, if the data reuse is limited because the advantage was calculated from $\pi_{old}$ and that advantage is stale (which leads to the issue #2 and #3) above why don't we find the advantage of the action-state pairs under the updated policy?

The reason we don't do this is because even though mechanically we could compute the advantages, since the continuation from a action-state pair came from the old policy, the $V_{new}$ is not a good estimate of the $V^{\pi_{new}}$ as $V_{new}$ did not get the correct training target (the policy was not sampled from $\pi_{new}$). To correctly estimate $V^{\pi_{new}}$ we would need to start from the stored prefix $s_t$, take the stored action $a_t$ and then generate the remainder using $\pi_{new}$ and score the new completion and use that to train the value function $V_{new}$ -- the entire motivation of the generation cost we want to avoid in the first place (For a length $T$ completion, generating one continuation from every prefix could require approximately: $T + T-1 + T-2 + \dots = \mathcal{O}(T^2)$ generated tokens), therefore we keep the advantages from the original policy $\hat{A}_{i,t}^{\pi_{old}}$ as fixed in the training.

How do we know when $\pi_{new}$ is substantially far from $\pi_{old}$?

We have established that (a) we want to reuse the sampled trajectories to avoid the cost of generating trajectories with every update step, and (b) we cannot forever reuse the sampled trajectories as they have become stale and the updates from them aren't meaningful. This begets the question, how do we know when to stop using the previously sampled trajectories? The answer to this question cannot be a fixed hyperparameter like 3 or 5, it depends on amount of policy movement which in turn depends on the learning rate, gradients, advantages, etc. We also cannot measure this using the change in parameter space $\lVert\theta-\theta_{\mathrm{old}}\rVert$ because a large parameter change might only slightly change the output distribution and vice-versa, therefore we need to measure how much the policy's action distribution at the states in the sampled trajectory has changed?

A natural measure of doing this is the KL-divergence between the distributions, i.e.

$$D_{\mathrm{KL}}(\pi_{old}(. \mid s) \Vert \pi_{new}(. \mid s)) = \sum_{a} \pi_{old}(a \mid s) \log \frac{\pi_{old}(a \mid s)}{\pi_{new}(a \mid s)}$$

Therefore the RL objective now becomes constrained

$$max_{\theta} \ E_{\pi_{old}} \ [\rho_{i,t}(\theta)\hat{A}_{i,t}] \ \text{such that } \ D_{KL}(\theta) \le \eta$$

This inteprets as "improve the policy using the old batch data, but remain inside a region where the old states and advantages are still reasonably representative under the updated policy". How should one solve this constrained optimization problem? There are several ways of solving such optimization problems like:

1. Projected gradient descent: Performs a step of parameter update and projects it back to the feasible set. However, here the constraint is not on the parameters but on the KL divergence of the output distribution which makes the projection part non-trivial.

2. Lagrange multipliers: This requires us to optimize: $L(\theta) - \beta (C(\theta) - \eta)$, and the problem moves to choosing the optimum parameter $\beta$. Too small, makes the KL grow beyond the desired region and too large makes the policy barely move. An adaptive primal-dual version can update both $\theta$ and $\beta$ alternatively, however that can also lead to oscillation.

3. Conjugate-gradients: $J_{surr}(\theta_{old} + \Delta \theta) \approx J_{surr}(\theta_{old}) + g^T \Delta\theta$ (where $g = \nabla_{\theta}J_{surr}(\theta_{old})$ )and $C(\theta_{old} + \Delta \theta) \approx \frac{1}{2}\Delta \theta^T F \Delta \theta$ (where F is the Hessian of KL divergence, which is the policy's Fisher information matrix) This leads to the optimization problem taking the form of $max_{\Delta \theta} \ g^T \Delta \theta$ subject to $\frac{1}{2}\Delta \theta^T F \Delta \theta \le \eta$. For a policy $\pi_{\theta}(a \mid s)$, the fisher matrix identifies which parameter direction most strongly change the policy's action probabilities and thus the above constraint can be interpreted as choosing a parameter update whose resulting policy distribution remains within a said distance (Fisher matrix is the local curvature of KL divergence, i.e., $D_{KL}(p_\theta || p_{\theta + \Delta \theta}) \approx \frac{1}{2} \Delta \theta^T F(\theta) \Delta \theta$. The KL is zero at $\Delta \theta = 0$ and its first derivative if zero there, the fisher matrix captures the second-order change, i.e., $\Delta \theta^T F \Delta \theta$ measures how much a parameter update changes the model's output distribution -- which is what we exactly need). The optimal direction satisfies: $\Delta \theta \propto F^{-1}g.$ Therefore we need to solve $Fx = g$. Constructing or inverting the massive fisher matrix F is very expensive, conjugate gradient approach approximately solve $Fx = g$ using only Fisher-vector products. After obtaining $x \approx F^{-1}g$, we scale it to the KL divergence limit, i.e., $\Delta \theta = \sqrt{\frac{2\eta}{x^TFx}}x$. A line search is then uses to check the satisfaction of the KL constraint under non-linear landscape of non-linear neural networks (note that conjugate-gradients assume linearity in local sense its formulation), so it might not satisfy the constraint after prediction, so we might need to reduce the magnitude of $\Delta \theta$ until satisfaction. This approach is what was proposed in the paper Trust-region policy optimization (TRPO).

However, as we have seen above the solving this constrained optimization problem using any of the popular three techniques is not easy. Recall the original goal, "how to know when you are too far from the original policy and stop using the sampled trajectories. This alludes to an interesting question which is "can we obtain a conservative update using the already computed IS ratio?"

So we want to move the improve the policy, but stop rewarding the change once it starts exploiting the fixed advantage loop. If $\hat{A}_{i,t} > 0$, it alludes that we should increase $\rho_{i,t}$ but only upto a threshold, similarly if $\hat{A}_{i,t} < 0$, it indicates we should decrease $\rho_{i,t}$, but only upto a threshold. A natural solution to the problem is we clip the $\rho_{i,t}$ update to an acceptable interval $1 - \epsilon < \rho_t < 1 + \epsilon$, and therefore the objective becomes $J_{i,t}(\theta) = clip(\rho_{i,t}, 1 - \epsilon, 1 + \epsilon) \hat{A}_{i,t}$. However this creates a small problem, dead gradients in the wrong regions. So say for a state-action pair, the stored advantage is $\hat{A}_{i,t} > 0$, but another minibatch update pushed $\rho_{i,t} < 1 - \epsilon$, if we simply clip $\rho_{i,t}$, then we won't be able to update the probability of this action-state pair which should be increased (since its $\hat{A}_{i,t} > 0$). Similarly, if for another action-state pair $\hat{A}_{i,t} < 0$ and a different minibatch has made $\rho_{i,t} > 1 + \epsilon$, then clipping the IS ratio will not give any gradients to this action-state pair whose probability should be decreased, and therefore important corrective updates would not apply. The best of both worlds (clipped ratios while allowed corrective updates) can be reached if we take the smaller of the claimed improvement, i.e., change the objective $J_t^{conservative} = min(\rho_t A_{i,t}, clip(\rho_t, 1 - \epsilon, 1 + \epsilon)A_{i,t})$. So now when $\hat{A}_{i,t} > 0$ and $\rho_t < 1 - \epsilon$, $J_t^{conservative} = \rho_t A_{i,t}$, and also when $\hat{A}_{i,t} < 0$ and $\rho_{i,t} > 1 + \epsilon$, $J_t^{conservative} = \rho_t A_{i,t}$ (so gradients can still update the actor to change $\rho_t$).

Here are the six possible cases:

The min operation is particularly important in the second and sixth cases because it preserves corrective gradients that naive clipping would remove.

This approach is called proximal policy optimization (PPO), one of the most popular RL optimization algorithms. It is important to note that clipping does not guarantee the new policy to satisfy KL divergence constraint, and therefore in practice the KL divergence is monitored and the old batches can be switched out.

Overall here's PPO's algorithm:

1. Collect rollouts using $\pi_{old}$ (which is frozen) and collect old log probabilities and old value predictions.
2. Calculate TD errors and generalized advantage estimates, and value targets.
3. Detach the advantages and value targets.
4. For several minibatch epochs: a. recompute current log probabilities and IS ratios. a. optimize the PPO-clip objective to update the actor b. optimize the critic using regression loss c. monitor the approximate KL which might require to stop early.
5. Discard the batch and set $\pi_{old} \leftarrow \pi_{\theta}$

How do we prevent the overall policy to drift too far?

The above clipping based approach prevents the current policy from drifting too much from the rollout generation policy, however, the individual small updates can move the model very far from its original behavior. This can be undesirable in certain cases like when the reward does not capture all aspects of the answer (it might capture verifier success, but not language quality, diversity, and factuality). This requires the updated policy to stay near the trusted language model across the entire training run (not just a single rollout update loop). The way this is usually achieved by adding a KL divergence from the original reference model as a penalty to the overall objective.

$$J_{regularized}(\theta) = J_{task}(\theta) - \beta E_{s,a}[D_{KL}(\pi_{\theta}(. \mid s) || \pi_{ref}(. \mid s))]$$

Note that the KL divergence problem is the same we had earlier at each step, however the solution deployed in this case is different, here its a fixed $\beta$ as a soft penalty coefficient. Why? The reason is that in the previous case, the KL constraint is a hard contraint, we do not want to use previous rollouts when using sufficiently different policy. In this case the constraint is softer, we can go away from the reference policy just not too much. The manner in which the KL constraint is operationalized is by computing the KL divergence on the output token distribution at each state and then averaged over all states. Another way to opertionalize is to directly subtract the KL divergence loss from the reward at each token position.

Can we replace the expensive critic with something else?

Reward-to-go provides a valid but high-variance policy-gradient signal. Subtracting state-dependent baseline reduces that variance. Using a learned critic to estimate this baseline is the usual way to go. However there are real disadvantages of using a critic:

1. a critic might be another similarly sized LLM (to the actor LLM) which requires significant resources to train it.
2. Training a reliable value predictor is hard because its target changes as soon as the actor is updated and sparse terminal rewards make prefix-level value estimation difficult.
3. Critic error affects GAE estimation and actor updates.

These issues begets us to ask, is there an alternate way to estimate the expected return of a initial prompt? Turns out we use the simple way of multiple sampling and averaging its reward as the expected return (kinda obvious, that's what "expected" means). So for a prompt $x$, we sample several completions from an LLM $y^{(1)}, y^{(2)}, y^{(3)}, \dots, y^{(G)} \sim \pi_{old}(. \mid x)$. And we get the rewards corresponding to each of them $R_1, R_2, R_3, \dots, R_{G}$ and the expected return can be estimated as the mean reward $\bar{R}_x = \frac{1}{G}\sum_{i=1}^{G}R_i \approx V^{\pi_{old}}(x)$. This acts as the value function without the need of the critic. Then the advantage for a particular completion can be computed as $\hat{A}_i = R_i - \bar{R}_x$. If $R_i$ was better than $\bar{R}_x$ then the completion $i$ was better than the average or worse in the other case. Note that unlike the value function case that provides us with the estimate value for each token, this case assumes all tokens in the completion to have the same advantage, $\hat{A}_{i,t} = \hat{A}_i$, this is because the group-level reward mean estimates the value of the original prompt $x$, not of the intermediate tokens. Therefore this avoids the critic, but at the cost of eliminating the token-level credits. This approach was introduced in Group-relative Policy Optimization (GRPO) paper. Note that in the original proposed formulation the advantage was also divided by the standard deviation of the group rewards, $\hat{A}_i = \frac{R_i - \bar{R}_x}{\sigma(R_x) + \epsilon}$.

For this approach to work, it is important that we only group completion of the same prompt and not completions of different prompts. This is because the value function can be underestimated or overestimated if we mix a prompt with other hard or easy prompts respectively.

So here's the modified objective function:

1. $\hat{A}_{i,t} = \hat{A}_i , t = 1, \dots, T_{i}$
2. $\rho_i(\theta) = \frac{\pi_{\theta}(y_t^{(i)} \mid x, y_{<t}^{(i)})}{\pi_{old}(y_t^{(i)} \mid x,y_{<t}^{(i)})}$
3. The clipped token contribution $J_{clipped}(\theta) = min(\rho_t(\theta) A_i, clip(\rho_t(\theta), 1-\epsilon, 1 + \epsilon)A_i)$
4. Average this objective over each token and each completion: $J(\theta) = \frac{1}{G}\sum_{i=1}^{G}\frac{1}{T_i}\sum_{t=1}^{T_i} min(\rho_t(\theta)A_i, clip(\rho_t(\theta), 1-\epsilon, 1+\epsilon)A_i)$
5. Optionally add the soft KL penalty to the reference model to the objective $-\frac{\beta}{G}\sum_{i=1}^{G}\frac{1}{T_i}\sum_{t=1}^{T_i} D_{KL}(\pi_{\theta}(. \mid s) \Vert \pi_{ref}(. \mid s) )$
6. $\mathcal{L}_{actor} = -J(\theta)$

Note that the factor $\frac{1}{T_i}$ gives each completion equal weight, without, longer completions by the virtue of having more tokens will get more gradient. This is an implementation choice. There are 3 valid choices, (a) sequence level sum $\frac{1}{G}\sum_i\sum_t J_{i,t}$, (b) all-token mean $\frac{1}{M}\sum_i\sum_t J_{i,t}, \text{ and } M = \sum_i T_i$, and (c) per-completion mean: $\frac{1}{G}\sum_i\frac{1}{T_i}\sum_t J_{i,t}$.

Now lets compare the two very popular approaches for doing RL on LLMs, PPO and GRPO

The two most popular application of RL in LLMs is RLHF and RLVR. RLHF stands for RL from human feedback and RLVR stands for RL from verified reward. RLHF works by training a preference reward model by providing humans with two responses to a prompt and asking them to choose the prefered one. And then optimizing the LLM to increase the reward from this preference model in order to tune its responses to human preferences. This was introduced in the InstructGPT paper Training Language Models to Follow Instructions with Human Feedback (InstructGPT) in 2022.

RLVR on the other hand uses verifiers, so for example a math problem which has one correct answer. There is usually no ambiguity in this reward, if the completion gets correct answer we give 1, otherwise 0. RLVR had led to the increased reasoning ability of LLMs where they can think for longer and can produce better answers.

Discussion

Name is optional. Sign in to edit or delete your own comments.

What did you think?

0 Helpful 0 Loved it