RLHF: Reinforcement Learning from Human Feedback

# RLHF Training Pipeline Overview
import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel, GPT2Tokenizer
from typing import List, Dict, Any

class RLHFPipeline:
    """Complete RLHF training pipeline for language models"""
    
    def __init__(self, base_model_name: str):
        self.tokenizer = GPT2Tokenizer.from_pretrained(base_model_name)
        self.tokenizer.pad_token = self.tokenizer.eos_token
        
        # Initialize models for each phase
        self.sft_model = None
        self.reward_model = None
        self.rl_model = None
        self.reference_model = None
    
    def phase_1_supervised_finetuning(self, instruction_data: List[Dict]):
        """
        Phase 1: Supervised fine-tuning on high-quality demonstrations
        Goal: Teach the model to follow instructions well
        """
        print("Phase 1: Supervised Fine-tuning")
        
        # Load pre-trained model
        self.sft_model = GPT2LMHeadModel.from_pretrained('gpt2')
        
        # Prepare instruction-following data
        # Format: {"prompt": "...", "completion": "..."}
        training_examples = []
        for example in instruction_data:
            formatted_input = f"Human: {example['prompt']}\n\nAssistant: {example['completion']}"
            training_examples.append(formatted_input)
        
        # Fine-tune on demonstrations
        # This teaches the model to follow instructions and be helpful
        print(f"Training on {len(training_examples)} demonstration examples")
        
        return self.sft_model
    
    def phase_2_reward_modeling(self, preference_data: List[Dict]):
        """
        Phase 2: Train reward model on human preferences
        Goal: Learn what humans prefer in responses
        """
        print("Phase 2: Reward Model Training")
        
        # Create reward model by adding scalar head to base model
        class RewardModel(nn.Module):
            def __init__(self, base_model):
                super().__init__()
                self.base_model = base_model
                self.reward_head = nn.Linear(base_model.config.hidden_size, 1)
                
            def forward(self, input_ids, attention_mask=None):
                outputs = self.base_model(input_ids, attention_mask=attention_mask, output_hidden_states=True)
                # Use last token's hidden state
                last_hidden_state = outputs.hidden_states[-1][:, -1, :]
                reward = self.reward_head(last_hidden_state)
                return reward
        
        self.reward_model = RewardModel(GPT2LMHeadModel.from_pretrained('gpt2'))
        
        # Train on preference comparisons
        # Format: {"prompt": "...", "chosen": "...", "rejected": "..."}
        print(f"Training reward model on {len(preference_data)} preference pairs")
        
        # Bradley-Terry model for preference learning
        def preference_loss(chosen_rewards, rejected_rewards):
            return -torch.log(torch.sigmoid(chosen_rewards - rejected_rewards)).mean()
        
        return self.reward_model
    
    def phase_3_reinforcement_learning(self, prompts: List[str]):
        """
        Phase 3: Optimize policy using reinforcement learning
        Goal: Maximize reward while staying close to original model
        """
        print("Phase 3: Reinforcement Learning with PPO")
        
        # Initialize RL models
        self.rl_model = GPT2LMHeadModel.from_pretrained('gpt2')  # Policy to optimize
        self.reference_model = GPT2LMHeadModel.from_pretrained('gpt2')  # Frozen reference
        
        # PPO hyperparameters
        kl_coeff = 0.1  # KL divergence penalty
        clip_ratio = 0.2  # PPO clip parameter
        
        def compute_rewards(prompt_ids, response_ids):
            """Compute rewards including KL penalty"""
            full_ids = torch.cat([prompt_ids, response_ids], dim=1)
            
            # Get reward from reward model
            reward_score = self.reward_model(full_ids)
            
            # Compute KL divergence penalty
            with torch.no_grad():
                ref_logits = self.reference_model(full_ids).logits
                ref_log_probs = torch.log_softmax(ref_logits, dim=-1)
            
            policy_logits = self.rl_model(full_ids).logits
            policy_log_probs = torch.log_softmax(policy_logits, dim=-1)
            
            # KL divergence between policy and reference
            kl_div = torch.sum(torch.exp(ref_log_probs) * (ref_log_probs - policy_log_probs), dim=-1)
            
            # Final reward = base reward - KL penalty
            final_reward = reward_score.squeeze() - kl_coeff * kl_div.mean(dim=-1)
            
            return final_reward
        
        print(f"Optimizing policy on {len(prompts)} prompts with PPO")
        
        return self.rl_model
    
    def generate_rlhf_response(self, prompt: str, max_length: int = 100):
        """Generate response using RLHF-trained model"""
        if self.rl_model is None:
            raise ValueError("Must complete RLHF training first")
        
        inputs = self.tokenizer.encode(prompt, return_tensors='pt')
        
        with torch.no_grad():
            outputs = self.rl_model.generate(
                inputs,
                max_length=max_length,
                do_sample=True,
                temperature=0.7,
                pad_token_id=self.tokenizer.eos_token_id
            )
        
        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return response[len(prompt):]  # Return only the generated part

# Example usage
def run_rlhf_pipeline():
    pipeline = RLHFPipeline('gpt2')
    
    # Phase 1: SFT data
    sft_data = [
        {"prompt": "Explain quantum computing", "completion": "Quantum computing uses quantum mechanics..."},
        {"prompt": "Write a poem about AI", "completion": "In circuits deep and neural vast..."}
    ]
    
    # Phase 2: Preference data  
    preference_data = [
        {
            "prompt": "How do I cook pasta?",
            "chosen": "Boil water, add salt, cook pasta for 8-12 minutes until al dente.",
            "rejected": "Just put it in water and hope for the best."
        }
    ]
    
    # Phase 3: RL prompts
    rl_prompts = ["Explain machine learning", "Write a helpful response about..."]
    
    # Run pipeline
    sft_model = pipeline.phase_1_supervised_finetuning(sft_data)
    reward_model = pipeline.phase_2_reward_modeling(preference_data)
    rl_model = pipeline.phase_3_reinforcement_learning(rl_prompts)
    
    # Generate RLHF response
    response = pipeline.generate_rlhf_response("How can I help you today?")
    print(f"RLHF Response: {response}")
    
    return pipeline