3.1 RLHF Pipeline: Reward Modeling and Preference Data

In Chapters 1 and 2, you learned how to teach a model to follow instructions using supervised examples. You curated datasets, fine-tuned models, and adapted them efficiently with LoRA and QLoRA. That process shaped behavior directly through imitation learning: "Here is the correct answer. Learn to reproduce it." The model observed input-output pairs and learned to predict the next token that matches the training distribution.

But what if there isn't a single correct answer?

What if quality depends on subtle human judgment — clarity, harmlessness, politeness, usefulness, depth of reasoning, or even stylistic preferences? In real-world applications, especially conversational AI, there are often multiple valid responses to the same prompt, each with different trade-offs. One response might be technically accurate but overly verbose. Another might be concise but lack important nuance. A third might be clear and helpful but use an inappropriate tone.

Supervised fine-tuning can only go so far. It teaches imitation. It does not teach preference. When you provide a single "correct" example, you're making an implicit claim that this is the best possible response — but human preferences are far more nuanced than binary correctness. We don't just want models that can mimic examples; we want models that understand what makes one response better than another.

This is where Reinforcement Learning with Human Feedback (RLHF) enters.

Instead of telling the model exactly what to say, RLHF teaches the model what humans prefer. It aligns model behavior with human values by rewarding better responses and discouraging worse ones. Rather than optimizing for likelihood of exact token sequences, RLHF optimizes for human-judged quality. This paradigm shift — from imitation to preference optimization — is what enabled models like ChatGPT and Claude to feel more helpful, harmless, and aligned with user intent.

The RLHF approach introduces several key innovations: it replaces ground-truth labels with comparative rankings, it trains a separate reward model to capture human preferences, and it uses reinforcement learning to optimize policy behavior based on learned rewards rather than supervised targets. This multi-stage pipeline is more complex than standard fine-tuning, but it unlocks capabilities that supervised learning alone cannot achieve.

In this chapter, you will learn:

• How the RLHF pipeline works — from preference collection through policy optimization
• How reward models are trained to score response quality
• How preference datasets are constructed from human annotations
• How optimization is performed using algorithms like PPO
• How newer alternatives (like DPO) simplify the process by eliminating the need for explicit reward modeling

We begin with the core foundation: the RLHF pipeline itself. Understanding the architecture of this multi-stage system — and why each stage is necessary — is essential before diving into implementation details. By the end of this chapter, you'll not only understand how RLHF works theoretically, but you'll have practical knowledge of how to implement each component and understand the trade-offs involved in different approaches to preference learning.

Reinforcement Learning with Human Feedback is not a single training step. It is a multi-stage pipeline. Each stage has a specific purpose, and understanding this flow is critical before writing any code. Unlike supervised fine-tuning, where you prepare data, configure a trainer, and run a single training loop, RLHF requires careful orchestration of multiple models, datasets, and optimization procedures. Each stage builds on the previous one, and skipping or misunderstanding any component can lead to suboptimal alignment or even model degradation.

At a high level, RLHF consists of three stages:

1. Supervised Fine-Tuning (SFT) — This creates your initial policy, a model capable of following instructions and generating coherent responses
2. Reward Model Training — This teaches a separate model to evaluate response quality based on human preferences
3. Reinforcement Learning Optimization — This refines the policy to maximize rewards while maintaining stability

You have already mastered Stage 1 through the instruction-tuning techniques covered in previous chapters.

Now we move into Stages 2 and 3, where the paradigm shifts fundamentally from imitation to preference learning.

The Big Picture

Imagine you prompt a model:

"Explain why gradient accumulation is useful."

The model might produce several valid answers. Some are clearer. Some are safer. Some are more concise. Some are verbose but technically accurate. Each response might be factually correct, yet they vary significantly in quality dimensions like clarity, helpfulness, depth, and tone.

Instead of labeling a single "correct" output — which would require you to arbitrarily choose one valid response over others — we ask humans a different question:

Which response do you prefer?

That simple question changes everything. It acknowledges that language generation is not a classification problem with a single ground truth. It's a preference optimization problem where quality exists on a spectrum.

Instead of predicting text tokens through maximum likelihood estimation, we now predict preference through comparative ranking. Rather than asking "What would a human write?" we ask "What would a human prefer?" This reframing allows us to capture nuanced human values — helpfulness, harmlessness, conciseness, clarity, safety — that cannot be easily encoded in supervised labels.

That is the heart of RLHF: replacing direct behavior cloning with preference-guided optimization.

3.1.1 Stage 1: Supervised Fine-Tuning (Recap)

