simple_ppo.py

import math
import random

def main():
    print("Ultra-Simple PPO (Proximal Policy Optimization) Example")
    print("=" * 60)
    print("Task: Learn to output large numbers (1-10) to get high rewards")
    print()
    
    # Our "environment": numbers 1 to 10, reward = number/10
    actions = list(range(1, 11))  # [1, 2, 3, ..., 10]
    
    def get_reward(action):
        """Simple reward: larger numbers = better rewards"""
        return action / 10.0  # So action 10 gives reward 1.0
    
    # PPO Agent: Just probability distribution over actions
    policy_probs = [0.1] * 10  # [0.1, 0.1, 0.1, ...] for actions 1-10
    
    # Value function: estimates "how good is my current situation?"
    value_estimate = 0.55  # Start with reasonable guess (average of 1-10 scaled)
    
    print("Starting policy probabilities:")
    for i, prob in enumerate(policy_probs):
        print(f"  Action {i+1}: {prob:.3f}")
    print(f"Value estimate: {value_estimate:.3f}")
    print()
    
    # PPO Training parameters
    learning_rate = 0.1
    clip_epsilon = 0.2  # PPO clipping parameter
    kl_penalty = 0.01   # Penalty for changing policy too much
    
    print("Training with PPO...")
    print("Step | Action | Reward | Advantage | Policy Loss | Value Loss")
    print("-" * 65)
    
    for step in range(15):
        # === 1. COLLECT EXPERIENCE ===
        # Sample an action from current policy
        rand = random.random()
        cumulative = 0
        action_idx = 0
        for i, prob in enumerate(policy_probs):
            cumulative += prob
            if rand <= cumulative:
                action_idx = i
                break
        
        action = actions[action_idx]  # The actual number (1-10)
        reward = get_reward(action)
        
        # === 2. CALCULATE ADVANTAGE ===
        # Advantage = "How much better was this action than expected?"
        # If advantage > 0: this action was better than average
        # If advantage < 0: this action was worse than average
        advantage = reward - value_estimate
        
        # === 3. PPO POLICY UPDATE ===
        old_prob = policy_probs[action_idx]
        
        # Update the probability for the chosen action
        # If advantage > 0: increase probability (good action!)
        # If advantage < 0: decrease probability (bad action!)
        new_prob_unclipped = old_prob + learning_rate * advantage
        
        # PPO CLIPPING: Prevent huge changes that could destabilize learning
        # Clip the ratio new_prob/old_prob to be between (1-ε, 1+ε)
        ratio = new_prob_unclipped / old_prob
        clipped_ratio = max(1 - clip_epsilon, min(1 + clip_epsilon, ratio))
        new_prob = old_prob * clipped_ratio
        
        # Make sure probability stays positive and reasonable
        new_prob = max(0.01, min(0.9, new_prob))
        
        # === 4. KL DIVERGENCE PENALTY ===
        # KL measures "how much did the policy change?"
        # We want some change (to learn) but not too much (to stay stable)
        kl_div = old_prob * math.log(old_prob / new_prob + 1e-8)
        policy_loss = -(advantage * math.log(new_prob + 1e-8)) + kl_penalty * kl_div
        
        # Update the policy
        policy_probs[action_idx] = new_prob
        
        # Normalize probabilities to sum to 1
        total = sum(policy_probs)
        policy_probs = [p / total for p in policy_probs]
        
        # === 5. VALUE FUNCTION UPDATE ===
        # Update our estimate of "how good is the current situation?"
        # We use the actual reward to improve our estimate
        value_error = reward - value_estimate
        value_loss = value_error ** 2
        value_estimate += 0.05 * value_error  # Slow update for stability
        
        print(f"{step:4d} | {action:6d} | {reward:6.3f} | {advantage:9.3f} | {policy_loss:11.3f} | {value_loss:10.3f}")
    
    print()
    print("Training complete!")
    print()
    print("Final policy probabilities:")
    for i, prob in enumerate(policy_probs):
        print(f"  Action {i+1}: {prob:.3f}")
    print(f"Final value estimate: {value_estimate:.3f}")
    print()
    
    # Test the learned policy
    print("Testing learned policy (10 samples):")
    test_actions = []
    for _ in range(10):
        rand = random.random()
        cumulative = 0
        for i, prob in enumerate(policy_probs):
            cumulative += prob
            if rand <= cumulative:
                test_actions.append(i + 1)
                break
    
    print("Actions chosen:", test_actions)
    print(f"Average action: {sum(test_actions)/len(test_actions):.1f}")
    print()
    print("Notice: The agent learned to prefer larger numbers!")
    print()
    print("=== PPO KEY CONCEPTS EXPLAINED ===")
    print("• ADVANTAGE: How much better was this action than expected?")
    print("• CLIPPING: Prevents policy from changing too drastically")
    print("• KL PENALTY: Keeps new policy close to old policy for stability") 
    print("• VALUE FUNCTION: Estimates expected future reward")
    print("PPO balances learning (improve good actions) with stability (don't change too fast)!")

if __name__ == "__main__":
    main()