# The Three Stages of the RLHF Workflow with PPO: Implementation Guide

> Master the three stages of RLHF with PPO: Rollout, Evaluation, and Optimization. This guide details implementing the RLHF workflow to fine-tune LLMs effectively.

- Repository: [Tongxin Yuan/dive-into-llms](https://github.com/Lordog/dive-into-llms)
- Tags: how-to-guide
- Published: 2026-04-16

---

**The RLHF workflow with PPO consists of three sequential stages: Rollout (generating responses from the language model policy), Evaluation (scoring query-response pairs with a reward model to produce scalar rewards), and Optimization (updating the policy using Proximal Policy Optimization while regularizing against a reference model via KL divergence penalties).**

Reinforcement Learning from Human Feedback (RLHF) using Proximal Policy Optimization (PPO) forms the backbone of modern LLM alignment techniques. According to the `Lordog/dive-into-llms` repository, this **RLHF workflow with PPO** follows a precise iterative loop described in [`documents/chapter11/README.md`](https://github.com/Lordog/dive-into-llms/blob/main/documents/chapter11/README.md), where each training iteration cycles through generation, reward computation, and policy refinement.

## Stage 1: Rollout – Generating Responses from the Policy

The first stage initiates the loop by producing candidate outputs from the active language model. During **Rollout**, the policy (the trainable LLM) receives input queries and generates corresponding responses.

As documented in [`documents/chapter11/README.md`](https://github.com/Lordog/dive-into-llms/blob/main/documents/chapter11/README.md) (lines 12-13), this stage involves the language model generating responses based on the incoming query. In practical implementations using the 🤗 TRL library, this corresponds to calling `ppo_trainer.generate()`, which samples token continuations from the current policy distribution.

The output of this stage is a batch of query-response pairs ready for assessment.

## Stage 2: Evaluation – Computing Scalar Rewards

The second stage assigns quantitative values to the generated outputs. During **Evaluation**, the system assesses each query-response pair using a reward model—this can be a classifier, a human feedback function, or any combination of scoring mechanisms—to produce a **scalar reward** for each sample.

According to the source documentation in [`documents/chapter11/README.md`](https://github.com/Lordog/dive-into-llms/blob/main/documents/chapter11/README.md) (lines 13-14), this evaluation generates a scalar value that indicates the quality of the response. In implementation terms, this often involves a sentiment classifier or similar reward model that outputs numerical scores, as demonstrated in the repository's `RLHF.ipynb` notebook where a BERT-based pipeline returns positive-sentiment scores as rewards.

## Stage 3: Optimization – PPO Policy Updates with KL Penalty

The final stage updates the policy parameters to maximize expected rewards while preventing catastrophic drift from the original model. During **Optimization**, the PPO algorithm adjusts the policy weights based on the collected rewards, incorporating a KL divergence penalty computed against a frozen **reference model**.

As specified in [`documents/chapter11/README.md`](https://github.com/Lordog/dive-into-llms/blob/main/documents/chapter11/README.md) (lines 14-15), the KL divergence between the active policy and the reference model serves as an additional reward signal to ensure the updated model remains close to the original. The `ppo_trainer.step()` function executes this update, processing the queries, responses, and rewards to perform a single PPO optimization step.

## Implementation Example: RLHF with PPO in Practice

The `documents/chapter11/RLHF.ipynb` file in the `Lordog/dive-into-llms` repository provides a complete implementation of these three stages. Below is the executable workflow demonstrating the full RLHF loop:

```python
import torch
from transformers import AutoTokenizer, pipeline
from datasets import load_dataset
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead
from trl.core import LengthSampler

# ---- Stage 1 Configuration: Rollout Setup ------------------------------------

config = PPOConfig(
    model_name="model/gpt2-imdb", 
    learning_rate=1.41e-5, 
    log_with="wandb"
)
tokenizer = AutoTokenizer.from_pretrained(config.model_name)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)
ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(config.model_name)

def build_dataset():
    ds = load_dataset("imdb", split="train")
    ds = ds.filter(lambda x: len(x["text"]) > 200)
    sampler = LengthSampler(2, 8)
    def tokenize(sample):
        ids = tokenizer.encode(sample["text"])[: sampler()]
        sample["input_ids"] = ids
        sample["query"] = tokenizer.decode(ids)
        return sample
    return ds.map(tokenize, batched=False).set_format(type="torch")

dataset = build_dataset()
ppo_trainer = PPOTrainer(config, model, ref_model, tokenizer, dataset=dataset)

# ---- Stage 2 Setup: Reward Model (Evaluation) --------------------------------

sentiment_pipe = pipeline(
    "sentiment-analysis",
    model="model/distilbert-imdb",
    device=ppo_trainer.accelerator.device,
)

def get_reward(texts):
    pipe_out = sentiment_pipe(
        texts, 
        top_k=None, 
        function_to_apply="none", 
        batch_size=16
    )
    return [
        item["score"]
        for out in pipe_out
        for item in out
        if item["label"] == "POSITIVE"
    ]

# ---- Stage 3 Execution: Optimization Loop ----------------------------------

output_len_sampler = LengthSampler(4, 16)

for batch in ppo_trainer.dataloader:
    # Rollout: Generate responses

    queries = batch["input_ids"]
    responses = []
    for q in queries:
        gen_len = output_len_sampler()
        resp = ppo_trainer.generate(q, max_new_tokens=gen_len).squeeze()
        responses.append(resp[len(q):])
    
    # Evaluation: Compute scalar rewards

    full_texts = [
        tokenizer.decode(q) + tokenizer.decode(r) 
        for q, r in zip(queries, responses)
    ]
    rewards = torch.tensor(get_reward(full_texts))
    
    # Optimization: PPO update with KL penalty

    stats = ppo_trainer.step(queries, responses, rewards)
    ppo_trainer.log_stats(stats, batch, rewards)

```

This implementation follows the exact **three-stage RLHF workflow** documented in the repository:

- **Rollout**: `ppo_trainer.generate()` creates responses for each query
- **Evaluation**: The sentiment pipeline computes scalar rewards indicating positive sentiment
- **Optimization**: `ppo_trainer.step()` executes the PPO update, automatically applying the KL penalty against the `ref_model`

## Summary

The **RLHF workflow with PPO** implemented in the `Lordog/dive-into-llms` repository follows a rigorous three-stage iterative process:

- **Rollout** generates candidate responses from the active policy model using `ppo_trainer.generate()`
- **Evaluation** produces scalar rewards via reward models or classifiers, as documented in [`documents/chapter11/README.md`](https://github.com/Lordog/dive-into-llms/blob/main/documents/chapter11/README.md)
- **Optimization** updates policy parameters using PPO with KL divergence regularization against a reference model, executed through `ppo_trainer.step()`

The [`requirements.txt`](https://github.com/Lordog/dive-into-llms/blob/main/requirements.txt) file in the repository root lists necessary dependencies including `torch`, `transformers`, `datasets`, and `trl` to execute this workflow.

## Frequently Asked Questions

### What is the role of the reference model in PPO-based RLHF?

The **reference model** serves as a frozen copy of the original language model that provides the baseline distribution for computing the KL divergence penalty. According to the implementation in [`documents/chapter11/README.md`](https://github.com/Lordog/dive-into-llms/blob/main/documents/chapter11/README.md), this KL term acts as an additional reward signal that prevents the optimized policy from drifting too far from the pre-trained model's behavior, ensuring stable training.

### How does the reward model generate scalar values during the Evaluation stage?

The reward model accepts query-response pairs and outputs a single numerical value representing response quality. In the `Lordog/dive-into-llms` implementation, this can involve any scoring function, model, or human feedback mechanism that produces a **scalar reward**. The provided example uses a BERT-based sentiment classifier that returns positive sentiment scores as the reward values.

### Why is a KL penalty necessary in the Optimization stage?

The **KL penalty** ensures the policy updates remain conservative and proximate to the reference model. Without this constraint, PPO might exploit the reward model by producing high-scoring but incoherent or out-of-distribution text. As documented in the repository, the KL divergence between the active policy and reference model outputs serves as a regularization term that maintains linguistic coherence while optimizing for human preferences.

### What dependencies are required to implement this RLHF workflow?

The implementation requires `torch`, `transformers`, `datasets`, and `trl` (Transformer Reinforcement Learning), as specified in the repository's [`requirements.txt`](https://github.com/Lordog/dive-into-llms/blob/main/requirements.txt). These libraries provide the `PPOTrainer`, `AutoModelForCausalLMWithValueHead`, and pipeline utilities necessary to execute the three-stage rollout, evaluation, and optimization loop.