Calculating Advantages and Returns

Optimizing the policy with the PPO objective requires a reliable estimate of how much better a chosen action (generating a specific token) performs compared to the average action the policy would take in that state (the current generated sequence). The advantage function, $A(s_t, a_t)$, provides this estimate. Relying solely on the immediate reward $r_t$, which typically combines a reward model signal and a KL penalty, proves insufficient for this optimization. This is because immediate rewards neglect the long-term consequences of an action. Instead, the total accumulated reward, or return, must be examined and compared against a baseline. This baseline is established by a value function $V(s_t)$, which estimates the expected return from state $s_t$.

Returns Calculation

The return $G_t$ is the total discounted reward received starting from time step $t$ until the end of the episode (the generated sequence). For a sequence of length $T$, it's defined as:

$$G_t = \sum_{k=0}^{T-t-1} \gamma^k r_{t+k+1}$$

Here, $r_{t+k+1}$ is the reward received after taking an action in state $s_{t+k}$. In the RLHF context for LLMs, the reward $r_t$ at each step typically consists of two components: a penalty based on the KL divergence between the current policy and the reference (SFT) policy, and potentially a contribution from the final reward model score $R(x)$ assigned to the complete sequence $x$. A common practice is to apply the KL penalty $r_{KL, t} = -\beta D_{KL}(\pi_{\theta}(\cdot|s_t) || \pi_{ref}(\cdot|s_t))$ at each token generation step $t$, and add the final reward model score $R(x)$ only at the last step, $T$. So, $r_t = r_{KL, t}$ for $t < T$, and $r_T = r_{KL, T} + R(x)$.

The discount factor $\gamma \in [0, 1]$ determines the present value of future rewards. A $\gamma$ closer to 1 gives more weight to future rewards, while a $\gamma$ closer to 0 prioritizes immediate rewards. For text generation tasks, $\gamma$ is often set close to 1 (e.g., 0.99 or 1.0) because the quality assessment (via the reward model) often depends on the complete sequence.

The Advantage Function

The advantage function $A(s_t, a_t)$ measures the relative value of taking action $a_t$ in state $s_t$ compared to the expected value of the state $V(s_t)$ under the current policy $\pi_\theta$. It's formally defined as:

$$A(s_t, a_t) = Q(s_t, a_t) - V(s_t)$$

where $Q(s_t, a_t)$ is the action-value function, representing the expected return after taking action $a_t$ in state $s_t$ and following the policy $\pi_\theta$ thereafter. Since we typically learn $V(s_t)$ directly using a value network (the critic), we can estimate $Q(s_t, a_t)$ using the immediate reward $r_{t+1}$ and the value of the next state $V(s_{t+1})$:

$$Q(s_t, a_t) \approx r_{t+1} + \gamma V(s_{t+1})$$

Substituting this into the advantage definition gives the one-step Temporal Difference (TD) error, $\delta_t$, often used as a basic advantage estimator:

$$\hat{A}_t \approx \delta_t = r_{t+1} + \gamma V(s_{t+1}) - V(s_t)$$

This estimate tells us whether the observed outcome ($r_{t+1} + \gamma V(s_{t+1})$) was better or worse than what was expected ($V(s_t)$).

Addressing High Variance with Generalized Advantage Estimation (GAE)

While the TD error $\delta_t$ is an unbiased estimate of the advantage (if $V$ is accurate), it can suffer from high variance because it relies heavily on the single-step reward $r_{t+1}$ and the next-state value estimate $V(s_{t+1})$. This variance can make the PPO updates unstable, especially in complex tasks like language generation.

Generalized Advantage Estimation (GAE) is a technique designed to reduce this variance by incorporating information from multiple time steps, effectively blending the single-step TD error with longer-term Monte Carlo returns. GAE introduces a parameter $\lambda \in [0, 1]$ (often called the GAE lambda) to control this trade-off.

The GAE advantage estimator is calculated as an exponentially weighted sum of TD errors:

$$\hat{A}_t^{GAE(\gamma, \lambda)} = \sum_{k=0}^{T-t-1} (\gamma \lambda)^k \delta_{t+k}$$