Before reinforcement learning begins, the base model must first undergo supervised fine-tuning (SFT). This initial stage is not optional — it's a prerequisite for effective RLHF. Without it, the model would lack the basic instruction-following capabilities needed to generate coherent candidate responses during preference learning.

During SFT, the model learns to map prompts to reasonable completions by training on curated instruction-response pairs. For example, given the prompt "Explain gradient descent," a well-instruction-tuned model will produce a coherent explanation rather than random tokens or off-topic text. This capability is essential because the subsequent stages of RLHF depend on the model's ability to generate plausible responses that humans can then compare and rank.

This stage gives us what reinforcement learning practitioners call a competent starting policy — a model that already demonstrates baseline competence at the task we want to improve through preference optimization.

In reinforcement learning terminology, we use specific vocabulary to describe the components of this system:

• The model itself is called the policy (often denoted as π). In RL terms, a policy is any function that maps states to actions.
• In our case, the policy maps prompts (states) to generated responses (actions). More precisely: π(response | prompt).
• The goal of RLHF is to transform this initial policy πSFT into an improved policy πRLHF that better aligns with human preferences.

We now want to improve this policy based on preference signals rather than direct supervision. Instead of showing the model more labeled examples of "correct" outputs, we will guide it using comparative feedback: "Response A is better than Response B." This shift from imitation to preference optimization is what distinguishes RLHF from standard fine-tuning and enables the nuanced alignment capabilities that make modern conversational AI systems feel more helpful and human-aligned.

3.1.2 Stage 2: Collecting Preference Data

Instead of labeled outputs, we collect ranked responses. This fundamental shift in data structure is what distinguishes RLHF from traditional supervised learning. Rather than asking annotators to produce a single "correct" response—which would force an arbitrary decision when multiple valid responses exist—we acknowledge that language quality is comparative in nature.

The preference collection process follows a systematic workflow:

1. Generate multiple candidate responses from the model (typically 2-4 responses per prompt)
2. Present these responses to human annotators who evaluate them based on criteria like helpfulness, harmlessness, accuracy, clarity, and tone
3. Record which response is preferred, creating pairwise comparisons that capture relative quality rather than absolute correctness

This approach has several advantages. First, it's often easier for humans to judge "Which response is better?" than to produce or evaluate a perfect ground-truth response. Second, it allows us to capture subjective quality dimensions that vary across contexts—what counts as "helpful" may differ between a technical explanation and casual conversation. Third, by collecting multiple annotations per prompt pair, we can measure annotator agreement and filter out low-quality or ambiguous comparisons.

The dataset format typically looks like this:

{  "prompt": "Explain why gradient accumulation is useful.",  "chosen": "Gradient accumulation allows small batches to simulate larger batch training without increasing memory usage.",  "rejected": "Gradient accumulation is a thing used in neural networks sometimes."}

COPY

Notice the critical differences between these two responses:

• We are not providing a single perfect answer with an absolute quality label
• We are providing a preference pair that establishes relative ordering
• The "chosen" response is clearer, more informative, and more helpful—but it doesn't need to be perfect
• The "rejected" response isn't necessarily wrong—it's just demonstrably worse in terms of helpfulness and informativeness

The model will later learn that "chosen" > "rejected" through the reward modeling process. Importantly, this preference signal is transitive: if we collect enough comparisons showing A > B and B > C, the reward model can infer that A > C, allowing it to generalize beyond the exact pairs seen during training.

This preference-based approach also enables us to collect feedback on multiple quality dimensions simultaneously. A single comparison might reflect judgments about factual accuracy, clarity, safety, conciseness, and tone all at once. The reward model learns to compress these multidimensional human values into a single scalar score, effectively learning a implicit representation of what humans consider "good" responses in a given context.d.”

3.1.3 Stage 3: Training a Reward Model

Now we train a separate model called a reward model (RM). This is a crucial component that bridges the gap between human preferences and machine optimization. The reward model is not the same as the policy model we're trying to improve—it's a distinct neural network whose sole purpose is to evaluate the quality of responses.

The Reward Model's Role

The reward model's job is conceptually simple but practically powerful: given a prompt and a candidate response, it outputs a scalar score representing how good that response is according to human preferences. This score becomes the optimization signal that guides policy improvement in the reinforcement learning stage.

Formally, we can express the reward model as:

Reward Model: (prompt, response) → reward score

Or in mathematical notation: r_{\theta}(x, y) where x is the prompt, y is the response, and \theta represents the model's learned parameters. The output is a real-valued scalar that quantifies response quality.

Training Through Pairwise Comparison

