RLHF Explained: Reinforcement Learning from Human Feedback

Hey everyone, Khadervali here! If you’ve been following the whirlwind of advancements in AI, especially with Large Language Models (LLMs) like ChatGPT, you’ve undoubtedly encountered the term “alignment.” It’s the grand challenge of making powerful AI systems not just intelligent, but also helpful, harmless, and honest. And at the heart of much of this alignment success lies a fascinating and powerful technique: Reinforcement Learning from Human Feedback, or RLHF.

As a senior engineer, I’ve seen firsthand how traditional Machine Learning often struggles when the “right” answer isn’t clear-cut, or when human preferences are nuanced. RLHF is a game-changer because it directly tackles this by bringing human intuition into the loop, allowing us to train AI models that genuinely understand and embody our values. It’s a complex dance of supervised learning, reward modeling, and reinforcement learning, all orchestrated to shape AI behavior. In this deep dive, we’ll unpack RLHF piece by piece, from its foundational concepts to its real-world impact, complete with architectural insights and conceptual code examples.

The Problem with Traditional Reinforcement Learning for Complex Tasks

Before we dissect RLHF, let’s briefly revisit traditional Reinforcement Learning (RL). In classic RL, an agent learns to make decisions by interacting with an environment. It receives a numerical “reward” signal for its actions, and its goal is to maximize the cumulative reward over time. Think AlphaGo mastering Go or a robot learning to walk – the reward is clear: win the game, or don’t fall over.

However, many real-world tasks, especially those involving human interaction or creativity, don’t come with easily quantifiable reward functions. How do you assign a numerical score to a “good” story, a “helpful” chatbot response, or a “safe” robot manipulation? Crafting such a reward function by hand is often:

• Extremely difficult: Defining objective metrics for subjective qualities is a Sisyphean task.
• Prone to incompleteness: A hand-engineered reward function might miss crucial aspects of what makes a behavior truly desirable.
• Vulnerable to “reward hacking”: The agent might find clever ways to maximize the reward function without actually achieving the intended goal (e.g., a chatbot generating overly positive but unhelpful responses just because “positive” was rewarded).

This is the fundamental “alignment problem”: how do we ensure an AI system’s objectives are truly aligned with human intentions and preferences, rather than just optimizing a proxy metric? This is precisely where RLHF steps in.

What is Reinforcement Learning from Human Feedback (RLHF)?

RLHF is a methodology that leverages human preferences to train a reward model, which then serves as the reward signal for a reinforcement learning agent. Instead of designing a perfect reward function ourselves, we ask humans to judge the quality of an AI’s output. These human judgments are then used to train a separate model, the “reward model,” which learns to approximate human preferences. Finally, a policy model (the AI we want to improve) is fine-tuned using reinforcement learning, where the reward signal comes from this learned reward model.

The core intuition is simple yet powerful: humans might struggle to articulate a perfect rule set for “good” behavior, but they are remarkably good at judging and comparing different examples of behavior. RLHF capitalizes on this comparative judgment to implicitly capture complex human preferences.

The Architecture of RLHF: Three Core Phases

RLHF is typically broken down into three sequential, yet often iterative, phases. Let’s walk through each one, understanding its purpose and how it contributes to the overall goal of alignment.

Phase 1: Supervised Fine-Tuning (SFT) of a Pre-trained Model

The journey begins with a powerful, pre-trained base model. For LLMs, this would be a large transformer model (like GPT-3, LLaMA, etc.) that has been trained on a massive corpus of text to predict the next token. While incredibly capable, these models are often very general-purpose and may not be optimally aligned with specific human instructions or safety guidelines.

The first step in RLHF is to further fine-tune this base model using a smaller dataset of high-quality, human-curated demonstrations. This phase is called Supervised Fine-Tuning (SFT). The goal here is to give the model an initial understanding of the desired behavior, making it more helpful and instruction-following from the get-go. This is essentially supervised learning, where the input is a prompt and the output is a desired response generated by a human expert.

Why SFT?

• Initial Alignment: It provides a baseline of desirable behavior, ensuring the model doesn’t start from a completely unaligned state.
• Instruction Following: It teaches the model to follow instructions more effectively, which is crucial for subsequent phases.
• Reduces Search Space: By making the model generally more useful, it narrows down the space of possible outputs, making the reward model’s job easier.

