PPO for Preference Alignment

Your model can write fluent prose, but ask it a politically neutral question and it might dodge, prevaricate, or give a biased answer. Supervised fine-tuning alone won't fix that — it only imitates examples and can't learn from the preferences that shape human judgment. Reinforcement Learning from Human Feedback (RLHF) with Proximal Policy Optimization (PPO) is the missing layer: it uses human preference data to actively optimize your model's behavior, turning a merely competent generator into one that genuinely aligns with your values and standards. This lesson shows you exactly how to apply RLHF with PPO for preference alignment, step by step, with working code you can adapt immediately.

What problem does RLHF with PPO solve?

Before RLHF, aligning a language model meant one of two things:

• Supervised fine-tuning (SFT) on curated examples of good responses
• Prompt engineering to guide the model's output

Both have fundamental limits. SFT teaches the model to mimic patterns but never tells it which response is better among multiple valid ones — it treats every training pair as equally correct. Prompt engineering can nudge behavior but can't guarantee consistent alignment across millions of generations.

RLHF with PPO bridges this gap by converting human preferences into a reward signal the model can optimize. Instead of learning from static labels, the model explores different responses, receives a reward based on how much humans prefer them, and updates its policy to maximize that reward. This is the technology behind ChatGPT's alignment, and it's what turns a raw base model into a helpful, harmless assistant.

Without this step, even a fine-tuned model can behave unpredictably when faced with edge cases, sensitive topics, or ambiguous requests. RLHF with PPO gives you a principled way to encode your definition of good behavior.

Core concept / mental model: the reward cloud and the policy ship

Think of your language model as a ship navigating a vast ocean of possible responses. SFT gives the ship a map of examples, but RLHF gives it a compass tuned to human preferences — a reward function that tells it which direction to steer.

Here's the three-stage pipeline:

1. Supervised Fine-Tuning (SFT): Train a base model on high-quality demonstrations to get a reasonable starting policy.
2. Reward Model Training: Collect human comparisons (response A vs. response B), train a second model to predict which one humans prefer.
3. PPO Optimization: Use the reward model as a critic to guide the SFT model's updates, reinforcing responses that score high rewards and penalizing those that don't.

Key definitions you'll encounter:

• Policy: Your language model that generates responses
• Reward model: A scoring function that estimates human preference
• PPO: A policy-gradient algorithm that updates the policy while keeping updates stable
• KL divergence: A penalty that stops the policy from drifting too far from the original SFT model

Pro tip: PPO is called proximal because it clips updates to stay close to the current policy, preventing catastrophic changes. This is what makes it safe for large language models.

The mental model: you're not just fitting data; you're reinforcing good behavior through a trial-and-error loop guided by a learned reward signal — just like training a dog with treats, where treats are the reward score.

How does it work step by step?

PPO for RLHF isn't a single function call — it's a carefully orchestrated sequence. Here's the high-level flow:

1. Generate responses from the current policy (the SFT model).
2. Score them with the reward model to get a reward per response.
3. Compute advantages by comparing the reward to a baseline (the value model). This tells you if a response was better or worse than expected.
4. Update the policy using PPO's clipped objective, which adjusts the model to increase the probability of high-advantage responses and decrease low-advantage ones.
5. Penalize drift with a KL divergence term, keeping the policy close to the reference model.
6. Repeat for multiple epochs, gradually aligning the model.

Mathematically, the PPO objective looks like:

L = E[ min( ratio * A, clip(ratio, 1-ε, 1+ε) * A ) ] - β * KL(π_θ || π_ref)

Where ratio is the probability ratio between old and new policies, A is the advantage, and β controls how much we penalize divergence.

Each step is designed to maximize reward while avoiding the "reward hacking" where the model finds loopholes in the reward function.

Hands-on walkthrough: applying RLHF with PPO using TRL

The easiest way to apply RLHF with PPO is with Hugging Face's trl library. Let's build a complete example that aligns a model to be more helpful.

Setup

First, install the necessary packages:

pip install transformers trl datasets accelerate bitsandbytes

Step 1: Prepare your SFT model, reward model, and dataset

You'll need three components: a base policy (usually SFT'd), a reward model (already trained on preference data), and a dataset of prompts.

from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import PPOConfig, PPOTrainer, create_reference_model
from trl.core import respond_to_batch
import torch

# Load your SFT model (in practice, you'd fine-tune a base model first)
model_name = "gpt2"  # replace with your SFT'd model
model = AutoModelForCausalLM.from_pretrained(model_name)
model_ref = create_reference_model(model)  # copy for KL penalty
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

# Load a reward model (e.g., a classifier that scores helpfulness)
reward_model_name = "your-reward-model"  # load from Hugging Face Hub or local
reward_model = AutoModelForSequenceClassification.from_pretrained(reward_model_name)

Step 2: Set up PPO training

config = PPOConfig(
    batch_size=4,
    learning_rate=1.41e-5,
    ppo_epochs=4,
    cliprange=0.2,
    kl_penalty="kl",  # use KL penalty to prevent drift
    log_with="wandb",  # or "tensorboard"
)

ppo_trainer = PPOTrainer(
    config=config,
    model=model,
    ref_model=model_ref,
    tokenizer=tokenizer,
)