We train the reward model using pairwise comparison loss, which directly reflects how we collected our preference data. Rather than trying to predict absolute quality scores—which would require us to define what "a score of 7.5" means in absolute terms—we train the model to correctly rank pairs of responses.

The training process works as follows: for each prompt x in our preference dataset, we have a chosen response y{chosen} that humans preferred and a rejected response y{rejected} that humans dispreferred. We want our reward model to assign a higher score to the chosen response than to the rejected one.

Let's denote:

• r{\theta}(x, y{chosen}) as the reward score for the preferred response
• r{\theta}(x, y{rejected}) as the reward score for the dispreferred response

The training objective encourages the model to maximize the difference between these two scores. We optimize the following loss function:

\mathcal{L}(\theta) = -\mathbb{E}{(x, y{chosen}, y{rejected})} \left[ \log \sigma(r{\theta}(x, y{chosen}) - r{\theta}(x, y_{rejected})) \right]

Breaking this down:

• The difference r{\theta}(x, y{chosen}) - r{\theta}(x, y{rejected}) measures how much higher the reward model scores the chosen response compared to the rejected one
• The sigmoid function \sigma(\cdot) converts this difference into a probability between 0 and 1, representing the model's confidence that the chosen response is indeed better
• Taking the log-sigmoid and negating it creates a loss that is minimized when the reward model confidently assigns higher scores to preferred responses
• The expectation \mathbb{E} indicates we average this loss across all preference pairs in our dataset

What This Loss Achieves

This pairwise ranking loss encourages the reward model to assign higher scores to preferred responses while maintaining relative ordering. Importantly, it doesn't force the model to predict specific numeric values—only to maintain the correct ranking. This is more robust than regression-based approaches because it focuses on what matters: relative quality rather than absolute scores.

As training progresses, the reward model learns to internalize the patterns of human preference present in the dataset. It learns that clear explanations score higher than vague ones, that helpful responses score higher than dismissive ones, that safe responses score higher than potentially harmful ones—all without being explicitly programmed with these rules. Instead, these preferences emerge from the patterns in the comparative rankings provided by human annotators.

The trained reward model becomes a learned proxy for human judgment, capable of evaluating novel responses it has never seen before by generalizing from the preference patterns in its training data. This generalization capability is what makes the reward model so powerful: it can provide feedback signals for the billions of possible responses the policy might generate during reinforcement learning, even though it was trained on only thousands or millions of preference pairs.

3.1.4 Implementing a Simple Reward Model

In practice, the reward model is often:

• A copy of the base model (or the SFT model from Stage 1)
• With a small classification head added on top
• Outputting a single scalar value representing response quality

What the reward model consumes and produces

• Input: a prompt and a candidate response (usually concatenated into one sequence)
• Output: a single scalar score that represents “how preferred” the response is

Why start from the base model?

The reward model inherits the base model’s language understanding. We are not teaching it new knowledge. We are teaching it to map language representations to a quality judgment.

Minimal architecture (base model + scalar head)

Here is a simplified example using Hugging Face:

import torchimport torch.nn as nnfrom transformers import AutoModel class RewardModel(nn.Module):    def __init__(self, base_model_name):        super().__init__()        self.model = AutoModel.from_pretrained(base_model_name)        hidden_size = self.model.config.hidden_size        self.reward_head = nn.Linear(hidden_size, 1)     def forward(self, input_ids, attention_mask):        outputs = self.model(            input_ids=input_ids,            attention_mask=attention_mask        )        pooled = outputs.last_hidden_state[:, -1, :]  # simple pooling choice        reward = self.reward_head(pooled)             # (batch, 1)        return reward

COPY

Training objective (pairwise ranking)

We do not need absolute “quality scores.” We only need the model to rank chosen > rejected.

import torch.nn.functional as F def preference_loss(chosen_reward, rejected_reward):    return -F.logsigmoid(chosen_reward - rejected_reward).mean()

COPY

How this loss works (intuition)

• If chosen_reward is much larger than rejected_reward, the loss becomes small.
• If the reward model ranks them incorrectly, the loss becomes large and pushes the model to flip the ordering.

End-to-End Example: Training a Reward Model on Preference Pairs

Below is a complete, runnable-style example that shows the full training path:

• load a preference dataset with prompt/chosen/rejected
• tokenize prompt+response pairs
• compute the pairwise ranking loss
• update the reward model

