simple_grpo.py

"""
Ultra-Simple GRPO (Group Relative Policy Optimization) Example
Used by DeepSeek for efficient reasoning model training
"""

import math
import random

def main():
    print("Ultra-Simple GRPO (Group Relative Policy Optimization) Example")
    print("=" * 65)
    print("Task: Learn to generate high-scoring math answers through group comparison")
    print()
    
    # Our "environment": Simple math problems with scoring
    # Higher scores = better mathematical reasoning
    problems = [
        "What is 2 + 3?",
        "What is 7 * 6?", 
        "What is 15 / 3?",
        "What is 8 - 2?"
    ]
    
    # Possible answers with different quality scores
    answer_pool = {
        "What is 2 + 3?": [
            ("5", 10.0),           # Perfect answer
            ("4", 2.0),            # Close but wrong  
            ("6", 2.0),            # Close but wrong
            ("I don't know", 0.5), # Honest but unhelpful
            ("Purple", 0.1)        # Nonsensical
        ],
        "What is 7 * 6?": [
            ("42", 10.0),          # Perfect
            ("41", 2.0),           # Close
            ("43", 2.0),           # Close  
            ("I think it's 40", 1.0), # Roughly right
            ("Seven times six", 0.5)  # Restates question
        ]
    }
    
    # GRPO Policy: probability distribution over answer types
    # [perfect, close_wrong, honest, nonsense] 
    policy_probs = [0.25, 0.25, 0.25, 0.25]  # Start uniform
    
    print("Starting policy probabilities:")
    print(f"  Perfect answers: {policy_probs[0]:.3f}")  
    print(f"  Close but wrong: {policy_probs[1]:.3f}")
    print(f"  Honest 'don't know': {policy_probs[2]:.3f}")
    print(f"  Nonsensical: {policy_probs[3]:.3f}")
    print()
    
    # GRPO parameters
    learning_rate = 0.15
    group_size = 4  # Key GRPO parameter: how many responses to compare
    clip_epsilon = 0.2
    
    print("Training with GRPO...")
    print("Step | Problem | Group Responses | Rewards | Advantages | Policy Update")
    print("-" * 85)
    
    for step in range(12):
        # === 1. SELECT PROBLEM ===
        problem = random.choice(list(answer_pool.keys()))
        available_answers = answer_pool[problem]
        
        # === 2. GENERATE GROUP OF RESPONSES ===
        # This is the key GRPO innovation: sample multiple responses at once
        group_responses = []
        group_rewards = []
        
        for _ in range(group_size):
            # Sample answer type from current policy
            rand = random.random()
            cumulative = 0
            answer_type_idx = 0
            
            for i, prob in enumerate(policy_probs):
                cumulative += prob
                if rand <= cumulative:
                    answer_type_idx = i
                    break
            
            # Get actual answer and reward for this type
            # Map policy indices to answer pool indices
            if answer_type_idx == 0:  # Perfect
                answer, reward = available_answers[0]
            elif answer_type_idx == 1:  # Close wrong
                answer, reward = random.choice(available_answers[1:3])
            elif answer_type_idx == 2:  # Honest
                answer, reward = available_answers[3]  
            else:  # Nonsensical
                answer, reward = available_answers[4]
            
            group_responses.append((answer, answer_type_idx))
            group_rewards.append(reward)
        
        # === 3. GRPO ADVANTAGE CALCULATION ===
        # Key difference from PPO: use group statistics instead of value network
        group_mean = sum(group_rewards) / len(group_rewards)
        group_std = math.sqrt(sum((r - group_mean)**2 for r in group_rewards) / len(group_rewards))
        group_std = max(group_std, 0.1)  # Avoid division by zero
        
        # Calculate normalized advantages for each response in group
        advantages = [(reward - group_mean) / group_std for reward in group_rewards]
        
        # === 4. POLICY UPDATE ===
        # Update policy based on which answer types performed well/poorly
        for response_info, advantage in zip(group_responses, advantages):
            answer, answer_type_idx = response_info
            old_prob = policy_probs[answer_type_idx]
            
            # GRPO policy gradient update (simplified)
            new_prob_unclipped = old_prob + learning_rate * advantage * old_prob
            
            # PPO-style clipping for stability  
            ratio = new_prob_unclipped / old_prob
            clipped_ratio = max(1 - clip_epsilon, min(1 + clip_epsilon, ratio))
            new_prob = old_prob * clipped_ratio
            
            # Keep probabilities reasonable
            new_prob = max(0.05, min(0.7, new_prob))
            policy_probs[answer_type_idx] = new_prob
        
        # Normalize probabilities to sum to 1
        total = sum(policy_probs)
        policy_probs = [p / total for p in policy_probs]
        
        # === 5. DISPLAY RESULTS ===
        prob_text = f"2+3?" if "2 + 3" in problem else f"7*6?"
        responses_text = ", ".join([f"{ans}({typ})" for ans, typ in group_responses])
        rewards_text = ", ".join([f"{r:.1f}" for r in group_rewards])
        advantages_text = ", ".join([f"{a:+.2f}" for a in advantages])
        
        print(f"{step:4d} | {prob_text:7} | {responses_text[:20]:20} | {rewards_text[:15]:15} | {advantages_text[:15]:15} | Updated")
    
    print()
    print("Training complete!")
    print()
    print("Final policy probabilities:")
    print(f"  Perfect answers: {policy_probs[0]:.3f}")  
    print(f"  Close but wrong: {policy_probs[1]:.3f}")
    print(f"  Honest 'don't know': {policy_probs[2]:.3f}")
    print(f"  Nonsensical: {policy_probs[3]:.3f}")
    print()
    
    # Test the learned policy
    print("Testing learned policy (6 samples):")
    test_problem = "What is 2 + 3?"
    test_answers = []
    
    for _ in range(6):
        rand = random.random()
        cumulative = 0
        for i, prob in enumerate(policy_probs):
            cumulative += prob
            if rand <= cumulative:
                if i == 0:
                    test_answers.append("5 (perfect)")
                elif i == 1:  
                    test_answers.append("4/6 (close)")
                elif i == 2:
                    test_answers.append("don't know (honest)")
                else:
                    test_answers.append("purple (nonsense)")
                break
    
    print("Sample responses:", test_answers)
    print()
    print("=== GRPO KEY CONCEPTS EXPLAINED ===")
    print("• GROUP SAMPLING: Generate multiple responses simultaneously")
    print("• RELATIVE ADVANTAGE: Compare each response to group mean (no critic needed!)")  
    print("• NORMALIZATION: advantage = (reward - group_mean) / group_std")
    print("• EFFICIENCY: No separate value network, ~50% compute reduction vs PPO")
    print("GRPO learns by comparing responses within each group - brilliant simplicity!")
    
if __name__ == "__main__":
    main()