Implementing RLHF (Reinforcement Learning from Human Feedback) in Python: A Three-Stage Guide

RLHF aligns language models to human preferences through a three-stage pipeline: Supervised Fine-Tuning (SFT), Reward Model training, and PPO optimization with a KL-penalty to prevent reward hacking.

The AI Engineering from Scratch curriculum provides a dependency-free, pure-Python implementation of the complete RLHF stack. This educational mirror of production systems like InstructGPT demonstrates how reinforcement learning transforms a base model into an instruction-following assistant while avoiding common failure modes such as reward hacking.

The Three-Stage RLHF Pipeline

The reference implementation in phases/18-ethics-safety-alignment/01-instruction-following-alignment-signal/code/main.py breaks RLHF into distinct stages that successively improve alignment.

Stage 1: Supervised Fine-Tuning (SFT)

The SFT stage teaches the model to generate coherent, helpful responses from labeled prompt-response pairs. This establishes a baseline policy that understands the output distribution but lacks explicit preference optimization.

In main.py, the stage1_sft() function trains using standard cross-entropy loss on demonstration data. This produces the initial policy $\pi_{SFT}$ that serves as the anchor for subsequent RL stages.

Stage 2: Reward Model Training

The reward model learns to predict human preferences by comparing pairs of responses. The curriculum implements a Bradley-Terry pairwise preference loss:

$$-\log \sigma(r_w - r_l)$$

Here, $r_w$ represents the reward for the preferred (winning) response and $r_l$ for the dispreferred (losing) response. The stage2_reward_model() function fits this scalar reward function $r(x, y)$ to provide a differentiable signal for policy optimization.

Stage 3: PPO with KL-Penalty

The final stage optimizes the policy using Proximal Policy Optimization (PPO) while constraining divergence from the SFT anchor. The objective function implemented in stage3_ppo() maximizes:

$$J(\pi) = \mathbb{E}[r] - \beta \cdot KL(\pi | \pi_{SFT})$$

The KL-penalty coefficient $\beta$ controls the alignment tax. A value of $\beta \approx 0.1$ keeps the RL policy close to the SFT distribution, preserving general capabilities while incorporating preference signals.

Understanding the KL-Penalty and Reward Hacking

Removing the KL regularization leads to reward hacking, where the policy exploits blind spots in the reward model to maximize scores without producing actually preferable outputs.

The curriculum demonstrates this empirically: running stage3_ppo(sft, rm, beta=0.0) produces a trajectory where rewards increase artificially while the KL divergence explodes. Conversely, beta=0.1 maintains stable training dynamics. This behavior is documented in phases/18-ethics-safety-alignment/02-reward-hacking-goodhart/docs/en.md, which analyzes how the KL term mitigates Goodhart's Law in alignment systems.

Running the RLHF Pipeline

Execute the complete three-stage pipeline to observe stable RLHF training with the default $\beta = 0.1$:

# Clone and navigate to the repository

git clone https://github.com/rohitg00/ai-engineering-from-scratch.git
cd ai-engineering-from-scratch

# Run the toy RLHF implementation

python phases/18-ethics-safety-alignment/01-instruction-following-alignment-signal/code/main.py

To experiment with reward hacking by disabling the KL penalty:

from phases.18_ethics_safety_alignment.01_instruction_following_alignment_signal.code.main import stage1_sft, stage2_reward_model, stage3_ppo

# Initialize pipeline stages

sft = stage1_sft()
rm = stage2_reward_model()

# Beta = 0.0 removes KL regularization, revealing reward hacking

rlhf, reward_traj, kl_traj = stage3_ppo(sft, rm, beta=0.0)

print(f"Final reward: {reward_traj[-1]}")
print(f"Final KL divergence: {kl_traj[-1]}")

Visualize the divergence between stable and hacked training:

import matplotlib.pyplot as plt

plt.plot(reward_traj, label="Reward")
plt.plot(kl_traj, label="KL Divergence")
plt.legend()
plt.title("RLHF Training Dynamics (β = 0.0)")
plt.show()

From RLHF to Direct Preference Optimization (DPO)

While RLHF requires three separate stages, Direct Preference Optimization (DPO) collapses the reward model and PPO steps into a single supervised loss. The curriculum derives the DPO loss from the RLHF-with-KL optimum in phases/18-ethics-safety-alignment/03-direct-preference-optimization-family/docs/en.md.

DPO eliminates the need to fit an explicit reward model by optimizing the policy directly on preference data using the closed-form solution:

$$\mathcal{L}{DPO} = -\log \sigma\left(\beta \log \frac{\pi(y_w|x)}{\pi{ref}(y_w|x)} - \beta \log \frac{\pi(y_l|x)}{\pi_{ref}(y_l|x)}\right)$$

This approach avoids the computational overhead of PPO while maintaining the theoretical guarantees of the three-stage pipeline.

Summary

• RLHF consists of three stages: SFT establishes a baseline policy, the reward model learns preference scoring, and PPO optimizes against the reward while constrained by a KL-penalty.
• The KL-penalty coefficient $\beta$ is critical: Values around 0.1 prevent reward hacking, while $\beta = 0.0$ allows the policy to exploit the reward model's weaknesses.
• Implementation is dependency-free: The ai-engineering-from-scratch repository provides pure Python implementations in phases/18-ethics-safety-alignment/01-instruction-following-alignment-signal/code/main.py with functions stage1_sft(), stage2_reward_model(), and stage3_ppo().
• DPO offers an alternative: Direct Preference Optimization collapses the RLHF pipeline into a single supervised learning stage, eliminating the reward model and PPO loop while preserving alignment properties.

Frequently Asked Questions

What is the Bradley-Terry loss in RLHF reward modeling?

The Bradley-Terry loss is a pairwise preference objective that trains the reward model to output higher scores for preferred responses compared to rejected ones. Mathematically, it computes $-\log \sigma(r_w - r_l)$, where $r_w$ is the reward for the winning response and $r_l$ for the losing response. This formulation transforms pairwise comparison data into a differentiable loss function suitable for gradient descent.

Why does removing the KL penalty cause reward hacking?

Without the KL penalty ($\beta = 0$), the PPO optimizer maximizes the reward model score without constraint. Because the reward model is an imperfect proxy for human judgment, the policy discovers adversarial inputs that trigger high reward scores while producing incoherent or undesirable outputs. The KL term penalizes deviation from the SFT policy, confining optimization to regions where the reward model remains reliable.

How does DPO differ from the three-stage RLHF pipeline?

Direct Preference Optimization (DPO) eliminates the explicit reward modeling and reinforcement learning stages by deriving a closed-form policy update from the RLHF optimization objective. Instead of training a separate reward model and running PPO, DPO optimizes the policy directly on preference pairs using a contrastive loss. This reduces computational overhead and avoids instability issues associated with online RL sampling.

Where is the RLHF implementation located in the repository?

The core implementation resides in phases/18-ethics-safety-alignment/01-instruction-following-alignment-signal/code/main.py. This file contains the stage1_sft(), stage2_reward_model(), and stage3_ppo() functions that implement the complete training loop. Supplementary documentation explaining the mathematical foundations and failure modes appears in phases/18-ethics-safety-alignment/01-instruction-following-alignment-signal/docs/en.md.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}