The data for SFT typically consists of prompts and their corresponding ideal responses, handcrafted or carefully filtered by humans. For example, a human might write a prompt like “Explain quantum entanglement simply” and then provide a clear, concise, and accurate explanation.

Conceptually, the SFT process looks like this:

# Conceptual SFT Training Loop

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
from datasets import Dataset

# 1. Load a pre-trained base LLM and its tokenizer
model_name = "EleutherAI/gpt-neo-125M" # Example, typically much larger models are used
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Add a padding token if the tokenizer doesn't have one (common for GPT-like models)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 2. Prepare a dataset of human-curated (prompt, ideal_response) pairs
sft_data = [
    {"prompt": "Tell me a short story about a brave knight.", "response": "Once upon a time, in a land of dragons and despair, Sir Reginald, known for his gleaming armor and even brighter courage, embarked on a quest to save the princess..."},
    {"prompt": "Explain the concept of photosynthesis.", "response": "Photosynthesis is the process by which green plants and some other organisms use sunlight to synthesize foods with the help of chlorophyll. It converts light energy into chemical energy..."},
    # ... many more human-curated examples
]

# Format data for fine-tuning (e.g., as 'prompt\nresponse' or similar, depending on model)
def format_sft_example(example):
    # For causal LMs, we concatenate prompt and response and train to predict the next token
    return {"text": example["prompt"] + "\n" + example["response"]}

sft_dataset = Dataset.from_list(sft_data)
sft_dataset = sft_dataset.map(format_sft_example)

# 3. Tokenize the dataset
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True, max_length=512)

tokenized_sft_dataset = sft_dataset.map(tokenize_function, batched=True, remove_columns=["text"])

# 4. Define training arguments and train the model
training_args = TrainingArguments(
    output_dir="./sft_results",
    num_train_epochs=3,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8, # Simulate larger batch size
    learning_rate=2e-5,
    logging_dir="./sft_logs",
    logging_steps=100,
    save_strategy="epoch",
    report_to="none", # Disable logging to external services for this example
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_sft_dataset,
    tokenizer=tokenizer,
)

print("Starting Supervised Fine-Tuning (SFT)...")
trainer.train()
print("SFT complete. Model saved to ./sft_results")

# The SFT model is now ready for the next phase, or to be used as the "policy model" base.
sft_model = model # This is our initial policy model for RL

After this phase, we have a `sft_model` that is reasonably good at generating coherent and relevant responses to prompts, acting as a solid foundation for the subsequent RL steps.

Phase 2: Training a Reward Model (RM)

This is arguably the most crucial and innovative phase of RLHF. Instead of directly using human feedback to train the policy model (which would be slow and inefficient due to the need for continuous human input), we train a separate model – the Reward Model (RM) – to *predict* human preferences.

How it Works:

1. Generate Responses: We take the SFT model (or even the base pre-trained model) and generate multiple diverse responses to a variety of prompts. For a given prompt, we might generate 4-9 different completions.
2. Human Ranking: Human annotators are presented with these prompts and their corresponding generated responses. Crucially, instead of assigning a numerical score (which can be inconsistent between annotators), humans are asked to *rank* or *compare* the responses from best to worst. For example, given options A, B, C, they might say B > A > C. This comparative judgment is much easier and more consistent for humans than absolute scoring.
3. Train the Reward Model: This collected dataset of human preferences is then used to train the Reward Model. The RM is typically another neural network (often a slightly modified version of the SFT model, but with a scalar output head instead of a language modeling head). Its input is a prompt and a response, and its output is a single scalar value representing the “reward” or preference score for that response, given the prompt.

Reward Model Architecture (in words):

Imagine the SFT model is a large transformer. The Reward Model can be constructed by taking a copy of this SFT model, removing its language modeling head (the part that predicts the next token), and adding a new small linear layer on top. This linear layer projects the final hidden state of the transformer (corresponding to the entire sequence or just the last token) into a single scalar value. This scalar value is the predicted reward.

Diagrammatically:

    [Prompt + Response]
           |
           V
    (Transformer Encoder/Decoder Blocks - like the SFT model)
           |
           V
    (Final Hidden State Representation of the sequence)
           |
           V
    (Linear Layer: Maps hidden state to a single scalar)
           |
           V
    [Reward Score]