where $\delta_{t+k} = r_{t+k+1} + \gamma V(s_{t+k+1}) - V(s_{t+k})$ is the TD error at time step $t+k$. (Note: $V(s_T)$ is typically defined as 0 if $s_T$ is a terminal state).

• If $\lambda = 0$: $\hat{A}_t^{GAE(\gamma, 0)} = \delta_t$. This recovers the high-variance, low-bias one-step TD error.
• If $\lambda = 1$: $\hat{A}_t^{GAE(\gamma, 1)} = \sum_{k=0}^{T-t-1} \gamma^k \delta_{t+k}$. This expands to $\sum_{k=0}^{T-t-1} \gamma^k (r_{t+k+1} + \gamma V(s_{t+k+1}) - V(s_{t+k}))$. This sum telescopes and approximates the Monte Carlo advantage estimate $G_t - V(s_t)$, which generally has lower variance but can be biased if the value function $V$ is inaccurate, especially for long sequences (large $T$).
• For $0 < \lambda < 1$: GAE provides an intermediate estimate that balances the bias-variance trade-off. Values like $\lambda = 0.95$ are common in practice, often yielding more stable and efficient learning than the endpoints.

The diagram below illustrates how TD errors over multiple steps contribute to the GAE calculation for $\hat{A}_t$.

Calculation flow for Generalized Advantage Estimation (GAE). Rewards ($r$) and value function estimates ($V$) for subsequent states ($s$) are used to compute Temporal Difference (TD) errors ($\delta$). These TD errors are then combined using weights based on $\gamma$ and $\lambda$ to form the final GAE advantage estimate $\hat{A}_t^{GAE}$.

Implementation Details

In a typical RLHF implementation using libraries like TRL (Transformer Reinforcement Learning), GAE is computed efficiently. During the PPO rollout phase, the policy generates sequences token by token. For each token generated up to the end of the sequence (or a maximum length), the following are stored:

• The state (sequence generated so far) $s_t$.
• The log probability of the chosen action (token) $\log \pi_\theta(a_t|s_t)$.
• The value estimate $V(s_t)$ from the critic network.
• The KL divergence penalty component of the reward $r_{KL, t}$.

Once a batch of complete sequences is generated:

1. The final reward model score $R(x)$ is computed for each sequence $x$ and added to the reward for the final step $T$.
2. The value function $V(s_t)$ is used to compute the TD errors $\delta_t = r_{t+1} + \gamma V(s_{t+1}) - V(s_t)$ for all steps $t$ in each sequence (working backwards from $T-1$, setting $V(s_T)=0$). Remember $r_{t+1}$ includes $r_{KL, t+1}$ and potentially $R(x)$ if $t+1=T$.
3. The GAE advantages $\hat{A}_t^{GAE}$ are computed for all steps $t$ using the calculated TD errors $\delta_t$, $\delta_{t+1}$, ..., $\delta_{T-1}$ and the chosen $\gamma$ and $\lambda$. This is typically done with another backward pass through the sequence:
   • Initialize $\hat{A}_T^{GAE} = 0$.
   • For $t = T-1$ down to $0$: $$\hat{A}_t^{GAE} = \delta_t + \gamma \lambda \hat{A}_{t+1}^{GAE}$$
4. The calculated returns $G_t$ (often estimated using $G_t \approx \hat{A}_t^{GAE} + V(s_t)$) are used as targets for training the value function $V(s_t)$ via minimizing a loss like Mean Squared Error (MSE): $L_{VF} = \mathbb{E}[(V_{\phi}(s_t) - G_t^{target})^2]$.
5. The computed GAE advantages $\hat{A}_t^{GAE}$ are used directly in the PPO policy loss function to update the policy parameters $\theta$.

It is a standard practice to normalize the advantages across a batch before using them in the PPO loss calculation. This involves subtracting the mean and dividing by the standard deviation of the advantages within the batch, which helps stabilize training by preventing excessively large policy updates.

By carefully calculating returns and using GAE to estimate advantages, we provide the PPO algorithm with a stable and informative signal to guide the LLM policy towards generating responses that align better with the preferences captured by the reward model, while mitigating the instability associated with high-variance gradient estimates.