import torchimport torch.nn as nnimport torch.nn.functional as Ffrom torch.utils.data import DataLoaderfrom transformers import AutoTokenizer, AutoModel # -----------------------------# 1) Reward model definition# -----------------------------class RewardModel(nn.Module):    def __init__(self, base_model_name: str):        super().__init__()        self.model = AutoModel.from_pretrained(base_model_name)        hidden_size = self.model.config.hidden_size        self.reward_head = nn.Linear(hidden_size, 1)     def forward(self, input_ids, attention_mask):        outputs = self.model(            input_ids=input_ids,            attention_mask=attention_mask        )        pooled = outputs.last_hidden_state[:, -1, :]        reward = self.reward_head(pooled)        return reward def preference_loss(chosen_reward, rejected_reward):    return -F.logsigmoid(chosen_reward - rejected_reward).mean() # -----------------------------# 2) Dataset formatting# -----------------------------# Each item has: {"prompt": ..., "chosen": ..., "rejected": ...} def format_pair(prompt: str, response: str) -> str:    # Simple formatting. In a production chat setup, use your chat template.    return f"Prompt:\n{prompt}\n\nResponse:\n{response}" def collate_fn(batch, tokenizer, max_length=512):    prompts = [b["prompt"] for b in batch]    chosen = [b["chosen"] for b in batch]    rejected = [b["rejected"] for b in batch]     chosen_text = [format_pair(p, c) for p, c in zip(prompts, chosen)]    rejected_text = [format_pair(p, r) for p, r in zip(prompts, rejected)]     chosen_tok = tokenizer(        chosen_text,        padding=True,        truncation=True,        max_length=max_length,        return_tensors="pt",    )    rejected_tok = tokenizer(        rejected_text,        padding=True,        truncation=True,        max_length=max_length,        return_tensors="pt",    )     return {        "chosen_input_ids": chosen_tok["input_ids"],        "chosen_attention_mask": chosen_tok["attention_mask"],        "rejected_input_ids": rejected_tok["input_ids"],        "rejected_attention_mask": rejected_tok["attention_mask"],    } # -----------------------------# 3) Training loop# -----------------------------# NOTE: We use an encoder backbone here only to keep the example lightweight.# In LLM RLHF, reward models are commonly decoder-only backbones + a scalar head.base_model_name = "distilbert-base-uncased" tokenizer = AutoTokenizer.from_pretrained(base_model_name)model = RewardModel(base_model_name) device = torch.device("cpu")model.to(device) optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5) # Placeholder datasetdataset = [    {        "prompt": "Explain why gradient accumulation is useful.",        "chosen": "It lets you simulate a larger batch size by accumulating gradients across steps, without needing extra GPU memory.",        "rejected": "Gradient accumulation is a thing used sometimes in training."    },    {        "prompt": "What is KL divergence used for in RLHF?",        "chosen": "It acts as a constraint that penalizes the policy for drifting too far from a reference model, helping prevent reward hacking.",        "rejected": "It is a divergence that is used for math stuff."    },] loader = DataLoader(    dataset,    batch_size=2,    shuffle=True,    collate_fn=lambda b: collate_fn(b, tokenizer),) model.train()for epoch in range(3):    for batch in loader:        batch = {k: v.to(device) for k, v in batch.items()}         chosen_reward = model(batch["chosen_input_ids"], batch["chosen_attention_mask"])        rejected_reward = model(batch["rejected_input_ids"], batch["rejected_attention_mask"])         loss = preference_loss(chosen_reward, rejected_reward)         optimizer.zero_grad()        loss.backward()        optimizer.step()     print(f"epoch={epoch} loss={loss.item():.4f}")

COPY

Code breakdown (what each part is doing)

• RewardModel class
• Loads a pretrained backbone.
• Adds a linear head that outputs a single number.
• format_pair
• Ensures the reward model judges the response in context.
• collate_fn
• Builds two batches: chosen and rejected.
• Pads and truncates so tensors align.
• preference_loss
• Trains the model to rank chosen higher than rejected.
• Training loop
• Forward pass chosen and rejected.
• Backprop ranking loss.
• Repeat until the reward model consistently ranks preferred responses higher.

Training considerations (what usually matters in practice)

• Pooling choice: last-token pooling is simple, but you may want to pool over the response tokens only.
• Batch composition: keep prompts diverse to avoid prompt-specific shortcuts.
• Score scale: absolute reward values do not matter, but extreme magnitudes can destabilize PPO later.
• Learning rate: the head may need a higher LR than the backbone.

Practical note: In decoder-only RLHF, reward models often use a causal LM backbone plus a scalar head. Pooling frequently targets the end of the response segment (not necessarily the end of the full concatenated sequence).