Loss Function for RM:

The RM is trained using a pairwise ranking loss. Given two responses, R1 and R2, for the same prompt, where humans preferred R1 over R2, the RM is trained to output a higher score for R1 than for R2. A common loss function is based on the Bradley-Terry model or a variant of it, which aims to maximize the log-probability that the preferred response is scored higher than the dispreferred one.

Mathematically, if r_theta(x, y) is the reward score given by the RM for a prompt x and response y, and humans preferred y_w (winner) over y_l (loser) for a prompt x, the loss for that pair would be:

Loss = -log(sigmoid(r_theta(x, y_w) - r_theta(x, y_l)))

This loss encourages r_theta(x, y_w) to be significantly higher than r_theta(x, y_l).

Conceptual RM training loop:

# Conceptual Reward Model Training Loop

import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel
from datasets import Dataset

# Assuming tokenizer is already loaded from SFT phase
# sft_model is the model from Phase 1, its backbone is used for RM

# 1. Prepare human preference data
# Each entry: {"prompt": "...", "chosen": "preferred_response", "rejected": "less_preferred_response"}
preference_data = [
    {"prompt": "Explain AI.", "chosen": "AI is the simulation of human intelligence processes by machines.", "rejected": "Robots are AI."},
    {"prompt": "Write a poem.", "chosen": "In fields of green, where whispers softly sway, A gentle breeze, a brand new day.", "rejected": "Poem go brrr, words are hard."},
    # ... many more human-ranked pairs
]

pref_dataset = Dataset.from_list(preference_data)

