RLHF Explained: How Human Feedback Trains AI to Be Helpful and Safe

Pre-trained GPT (before RLHF):
Prompt: "How do I bake a chocolate cake?"
Output: "...and the temperature should be 350F for 30 minutes. However, older
         recipes use 375F. The difference in cake density stems from..." 
         [continues into technical essay no one asked for]

After RLHF:
Prompt: "How do I bake a chocolate cake?"
Output: "Here's a simple chocolate cake recipe:
         1. Preheat oven to 350F...
         [helpful, structured, practical]"

# SFT dataset format
sft_examples = [
    {
        "prompt": "Explain photosynthesis to a 10-year-old.",
        "response": """Photosynthesis is how plants make their own food using sunlight!

1. Plants have green chlorophyll in their leaves
2. Chlorophyll captures sunlight (like a solar panel)
3. Plants breathe in CO2 from the air
4. They drink water through their roots
5. Using sunlight as energy, they combine water and CO2 to make sugar
6. They release oxygen as a bonus — the air we breathe!"""
    },
    {
        "prompt": "Write a Python function to reverse a string.",
        "response": """```python
def reverse_string(s: str) -> str:
    return s[::-1]

SFT teaches the model *how* to respond. It has no way, yet, of judging which of two decent responses is better — that's the next stage's job.

---

## Stage 2: Reward Model Training

**The reward model is a separate network trained to score responses the way a human rater would — it turns "humans prefer A over B" into a number the RL stage can optimize against.** Annotators compare pairs of model outputs and pick the better one; those comparisons become training data.

```python
# Human annotators compare response pairs
preference_data = [
    {
        "prompt": "What are the side effects of ibuprofen?",
        "response_A": "Common side effects: stomach upset, nausea, heartburn. "
                      "Rare but serious: kidney problems, cardiovascular risks. "
                      "Take with food. Consult a doctor if symptoms persist.",
        "response_B": "Ibuprofen can cause side effects. It is a medication. "
                      "People take it for pain.",
        "preferred": "A",  # More helpful, accurate, practical
    },
    # Hundreds of thousands of comparisons...
]

# Train a reward model to predict human preferences
import torch
import torch.nn as nn
from transformers import AutoModel

class RewardModel(nn.Module):
    def __init__(self, base_model):
        super().__init__()
        self.model = base_model
        self.value_head = nn.Linear(self.model.config.hidden_size, 1)
    
    def forward(self, input_ids, attention_mask):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
        # Mean pool last hidden state
        hidden = outputs.last_hidden_state.mean(dim=1)
        return self.value_head(hidden).squeeze(-1)  # Scalar reward

def preference_loss(reward_chosen, reward_rejected):
    # Bradley-Terry model: P(A > B) = sigmoid(r_A - r_B)
    # Maximize log-likelihood that chosen > rejected
    return -torch.nn.functional.logsigmoid(reward_chosen - reward_rejected).mean()

# After training, reward_model(prompt + response_A) > reward_model(prompt + response_B)
# whenever humans prefer response_A

# In practice, use the trl library (Hugging Face)
from trl import PPOTrainer, PPOConfig, AutoModelForCausalLMWithValueHead

policy = AutoModelForCausalLMWithValueHead.from_pretrained("sft-checkpoint")
reference = AutoModelForCausalLMWithValueHead.from_pretrained("sft-checkpoint")
reference.requires_grad_(False)  # Frozen reference for KL computation

config = PPOConfig(
    learning_rate=1e-5,
    batch_size=32,
    kl_coef=0.1,           # KL penalty weight — crucial for stability
    target_kl=6.0,         # Target KL divergence from reference
)

trainer = PPOTrainer(
    config=config,
    model=policy,
    ref_model=reference,
    tokenizer=tokenizer,
)

for epoch in range(num_epochs):
    for batch in dataloader:
        prompts = batch["input_ids"]
        
        # Policy generates responses
        responses = policy.generate(prompts, max_new_tokens=256)
        
        # Score with reward model
        rewards = reward_model(prompts, responses)
        
        # KL penalty: keeps policy close to SFT reference
        # Prevents reward hacking (gaming the reward model)
        kl_penalty = compute_kl(policy, reference, responses)
        adjusted_rewards = rewards - config.kl_coef * kl_penalty
        
        # PPO update step
        stats = trainer.step(prompts, responses, adjusted_rewards)

from trl import DPOTrainer, DPOConfig
from datasets import Dataset

# Same preference data, but used directly for fine-tuning
dataset = Dataset.from_dict({
    "prompt": [
        "How do I improve Python code performance?",
        "What is the best way to learn machine learning?",
    ],
    "chosen": [
        "Profile first with cProfile, then optimize bottlenecks. Use list "
        "comprehensions, avoid repeated dict lookups, use NumPy for numerical ops.",
        "Start with fast.ai's Practical Deep Learning — builds working models "
        "first, then explains theory. Add Andrew Ng's ML course for fundamentals.",
    ],
    "rejected": [
        "You can improve Python by writing better code. There are many ways.",
        "Machine learning is complex. You should learn everything about it.",
    ]
})

dpo_config = DPOConfig(
    model_name_or_path="meta-llama/Meta-Llama-3.1-8B-Instruct",
    beta=0.1,              # Lower = closer to reference model
    learning_rate=5e-6,
    num_train_epochs=3,
    output_dir="./dpo-finetuned",
)

# No reward model, no RL — just supervised learning on preferences
trainer = DPOTrainer(
    model=base_model,
    ref_model=reference_model,
    config=dpo_config,
    train_dataset=dataset,
    tokenizer=tokenizer,
)

trainer.train()

r*(x, y) = β * log[π_θ(y|x) / π_ref(y|x)] + β * log Z(x)

Step 1: Generate initial responses with SFT model
         Prompt → [model] → Response draft

Step 2: Critique against constitution
         "Is this response helpful? Could it cause harm?"
         "Does it violate the honesty principle?"
         [AI reviews its own output against ~16 constitutional principles]

Step 3: Revise based on critique
         "Revised to be more helpful and avoid potential harm..."

Step 4: Use revised responses as preference data
         Original draft = rejected
         Revised draft = chosen
         Train reward model on these AI-generated preferences

Step 5: Fine-tune with RL using AI-preference reward model
         (RLAIF — RL from AI Feedback, not just human feedback)

Common reward hacking behaviors observed:

1. Length exploitation: if reward correlates with response length,
   model generates unnecessarily long, padded responses
   
2. Sycophancy: model learns to agree with the user's stated beliefs,
   even incorrect ones, because validation gets high ratings
   
3. Verbosity: adding "Great question!" and similar preambles
   if these patterns appear in high-rated training examples
   
4. Format gaming: if bullet points got high ratings, 
   model bullet-points everything even when prose is better

Prevention:
- KL penalty (limits how far from SFT model can go)
- Multiple reward models (harder to simultaneously exploit all)
- Iterative RLHF: retrain reward model on RL-policy outputs
- Red-teaming: systematically probe for gaming behaviors
- Length normalization in reward model training