3.1.5 Stage 4: Reinforcement Learning Optimization

Once the reward model is trained, we enter the final stage: optimizing the policy model to generate responses that maximize the reward signal we've carefully constructed.

The Core Objective

The fundamental goal is straightforward: generate responses that maximize the reward model's score.However, the implementation involves sophisticated machinery to achieve this safely and effectively. We're essentially teaching the policy to satisfy the preferences encoded in our reward model while maintaining the linguistic capabilities it acquired during pretraining and supervised fine-tuning.

Why Proximal Policy Optimization (PPO)?

This optimization is typically performed using Proximal Policy Optimization (PPO),a reinforcement learning algorithm specifically designed for stable policy updates. PPO has become the de facto standard for RLHF because it balances two competing needs: making meaningful progress toward better responses while preventing catastrophic failures that could occur from overly aggressive updates.

The Optimization Loop

The conceptual flow operates as follows:

• Generation: The policy model receives a prompt and generates a candidate response.This is standard autoregressive sampling—the model predicts tokens one at a time, building a complete response.
• Evaluation: The reward model evaluates the (prompt, response) pair and outputs a scalar reward score.This score represents how well the response aligns with learned human preferences—higher scores indicate better alignment.
• Policy Update: Using the reward signal, we adjust the policy's parameters to increase the expected reward for similar future prompts.This is where reinforcement learning theory comes into play: we're performing gradient ascent on expected reward, making the policy more likely to generate high-scoring responses.
• Constraint Enforcement: Critically, we constrain how much the policy can change in a single update, ensuring behavior doesn't drift too far from the reference policy.This is PPO's defining characteristic—the "proximal" constraint that keeps updates within a trust region.

The Critical Role of Constraints

The constraint mechanism deserves special attention because it addresses one of RLHF's most fundamental challenges.Without constraints, the policy might exploit weaknesses in the reward model—a phenomenon called reward hacking.

Consider what could go wrong: the reward model is imperfect. It's a learned approximation of human preferences, trained on limited data. If we allow the policy to change arbitrarily, it might discover adversarial patterns—responses that score highly according to the reward model but would be judged poorly by actual humans. For example, it might learn to generate verbose, repetitive text that exploits quirks in how the reward model processes length, or it might discover that certain phrases reliably trigger high scores regardless of whether they're actually appropriate.

The constraint prevents this by anchoring the policy to a reference model—typically a copy of the policy before RL training begins, or the supervised fine-tuned model from Stage 1. We add a KL divergence penalty term to the optimization objective that penalizes the policy for generating responses whose token probability distribution differs too much from the reference model. This keeps the policy "grounded" in sensible language generation while still allowing it to improve according to the reward signal.

Mathematical Formulation

The complete objective that PPO optimizes can be expressed as:

maximize: E[reward(x, y)] - β × KL(πθ || πref)

Where:

• π_θ is the policy we're training
• π_ref is the reference policy (frozen)
• β controls the strength of the KL penalty
• The expectation is over prompts x and generated responses y

This formulation makes the trade-off explicit: we want high rewards, but not at the cost of deviating too far from the reference distribution. The β hyperparameter determines this balance—higher values enforce stronger constraints, while lower values allow more aggressive optimization.

The Iterative Nature of Training

Unlike supervised training where each example has a fixed target, RL training is inherently dynamic. As the policy improves and generates better responses, the training distribution shifts. Early in training, the policy might generate low-quality responses that receive poor rewards, providing strong learning signals about what to avoid. Later, as the policy improves, the responses become better on average, and the learning signal becomes more subtle—distinguishing between "good" and "very good" rather than "bad" and "good."

This creates a moving target that requires careful curriculum design and hyperparameter scheduling. Too aggressive early updates can destabilize training, while too conservative late updates can prevent the policy from reaching its full potential.

Connection to Earlier Stages

Stage 4 builds directly on the foundation established in earlier stages. The supervised fine-tuning from Stage 1 provides a strong initialization—the policy already knows how to follow instructions and generate coherent responses. The preference data from Stage 2 and reward model from Stage 3 provide the optimization signal. Without these foundations, RL optimization would be intractable—the search space of possible responses is too vast to explore from scratch.

The result is a model that maintains the knowledge from pretraining, the instruction-following capability from supervised fine-tuning, and the preference alignment from reward-guided optimization—creating an AI system that is simultaneously capable, controllable, and aligned with human values.

3.1.6 Why This Pipeline Works

Supervised fine-tuning teaches imitation.

