Proximal Policy Optimization (PPO)

The core innovation of PPO is its clipped surrogate objective function. Instead of directly maximizing the expected reward, it constrains the policy update by clipping the probability ratio between the new and old policies. This ratio, ( r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)} ), is clipped within the interval ( [1 - \epsilon, 1 + \epsilon] ), where ( \epsilon ) is a small hyperparameter (typically 0.1 or 0.2). The objective becomes:

( L^{CLIP}(\theta) = \mathbb{E}_t [ \min( r_t(\theta) A_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t ) ] )

• Prevents Destructive Updates: The min() operator ensures the update is the pessimistic (lower) bound of the clipped and unclipped objectives.
• Enforces a Trust Region: The clipping implicitly prevents the policy from changing too much in a single update, stabilizing training.
• This mechanism is the primary reason for PPO's robustness compared to vanilla policy gradient methods.

PPO typically uses Generalized Advantage Estimation (GAE) to compute a low-variance, bias-controlled advantage function ( A_t ). The advantage measures how much better a specific action is compared to the average action in a state.

• Reduces Variance: GAE creates a weighted average of k-step Temporal Difference (TD) error estimates, balancing bias and variance.
• Key Parameter: The discount factor ( \gamma ) and an additional smoothing parameter ( \lambda ) (between 0 and 1) control the trade-off. ( \lambda=1 ) is high-variance Monte Carlo, ( \lambda=0 ) is biased 1-step TD.
• Formula: ( \hat{A}t^{GAE(\gamma, \lambda)} = \sum{l=0}^{\infty} (\gamma \lambda)^l \delta_{t+l} ), where ( \delta_t = r_t + \gamma V(s_{t+1}) - V(s_t) ) is the TD residual. This provides a smooth, informative signal for the policy update, crucial for learning in environments with sparse or delayed rewards.

PPO employs an actor-critic architecture, utilizing two neural networks (often with shared base layers):

• Actor Network (Policy): Parameterized by ( \pi_\theta(a|s) ), this network selects actions. It is optimized using the clipped objective and advantage estimates.
• Critic Network (Value Function): Parameterized by ( V_\phi(s) ), this network estimates the expected cumulative reward from a given state. It is trained via mean-squared error against target values (e.g., reward + discounted next value).

Shared Benefits:

• The critic reduces variance in policy updates by providing a baseline (the value function).
• The actor focuses on improving action selection relative to that baseline.
• This separation allows for more stable and efficient learning than pure policy gradient methods.

Unlike traditional policy gradient methods that use a trajectory once, PPO collects a batch of trajectories and then performs multiple epochs of stochastic gradient descent on that data.

• Improved Sample Efficiency: By re-using sampled data for several updates, PPO extracts more learning signal from each interaction with the environment.
• Requires Importance Sampling: Since the policy is updated between epochs, the data becomes off-policy relative to the current policy. The probability ratio ( r_t(\theta) ) acts as an importance weight to correct for this distribution shift.
• Early Stopping: Training is often halted early if the mean KL-divergence between the new and old policy grows beyond a threshold, providing an additional safeguard against large updates.

The original PPO paper proposed two main variants: the clipped objective (standard) and an adaptive KL penalty objective.

• Objective: ( L^{KLPEN}(\theta) = \mathbb{E}t [ \frac{\pi\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)} A_t - \beta , KL[\pi_{\theta_{old}}(\cdot|s_t), \pi_\theta(\cdot|s_t)] ] )
• Adaptive ( \beta ): The coefficient ( \beta ) is adjusted each update:
  • Increase ( \beta ) if the actual KL divergence is above a target threshold.
  • Decrease ( \beta ) if the KL divergence is below the target.
• Comparison: While theoretically sound, the clipped objective is simpler to tune and implement, making it the more widely adopted variant in practice. The KL penalty is less sensitive to the choice of ( \epsilon ) but introduces the complexity of adapting ( \beta ).