# 2. Define the Reward Model architecture
# Often, this is a language model with a single linear layer on top for scalar output.
# We'll use AutoModelForSequenceClassification for simplicity, adapting its head for 1 output.
# In practice, you might take the SFT model's body and add a new head.
class RewardModel(nn.Module):
    def __init__(self, base_model, tokenizer):
        super().__init__()
        self.tokenizer = tokenizer
        self.base_model = base_model # e.g., AutoModel.from_pretrained(model_name)
        # Assuming base_model outputs a sequence of hidden states, we take the last token's representation
        # and project it to a single scalar.
        self.reward_head = nn.Linear(base_model.config.hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        outputs = self.base_model(input_ids=input_ids, attention_mask=attention_mask)
        # Use the representation of the last token (or mean pool, etc.) for sequence classification
        last_hidden_state = outputs.last_hidden_state
        # For simplicity, taking the last token's representation before padding
        # In practice, need to handle variable lengths and padding correctly
        sequence_output = last_hidden_state[:, -1, :] # Take last non-padded token

        reward = self.reward_head(sequence_output)
        return reward

# Load a base model (e.g., the same as SFT or a smaller one) for the RM backbone
rm_base_model = AutoModel.from_pretrained(model_name)
reward_model = RewardModel(rm_base_model, tokenizer)

# 3. Define the preference loss function
class PreferenceLoss(nn.Module):
    def forward(self, chosen_rewards, rejected_rewards):
        # chosen_rewards and rejected_rewards are scalar outputs from the RM
        # We want chosen_rewards to be higher than rejected_rewards
        # Using a sigmoid difference to calculate the probability of chosen > rejected
        loss = -torch.log(torch.sigmoid(chosen_rewards - rejected_rewards)).mean()
        return loss

loss_fn = PreferenceLoss()
optimizer = torch.optim.Adam(reward_model.parameters(), lr=1e-5)

# 4. Tokenize and prepare data for RM training
def tokenize_preference_example(example):
    # Tokenize chosen and rejected responses separately for direct comparison
    chosen_tokens = tokenizer(example["prompt"] + "\n" + example["chosen"], truncation=True, max_length=512, return_tensors="pt")
    rejected_tokens = tokenizer(example["prompt"] + "\n" + example["rejected"], truncation=True, max_length=512, return_tensors="pt")
    return {
        "chosen_input_ids": chosen_tokens["input_ids"].squeeze(0),
        "chosen_attention_mask": chosen_tokens["attention_mask"].squeeze(0),
        "rejected_input_ids": rejected_tokens["input_ids"].squeeze(0),
        "rejected_attention_mask": rejected_tokens["attention_mask"].squeeze(0),
    }

tokenized_pref_dataset = pref_dataset.map(tokenize_preference_example)

# Convert to PyTorch tensors for DataLoader
tokenized_pref_dataset.set_format(type="torch", columns=[
    "chosen_input_ids", "chosen_attention_mask",
    "rejected_input_ids", "rejected_attention_mask"
])

from torch.utils.data import DataLoader
train_dataloader = DataLoader(tokenized_pref_dataset, batch_size=4)

print("Starting Reward Model Training...")
reward_model.train()
for epoch in range(2): # Train for a few epochs
    for batch_idx, batch in enumerate(train_dataloader):
        optimizer.zero_grad()

        # Get rewards for chosen and rejected responses
        chosen_rewards = reward_model(
            input_ids=batch["chosen_input_ids"],
            attention_mask=batch["chosen_attention_mask"]
        )
        rejected_rewards = reward_model(
            input_ids=batch["rejected_input_ids"],
            attention_mask=batch["rejected_attention_mask"]
        )

        loss = loss_fn(chosen_rewards, rejected_rewards)
        loss.backward()
        optimizer.step()

        if batch_idx % 50 == 0:
            print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {loss.item():.4f}")

print("Reward Model training complete.")
torch.save(reward_model.state_dict(), "./reward_model.pth")

The trained Reward Model (RM) is a critical asset. It has learned to mimic human preferences and can now provide a continuous, scalable reward signal for any generated response without requiring further human intervention for every decision.

Phase 3: Reinforcement Learning (RL) with the Reward Model

Now that we have a Reward Model that can provide a “score” for how much a human would prefer a given response, we can use this score as the reward signal for a standard reinforcement learning algorithm. The goal in this phase is to fine-tune the SFT model (which we now call the “policy model”) to maximize the rewards given by the RM.

The most commonly used algorithm for this phase is Proximal Policy Optimization (PPO). PPO is an on-policy algorithm that optimizes a “policy” (our SFT model’s generative capabilities) by taking small, iterative steps to improve its expected reward. It’s known for its stability and performance in complex environments.

How PPO Works in RLHF:

1. Initialize Policy: We start with the SFT model from Phase 1. This is our “policy” – the model that generates responses.
2. Generate Data (Rollouts): For a given set of prompts, the current policy model generates responses. These (prompt, response) pairs form the “trajectories” or “rollouts.”
3. Calculate Rewards: Each generated (prompt, response) pair is fed into the trained Reward Model (RM) from Phase 2, which outputs a scalar reward score.
4. Calculate Loss and Update Policy: PPO calculates a loss function that aims to increase the probability of actions (tokens) that lead to higher rewards. It then updates the policy model’s weights to maximize these rewards.
5. KL Divergence Penalty: A crucial addition in RLHF is a KL divergence penalty. This penalty term is added to the PPO loss to prevent the policy model from deviating too far from its original SFT version. This is important for several reasons:
   • Prevents Reward Hacking: It discourages the policy from finding “shortcuts” to high RM scores that might lead to nonsensical or undesired outputs (e.g., generating repetitive phrases that the RM accidentally overvalues).
   • Maintains Coherence: It ensures the model retains its general language understanding and generation capabilities learned during pre-training and SFT.
   • Stability: It helps stabilize the RL training process, preventing catastrophic forgetting or unstable behavior.

   The KL divergence term essentially measures how much the new policy’s distribution of outputs has shifted from the original SFT policy’s distribution.

Architecture Diagram (in words) for RL Phase:

                                  [Human Prompts]
                                        |
                                        V
                      +---------------------------------+
                      |     Policy Model (SFT Model)    |
                      | (Generates responses token by token) |
                      +---------------------------------+
                                        |
                                        V
                      [Generated Response] + [Original SFT Model (for KL divergence)]
                                        |
                                        V
                      +---------------------------------+
                      |     Reward Model (from Phase 2)   |
                      |     (Scores the generated response) |
                      +---------------------------------+
                                        |
                                        V
                      [Reward Score] + [KL Divergence Penalty]
                                        |
                                        V
                      +---------------------------------+
                      |         PPO Optimization        |
                      |   (Updates Policy Model weights)  |
                      +---------------------------------+
                                        |
                                        V
                                [Improved Policy Model]

The PPO loop continues for many iterations, iteratively refining the policy model. Each iteration involves generating responses, getting rewards from the RM, and updating the policy. The KL divergence penalty ensures that while the model learns to maximize the RM’s reward, it doesn’t forget how to be a generally good language model.

Conceptual PPO training step:

# Conceptual PPO Training Loop with Reward Model and KL Penalty

import torch

import torch.nn as nn

from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

import trl # A library like TRL (Transformer Reinforcement Learning) simplifies this
# Assuming SFT model and Reward Model are already trained and loaded

# sft_model is our policy model that we want to fine-tune

# reward_model is the model from Phase 2
# For RLHF, we usually need two copies of the policy model:

# 1. The 'active' policy model that is being updated (policy_model)

# 2. A 'reference' policy model that stays fixed to calculate KL divergence (ref_policy_model)
# 1. Load the SFT model as the initial policy model and a reference model

policy_model = AutoModelForCausalLM.from_pretrained("./sft_results") # Load SFT model

ref_policy_model = AutoModelForCausalLM.from_pretrained("./sft_results") # A copy for KL divergence

# Ensure tokenizer is correctly loaded and configured

tokenizer = AutoTokenizer.from_pretrained(model_name)

if tokenizer.pad_token is None:

    tokenizer.pad_token = tokenizer.eos_token
# Load the trained Reward Model

reward_model_state_dict = torch.load("./reward_model.pth")

rm_base_model_for_ppo = AutoModel.from_pretrained(model_name)

reward_model = RewardModel(rm_base_model_for_ppo, tokenizer)

reward_model.load_state_dict(reward_model_state_dict)

reward_model.eval() # Set RM to evaluation mode
# For actual PPO, you'd typically use a library like TRL or RLHF frameworks

# Here's a very high-level conceptual loop without using TRL directly to show steps
ppo_epochs = 1 # Number of PPO epochs (iterations over collected data)

kl_coef = 0.1 # Coefficient for KL divergence penalty
optimizer = torch.optim.Adam(policy_model.parameters(), lr=1e-6)
# Example prompts for RL training

rl_prompts = [

    "Write a short, engaging description of a new AI assistant.",

    "Draft an email to a colleague about a project update.",

    "Explain the concept of 'black holes' in simple terms.",

    # ... more prompts

]
print("Starting Reinforcement Learning (PPO) with Reward Model...")
for _ in range(ppo_epochs):

    # 1. Generate responses from the current policy model

    generated_responses = []

    for prompt_text in rl_prompts:

        input_ids = tokenizer(prompt_text, return_tensors="pt").input_ids

        # Generate a response

        output = policy_model.generate(

            input_ids,

            max_new_tokens=64,

            do_sample=True,

            top_k=50,

            top_p=0.95,

            num_return_sequences=1,

            pad_token_id=tokenizer.eos_token_id

        )

        response_text = tokenizer.decode(output[0], skip_special_tokens=True)

        generated_responses.append({"prompt": prompt_text, "response": response_text})
    # 2. Collect rewards from the Reward Model

    rewards = []

    for gen_pair in generated_responses:

        full_text = gen_pair["prompt"] + "\n" + gen_pair["response"]

        input_ids = tokenizer(full_text, return_tensors="pt", truncation=True, max_length=512).input_ids

        attention_mask = tokenizer(full_text, return_tensors="pt", truncation=True, max_length=512).attention_mask

        with torch.no_grad():

            reward = reward_model(input_ids=input_ids, attention_mask=attention_mask).item()

        rewards.append(reward)
    # 3. Calculate KL divergence (conceptual for this example, usually done with log probabilities)

    # This part is highly simplified. In a real PPO implementation, you would:

    # a. Get log probabilities of generated tokens from policy_model

    # b. Get log probabilities of generated tokens from ref_policy_model

    # c. Compute KL divergence for each token/sequence

    # For this example, we'll just simulate a placeholder KL penalty.

    kl_penalties = [kl_coef * (torch.rand(1).item() * 0.1) for _ in rewards] # Placeholder
    # 4. PPO Update (highly simplified conceptual step)

    # In a real PPO, this involves computing advantage estimates, ratio of new vs old policy probs,

    # clipping, and optimizing the PPO objective.

    # Here