In Stage 1, the model learns by observing input-output pairs where correct answers are explicitly provided. This is direct behavioral cloning—the model sees "here is a question, here is the right response" and learns to reproduce similar patterns. It's learning to follow instructions through demonstration, building the foundational capability to understand what humans want and how to structure appropriate responses. However, this approach is limited to scenarios where there exists a clear, demonstrable correct answer.

Reward modeling teaches judgment.

Stage 3 introduces a fundamentally different capability: the ability to evaluate quality. Rather than learning what to say, the reward model learns what makes one response better than another. By training on preference pairs where humans have indicated "this response is better than that one," the model develops a nuanced understanding of qualities like helpfulness, clarity, harmlessness, and appropriateness. This is judgment—the ability to score and rank responses according to learned human values. The reward model becomes a differentiable, learned proxy for human judgment, enabling automated evaluation at scale.

Reinforcement learning teaches optimization under preference. Stage 4 completes the pipeline by teaching the policy model to actively maximize the rewards defined by human preferences. Unlike supervised learning where targets are fixed, or reward modeling where the goal is evaluation, RL training is about optimization—the policy learns to generate responses that score highly according to the reward model while maintaining its linguistic capabilities through KL divergence constraints. This creates a dynamic, iterative process where the model doesn't just imitate or judge, but actively seeks to produce outputs that satisfy learned preferences.

This layered process produces models that:

• Follow instructions: Through supervised fine-tuning, models gain the foundational ability to understand and execute user requests in a structured, coherent manner.
• Prefer helpful responses: The reward model encodes what makes responses valuable—depth, clarity, usefulness—and the policy learns to optimize for these qualities through the RL loop.
• Avoid harmful content: Preference data explicitly captures safety considerations, teaching the model to recognize and avoid generating responses that could be harmful, biased, or inappropriate.
• Align more closely with human expectations: The complete pipeline creates alignment—the model's behavior increasingly reflects human values and preferences, not just linguistic patterns from training data. This alignment comes from preference shaping, not from increased knowledge.

Each stage builds on the previous one: supervised fine-tuning provides strong initialization, preference data and reward modeling provide the optimization signal, and RL optimization brings them together to create models that are simultaneously capable, controllable, and aligned with human values.

3.1.7 Practical Example with TRL (Conceptual Outline)

The TRL (Transformer Reinforcement Learning) library provides high-level utilities for implementing PPO training, abstracting away much of the complexity while maintaining the flexibility needed for effective RLHF implementation.

End-to-End Example: Minimal PPO-Style RLHF with TRL (Skeleton)

This example shows the moving pieces of a PPO-based RLHF run:

• a policy you update
• a reference policy you keep frozen (for KL control)
• a reward model that scores generations

This is a minimal skeleton meant to make the pipeline concrete. Real training needs careful hyperparameters, batching, and stability tricks.

import torchfrom transformers import AutoTokenizer, AutoModelForCausalLMfrom trl import PPOTrainer, PPOConfig # -----------------------------# 1) Load models and tokenizer# -----------------------------policy_name = "gpt2"  # placeholder; use an instruction-tuned causal LM in practice tokenizer = AutoTokenizer.from_pretrained(policy_name)if tokenizer.pad_token is None:    tokenizer.pad_token = tokenizer.eos_token policy_model = AutoModelForCausalLM.from_pretrained(policy_name)ref_model = AutoModelForCausalLM.from_pretrained(policy_name)ref_model.eval()  # frozen reference # Suppose you already trained/loaded a reward model:# reward_model = ... # -----------------------------# 2) PPO config and trainer# -----------------------------config = PPOConfig(    batch_size=4,    mini_batch_size=2,    learning_rate=1e-5,) ppo_trainer = PPOTrainer(    config=config,    model=policy_model,    ref_model=ref_model,    tokenizer=tokenizer,) # -----------------------------# 3) Prompts# -----------------------------prompts = [    "Explain gradient accumulation in 3 bullet points.",    "What does the KL penalty do in RLHF?",    "Give a safe, concise answer: what is PPO?",    "Explain preference datasets with an example.",] # Tokenize promptsquery_tensors = [tokenizer(p, return_tensors="pt").input_ids.squeeze(0) for p in prompts] # -----------------------------# 4) RLHF loop: generate -> score -> update# -----------------------------policy_model.train()for step in range(10):    # Generate responses from the current policy    response_tensors = ppo_trainer.generate(        query_tensors,        max_new_tokens=64,        do_sample=True,        top_p=0.9,        temperature=0.8,    )     # Decode for reward scoring    queries = [tokenizer.decode(q, skip_special_tokens=True) for q in query_tensors]    responses = [tokenizer.decode(r, skip_special_tokens=True) for r in response_tensors]     # Compute rewards (placeholder)    # In practice, your reward model scores (prompt, response) pairs.    rewards = []    for q, r in zip(queries, responses):        # score = reward_model.score(q, r)        score = 0.0  # replace with real reward model output        rewards.append(torch.tensor(score))     # PPO update step    stats = ppo_trainer.step(query_tensors, response_tensors, rewards)     if step % 2 == 0:        print(f"step={step} stats_keys={list(stats.keys())[:5]}")