In Reinforcement Learning from Human Feedback (RLHF), PPO is the workhorse algorithm for the final reinforcement learning stage.

• Policy (Actor): The language model to be aligned (initialized from the SFT model).
• Reward Signal: Provided by the Reward Model (RM) trained on human preference data.
• Value Function (Critic): Trained to estimate the expected future reward from a given state (prompt + partial response).
• KL Penalty: A critical addition in RLHF is an extra penalty term based on the KL divergence between the current policy and the original SFT model. This prevents the policy from deviating too far into "reward hacking" and losing its general language capabilities (mitigating alignment tax).
• Efficiency via PEFT: To reduce massive computational cost, Parameter-Efficient Fine-Tuning (PEFT) methods like LoRA for RLHF are applied, where only low-rank adapter weights for the actor and critic are trained, while the base model remains frozen.

RLHF is the overarching alignment framework where PPO operates. It is a three-stage process:

• Supervised Fine-Tuning (SFT): A base model is first fine-tuned on high-quality demonstration data.
• Reward Model Training: A separate neural network is trained to predict a scalar reward from human preference data (e.g., choosing between two model outputs).
• Policy Optimization: The SFT model (the policy) is fine-tuned using reinforcement learning (typically PPO) to maximize the reward from the reward model, while staying close to the original policy via a KL divergence penalty. PPO provides the stable, on-policy RL algorithm that performs this final optimization step.

The reward model is a crucial component that provides the training signal for PPO in RLHF.

• It is typically a neural network that takes a prompt and a model-generated response as input and outputs a scalar reward value.
• It is trained via supervised learning on a preference dataset, learning to predict which of two responses a human would prefer.
• During PPO training, this model acts as the environment, scoring the policy's outputs. A poorly trained or overfitted reward model can lead to reward hacking, where PPO exploits flaws to generate high-reward but low-quality text.

The KL divergence penalty is a regularization term central to PPO's stability in RLHF.

• In the RLHF objective, the total reward is: Reward from RM - β * KL(π_θ || π_ref).
• π_θ is the current policy being optimized by PPO, and π_ref is typically the initial SFT model.
• The term β * KL(...) penalizes the policy for deviating too far from the reference model. This prevents catastrophic forgetting of general language abilities and mitigates reward overoptimization by constraining updates within a trust region. PPO's clipping mechanism works in concert with this penalty to ensure stable training.

TRPO is the direct predecessor to PPO. Both are on-policy, actor-critic algorithms designed for stable policy updates.

• Core Idea: TRPO explicitly enforces a trust region constraint using a complex second-order optimization (conjugate gradient) to ensure the new policy does not deviate too far from the old policy, as measured by KL divergence.
• PPO's Innovation: PPO achieves similar stability but with a first-order method that is simpler and more computationally efficient. It uses a clipped surrogate objective or an adaptive KL penalty to approximate the trust region constraint, avoiding TRPO's expensive calculations.

PPO is an actor-critic algorithm. This architecture separates two functions:

• The Actor: This is the policy network (π_θ). It takes the state (the current prompt and generated tokens) and selects the next action (the next token).
• The Critic: This is the value function network (V_φ). It estimates the expected cumulative reward from a given state, judging the quality of the actor's decisions. During PPO training, the actor is updated to maximize reward, guided by the advantage function (estimated reward minus the critic's baseline). The critic is updated to better predict rewards, reducing variance. This separation leads to more stable and efficient learning than pure policy gradient methods.

GAE is a technique commonly used with PPO to compute a low-variance, bias-controlled estimate of the advantage function.

• The advantage A(s, a) measures how much better a specific action is compared to the average action in that state.
• GAE computes this by taking an exponentially weighted average of k-step Temporal Difference (TD) errors, controlled by parameters γ (discount) and λ (bias-variance trade-off).
• Using GAE(λ) provides PPO's actor with a smoother, more informative training signal than a simple Monte Carlo return, which is crucial for effective policy updates in environments with delayed rewards, like text generation.