# Your prompts
prompts = [
    "What is the capital of France?",
    "Explain how photosynthesis works.",
    "Write a short story about a robot.",
]

# Encode each prompt and generate response
for prompt in prompts:
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)
    # Generate a response using the current policy
    response_ids = respond_to_batch(model, input_ids, max_new_tokens=50)

    # Compute reward (mock — you'll use your reward model)
    # In reality: score response with reward_model
    reward = torch.tensor([1.0])  # placeholder

    # Create query-response pair
    query_tensor = input_ids
    response_tensor = response_ids[:, input_ids.shape[1]:]  # trim prompt

    # Train PPO on this batch
    train_stats = ppo_trainer.step([query_tensor], [response_tensor], [reward])
    print(f"Objective: {train_stats['objective/kl']:.4f}")

Step 3: Run the loop and evaluate

# After training, generate a response and check if it's more aligned
input_ids = tokenizer("What is the capital of France?", return_tensors="pt").input_ids.to(model.device)
response = model.generate(input_ids, max_new_tokens=50)
print(tokenizer.decode(response[0], skip_special_tokens=True))

Expected output (your model should give a direct, factual answer instead of a hedge):

The capital of France is Paris.

Before PPO, the same model might have said: "I'm not sure; some people say Paris, but others claim Lyon..." — now it gives the preferred, concise answer.

Pro tip: Always keep a reference model to compute KL divergence. Without it, your model can over-optimize the reward and produce unnatural or degenerate text.

The full loop using trl's PPOTrainer

Here's a more realistic training loop that processes a dataset:

from datasets import load_dataset

dataset = load_dataset("your-preference-prompts")

for epoch in range(2):
    for i in range(0, len(dataset), config.batch_size):
        batch = dataset[i:i+config.batch_size]
        query_tensors = [tokenizer(p['prompt'], return_tensors="pt").input_ids[0] for p in batch]
        # Generate responses in batch
        response_tensors = [respond_to_batch(model, q.unsqueeze(0), max_new_tokens=50)[0] for q in query_tensors]

        # Score with reward model
        rewards = [compute_reward(reward_model, tokenizer, q, r) for q, r in zip(query_tensors, response_tensors)]

        # PPO step
        stats = ppo_trainer.step(query_tensors, response_tensors, rewards)

        # Log metrics
        if i % 100 == 0:
            print(f"Epoch {epoch} Batch {i}: mean_reward={stats['objective/rewards'].mean():.4f}, kl={stats['objective/kl']:.4f}")

Compare options: PPO vs. other alignment methods

PPO isn't the only way to align models. Here’s how it stacks up against alternatives:

Choose PPO when: you need fine-grained control over alignment (e.g., content moderation, safety-critical applications) and you have the GPU budget. Choose DPO when: you want a faster, cheaper alternative and your preference data is clean.

Pro tip: Many teams use PPO as a final polish after DPO or SFT because it can push behavior to the edge of the reward boundary.

Troubleshooting & edge cases

Even experienced practitioners hit these common issues. Here's how to fix them:

Problem: Reward hacking — the model finds loopholes to score high without being truly aligned.

Example: A model learns to output "I don't know" for every question because that gets a higher reward on safety prompts.

Fix: Keep KL penalty strong (increase β), and monitor reward distribution. If rewards spike suddenly, your reward model might be exploitable. Use diverse prompt sets and avoid single-reward models.

Problem: KL divergence explodes, causing the model to produce gibberish.

Example: During training, loss/kl climbs to thousands — the model drifts far from the reference.

Fix: Lower the learning rate, reduce ppo_epochs, or increase cliprange to 0.3. Also, ensure your reference model is a frozen copy of the SFT model, not a base model.

Problem: PPO doesn't converge — reward stays flat.

Example: After many steps, the training reward remains at 0.5 with no improvement.

Fix: Check your reward model's quality — if it's poorly trained, there's no signal. Also, verify that the prompts represent a balanced distribution. Sometimes the batch size is too small; increase batch_size to 16 or 32.

Problem: NaN loss after a few steps.

Fix: Gradient explosion. Add gradient clipping (config.max_grad_norm = 0.5), use a lower learning rate, and ensure your reward values are normalized (e.g., z-scores).

Problem: The model becomes overly verbose to game the reward model.

Fix: Add a token length penalty to your reward function, or train your reward model to penalize verbosity. Also, use a KL penalty that discourages drastic policy changes.

What you learned & what's next

You now understand how to apply RLHF with PPO for preference alignment — from the core concepts to a working implementation with trl. You've seen how to:

• Set up a PPO trainer with a reference model
• Generate responses, compute rewards, and update the policy
• Handle common pitfalls like reward hacking and KL divergence explosion

This hands-on experience prepares you for the next lesson in the LLM Finetuning track: Evaluating aligned models — where you'll learn to measure whether your RLHF actually improved alignment, using automated metrics and human eval. You'll also explore advanced variants like DPO and RLAIF to broaden your toolkit.

Next step: Take the code from this lesson, load a real reward model (like OpenAssistant/reward-model-deberta-v3-large), and run PPO on a small dataset like Anthropic/hh-rlhf. Monitor reward and KL curves to observe alignment in action.

Remember, RLHF with PPO is a tool, not a magic bullet. Use it when you have clear preference data and the compute budget. Now go align your models!