COPY

Code breakdown (what to pay attention to)

• Two policies, not one
• policy_model is trainable.
• ref_model is frozen.
• PPO uses the reference to compute a KL penalty that discourages the policy from drifting too far.
• The generation call is part of training
• ppo_trainer.generate(...) is not just “inference.” The generated samples become training data for the update step.
• Sampling settings matter. If you sample too randomly, training becomes noisy.
• Rewards come from the reward model (your learned proxy for humans)
• In a real setup, you build a function that takes (prompt, response) and returns a scalar.
• You often normalize rewards (for stability) and clip extremes.
• ppo_trainer.step(...) is where learning happens
• It updates the policy to increase expected reward.
• It also applies PPO-specific clipping and KL regularization.

Common beginner trap: if rewards are always near-zero, always positive, or extremely large, PPO can either fail to learn or drift. Reward scaling and KL control are not optional details.

Minimal DPO Training Loop (Why Many Teams Prefer It)

Direct Preference Optimization (DPO) often feels more approachable because you can train directly on preference pairs without a separate reward model and PPO loop.

import torchfrom transformers import AutoTokenizer, AutoModelForCausalLM # Placeholder: in practice you would use TRL's DPOTrainer, but the idea is simple:# maximize logprob(chosen) - logprob(rejected) with a reference model term. policy_name = "gpt2" tokenizer = AutoTokenizer.from_pretrained(policy_name)if tokenizer.pad_token is None:    tokenizer.pad_token = tokenizer.eos_token policy = AutoModelForCausalLM.from_pretrained(policy_name)ref = AutoModelForCausalLM.from_pretrained(policy_name)ref.eval() def logprob(model, input_ids, attention_mask):    # Computes token-level logprobs for the sequence (simplified)    out = model(input_ids=input_ids, attention_mask=attention_mask)    logits = out.logits[:, :-1, :]    labels = input_ids[:, 1:]    logp = torch.log_softmax(logits, dim=-1)    token_logp = logp.gather(-1, labels.unsqueeze(-1)).squeeze(-1)    # sum over tokens    return (token_logp * attention_mask[:, 1:]).sum(dim=-1) # One preference pairprompt = "Explain KL penalty in RLHF in 1–2 sentences."chosen = "It penalizes the policy for moving too far from a reference model, helping keep updates stable and preventing reward hacking."rejected = "It is a penalty about KL." chosen_text = prompt + "\n" + chosenrejected_text = prompt + "\n" + rejected chosen_tok = tokenizer(chosen_text, return_tensors="pt", padding=True)rejected_tok = tokenizer(rejected_text, return_tensors="pt", padding=True) pi_chosen = logprob(policy, chosen_tok["input_ids"], chosen_tok["attention_mask"])pi_rejected = logprob(policy, rejected_tok["input_ids"], rejected_tok["attention_mask"]) ref_chosen = logprob(ref, chosen_tok["input_ids"], chosen_tok["attention_mask"])ref_rejected = logprob(ref, rejected_tok["input_ids"], rejected_tok["attention_mask"]) beta = 0.1# Conceptual DPO objective (simplified):# prefer chosen over rejected, corrected by a reference.advantage = (pi_chosen - pi_rejected) - beta * (ref_chosen - ref_rejected)loss = -torch.log(torch.sigmoid(advantage)).mean() loss.backward()

COPY

Code breakdown

• You still train on chosen vs rejected.
• You still keep a reference model to anchor behavior.
• You optimize a direct preference objective instead of fitting an explicit reward model and running PPO.

In the next section (DPO), we will formalize this properly and show a clean, library-level implementation.

The Basic Training Loop

A simplified conceptual implementation looks like this:

from trl import PPOTrainer ppo_trainer = PPOTrainer(    model=policy_model,    ref_model=reference_model,    tokenizer=tokenizer) for batch in prompts:    responses = ppo_trainer.generate(batch)    rewards = reward_model(batch, responses)    ppo_trainer.step(batch, responses, rewards)

COPY

This deceptively simple loop encapsulates the entire reinforcement learning optimization process described in Stage 4. Let's break down what's happening at each step:

Initialization: Setting Up the Training Components

The PPOTrainer initialization requires three core components, each serving a distinct purpose in the optimization pipeline:

• policy_model: This is the model we're actively training—the one whose parameters will be updated to maximize reward. It typically starts as the supervised fine-tuned model from Stage 1, already capable of following instructions and generating coherent responses.
• ref_model: The reference model is a frozen copy of the policy model at the start of RL training. It serves as the anchor point for the KL divergence constraint, preventing the policy from deviating too far from sensible language generation. This is the mechanism that prevents reward hacking—without it, the policy might exploit weaknesses in the reward model.
• tokenizer: Handles the conversion between text and token representations, ensuring consistency across generation and evaluation steps.

The Generation Step

When ppo_trainer.generate(batch) executes, the policy model receives a batch of prompts and generates complete responses through autoregressive sampling. This is standard language model generation—predicting one token at a time—but with a crucial difference: these responses will be used to compute gradients and update the policy. The generation process must balance exploration (trying diverse responses to discover what works) with exploitation (leveraging what the model has already learned).

The Evaluation Step

The reward model evaluates each (prompt, response) pair, outputting scalar scores that represent alignment with learned human preferences. These scores are the optimization signal—they tell the policy which directions in parameter space lead to better behavior. The reward model here is the one trained in Stage 3 on preference data, serving as a differentiable proxy for human judgment.

The Policy Update Step

ppo_trainer.step(batch, responses, rewards) is where the actual learning happens. This step computes gradients and updates the policy's parameters to increase the expected reward for similar future prompts. Critically, it also enforces the KL divergence constraint relative to the reference model, ensuring updates remain within a trust region. This is PPO's defining characteristic—making meaningful progress while maintaining stability.

Critical Hyperparameters Requiring Tuning

While the loop structure is straightforward, successful training depends on carefully tuning several hyperparameters:

• KL penalty strength: Controls the trade-off between reward maximization and staying close to the reference model. Too high, and the policy barely improves; too low, and it may drift into reward hacking territory. This parameter (β in the mathematical formulation) is perhaps the most critical tuning knob in RLHF.
• Reward scaling: Normalizes reward magnitudes to a range that works well with the optimization algorithm. Without proper scaling, extreme reward values can cause numerical instability or make the KL penalty ineffective.
• Batch size: Affects both the variance of gradient estimates and computational efficiency. Larger batches provide more stable updates but require more memory and computation per step. The batch should also contain diverse prompts to prevent learning prompt-specific biases.
• Clipping thresholds: PPO uses probability ratio clipping to prevent excessively large policy updates. The clipping range determines how much the policy can change in a single step, directly impacting training stability.

The Iterative Nature

This loop runs for many iterations, and the training dynamics evolve as the policy improves. Early iterations provide strong learning signals as the policy learns to avoid obviously bad responses. Later iterations become more nuanced, distinguishing between good and excellent responses. The hyperparameters often need to be adjusted throughout training—for example, reducing the learning rate or increasing the KL penalty as training progresses to maintain stability.

Connection to the Complete Pipeline

This training loop represents the culmination of all four stages: it takes the supervised fine-tuned initialization from Stage 1, uses prompts that might come from the same distribution as Stage 2's preference data, relies on the reward model from Stage 3 for evaluation, and implements the constrained optimization strategy described in Stage 4. Each component is essential—remove any one, and the system fails.

But structurally, that is the loop: generate, evaluate, update, repeat.The elegance lies in how this simple iteration, when executed with proper constraints and careful hyperparameter tuning, can transform a model from merely capable to genuinely aligned with human values and preferences.

3.1.8 Why Preference Data Is Powerful

Preference data captures nuance.

Instead of asking:

“What is the correct answer?”

We ask:

“Which answer is better?”

That allows us to encode:

• Helpfulness
• Clarity
• Harmlessness
• Tone
• Safety
• Conciseness

And because humans are imperfect, we often collect multiple annotations per prompt to reduce noise.

3.1.9 Key Insight

RLHF does not make the model more knowledgeable.

It makes the model more aligned.

That distinction matters.

Knowledge comes from pretraining.

Alignment comes from preference shaping.

In the next section, we will explore modern alternatives that simplify RLHF — particularly Direct Preference Optimization (DPO), which removes the need for a separate reward model and PPO loop.

Before moving on, reflect:

Can you explain the difference between:

• Supervised fine-tuning
• Reward modeling
• Reinforcement optimization?

If you can, you are ready to move deeper.