🧠 Policy Gradient Fundamentals
🎯 What are Policy Gradient Methods?
Policy Gradient methods are a family of reinforcement learning algorithms that directly learn a parameterized policy π_θ(a|s) by optimizing the expected return via gradient ascent. Unlike value-based methods (Q-learning), they can handle continuous action spaces, learn stochastic policies naturally, and have strong convergence guarantees.
Why Policy Gradients Matter
Continuous actions: Natural for robotics and control. Stochastic policies: Optimal in partially observable settings. Direct optimization: Optimize what you care about. Convergence: Theoretical guarantees. Foundation: Basis for PPO, TRPO, RLHF. Modern AI: Powers ChatGPT's alignment.
📌 Policy π_θ📌 Gradient Ascent📌 Continuous Actions📌 RLHF
# Policy gradient — the big picture
import numpy as np
import torch
import torch.nn as nn
class PolicyNetwork(nn.Module):
"""Neural network that outputs action probabilities."""
def __init__(self, state_dim, action_dim):
super().__init__()
self.fc1 = nn.Linear(state_dim, 64)
self.fc2 = nn.Linear(64, 64)
self.fc3 = nn.Linear(64, action_dim)
def forward(self, state):
x = torch.relu(self.fc1(state))
x = torch.relu(self.fc2(x))
logits = self.fc3(x)
return torch.softmax(logits, dim=-1)
# The core idea:
# 1. Have a policy π_θ(a|s)
# 2. Sample trajectories by running the policy
# 3. Compute gradient of expected return w.r.t. θ
# 4. Update θ in direction that increases return
# 5. Repeat
def policy_gradient_concept():
"""
θ ← θ + α ∇_θ J(θ)
where J(θ) = E[return] under policy π_θ
Key insight: We can estimate ∇_θ J(θ) from samples!
"""
pass
# Comparison with value-based methods
comparison = {
'Policy Gradient': {
'learns': 'Policy π(a|s) directly',
'action_space': 'Discrete or continuous',
'exploration': 'Stochastic policy',
'convergence': 'Smooth, local optima',
'examples': 'REINFORCE, PPO, TRPO'
},
'Value-Based': {
'learns': 'Value Q(s,a) or V(s)',
'action_space': 'Discrete only',
'exploration': 'ε-greedy',
'convergence': 'Can oscillate',
'examples': 'Q-Learning, DQN'
}
}
for method, info in comparison.items():
print(f"\n{method}:")
for k, v in info.items():
print(f" {k}: {v}")📊 Policy Function
🎯 The Learning Target
The policy π(a|s) defines the probability of taking action a in state s. It can be:
- Deterministic: π(s) = a
- Stochastic: π(a|s) = P(a|s)
- Tabular: Table of probabilities
- Parameterized: π_θ(a|s) — neural network
Policy Representations
Discrete actions: Softmax over action logits. Continuous actions: Gaussian with mean and std. Mixed: Different heads for different action types. Recurrent: For partially observable MDPs. Attention: For multi-agent or structured tasks.
📊 Categorical📈 Gaussian🎯 Deterministic🔊 Stochastic
# Policy representations
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical, Normal
# 1. DISCRETE ACTION POLICY (Categorical)
class DiscretePolicy(nn.Module):
"""Policy for discrete action spaces."""
def __init__(self, state_dim, action_dim):
super().__init__()
self.fc1 = nn.Linear(state_dim, 128)
self.fc2 = nn.Linear(128, 128)
self.fc3 = nn.Linear(128, action_dim)
def forward(self, state):
x = F.relu(self.fc1(state))
x = F.relu(self.fc2(x))
logits = self.fc3(x)
return logits
def get_distribution(self, state):
"""Return action distribution."""
logits = self.forward(state)
return Categorical(logits=logits)
def sample(self, state):
"""Sample action from policy."""
dist = self.get_distribution(state)
action = dist.sample()
log_prob = dist.log_prob(action)
return action, log_prob
def log_prob(self, state, action):
"""Compute log probability of action."""
dist = self.get_distribution(state)
return dist.log_prob(action)
# 2. CONTINUOUS ACTION POLICY (Gaussian)
class ContinuousPolicy(nn.Module):
"""Policy for continuous action spaces."""
def __init__(self, state_dim, action_dim):
super().__init__()
self.fc1 = nn.Linear(state_dim, 128)
self.fc2 = nn.Linear(128, 128)
self.mean_head = nn.Linear(128, action_dim)
self.log_std = nn.Parameter(torch.zeros(action_dim))
def forward(self, state):
x = F.relu(self.fc1(state))
x = F.relu(self.fc2(x))
mean = self.mean_head(x)
std = torch.exp(self.log_std.clamp(-20, 2))
return mean, std
def get_distribution(self, state):
"""Return Gaussian distribution."""
mean, std = self.forward(state)
return Normal(mean, std)
def sample(self, state):
dist = self.get_distribution(state)
action = dist.rsample() # Reparameterized sample
log_prob = dist.log_prob(action).sum(-1)
return action, log_prob
# 3. USING THE POLICY
state_dim = 4
action_dim = 2
# Discrete
policy_discrete = DiscretePolicy(state_dim, action_dim)
state = torch.randn(1, state_dim)
action, log_prob = policy_discrete.sample(state)
print(f"Discrete action: {action.item()}")
print(f"Log prob: {log_prob.item():.4f}")
# Continuous
policy_continuous = ContinuousPolicy(state_dim, action_dim)
action, log_prob = policy_continuous.sample(state)
print(f"\nContinuous action: {action.detach().numpy().round(3)}")
print(f"Log prob: {log_prob.item():.4f}")
# 4. KEY PROPERTIES
print("\n=== Policy Properties ===")
print("""
1. SUPPORTS CONTINUOUS ACTIONS
- Output mean and std of Gaussian
- Natural for robotics and control
2. STOCHASTIC BY DESIGN
- Built-in exploration
- Optimal for POMDPs
- Mix strategies
3. SMOOTH POLICY CHANGES
- Small θ change → small π change
- Better for continuous control
4. NATURAL FOR MULTI-ACTION
- Can have different heads
- Handles different action types
5. COMPATIBLE WITH FUNCTION APPROX
- Neural networks
- Can use any architecture
""")
# 5. ACTION SAMPLING
print("=== Action Sampling ===")
# Discrete: sample from categorical
state = torch.randn(1, state_dim)
logits = policy_discrete(state)
probs = F.softmax(logits, dim=-1)
print(f"Action probabilities: {probs.detach().numpy().round(3)}")
# Sample multiple times
samples = [policy_discrete.sample(state)[0].item() for _ in range(10)]
print(f"10 samples: {samples}")
# Continuous: sample from Gaussian
mean, std = policy_continuous(state)
print(f"\nContinuous mean: {mean.detach().numpy().round(3)}")
print(f"Continuous std: {std.detach().numpy().round(3)}")
# Sample multiple times
samples = [policy_continuous.sample(state)[0].detach().numpy() for _ in range(5)]
print(f"5 samples:")
for s in samples:
print(f" {s.round(3)}")🎯 Objective Function
🎯 What We Optimize
The goal of policy gradient is to maximize the expected return J(θ) = E_τ~π_θ[R(τ)], where τ is a trajectory and R(τ) is the total reward. We optimize θ via gradient ascent.
Objective Variations
Total reward: J(θ) = E[Σ r_t]. Discounted: J(θ) = E[Σ γ^t r_t]. Average reward: For continuing tasks. Regularized: Add entropy for exploration. Constrained: Subject to safety constraints. Multi-objective: Balance multiple rewards.
📊 Expected Return📉 Discounted🎯 Average🛡️ Regularized
# Objective functions for policy gradient
import numpy as np
import torch
# 1. TOTAL REWARD OBJECTIVE
def total_reward_objective(trajectories):
"""J(θ) = E[Σ_t r_t]"""
returns = []
for trajectory in trajectories:
total = sum(trajectory['rewards'])
returns.append(total)
return np.mean(returns)
# 2. DISCOUNTED RETURN OBJECTIVE
def discounted_objective(trajectories, gamma=0.99):
"""J(θ) = E[Σ_t γ^t r_t]"""
returns = []
for trajectory in trajectories:
discounted = sum(
gamma ** t * r
for t, r in enumerate(trajectory['rewards'])
)
returns.append(discounted)
return np.mean(returns)
# 3. AVERAGE REWARD OBJECTIVE
def average_reward_objective(trajectories):
"""For continuing tasks."""
rates = []
for trajectory in trajectories:
rate = np.mean(trajectory['rewards'])
rates.append(rate)
return np.mean(rates)
# 4. REGULARIZED OBJECTIVE (with entropy)
def regularized_objective(trajectories, policy, beta=0.01):
"""J(θ) = E[R] + β * E[H(π)]"""
returns = [sum(t['rewards']) for t in trajectories]
mean_return = np.mean(returns)
# Entropy bonus for exploration
entropy_bonus = 0
for trajectory in trajectories:
for state in trajectory['states']:
probs = policy(torch.FloatTensor(state))
entropy = -(probs * torch.log(probs + 1e-8)).sum()
entropy_bonus += entropy.item()
return mean_return + beta * entropy_bonus / len(trajectories)
# 5. CONSTRAINED OBJECTIVE
def constrained_objective(trajectories, cost_limit=0.1):
"""
maximize E[R]
subject to E[C] ≤ cost_limit
"""
returns = [sum(t['rewards']) for t in trajectories]
costs = [sum(t.get('costs', [0])) for t in trajectories]
mean_return = np.mean(returns)
mean_cost = np.mean(costs)
if mean_cost > cost_limit:
# Penalize constraint violation
return mean_return - 10 * (mean_cost - cost_limit)
return mean_return
# 6. GENERAL ADVANTAGE ESTIMATION (GAE)
def gae_objective(trajectories, gamma=0.99, lam=0.95):
"""
Use GAE for advantage estimation.
A_t = Σ_l (γλ)^l δ_{t+l}
where δ_t = r_t + γV(s_{t+1}) - V(s_t)
"""
advantages = []
for trajectory in trajectories:
rewards = trajectory['rewards']
values = trajectory['values']
advantages_t = []
gae = 0
for t in reversed(range(len(rewards))):
if t == len(rewards) - 1:
next_value = 0
else:
next_value = values[t + 1]
delta = rewards[t] + gamma * next_value - values[t]
gae = delta + gamma * lam * gae
advantages_t.insert(0, gae)
advantages.extend(advantages_t)
return np.array(advantages)
# 7. COMPARISON
print("=== Objective Function Comparison ===")
print("""
Total Reward:
+ Simple, intuitive
- Doesn't handle continuing tasks
- May not converge
Discounted Return:
+ Standard for episodic tasks
+ Prioritizes near-term rewards
- Bias-variance tradeoff with γ
Average Reward:
+ Good for continuing tasks
+ Natural for long-running systems
- Harder to estimate
Regularized:
+ Encourages exploration
+ Prevents premature convergence
- Tuning β is tricky
Constrained:
+ Safety-critical applications
+ Business constraints
- Requires careful formulation
""")
# 8. ENTROPY REGULARIZATION
print("=== Entropy Regularization ===")
def compute_entropy(probs):
"""H(π) = -Σ π log π"""
return -(probs * np.log(probs + 1e-8)).sum()
# Example probabilities
probs_confident = np.array([0.95, 0.03, 0.02])
probs_uniform = np.array([0.33, 0.33, 0.34])
print(f"Confident policy entropy: {compute_entropy(probs_confident):.4f}")
print(f"Uniform policy entropy: {compute_entropy(probs_uniform):.4f}")
print("\nHigher entropy = more exploration")
# Entropy in the objective
def policy_loss_with_entropy(log_probs, advantages, entropies, beta=0.01):
"""Loss = -E[A * log π] - β * H(π)"""
policy_loss = -(advantages * log_probs).mean()
entropy_bonus = -entropies.mean() # Negative for minimization
return policy_loss + beta * entropy_bonus
print("\nEntropy regularization encourages exploration!")📐 Policy Gradient Theorem
🎯 The Mathematical Foundation
The Policy Gradient Theorem provides an expression for the gradient of the expected return that doesn't require differentiating through the environment dynamics. This is what makes policy gradient methods practical.
The Theorem
∇_θ J(θ) = E_π[∇_θ log π_θ(a|s) · Q^π(s,a)]. Key insight: No environment gradient needed. Log-derivative trick: ∇_θ π = π · ∇_θ log π. Intuition: Increase prob of actions that led to high Q. Variants: With baseline, with advantage, with GAE.
📐 ∇_θ J(θ)🔍 Log Trick🎯 Q-value✅ No Dynamics
# Policy Gradient Theorem — derivation and implementation
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# THE THEOREM:
# ∇_θ J(θ) = E_π[∇_θ log π_θ(a|s) · Q^π(s,a)]
#
# Where:
# - J(θ) = expected return
# - π_θ(a|s) = policy
# - Q^π(s,a) = action-value function under π
# DERIVATION (simplified):
# 1. J(θ) = ∫ p_θ(τ) R(τ) dτ
# 2. ∇_θ J = ∫ ∇_θ p_θ(τ) R(τ) dτ
# 3. Log trick: ∇_θ p_θ(τ) = p_θ(τ) ∇_θ log p_θ(τ)
# 4. ∇_θ J = ∫ p_θ(τ) ∇_θ log p_θ(τ) R(τ) dτ
# 5. = E_τ[∇_θ log p_θ(τ) R(τ)]
#
# Now, log p_θ(τ) = log p(s_0) + Σ log π_θ(a_t|s_t) + Σ log p(s_{t+1}|s_t,a_t)
# Since p(s_0) and p(s_{t+1}|s_t,a_t) don't depend on θ:
# ∇_θ log p_θ(τ) = Σ ∇_θ log π_θ(a_t|s_t)
#
# Therefore: ∇_θ J(θ) = E[Σ_t ∇_θ log π_θ(a_t|s_t) · R(τ)]
def policy_gradient_estimate(policy, trajectories):
"""
Estimate policy gradient from trajectories.
∇_θ J(θ) ≈ (1/N) Σ_i Σ_t ∇_θ log π_θ(a_t|s_t) · R(τ_i)
"""
all_log_probs = []
all_returns = []
for trajectory in trajectories:
# Compute return
returns = compute_returns(trajectory['rewards'])
for t, (state, action) in enumerate(zip(trajectory['states'],
trajectory['actions'])):
# log π_θ(a|s)
state_tensor = torch.FloatTensor(state).unsqueeze(0)
action_tensor = torch.LongTensor([action])
log_prob = policy.log_prob(state_tensor, action_tensor)
all_log_probs.append(log_prob)
all_returns.append(returns[t])
# Compute gradient
all_returns = torch.FloatTensor(all_returns)
log_probs = torch.stack(all_log_probs)
# Weight log probs by returns
weighted_log_probs = log_probs * all_returns
loss = -weighted_log_probs.mean() # Negative for gradient ascent
return loss
def compute_returns(rewards, gamma=0.99):
"""Compute discounted returns-to-go."""
returns = []
R = 0
for r in reversed(rewards):
R = r + gamma * R
returns.insert(0, R)
return returns
# IMPLEMENTATION: Simple REINFORCE
class PolicyNet(nn.Module):
def __init__(self, state_dim, action_dim):
super().__init__()
self.fc1 = nn.Linear(state_dim, 128)
self.fc2 = nn.Linear(128, 128)
self.fc3 = nn.Linear(128, action_dim)
def forward(self, state):
x = F.relu(self.fc1(state))
x = F.relu(self.fc2(x))
return F.softmax(self.fc3(x), dim=-1)
def log_prob(self, state, action):
probs = self.forward(state)
return torch.log(probs.gather(1, action.unsqueeze(1)) + 1e-8).squeeze()
# TRAINING LOOP
def train_policy_gradient(env, policy, optimizer, episodes=1000):
"""Train policy using policy gradient."""
for episode in range(episodes):
# Collect trajectory
state = env.reset()
log_probs = []
rewards = []
done = False
while not done:
state_tensor = torch.FloatTensor(state).unsqueeze(0)
probs = policy(state_tensor)
# Sample action
dist = torch.distributions.Categorical(probs)
action = dist.sample()
# Step environment
next_state, reward, done, _ = env.step(action.item())
log_probs.append(dist.log_prob(action))
rewards.append(reward)
state = next_state
# Compute returns
returns = compute_returns(rewards)
returns = torch.FloatTensor(returns)
# Normalize returns (variance reduction)
returns = (returns - returns.mean()) / (returns.std() + 1e-8)
# Compute loss
log_probs = torch.stack(log_probs)
loss = -(log_probs * returns).mean()
# Update
optimizer.zero_grad()
loss.backward()
optimizer.step()
if episode % 100 == 0:
total_reward = sum(rewards)
print(f"Episode {episode}, Return: {total_reward:.2f}")
# KEY INSIGHTS
print("=== Policy Gradient Theorem — Key Insights ===")
print("""
1. NO ENVIRONMENT GRADIENT NEEDED
- Works with black-box environments
- Only requires sampling
- Can use real-world systems
2. LOG-DERIVATIVE TRICK
- ∇_θ π = π ∇_θ log π
- Enables Monte Carlo estimation
- Foundation of all PG methods
3. INTERPRETATION
- Increase prob of good actions
- Decrease prob of bad actions
- "Good" measured by Q-value
4. VARIANTS
- REINFORCE: uses return R(τ)
- Actor-Critic: uses Q(s,a)
- PPO: uses advantage A(s,a)
- All derived from the theorem
5. LIMITATIONS
- High variance
- Slow convergence
- Local optima
- Sample inefficient
""")
# ADVANTAGE-BASED VARIANT
def advantage_policy_gradient(policy, value_fn, trajectories):
"""
Use advantage A(s,a) = Q(s,a) - V(s) instead of return.
∇_θ J(θ) = E[∇_θ log π_θ(a|s) · A^π(s,a)]
Lower variance than using raw returns!
"""
log_probs = []
advantages = []
for trajectory in trajectories:
returns = compute_returns(trajectory['rewards'])
for t, (state, action) in enumerate(zip(trajectory['states'],
trajectory['actions'])):
# Get value estimate
state_tensor = torch.FloatTensor(state).unsqueeze(0)
value = value_fn(state_tensor).item()
# Advantage = return - baseline
advantage = returns[t] - value
# Log prob
action_tensor = torch.LongTensor([action])
log_prob = policy.log_prob(state_tensor, action_tensor)
log_probs.append(log_prob)
advantages.append(advantage)
# Weighted loss
log_probs = torch.stack(log_probs)
advantages = torch.FloatTensor(advantages)
return -(log_probs * advantages).mean()🔄 REINFORCE Algorithm
🎯 The Simplest Policy Gradient
REINFORCE (Williams, 1992) is the simplest policy gradient algorithm. It uses Monte Carlo returns to estimate the gradient. Simple but high variance.
How REINFORCE Works
1. Sample: Run policy to collect trajectory. 2. Compute returns: Discounted returns G_t. 3. Compute gradient: ∇_θ log π(a_t|s_t) · G_t. 4. Update: θ ← θ + α ∇_θ J. 5. Repeat. Key: Unbiased but high variance.
🎲 Monte Carlo📊 Unbiased📈 High Variance✅ Simple
# REINFORCE from scratch
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
import gym
class REINFORCEPolicy(nn.Module):
"""Policy network for REINFORCE."""
def __init__(self, state_dim, action_dim, hidden=128):
super().__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, action_dim)
)
def forward(self, state):
logits = self.net(state)
return F.softmax(logits, dim=-1)
class REINFORCE:
"""REINFORCE algorithm implementation."""
def __init__(self, state_dim, action_dim, lr=1e-3, gamma=0.99):
self.policy = REINFORCEPolicy(state_dim, action_dim)
self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)
self.gamma = gamma
def select_action(self, state):
"""Sample action from policy."""
state = torch.FloatTensor(state).unsqueeze(0)
probs = self.policy(state)
dist = Categorical(probs)
action = dist.sample()
return action.item(), dist.log_prob(action)
def compute_returns(self, rewards):
"""Compute discounted returns-to-go."""
returns = []
R = 0
for r in reversed(rewards):
R = r + self.gamma * R
returns.insert(0, R)
return returns
def update(self, log_probs, rewards):
"""Update policy using REINFORCE."""
# Compute returns
returns = self.compute_returns(rewards)
returns = torch.FloatTensor(returns)
# Normalize returns (variance reduction)
if len(returns) > 1:
returns = (returns - returns.mean()) / (returns.std() + 1e-8)
# Compute loss: -E[log π(a|s) * G_t]
log_probs = torch.stack(log_probs)
loss = -(log_probs * returns).mean()
# Update
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
def train_episode(self, env):
"""Train for one episode."""
state = env.reset()
if isinstance(state, tuple):
state = state[0]
log_probs = []
rewards = []
done = False
while not done:
action, log_prob = self.select_action(state)
next_state, reward, done, truncated, _ = env.step(action)
if isinstance(done, bool) and truncated:
done = True
log_probs.append(log_prob)
rewards.append(reward)
state = next_state
loss = self.update(log_probs, rewards)
return sum(rewards), loss
# TRAINING EXAMPLE
def train_reinforce():
"""Train REINFORCE on CartPole."""
try:
env = gym.make('CartPole-v1')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n
agent = REINFORCE(state_dim, action_dim, lr=1e-2)
episode_rewards = []
for episode in range(1000):
total_reward, loss = agent.train_episode(env)
episode_rewards.append(total_reward)
if episode % 50 == 0:
avg_reward = np.mean(episode_rewards[-50:])
print(f"Episode {episode}, "
f"Avg Reward: {avg_reward:.2f}, "
f"Loss: {loss:.4f}")
if np.mean(episode_rewards[-50:]) > 475:
print(f"Solved in {episode} episodes!")
break
env.close()
return agent, episode_rewards
except Exception as e:
print(f"Training demo: {e}")
return None, None
# REINFORCE KEY CONCEPTS
print("=== REINFORCE Algorithm ===")
print("""
Algorithm:
1. Initialize policy π_θ with random θ
2. For each episode:
a. Generate trajectory τ = (s_0, a_0, r_0, ..., s_T)
b. For each timestep t:
- Compute return G_t = Σ_{k=t}^{T} γ^{k-t} r_k
- Compute gradient ∇_θ log π_θ(a_t|s_t) · G_t
c. Update: θ ← θ + α Σ_t ∇_θ log π_θ(a_t|s_t) · G_t
Properties:
✓ Unbiased gradient estimate
✓ Simple to implement
✓ Works with any differentiable policy
✗ High variance
✗ Slow convergence
✗ Requires full episodes
Variants:
- REINFORCE with baseline
- REINFORCE with entropy
- REINFORCE with normalization
""")
# VARIANCE REDUCTION TECHNIQUES
print("\n=== Variance Reduction for REINFORCE ===")
# 1. Return normalization
def normalized_returns(rewards, gamma=0.99):
"""Normalize returns for lower variance."""
returns = []
R = 0
for r in reversed(rewards):
R = r + gamma * R
returns.insert(0, R)
returns = np.array(returns)
return (returns - returns.mean()) / (returns.std() + 1e-8)
# 2. Baseline subtraction
def with_baseline(log_probs, returns, baseline_values):
"""Use baseline for variance reduction."""
advantages = returns - baseline_values
return -(log_probs * advantages).mean()
# 3. Reward-to-go
def reward_to_go(rewards, gamma=0.99):
"""Use reward-to-go instead of full return."""
returns = []
R = 0
for r in reversed(rewards):
R = r + gamma * R
returns.insert(0, R)
return returns
# 4. Causality trick
# Only use rewards AFTER action was taken
# (already incorporated in reward-to-go)
# COMPLETE REINFORCE WITH ALL IMPROVEMENTS
class ImprovedREINFORCE(REINFORCE):
"""REINFORCE with variance reduction."""
def __init__(self, state_dim, action_dim, lr=1e-3, gamma=0.99):
super().__init__(state_dim, action_dim, lr, gamma)
# Value function for baseline
self.value_fn = nn.Sequential(
nn.Linear(state_dim, 128),
nn.ReLU(),
nn.Linear(128, 1)
)
self.value_optimizer = optim.Adam(self.value_fn.parameters(), lr=lr)
def update(self, states, actions, log_probs, rewards):
"""Update with baseline and all improvements."""
# Compute returns
returns = self.compute_returns(rewards)
returns = torch.FloatTensor(returns)
# Compute values (baseline)
states_tensor = torch.FloatTensor(np.array(states))
values = self.value_fn(states_tensor).squeeze()
# Advantages
advantages = returns - values.detach()
# Normalize advantages
if len(advantages) > 1:
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# Policy loss
log_probs = torch.stack(log_probs)
policy_loss = -(log_probs * advantages).mean()
# Value loss
value_loss = F.mse_loss(values, returns)
# Update policy
self.optimizer.zero_grad()
policy_loss.backward()
self.optimizer.step()
# Update value function
self.value_optimizer.zero_grad()
value_loss.backward()
self.value_optimizer.step()
return policy_loss.item(), value_loss.item()📉 Variance Reduction & Baseline
🎯 Lower Variance, Same Gradient
The baseline is a function b(s) subtracted from returns. It doesn't change the expected gradient but can dramatically reduce variance. The optimal baseline is the value function V(s).
Why Baselines Work
Key insight: E[∇log π · b(s)] = 0 for any b(s). Proof: Σ_a π(a|s) ∇log π(a|s) = ∇Σ_a π(a|s) = ∇1 = 0. Advantage: A(s,a) = Q(s,a) - V(s). Best baseline: V(s). Effect: Reduces variance without bias.
📉 Variance📊 Baseline🎯 Advantage✅ Unbiased
# Variance reduction techniques for policy gradient
import numpy as np
import torch
import torch.nn as nn
# 1. THE BASELINE TRICK
print("=== The Baseline Trick ===")
print("""
Original gradient: E[∇log π(a|s) · R(τ)]
Modified gradient: E[∇log π(a|s) · (R(τ) - b(s))]
Why is this valid?
E[∇log π(a|s) · b(s)]
= Σ_a π(a|s) ∇log π(a|s) · b(s)
= Σ_a ∇π(a|s) · b(s)
= b(s) ∇Σ_a π(a|s)
= b(s) ∇1
= 0
So subtracting b(s) doesn't change the expected gradient!
But it CAN reduce variance dramatically.
""")
# 2. COMMON BASELINES
print("=== Common Baselines ===")
# a) No baseline (vanilla REINFORCE)
def no_baseline(log_probs, returns):
return -(log_probs * returns).mean()
# b) Constant baseline (average return)
def constant_baseline(log_probs, returns):
baseline = returns.mean()
advantages = returns - baseline
return -(log_probs * advantages).mean()
# c) State-dependent baseline (value function)
def value_baseline(log_probs, returns, values):
advantages = returns - values
return -(log_probs * advantages).mean()
# d) Optimal baseline (weighted by gradient magnitude)
def optimal_baseline(log_probs, returns):
"""b* = E[||∇log π||² R] / E[||∇log π||²]"""
grad_norms_sq = (log_probs ** 2).detach()
numerator = (grad_norms_sq * returns).sum()
denominator = grad_norms_sq.sum() + 1e-8
baseline = numerator / denominator
advantages = returns - baseline
return -(log_probs * advantages).mean()
# 3. VALUE FUNCTION AS BASELINE
class ValueBaseline(nn.Module):
"""Value function for baseline."""
def __init__(self, state_dim, hidden=128):
super().__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, 1)
)
def forward(self, state):
return self.net(state).squeeze(-1)
# 4. ADVANTAGE FUNCTION
print("=== Advantage Function ===")
print("""
A^π(s,a) = Q^π(s,a) - V^π(s)
Interpretation:
- How much better is action a than average?
- Positive: better than expected
- Negative: worse than expected
Advantage estimation methods:
1. Monte Carlo: A_t = G_t - V(s_t)
2. TD(0): A_t = r_t + γV(s_{t+1}) - V(s_t)
3. TD(λ): A_t = Σ (γλ)^l δ_{t+l}
4. GAE: Generalized Advantage Estimation
""")
# 5. VARIANCE COMPARISON
def compare_variance():
"""Compare variance of different estimators."""
np.random.seed(42)
n_samples = 1000
# Simulate returns and log probs
returns = np.random.randn(n_samples) * 10 + 50
log_probs = np.random.randn(n_samples) * 0.5 - 1
# a) No baseline
gradients_no_baseline = -log_probs * returns
var_no_baseline = gradients_no_baseline.var()
# b) Constant baseline
advantages_const = returns - returns.mean()
gradients_const = -log_probs * advantages_const
var_const = gradients_const.var()
# c) Value function baseline (assume perfect)
values = returns + np.random.randn(n_samples) * 2 # noisy
advantages_value = returns - values
gradients_value = -log_probs * advantages_value
var_value = gradients_value.var()
print("\n=== Variance Comparison ===")
print(f"No baseline: {var_no_baseline:.2f}")
print(f"Constant baseline: {var_const:.2f}")
print(f"Value baseline: {var_value:.2f}")
print(f"\nVariance reduction: {var_no_baseline/var_value:.1f}x")
compare_variance()
# 6. ADVANTAGE NORMALIZATION
print("\n=== Advantage Normalization ===")
def normalize_advantages(advantages):
"""Normalize advantages for stability."""
return (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# 7. FULL IMPLEMENTATION WITH BASELINE
class PPOWithBaseline:
"""Policy gradient with value function baseline."""
def __init__(self, policy, value_fn, optimizer, gamma=0.99):
self.policy = policy
self.value_fn = value_fn
self.optimizer = optimizer
self.gamma = gamma
def compute_gae(self, rewards, values, dones, lam=0.95):
"""Compute Generalized Advantage Estimation."""
advantages = []
gae = 0
for t in reversed(range(len(rewards))):
if t == len(rewards) - 1:
next_value = 0
else:
next_value = values[t + 1]
# TD error
delta = rewards[t] + self.gamma * next_value * (1 - dones[t]) - values[t]
# GAE
gae = delta + self.gamma * lam * (1 - dones[t]) * gae
advantages.insert(0, gae)
return advantages
def update(self, trajectories):
"""Update policy with baseline."""
# Collect all data
all_returns = []
all_log_probs = []
for traj in trajectories:
returns = self.compute_returns(traj['rewards'])
values = self.value_fn(traj['states']).detach().numpy()
advantages = self.compute_gae(
traj['rewards'], values, traj['dones']
)
all_returns.extend(returns)
all_log_probs.extend(traj['log_probs'])
# Normalize advantages
advantages = torch.FloatTensor(advantages)
advantages = normalize_advantages(advantages)
# Policy loss
log_probs = torch.stack(all_log_probs)
policy_loss = -(log_probs * advantages).mean()
# Update
self.optimizer.zero_grad()
policy_loss.backward()
self.optimizer.step()
return policy_loss.item()
print("""
=== Best Practices for Baselines ===
1. Always use a baseline
- Reduces variance dramatically
- Doesn't introduce bias
- Standard in modern methods
2. Use value function baseline
- V(s) is optimal baseline
- Train alongside policy
- Share network for efficiency
3. Normalize advantages
- Zero mean, unit variance
- Stabilizes training
- Standard in PPO
4. Use GAE for advantage estimation
- Balances bias and variance
- λ parameter controls tradeoff
- Standard in modern methods
5. Watch for baseline bias
- Value function errors
- Distribution shift
- Use old values for stability
""")🎭 Actor-Critic
🎯 Combining Policy and Value
Actor-Critic methods combine policy gradient (actor) with value function estimation (critic). The actor decides actions; the critic evaluates them. This reduces variance and speeds up learning.
Two Networks, One Goal
Actor: Policy π_θ(a|s) — decides actions. Critic: Value V_w(s) — evaluates states. Advantage: A(s,a) = r + γV(s') - V(s). Update: Actor uses advantage, critic uses TD error. Benefit: Lower variance, faster learning. Variants: A2C, A3C, PPO, SAC.
🎭 Actor📊 Critic🎯 Advantage📉 Lower Variance
# Actor-Critic implementation
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
import gym
class ActorCritic(nn.Module):
"""Shared actor-critic network."""
def __init__(self, state_dim, action_dim, hidden=128):
super().__init__()
# Shared layers
self.shared = nn.Sequential(
nn.Linear(state_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU()
)
# Actor head
self.actor = nn.Linear(hidden, action_dim)
# Critic head
self.critic = nn.Linear(hidden, 1)
def forward(self, state):
features = self.shared(state)
logits = self.actor(features)
value = self.critic(features)
return logits, value
def get_action(self, state):
"""Sample action and get value."""
logits, value = self.forward(state)
dist = Categorical(logits=logits)
action = dist.sample()
log_prob = dist.log_prob(action)
return action, log_prob, value
def evaluate(self, states, actions):
"""Evaluate actions (for training)."""
logits, values = self.forward(states)
dist = Categorical(logits=logits)
log_probs = dist.log_prob(actions)
entropy = dist.entropy()
return log_probs, values.squeeze(-1), entropy
class ActorCriticAgent:
"""Actor-Critic agent."""
def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99):
self.model = ActorCritic(state_dim, action_dim)
self.optimizer = optim.Adam(self.model.parameters(), lr=lr)
self.gamma = gamma
def train_episode(self, env):
"""Train for one episode."""
state = env.reset()
if isinstance(state, tuple):
state = state[0]
log_probs = []
values = []
rewards = []
entropies = []
done = False
while not done:
state_tensor = torch.FloatTensor(state).unsqueeze(0)
action, log_prob, value = self.model.get_action(state_tensor)
next_state, reward, done, truncated, _ = env.step(action.item())
if isinstance(done, bool) and truncated:
done = True
log_probs.append(log_prob)
values.append(value)
rewards.append(reward)
state = next_state
# Compute returns
returns = self.compute_returns(rewards)
returns = torch.FloatTensor(returns)
# Values as tensor
values = torch.cat(values).squeeze(-1)
# Advantages
advantages = returns - values.detach()
# Normalize advantages
if len(advantages) > 1:
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# Policy loss (actor)
log_probs = torch.stack(log_probs)
actor_loss = -(log_probs * advantages).mean()
# Value loss (critic)
critic_loss = F.mse_loss(values, returns)
# Combined loss
total_loss = actor_loss + 0.5 * critic_loss
# Update
self.optimizer.zero_grad()
total_loss.backward()
self.optimizer.step()
return sum(rewards), total_loss.item()
def compute_returns(self, rewards):
returns = []
R = 0
for r in reversed(rewards):
R = r + self.gamma * R
returns.insert(0, R)
return returns
# TRAINING
def train_actor_critic():
"""Train Actor-Critic on CartPole."""
try:
env = gym.make('CartPole-v1')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n
agent = ActorCriticAgent(state_dim, action_dim)
episode_rewards = []
for episode in range(1000):
total_reward, loss = agent.train_episode(env)
episode_rewards.append(total_reward)
if episode % 50 == 0:
avg_reward = np.mean(episode_rewards[-50:])
print(f"Episode {episode}, "
f"Avg Reward: {avg_reward:.2f}")
if np.mean(episode_rewards[-50:]) > 475:
print(f"Solved in {episode} episodes!")
break
env.close()
except Exception as e:
print(f"Training demo: {e}")
# ACTOR-CRITIC VARIANTS
print("=== Actor-Critic Variants ===")
print("""
A2C (Advantage Actor-Critic):
- Synchronous version of A3C
- Multiple workers in parallel
- Update from collected batch
- Simple but effective
A3C (Asynchronous Advantage Actor-Critic):
- Multiple workers asynchronously
- Each worker has own env
- Global shared parameters
- No replay buffer needed
PPO (Proximal Policy Optimization):
- Clipped objective
- Multiple epochs per batch
- State-of-the-art
- Used in RLHF
SAC (Soft Actor-Critic):
- Maximum entropy RL
- Off-policy
- Sample efficient
- Continuous control
TD3 (Twin Delayed DDPG):
- Deterministic policy
- Twin critics
- Delayed updates
- Continuous control
""")
# ADVANTAGE FUNCTION ESTIMATION
print("\n=== Advantage Estimation ===")
# 1. Monte Carlo (high variance, low bias)
def mc_advantage(rewards, values, gamma=0.99):
"""A_t = G_t - V(s_t)"""
returns = []
R = 0
for r in reversed(rewards):
R = r + gamma * R
returns.insert(0, R)
returns = torch.FloatTensor(returns)
return returns - values
# 2. TD(0) (low variance, high bias)
def td_advantage(rewards, values, gamma=0.99):
"""A_t = r_t + γV(s_{t+1}) - V(s_t)"""
advantages = []
for t in range(len(rewards)):
if t == len(rewards) - 1:
next_value = 0
else:
next_value = values[t + 1]
advantages.append(rewards[t] + gamma * next_value - values[t])
return torch.FloatTensor(advantages)
# 3. GAE (balanced)
def gae_advantage(rewards, values, gamma=0.99, lam=0.95):
"""Generalized Advantage Estimation."""
advantages = []
gae = 0
for t in reversed(range(len(rewards))):
if t == len(rewards) - 1:
next_value = 0
else:
next_value = values[t + 1]
delta = rewards[t] + gamma * next_value - values[t]
gae = delta + gamma * lam * gae
advantages.insert(0, gae)
return torch.FloatTensor(advantages)
# COMPARISON
print("""
Bias-Variance Tradeoff:
Method Bias Variance λ
------- ---- -------- ---
MC Low High λ=1
TD(0) High Low λ=0
GAE(0.95) Medium Medium 0.95
GAE(0.99) Lower Higher 0.99
Generalized Advantage Estimation:
- λ=0: TD(0) — high bias, low variance
- λ=1: Monte Carlo — low bias, high variance
- λ=0.95: Sweet spot — balanced
""")
# SHARED VS SEPARATE NETWORKS
print("\n=== Shared vs Separate Networks ===")
print("""
Shared Network:
✓ Fewer parameters
✓ Faster inference
✓ Feature sharing
✗ Interference between tasks
✗ Harder to tune
Separate Networks:
✓ Independent optimization
✓ No interference
✓ Easier to tune
✗ More parameters
✗ Slower inference
Rule of thumb:
- Shared: Small tasks, similar features
- Separate: Complex tasks, different features
- Shared is default in most implementations
""")📈 Generalized Advantage Estimation
🎯 The Sweet Spot
GAE (Schulman et al., 2016) provides a way to estimate advantages with controllable bias-variance tradeoff. It's the standard in modern policy gradient methods.
The GAE Formula
A_t^GAE(γ,λ) = Σ_{l=0}^∞ (γλ)^l δ_{t+l}. δ_t = r_t + γV(s_{t+1}) - V(s_t) (TD error). λ = 0: TD(0) — low variance, high bias. λ = 1: Monte Carlo — high variance, low bias. λ = 0.95: Balanced — sweet spot. γλ: Discount factor for multi-step returns.
📊 Bias-Variance🎯 TD Error⚖️ λ Parameter📈 Multi-step
# Generalized Advantage Estimation (GAE)
import torch
import numpy as np
def compute_gae(rewards, values, next_value=0, gamma=0.99, lam=0.95, dones=None):
"""
Compute Generalized Advantage Estimation.
A_t = Σ_{l=0}^{∞} (γλ)^l δ_{t+l}
where δ_t = r_t + γV(s_{t+1}) - V(s_t)
Args:
rewards: list of rewards [r_0, r_1, ..., r_{T-1}]
values: list of values [V(s_0), V(s_1), ..., V(s_{T-1})]
next_value: V(s_T) — bootstrap value
gamma: discount factor
lam: GAE parameter
dones: list of done flags
Returns:
advantages: GAE advantages
returns: TD(λ) returns
"""
if dones is None:
dones = [False] * len(rewards)
advantages = []
gae = 0
# Compute backwards (dynamic programming)
for t in reversed(range(len(rewards))):
if t == len(rewards) - 1:
next_val = next_value
else:
next_val = values[t + 1]
# Mask by done (if done, don't bootstrap)
mask = 1.0 - dones[t]
# TD error
delta = rewards[t] + gamma * next_val * mask - values[t]
# GAE accumulation
gae = delta + gamma * lam * mask * gae
advantages.insert(0, gae)
# Compute returns for value function
returns = [adv + val for adv, val in zip(advantages, values)]
return torch.FloatTensor(advantages), torch.FloatTensor(returns)
# EXAMPLE
print("=== GAE Example ===")
# Simulated episode
rewards = [1.0, 1.0, 1.0, -1.0, 1.0]
values = [0.5, 0.8, 0.9, 0.6, 0.3]
dones = [False, False, False, False, True]
# Compute GAE with different lambda values
for lam in [0.0, 0.5, 0.95, 1.0]:
advantages, returns = compute_gae(
rewards, values, next_value=0,
gamma=0.99, lam=lam, dones=dones
)
print(f"\nλ = {lam}:")
print(f" Advantages: {advantages.numpy().round(3)}")
print(f" Returns: {returns.numpy().round(3)}")
# INTERPRETING λ
print("\n\n=== Interpreting λ ===")
print("""
λ = 0 (TD(0)):
A_t = δ_t = r_t + γV(s_{t+1}) - V(s_t)
- Uses only 1-step return
- Low variance, high bias
- Fast but inaccurate
λ = 1 (Monte Carlo):
A_t = Σ_l γ^l r_{t+l} - V(s_t)
- Uses full return
- High variance, low bias
- Accurate but noisy
λ = 0.95 (GAE):
- Balanced estimator
- Geometric weighting of multi-step returns
- Standard in PPO and modern RL
- Trade-off controlled by λ
""")
# VISUALIZING GAE
print("=== GAE Weights ===")
print("\nWeight on k-step return (γλ)^{k-1}:")
gamma = 0.99
for lam in [0.0, 0.5, 0.9, 0.95, 0.99]:
weights = []
for k in range(1, 6):
weight = (gamma * lam) ** (k - 1)
weights.append(weight)
print(f" λ={lam}: {[f'{w:.3f}' for w in weights]}")
print("""
Notice:
- λ=0: Only 1-step (weight = [1, 0, 0, 0, 0])
- λ=0.95: Long tail (weights decrease slowly)
- λ=1: Never decays (all weights = 1)
The effective horizon is ~1/(1-γλ).
""")
# GAE IN PRACTICE
print("\n=== GAE in PPO ===")
class PPOWithGAE:
"""PPO with GAE for advantage estimation."""
def __init__(self, policy, value_fn, optimizer,
gamma=0.99, lam=0.95, clip_epsilon=0.2):
self.policy = policy
self.value_fn = value_fn
self.optimizer = optimizer
self.gamma = gamma
self.lam = lam
self.clip_epsilon = clip_epsilon
def collect_rollout(self, env, n_steps=2048):
"""Collect rollout data."""
states, actions, rewards, values, log_probs, dones = [], [], [], [], [], []
state = env.reset()
if isinstance(state, tuple):
state = state[0]
for _ in range(n_steps):
state_tensor = torch.FloatTensor(state).unsqueeze(0)
with torch.no_grad():
action, log_prob, value = self.policy.get_action(state_tensor)
next_state, reward, done, truncated, _ = env.step(action.item())
if isinstance(done, bool) and truncated:
done = True
states.append(state)
actions.append(action.item())
rewards.append(reward)
values.append(value.item())
log_probs.append(log_prob.item())
dones.append(done)
state = next_state if not done else env.reset()[0]
return {
'states': states,
'actions': actions,
'rewards': rewards,
'values': values,
'log_probs': log_probs,
'dones': dones
}
def update(self, rollout):
"""Update using PPO with GAE."""
# Compute GAE
with torch.no_grad():
next_state = rollout['states'][-1]
# Bootstrap value for last state
next_value = 0 # or value_fn(next_state)
advantages, returns = compute_gae(
rollout['rewards'],
rollout['values'],
next_value=next_value,
gamma=self.gamma,
lam=self.lam,
dones=rollout['dones']
)
# Normalize advantages
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# Convert to tensors
states = torch.FloatTensor(np.array(rollout['states']))
actions = torch.LongTensor(rollout['actions'])
old_log_probs = torch.FloatTensor(rollout['log_probs'])
# Multiple epochs
for epoch in range(10):
# Get current log probs and values
log_probs, values, entropy = self.policy.evaluate(states, actions)
# Ratio
ratio = torch.exp(log_probs - old_log_probs)
# Clipped objective
surr1 = ratio * advantages
surr2 = torch.clamp(ratio,
1 - self.clip_epsilon,
1 + self.clip_epsilon) * advantages
policy_loss = -torch.min(surr1, surr2).mean()
# Value loss
value_loss = F.mse_loss(values, returns)
# Entropy bonus
entropy_loss = -entropy.mean()
# Total loss
loss = policy_loss + 0.5 * value_loss + 0.01 * entropy_loss
# Update
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.policy.parameters(), 0.5)
self.optimizer.step()
# BEST PRACTICES
print("\n\n=== GAE Best Practices ===")
print("""
1. Use λ = 0.95 as default
- Works well across most tasks
- Standard in PPO implementations
2. Normalize advantages
- Zero mean, unit variance
- Critical for stability
- Especially with GAE
3. Bootstrap properly
- Use V(s_T) for terminal states
- Mask with done flag
- Important for episodic tasks
4. Watch for numerical issues
- Use float32 or float64
- Clip advantages if needed
- Monitor value predictions
5. Tune γ and λ together
- γ controls horizon
- λ controls bias-variance
- Consider task requirements
6. Combine with value function training
- Critic learns V(s)
- GAE uses V(s)
- Joint optimization works well
7. Use parallel environments
- Reduces variance further
- Faster data collection
- Standard in modern implementations
""")🚀 PPO (Proximal Policy Optimization)
🎯 The Workhorse of Modern RL
PPO (Schulman et al., 2017) is the most popular policy gradient algorithm. It's simple, stable, and effective. It's used to train ChatGPT (RLHF), robotics policies, and game-playing agents.
PPO's Key Innovation
Clipped objective: Prevents large policy updates. Trust region: Approximated via clipping. Multiple epochs: Reuse data efficiently. Simple: First-order optimization. Robust: Works across domains. Standard: Default in many applications.
✂️ Clipping🔄 Multi-epoch📊 Stable🎯 Simple
# PPO from scratch
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
import gym
class PPONetwork(nn.Module):
"""Actor-Critic network for PPO."""
def __init__(self, state_dim, action_dim, hidden=64):
super().__init__()
# Shared feature extractor
self.shared = nn.Sequential(
nn.Linear(state_dim, hidden),
nn.Tanh(),
nn.Linear(hidden, hidden),
nn.Tanh()
)
# Actor head
self.actor = nn.Linear(hidden, action_dim)
# Critic head
self.critic = nn.Linear(hidden, 1)
def forward(self, state):
features = self.shared(state)
logits = self.actor(features)
value = self.critic(features)
return logits, value
def get_action(self, state, action=None):
"""Sample action and compute log prob."""
logits, value = self.forward(state)
dist = Categorical(logits=logits)
if action is None:
action = dist.sample()
log_prob = dist.log_prob(action)
entropy = dist.entropy()
return action, log_prob, entropy, value
class PPO:
"""PPO algorithm implementation."""
def __init__(
self,
state_dim,
action_dim,
lr=3e-4,
gamma=0.99,
lam=0.95,
clip_epsilon=0.2,
value_coef=0.5,
entropy_coef=0.01,
epochs=10,
batch_size=64,
rollout_steps=2048,
):
self.network = PPONetwork(state_dim, action_dim)
self.optimizer = optim.Adam(self.network.parameters(), lr=lr)
self.gamma = gamma
self.lam = lam
self.clip_epsilon = clip_epsilon
self.value_coef = value_coef
self.entropy_coef = entropy_coef
self.epochs = epochs
self.batch_size = batch_size
self.rollout_steps = rollout_steps
def collect_rollout(self, env):
"""Collect rollout data."""
states = []
actions = []
rewards = []
log_probs = []
values = []
dones = []
state = env.reset()
if isinstance(state, tuple):
state = state[0]
for _ in range(self.rollout_steps):
state_tensor = torch.FloatTensor(state).unsqueeze(0)
with torch.no_grad():
action, log_prob, _, value = self.network.get_action(state_tensor)
next_state, reward, done, truncated, _ = env.step(action.item())
if isinstance(done, bool) and truncated:
done = True
states.append(state)
actions.append(action.item())
rewards.append(reward)
log_probs.append(log_prob.item())
values.append(value.item())
dones.append(done)
state = next_state
if done:
state = env.reset()
if isinstance(state, tuple):
state = state[0]
# Compute GAE
advantages = self.compute_gae(rewards, values, dones)
return {
'states': np.array(states),
'actions': np.array(actions),
'rewards': np.array(rewards),
'log_probs': np.array(log_probs),
'values': np.array(values),
'dones': np.array(dones),
'advantages': advantages,
}
def compute_gae(self, rewards, values, dones, next_value=0):
"""Compute Generalized Advantage Estimation."""
advantages = []
gae = 0
for t in reversed(range(len(rewards))):
if t == len(rewards) - 1:
next_val = next_value
else:
next_val = values[t + 1]
mask = 1.0 - dones[t]
delta = rewards[t] + self.gamma * next_val * mask - values[t]
gae = delta + self.gamma * self.lam * mask * gae
advantages.insert(0, gae)
return np.array(advantages)
def update(self, rollout):
"""PPO update."""
# Convert to tensors
states = torch.FloatTensor(rollout['states'])
actions = torch.LongTensor(rollout['actions'])
old_log_probs = torch.FloatTensor(rollout['log_probs'])
returns = torch.FloatTensor(rollout['advantages'] + rollout['values'])
# Normalize advantages
advantages = torch.FloatTensor(rollout['advantages'])
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# Training
n_samples = len(states)
indices = np.arange(n_samples)
total_loss = 0
for epoch in range(self.epochs):
np.random.shuffle(indices)
for start in range(0, n_samples, self.batch_size):
end = start + self.batch_size
batch_idx = indices[start:end]
batch_states = states[batch_idx]
batch_actions = actions[batch_idx]
batch_old_log_probs = old_log_probs[batch_idx]
batch_advantages = advantages[batch_idx]
batch_returns = returns[batch_idx]
# Forward pass
_, new_log_probs, entropy, values = self.network.get_action(
batch_states, batch_actions
)
# Policy loss (clipped)
ratio = torch.exp(new_log_probs - batch_old_log_probs)
surr1 = ratio * batch_advantages
surr2 = torch.clamp(
ratio, 1 - self.clip_epsilon, 1 + self.clip_epsilon
) * batch_advantages
policy_loss = -torch.min(surr1, surr2).mean()
# Value loss
value_loss = F.mse_loss(values.squeeze(), batch_returns)
# Entropy bonus
entropy_loss = -entropy.mean()
# Total loss
loss = (
policy_loss
+ self.value_coef * value_loss
+ self.entropy_coef * entropy_loss
)
# Update
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.network.parameters(), 0.5)
self.optimizer.step()
total_loss += loss.item()
return total_loss / (self.epochs * (n_samples // self.batch_size))
# TRAINING
def train_ppo():
"""Train PPO on CartPole."""
try:
env = gym.make('CartPole-v1')
state_dim = env.observation_space.shape[0]
action_dim = env.action_space.n
agent = PPO(state_dim, action_dim)
episode_rewards = []
episode_reward = 0
for iteration in range(300):
rollout = agent.collect_rollout(env)
# Track rewards
for r, d in zip(rollout['rewards'], rollout['dones']):
episode_reward += r
if d:
episode_rewards.append(episode_reward)
episode_reward = 0
# Update
loss = agent.update(rollout)
if iteration % 10 == 0:
avg_reward = np.mean(episode_rewards[-20:]) if episode_rewards else 0
print(f"Iteration {iteration}, "
f"Avg Reward: {avg_reward:.2f}, "
f"Loss: {loss:.4f}")
if episode_rewards and np.mean(episode_rewards[-20:]) > 475:
print(f"Solved in {iteration} iterations!")
break
env.close()
except Exception as e:
print(f"Training demo: {e}")
# PPO VARIANTS
print("=== PPO Variants ===")
print("""
PPO-Clip (original):
- Clipped surrogate objective
- Standard implementation
PPO-Penalty:
- KL penalty in objective
- Adaptive β coefficient
- Earlier version, less common
PPO with GAE:
- GAE for advantage estimation
- Standard in practice
PPO with Shared Network:
- Shared features for actor & critic
- More efficient
PPO with Separate Networks:
- Independent networks
- More flexible
PPO with Recurrent Policy:
- LSTM/GRU for memory
- For POMDPs
PPO for LLMs (RLHF):
- Token-level actions
- KL penalty to reference model
- Used in ChatGPT
""")
# PPO BEST PRACTICES
print("\n=== PPO Best Practices ===")
print("""
Hyperparameters (standard):
- Learning rate: 3e-4
- γ (gamma): 0.99
- λ (lambda): 0.95
- Clip ε: 0.2
- Value coefficient: 0.5
- Entropy coefficient: 0.01
- Epochs: 10
- Batch size: 64
- Rollout steps: 2048
Implementation tips:
1. Normalize observations
2. Normalize advantages
3. Clip gradients
4. Use GAE for advantages
5. Multiple epochs per rollout
6. Check KL divergence
7. Monitor entropy
8. Early stopping on KL
9. Use vectorized environments
10. Reward scaling
Debugging:
- If KL too high: reduce learning rate
- If policy entropy collapses: increase entropy coefficient
- If returns don't improve: check advantage normalization
- If unstable: reduce clip epsilon
Common pitfalls:
- Not normalizing advantages
- Wrong KL sign
- Incorrect GAE implementation
- Not using done masks
- Too few rollout steps
""")⚡ A3C / A2C
🎯 Parallel Training
A3C (Asynchronous Advantage Actor-Critic) and A2C (Advantage Actor-Critic) use multiple workers to collect experience in parallel. This reduces correlation between samples and speeds up training.
A3C/A2C Key Ideas
A3C: Multiple async workers, global network. A2C: Synchronous version — simpler, often better. Parallel envs: Diverse experience. No replay: Fresh samples. Actor-critic: Value baseline. Standard: Foundation for modern RL.
⚡ Parallel🔄 Async📊 Sync🎭 Actor-Critic
# A2C (synchronous version of A3C)
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Categorical
import gym
class A2CNetwork(nn.Module):
"""Actor-Critic network for A2C."""
def __init__(self, state_dim, action_dim, hidden=128):
super().__init__()
self.shared = nn.Sequential(
nn.Linear(state_dim, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU()
)
self.actor = nn.Linear(hidden, action_dim)
self.critic = nn.Linear(hidden, 1)
def forward(self, state):
features = self.shared(state)
logits = self.actor(features)
value = self.critic(features)
return logits, value
class A2CAgent:
"""A2C agent with vectorized environments."""
def __init__(
self,
state_dim,
action_dim,
n_envs=8,
lr=7e-4,
gamma=0.99,
value_coef=0.5,
entropy_coef=0.01,
):
self.network = A2CNetwork(state_dim, action_dim)
self.optimizer = optim.Adam(self.network.parameters(), lr=lr)
self.gamma = gamma
self.value_coef = value_coef
self.entropy_coef = entropy_coef
self.n_envs = n_envs
def collect_rollout(self, envs, n_steps=5):
"""Collect rollout from n_envs environments."""
states = []
actions = []
rewards = []
dones = []
values = []
log_probs = []
# Store per-env
env_states = []
for env in envs:
s = env.reset()
if isinstance(s, tuple):
s = s[0]
env_states.append(s)
for _ in range(n_steps):
states.append(env_states)
# Batch states
state_tensor = torch.FloatTensor(np.array(env_states))
with torch.no_grad():
logits, value = self.network(state_tensor)
dist = Categorical(logits=logits)
action = dist.sample()
log_prob = dist.log_prob(action)
actions.append(action.numpy())
log_probs.append(log_prob.numpy())
values.append(value.squeeze(-1).numpy())
# Step all envs
next_states = []
step_rewards = []
step_dones = []
for i, env in enumerate(envs):
s, r, d, trunc, _ = env.step(action[i].item())
if isinstance(d, bool) and trunc:
d = True
step_rewards.append(r)
step_dones.append(d)
if d:
s = env.reset()
if isinstance(s, tuple):
s = s[0]
next_states.append(s)
rewards.append(step_rewards)
dones.append(step_dones)
env_states = next_states
# Bootstrap final values
with torch.no_grad():
final_state = torch.FloatTensor(np.array(env_states))
_, final_value = self.network(final_state)
return {
'states': np.array(states), # (T, N, S)
'actions': np.array(actions), # (T, N)
'rewards': np.array(rewards), # (T, N)
'dones': np.array(dones), # (T, N)
'values': np.array(values), # (T, N)
'log_probs': np.array(log_probs), # (T, N)
'final_value': final_value.squeeze(-1).numpy(), # (N,)
}
def compute_returns_and_advantages(self, rollout):
"""Compute returns and advantages for all envs."""
T, N = rollout['rewards'].shape
returns = np.zeros((T, N))
advantages = np.zeros((T, N))
# Bootstrap from final value
next_value = rollout['final_value']
for t in reversed(range(T)):
# Mask for done
mask = 1.0 - rollout['dones'][t]
# TD target
td_target = rollout['rewards'][t] + self.gamma * next_value * mask
# TD error
td_error = td_target - rollout['values'][t]
# Advantage (TD(0))
advantages[t] = td_error
# Return
returns[t] = td_target
# Next value for previous step
next_value = rollout['values'][t]
return returns, advantages
def update(self, rollout):
"""Update using A2C loss."""
returns, advantages = self.compute_returns_and_advantages(rollout)
# Flatten
T, N = rollout['rewards'].shape
states = rollout['states'].reshape(T * N, -1)
actions = rollout['actions'].flatten()
old_log_probs = rollout['log_probs'].flatten()
returns = returns.flatten()
advantages = advantages.flatten()
# Normalize advantages
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
# Convert to tensors
states = torch.FloatTensor(states)
actions = torch.LongTensor(actions)
returns = torch.FloatTensor(returns)
advantages = torch.FloatTensor(advantages)
# Forward
logits, values = self.network(states)
dist = Categorical(logits=logits)
log_probs = dist.log_prob(actions)
entropy = dist.entropy()
# Losses
policy_loss = -(log_probs * advantages).mean()
value_loss = F.mse_loss(values.squeeze(), returns)
entropy_loss = -entropy.mean()
loss = (
policy_loss
+ self.value_coef * value_loss
+ self.entropy_coef * entropy_loss
)
# Update
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.network.parameters(), 0.5)
self.optimizer.step()
return loss.item()
# TRAINING
def train_a2c():
"""Train A2C on CartPole."""
try:
n_envs = 8
envs = [gym.make('CartPole-v1') for _ in range(n_envs)]
state_dim = envs[0].observation_space.shape[0]
action_dim = envs[0].action_space.n
agent = A2CAgent(state_dim, action_dim, n_envs=n_envs)
episode_rewards = [[] for _ in range(n_envs)]
current_rewards = [0] * n_envs
for iteration in range(2000):
rollout = agent.collect_rollout(envs, n_steps=5)
loss = agent.update(rollout)
# Track rewards
for t in range(len(rollout['rewards'])):
for i in range(n_envs):
current_rewards[i] += rollout['rewards'][t][i]
if rollout['dones'][t][i]:
episode_rewards[i].append(current_rewards[i])
current_rewards[i] = 0
if iteration % 50 == 0:
all_rewards = [r for env_rs in episode_rewards for r in env_rs[-20:]]
if all_rewards:
avg_reward = np.mean(all_rewards)
print(f"Iteration {iteration}, "
f"Avg Reward: {avg_reward:.2f}")
for env in envs:
env.close()
except Exception as e:
print(f"Training demo: {e}")
# A3C vs A2C
print("=== A3C vs A2C ===")
print("""
A3C (Asynchronous):
- Multiple workers in parallel
- Each worker has own env
- Asynchronous parameter updates
- No synchronization needed
- Historically important
- Rarely used now
A2C (Synchronous):
- Multiple envs in parallel
- Synchronized updates
- Simpler implementation
- Often better performance
- Standard in practice
- Good baseline
Key differences:
A3C: async updates, potential stale gradients
A2C: sync updates, consistent gradients
Both use:
- Actor-Critic architecture
- Advantage estimation
- Parallel environments
- Entropy regularization
""")
# VECTORIZED ENVIRONMENTS
print("\n=== Vectorized Environments ===")
print("""
Benefits:
1. Faster data collection
2. Less correlated samples
3. Better gradient estimates
4. More stable training
Implementation options:
- Stable-Baselines3: DummyVecEnv, SubprocVecEnv
- Gym: make_vec_env
- Ray: Distributed environments
Example:
from stable_baselines3.common.vec_env import DummyVecEnv
def make_env():
return gym.make('CartPole-v1')
env = DummyVecEnv([make_env for _ in range(8)])
""")
# N-STEP RETURNS
print("\n=== N-Step Returns ===")
print("""
A2C typically uses n-step returns:
G_t^{(n)} = r_t + γr_{t+1} + ... + γ^{n-1}r_{t+n-1} + γ^n V(s_{t+n})
Trade-off:
n=1: TD(0), low variance, high bias
n=5: Common choice for A2C
n=∞: Monte Carlo, high variance, low bias
In practice:
- Use n=5 for A2C
- Use GAE (λ=0.95) for PPO
- Both work well
""")🔒 TRPO (Trust Region)
🎯 Safe Policy Updates
TRPO (Schulman et al., 2015) constrains policy updates to a trust region, ensuring monotonic improvement. It's theoretically sound but complex — PPO is its practical successor.
Trust Region Concept
Constraint: KL(π_old || π_new) ≤ δ. Surrogate: L(θ) = E[ratio · A]. Solve: Constrained optimization. Theory: Monotonic improvement guarantee. Practice: Conjugate gradient + line search. Successor: PPO approximates TRPO with clipping.
📊 KL Constraint🎯 Trust Region📈 Monotonic🔧 Complex
# TRPO (Trust Region Policy Optimization)
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Categorical
# THE TRPO OBJECTIVE
# maximize_θ E[π_θ(a|s)/π_old(a|s) · A^π_old(s,a)]
# subject to E[KL(π_old(·|s) || π_θ(·|s))] ≤ δ
class TRPOPolicy(nn.Module):
"""Policy for TRPO."""
def __init__(self, state_dim, action_dim, hidden=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, hidden),
nn.Tanh(),
nn.Linear(hidden, hidden),
nn.Tanh(),
nn.Linear(hidden, action_dim)
)
def forward(self, state):
logits = self.net(state)
return Categorical(logits=logits)
def get_distribution(self, states):
return self.forward(states)
class TRPOValue(nn.Module):
"""Value function for TRPO."""
def __init__(self, state_dim, hidden=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(state_dim, hidden),
nn.Tanh(),
nn.Linear(hidden, hidden),
nn.Tanh(),
nn.Linear(hidden, 1)
)
def forward(self, states):
return self.net(states).squeeze(-1)
class TRPO:
"""TRPO algorithm."""
def __init__(
self,
state_dim,
action_dim,
gamma=0.99,
lam=0.95,
max_kl=0.01,
damping=0.1,
cg_iters=10,
value_lr=1e-3,
value_iters=5,
):
self.policy = TRPOPolicy(state_dim, action_dim)
self.value = TRPOValue(state_dim)
self.value_optimizer = torch.optim.Adam(
self.value.parameters(), lr=value_lr
)
self.gamma = gamma
self.lam = lam
self.max_kl = max_kl
self.damping = damping
self.cg_iters = cg_iters
self.value_iters = value_iters
def collect_rollout(self, env, n_steps=2048):
"""Collect rollout."""
states, actions, rewards, values, log_probs, dones = [], [], [], [], [], []
state = env.reset()
if isinstance(state, tuple):
state = state[0]
for _ in range(n_steps):
state_tensor = torch.FloatTensor(state)
with torch.no_grad():
dist = self.policy.get_distribution(state_tensor)
action = dist.sample()
log_prob = dist.log_prob(action)
value = self.value(state_tensor)
next_state, reward, done, truncated, _ = env.step(action.item())
if isinstance(done, bool) and truncated:
done = True
states.append(state)
actions.append(action.item())
rewards.append(reward)
values.append(value.item())
log_probs.append(log_prob.item())
dones.append(done)
state = next_state if not done else env.reset()[0]
return {
'states': np.array(states),
'actions': np.array(actions),
'rewards': np.array(rewards),
'values': np.array(values),
'log_probs': np.array(log_probs),
'dones': np.array(dones),
}
def compute_gae(self, rewards, values, dones, next_value=0):
"""Compute GAE advantages."""
advantages = []
gae = 0
for t in reversed(range(len(rewards))):
if t == len(rewards) - 1:
next_val = next_value
else:
next_val = values[t + 1]
mask = 1.0 - dones[t]
delta = rewards[t] + self.gamma * next_val * mask - values[t]
gae = delta + self.gamma * self.lam * mask * gae
advantages.insert(0, gae)
return np.array(advantages)
def conjugate_gradient(self, Avp, b, n_iters=10):
"""Conjugate gradient for solving Ax = b."""
x = torch.zeros_like(b)
r = b.clone()
p = r.clone()
r_dot_r = (r * r).sum()
for _ in range(n_iters):
Ap = Avp(p)
alpha = r_dot_r / (p * Ap).sum()
x = x + alpha * p
r_new = r - alpha * Ap
r_dot_r_new = (r_new * r_new).sum()
if r_dot_r_new < 1e-10:
break
beta = r_dot_r_new / r_dot_r
p = r_new + beta * p
r = r_new
r_dot_r = r_dot_r_new
return x
def update(self, rollout):
"""TRPO update using conjugate gradient."""
states = torch.FloatTensor(rollout['states'])
actions = torch.LongTensor(rollout['actions'])
old_log_probs = torch.FloatTensor(rollout['log_probs'])
# Compute advantages
advantages = self.compute_gae(
rollout['rewards'], rollout['values'], rollout['dones']
)
advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)
advantages = torch.FloatTensor(advantages)
# Compute returns
returns = advantages + torch.FloatTensor(rollout['values'])
# 1. Compute policy gradient
dist = self.policy.get_distribution(states)
log_probs = dist.log_prob(actions)
ratio = torch.exp(log_probs - old_log_probs)
surrogate = (ratio * advantages).mean()
# Policy gradient (using autograd)
grads = torch.autograd.grad(
surrogate, self.policy.parameters(), retain_graph=True
)
grads = torch.cat([g.flatten() for g in grads])
# 2. Compute Fisher vector product
def fisher_vector_product(v):
# Flatten v to match gradients
v = v.detach()
# Compute KL divergence
new_dist = self.policy.get_distribution(states)
kl = torch.distributions.kl_divergence(dist, new_dist).mean()
# Compute gradient of KL
kl_grads = torch.autograd.grad(
kl, self.policy.parameters(), create_graph=True
)
kl_grads = torch.cat([g.flatten() for g in kl_grads])
# FVP = ∇(∇KL · v)
kl_v = (kl_grads * v).sum()
fvp = torch.autograd.grad(kl_v, self.policy.parameters())
fvp = torch.cat([g.flatten() for g in fvp])
return fvp + self.damping * v
# 3. Conjugate gradient to solve Fx = g
step_dir = self.conjugate_gradient(
fisher_vector_product, grads, n_iters=self.cg_iters
)
# 4. Compute step size (line search)
# max_step = sqrt(2 * max_kl / (step_dir^T F step_dir))
fvp_step = fisher_vector_product(step_dir)
shs = 0.5 * (step_dir * fvp_step).sum()
lm = torch.sqrt(shs / self.max_kl)
max_step = 1.0 / (lm + 1e-8)
# 5. Line search
old_params = torch.cat([p.flatten() for p in self.policy.parameters()])
step_size = max_step.item()
improved = False
for _ in range(10):
new_params = old_params + step_size * step_dir
# Set new params
idx = 0
for p in self.policy.parameters():
size = p.numel()
p.data = new_params[idx:idx+size].view(p.shape)
idx += size
# Check KL constraint
with torch.no_grad():
new_dist = self.policy.get_distribution(states)
new_kl = torch.distributions.kl_divergence(dist, new_dist).mean()
# Compute new surrogate
new_log_probs = new_dist.log_prob(actions)
new_ratio = torch.exp(new_log_probs - old_log_probs)
new_surrogate = (new_ratio * advantages).mean()
if new_kl <= self.max_kl * 1.5 and new_surrogate > surrogate:
improved = True
break
step_size *= 0.5
if not improved:
# Restore old params
idx = 0
for p in self.policy.parameters():
size = p.numel()
p.data = old_params[idx:idx+size].view(p.shape)
idx += size
# 6. Update value function
for _ in range(self.value_iters):
values = self.value(states)
value_loss = F.mse_loss(values, returns)
self.value_optimizer.zero_grad()
value_loss.backward()
self.value_optimizer.step()
return new_kl.item()
# TRPO KEY CONCEPTS
print("=== TRPO Key Concepts ===")
print("""
1. TRUST REGION
- Limits policy change per update
- Prevents performance collapse
- KL divergence constraint
2. SURROGATE OBJECTIVE
L(θ) = E[π_θ(a|s)/π_old(a|s) · A(s,a)]
- Importance sampling
- Approximates policy improvement
3. CONSTRAINED OPTIMIZATION
maximize L(θ)
subject to E[KL(π_old || π_θ)] ≤ δ
4. SOLUTION METHOD
- Compute policy gradient
- Compute Fisher vector product
- Solve using conjugate gradient
- Line search for step size
- Check KL constraint
5. GUARANTEES
- Monotonic improvement (theory)
- Safe policy updates
- Sample efficient
6. LIMITATIONS
- Complex implementation
- Computationally expensive
- Second-order optimization
- PPO is a simpler alternative
""")
# TRPO vs PPO
print("=== TRPO vs PPO ===")
print("""
TRPO:
+ Theoretical guarantees
+ Monotonic improvement
+ Trust region constraint
- Complex (conjugate gradient)
- Second-order
- Slow per update
- Hard to implement
PPO:
+ Simple (first-order)
+ Fast per update
+ Easy to implement
+ Works well in practice
- No monotonic guarantee
- Clipping heuristic
- Sensitive to hyperparameters
Key insight:
PPO clips the ratio to approximate TRPO's
trust region constraint — simpler and effective.
""")
# WHY TRPO MATTERS
print("\n=== Why TRPO Matters ===")
print("""
Historical significance:
1. First practical trust region method
2. Theoretical foundation for PPO
3. Demonstrated safe RL updates
4. Enabled complex tasks
Conceptual contributions:
1. Trust region concept
2. Importance sampling for policy
3. Constrained optimization
4. Value function as baseline
Modern relevance:
- Mostly replaced by PPO
- Concepts used in safe RL
- Foundation for research
- Educational value
""")🎮 Continuous Action Spaces
🎯 Beyond Discrete Actions
Continuous action spaces are essential for robotics, autonomous driving, and control. Policy gradient methods naturally handle them via Gaussian policies, while value-based methods struggle.
Continuous Control
Gaussian policy: Output mean μ(s) and std σ(s). Reparameterization: a = μ + σ·ε, ε ~ N(0,1). Log prob: Computed from Gaussian PDF. State-dependent std: Learned or fixed. Squashing: tanh for bounded actions. Algorithms: DDPG, TD3, SAC, PPO.
📈 Gaussian🔧 Reparam🎯 Bounded🤖 Robotics
# Continuous action spaces with policy gradient
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.distributions import Normal
class ContinuousPolicy(nn.Module):
"""Gaussian policy for continuous actions."""
def __init__(self, state_dim, action_dim, hidden=64):
super().__init__()
# Shared features
self.shared = nn.Sequential(
nn.Linear(state_dim, hidden),
nn.Tanh(),
nn.Linear(hidden, hidden),
nn.Tanh()
)
# Mean head
self.mean_head = nn.Linear(hidden, action_dim)
# Log std head (learned)
self.log_std_head = nn.Linear(hidden, action_dim)
# Alternative: fixed log_std as parameter
# self.log_std = nn.Parameter(torch.zeros(action_dim))
def forward(self, state):
features = self.shared(state)
mean = self.mean_head(features)
log_std = self.log_std_head(features)
log_std = torch.clamp(log_std, -20, 2) # Clamp for stability
return mean, log_std
def get_distribution(self, state):
mean, log_std = self.forward(state)
std = torch.exp(log_std)
return Normal(mean, std)
def get_action(self, state, deterministic=False):
"""Sample action from policy."""
with torch.no_grad():
mean, log_std = self.forward(state)
std = torch.exp(log_std)
if deterministic:
return mean
# Sample with reparameterization
dist = Normal(mean, std)
action = dist.rsample()
# Compute log prob
log_prob = dist.log_prob(action).sum(-1)
return action, log_prob
def evaluate(self, states, actions):
"""Evaluate log prob and entropy."""
mean, log_std = self.forward(states)
std = torch.exp(log_std)
dist = Normal(mean, std)
log_probs = dist.log_prob(actions).sum(-1)
entropy = dist.entropy().sum(-1)
return log_probs, entropy
# BOUNDED ACTIONS (tanh squashing)
class SquashedContinuousPolicy(ContinuousPolicy):
"""Squashed Gaussian for bounded actions."""
def get_action(self, state, deterministic=False):
with torch.no_grad():
mean, log_std = self.forward(state)
std = torch.exp(log_std)
if deterministic:
return torch.tanh(mean)
dist = Normal(mean, std)
action = dist.rsample()
squashed = torch.tanh(action)
# Correct log prob for tanh squashing
log_prob = dist.log_prob(action).sum(-1)
log_prob -= torch.log(1 - squashed.pow(2) + 1e-6).sum(-1)
return squashed, log_prob
def evaluate(self, states, actions):
mean, log_std = self.forward(states)
std = torch.exp(log_std)
dist = Normal(mean, std)
# Inverse tanh
actions_clamped = torch.clamp(actions, -0.999999, 0.999999)
raw_actions = 0.5 * torch.log((1 + actions_clamped) / (1 - actions_clamped))
log_probs = dist.log_prob(raw_actions).sum(-1)
log_probs -= torch.log(1 - actions_clamped.pow(2) + 1e-6).sum(-1)
entropy = dist.entropy().sum(-1)
return log_probs, entropy
# TRAINING WITH CONTINUOUS ACTIONS
class ContinuousPPO:
"""PPO for continuous actions."""
def __init__(self, state_dim, action_dim, lr=3e-4):
self.policy = SquashedContinuousPolicy(state_dim, action_dim)
self.value = nn.Sequential(
nn.Linear(state_dim, 64), nn.Tanh(),
nn.Linear(64, 64), nn.Tanh(),
nn.Linear(64, 1)
)
self.optimizer = optim.Adam(
list(self.policy.parameters()) + list(self.value.parameters()),
lr=lr
)
def select_action(self, state):
state_tensor = torch.FloatTensor(state).unsqueeze(0)
action, log_prob = self.policy.get_action(state_tensor)
return action.squeeze(0).numpy(), log_prob.item()
# ACTION SPACE EXAMPLES
print("=== Continuous Action Space Examples ===")
print("""
Robotics:
- Joint angles: [θ1, θ2, ..., θn] ∈ ℝⁿ
- End-effector position: [x, y, z]
- Gripper force: scalar
Autonomous Driving:
- Steering angle: [-1, 1]
- Acceleration: [-1, 1]
- Braking: [0, 1]
Quadcopter Control:
- Thrust: [0, 1]
- Roll, pitch, yaw rates: [-π, π]
Portfolio Management:
- Asset allocations: [w1, w2, ..., wn] with Σwi = 1
Resource Allocation:
- CPU time, memory, bandwidth
- Real-valued allocations
Algorithms for continuous control:
- DDPG (Deterministic Policy Gradient)
- TD3 (Twin Delayed DDPG)
- SAC (Soft Actor-Critic)
- PPO for continuous
- A2C/A3C for continuous
""")
# LOG PROBABILITY FOR GAUSSIAN
print("=== Gaussian Log Probability ===")
print("""
For a Gaussian N(μ, σ²):
log p(a|s) = -½ [(a - μ)² / σ² + log(2πσ²)]
For multi-dimensional:
log p(a|s) = Σ_i log p(a_i|s)
Key insight: log prob is DIFFERENTIABLE w.r.t. μ and σ
Reparameterization trick:
a = μ + σ · ε, ε ~ N(0, 1)
This allows:
∇_θ a = ∇_θ μ + ε · ∇_θ σ
Gradient flows through!
""")
# WHY CONTINUOUS IS IMPORTANT
print("\n=== Why Continuous Actions Matter ===")
print("""
1. REAL-WORLD CONTROL
- Most physical systems are continuous
- Motors, hydraulics, etc.
2. FINE-GRAINED CONTROL
- Small adjustments
- Not limited to discrete choices
3. NATURAL REPRESENTATION
- Directly matches physical quantities
4. POWERFUL ALGORITHMS
- SAC, TD3 are SOTA
- Sample efficient
5. APPLICATION DOMAINS
- Robotics
- Autonomous vehicles
- Finance
- Resource management
- Drug discovery
""")
# COMPARISON
print("=== Discrete vs Continuous ===")
print("""
DISCRETE:
Actions: {a₁, a₂, ..., aₙ}
Policy: Categorical distribution
Methods: DQN, PPO, A2C
Examples: Chess, Atari, Go
CONTINUOUS:
Actions: a ∈ ℝⁿ
Policy: Gaussian distribution
Methods: DDPG, TD3, SAC, PPO
Examples: Robotics, Control
HYBRID:
Mixed discrete/continuous
Separate policy heads
Complex action structures
Key insight: Policy gradient methods
naturally extend to continuous actions,
while value-based methods struggle.
""")
# ADVANCED TOPICS
print("=== Advanced Topics ===")
print("""
1. State-dependent std
- Learn σ(s) as function of state
- More expressive than fixed σ
2. Beta distribution
- For bounded actions [0, 1]
- Alternative to tanh-Gaussian
3. Mixture of Gaussians
- Multi-modal action distributions
- More flexible
4. Normalizing flows
- Complex action distributions
- Invertible transformations
5. Hierarchical policies
- High-level: discrete options
- Low-level: continuous actions
6. Action repeat
- Repeat action multiple times
- Simpler temporal abstraction
""")💼 Applications (Games, Robotics, LLMs)
🎯 Policy Gradients in the Real World
Policy gradient methods power some of the most impressive AI systems: game-playing agents, robotic controllers, and the LLMs you use every day. Their versatility is unmatched.
Application Domains
Games: Atari, Dota 2, StarCraft II. Robotics: Manipulation, locomotion, navigation. LLMs: RLHF for ChatGPT, Claude. Recommendation: Personalization, ads. Finance: Trading, portfolio. Healthcare: Treatment, drug discovery.
🎮 Games🤖 Robotics💬 LLMs📊 Business
# Applications of policy gradient methods
import numpy as np
# 1. GAME PLAYING
print("=== Games ===")
print("""
OpenAI Five (Dota 2):
- PPO with LSTM
- Beat world champions
- Complex team coordination
AlphaStar (StarCraft II):
- Actor-Critic
- Grandmaster level
- Real-time strategy
AlphaGo/AlphaZero:
- MCTS + Neural Networks
- Superhuman at Go, Chess, Shogi
Atari Games:
- DQN initially
- PPO/IMPALA better
- 57 games benchmark
Applications:
- Game AI opponents
- Playtesting
- Game content generation
""")
# 2. ROBOTICS
print("=== Robotics ===")
print("""
Locomotion:
- Quadruped robots (ANYmal, Spot)
- Bipedal walking
- Running, jumping, parkour
- Sim-to-real transfer
Manipulation:
- Grasping objects
- Assembly tasks
- Tool use
- In-hand manipulation
Navigation:
- Autonomous driving
- Drone flight
- Mobile robot navigation
- SLAM + RL
Key successes:
- OpenAI Rubik's Cube (PPO)
- DeepMind robot soccer
- ANYmal parkour
- Dexterous manipulation
Challenges:
- Sim-to-real gap
- Sample efficiency
- Safety constraints
- Real-world variability
""")
# 3. LLMs and RLHF
print("=== LLMs and RLHF ===")
print("""
RLHF (Reinforcement Learning from Human Feedback):
1. Pretrain LLM on large corpus
2. Train reward model on human preferences
3. Fine-tune with PPO using reward model
4. Result: Aligned LLM
Used in:
- ChatGPT (OpenAI)
- Claude (Anthropic)
- Gemini (Google)
- Llama 2/3 (Meta)
PPO details for LLMs:
- Token-level actions
- KL penalty to reference model
- Reward shaping
- Value function for advantages
Reward model training:
- Human comparisons
- Bradley-Terry model
- Preference learning
Modern variants:
- DPO (Direct Preference Optimization)
- RLAIF (RL from AI Feedback)
- Constitutional AI
- GRPO (Group Relative Policy Optimization)
""")
# RLHF pseudocode
print("--- RLHF PPO Pseudocode ---")
print("""
# 1. Initialize
policy = LLM()
reference = copy(policy) # frozen
reward_model = train_reward_model(human_preferences)
# 2. RLHF training loop
for iteration in range(N):
# Sample prompts
prompts = sample_prompts()
# Generate responses
responses = policy.generate(prompts)
# Compute rewards
rewards = reward_model(prompts, responses)
# Compute KL penalty
kl = kl_divergence(policy, reference, prompts, responses)
rewards = rewards - beta * kl
# PPO update
advantages = compute_gae(rewards, values)
ppo_update(policy, advantages)
# 3. Result: Aligned LLM
""")
# 4. RECOMMENDATION SYSTEMS
print("=== Recommendation Systems ===")
print("""
Applications:
- Content recommendation (YouTube, Netflix)
- Product recommendation (Amazon)
- Ad placement
- News ranking
RL approach:
- State: User features + context
- Action: Items to recommend
- Reward: Clicks, watch time, purchases
- Policy: Which items to show
Challenges:
- Off-policy evaluation
- Counterfactual reasoning
- Exploration vs exploitation
- Long-term rewards
Methods:
- PPO for ranking
- Actor-Critic for personalization
- Slate optimization
""")
# 5. FINANCE
print("=== Finance ===")
print("""
Applications:
- Algorithmic trading
- Portfolio management
- Risk management
- Market making
RL setup:
- State: Market features, portfolio
- Action: Buy/sell/hold quantities
- Reward: Returns - risk penalty
- Constraints: Budget, position limits
Challenges:
- Non-stationarity
- Low signal-to-noise
- Transaction costs
- Regulatory constraints
Methods:
- PPO for trading
- SAC for portfolio
- Multi-agent for market making
""")
# 6. HEALTHCARE
print("=== Healthcare ===")
print("""
Applications:
- Treatment optimization
- Drug discovery
- Clinical trial design
- Personalized medicine
RL setup:
- State: Patient features, history
- Action: Treatment decisions
- Reward: Health outcomes
- Constraints: Safety, cost
Challenges:
- Safety-critical
- Small datasets
- Causal inference
- Regulatory approval
Methods:
- Off-policy evaluation
- Safe RL
- Model-based RL
Examples:
- Sepsis treatment
- Diabetes management
- Chemotherapy dosing
- Ventilator settings
""")
# 7. OTHER DOMAINS
print("=== Other Domains ===")
print("""
Energy:
- Smart grid control
- Data center cooling
- Battery management
- HVAC optimization
Manufacturing:
- Process control
- Quality optimization
- Predictive maintenance
- Supply chain
Education:
- Adaptive learning
- Personalized tutoring
- Curriculum design
- Student engagement
Marketing:
- Campaign optimization
- Customer lifetime value
- Dynamic pricing
- Retention
Transportation:
- Traffic signal control
- Fleet management
- Route optimization
- Autonomous vehicles
""")
# 8. LANDMARK PAPERS
print("=== Landmark Papers ===")
print("""
REINFORCE (1992):
Williams — Policy gradient theorem
Actor-Critic (1999):
Sutton et al. — Actor-Critic methods
TRPO (2015):
Schulman et al. — Trust region policy optimization
GAE (2016):
Schulman et al. — Generalized advantage estimation
A3C (2016):
Mnih et al. — Asynchronous methods
PPO (2017):
Schulman et al. — Proximal policy optimization
SAC (2018):
Haarnoja et al. — Soft actor-critic
RLHF (2017, 2022):
Christiano et al., Ouyang et al. — Human feedback
Decision Transformer (2021):
Chen et al. — Sequence modeling for RL
DPO (2023):
Rafailov et al. — Direct preference optimization
""")
# 9. IMPACT
print("=== Impact on AI ===")
print("""
Policy gradient methods have enabled:
1. Game-playing AI
- Beating world champions
- Complex multi-agent scenarios
2. Real-world robotics
- Robust locomotion
- Dexterous manipulation
3. Aligned LLMs
- ChatGPT, Claude, Gemini
- Instruction following
- Safety alignment
4. Business applications
- Recommendation
- Advertising
- Trading
5. Scientific discovery
- Drug design
- Materials science
- Physics simulation
Key insight:
Policy gradients are the bridge between
supervised learning and interactive AI.
""")🚀 Future: RLHF & Beyond
🎯 The Next Frontier
Policy gradient methods are evolving rapidly: RLHF for LLM alignment, offline RL for real-world applications, model-based RL for sample efficiency, and multi-agent RL for complex coordination.
Future Directions
RLHF: Aligning LLMs with human values. Offline RL: Learning from fixed datasets. Model-based: Learning environment models. Multi-agent: Coordination and competition. Safe RL: Constrained optimization. Real-world: Deploying at scale.
🧠 RLHF📊 Offline🌍 Model-based👥 Multi-agent
# Future of policy gradient methods
import numpy as np
# 1. RLHF AND BEYOND
print("=== RLHF Evolution ===")
print("""
RLHF (2022):
- PPO for LLM alignment
- Reward modeling from preferences
- KL penalty to reference
DPO (2023):
- Direct preference optimization
- No reward model needed
- Simpler, more stable
- Closed-form objective
RLAIF (2023):
- RL from AI feedback
- Scale beyond human labels
- Constitutional AI
GRPO (2024):
- Group Relative Policy Optimization
- Used in DeepSeek
- No value function needed
- Efficient for LLMs
Future trends:
- More efficient alignment
- Better reward modeling
- Multi-modal alignment
- Long-context RL
- Online learning
""")
# DPO pseudocode
print("--- DPO Pseudocode ---")
print("""
# DPO: Direct Preference Optimization
# No reward model, no RL loop
# Loss function:
# L_DPO = -E[log σ(β log(π(y_w|x)/π_ref(y_w|x))
# - β log(π(y_l|x)/π_ref(y_l|x)))]
# where:
# y_w: preferred response
# y_l: dispreferred response
# π: policy being trained
# π_ref: reference policy
# β: temperature parameter
# Training:
for batch in data:
y_w, y_l, x = batch
# Compute log probabilities
log_pi_yw = compute_log_prob(policy, x, y_w)
log_pi_yl = compute_log_prob(policy, x, y_l)
log_ref_yw = compute_log_prob(reference, x, y_w)
log_ref_yl = compute_log_prob(reference, x, y_l)
# Compute DPO loss
logits = beta * (log_pi_yw - log_ref_yw - log_pi_yl + log_ref_yl)
loss = -F.logsigmoid(logits).mean()
# Update
loss.backward()
optimizer.step()
""")
# 2. OFFLINE RL
print("\n=== Offline RL ===")
print("""
Problem:
- Learn from fixed dataset
- No environment interaction
- Critical for safety, cost, time
Challenges:
- Distribution shift
- Extrapolation error
- Uncertainty quantification
Methods:
- CQL (Conservative Q-Learning)
- IQL (Implicit Q-Learning)
- TD3+BC (Behavior Cloning)
- Decision Transformer
- AWAC
Applications:
- Healthcare (treatment)
- Robotics (from demos)
- Recommendation
- Autonomous driving
- Finance
Future:
- Combining offline + online
- Better uncertainty
- Foundation models for RL
""")
# 3. MODEL-BASED RL
print("=== Model-Based RL ===")
print("""
Learn environment model, plan with it:
Components:
- Dynamics model: p(s'|s,a)
- Reward model: r(s,a)
- Policy/planner
Approaches:
- MBPO (Model-Based Policy Optimization)
- Dreamer / DreamerV3
- MuZero (learned model)
- PETS (Probabilistic Ensembles)
Benefits:
- Sample efficiency (10-100x)
- Planning capability
- Model reuse
Challenges:
- Model errors compound
- Uncertainty estimation
- Sim-to-real gap
Future:
- Foundation models for dynamics
- Better uncertainty quantification
- Hybrid model-free/model-based
""")
# DreamerV3 pseudocode
print("--- DreamerV3 Pseudocode ---")
print("""
# Dreamer: Learning in latent imagination
# 1. World model
encoder: s_t = enc(o_t) # Obs to latent
dynamics: s_{t+1} = dyn(s_t, a_t) # Latent dynamics
reward: r_t = rew(s_t) # Reward prediction
decoder: o_t = dec(s_t) # Reconstruction
# 2. Imagination training
for each iteration:
# Sample latents
s = sample_state_from_replay()
# Imagine trajectories
for t in range(horizon):
a = actor(s)
s = dynamics(s, a)
r = reward(s)
# Update actor using imagined returns
update_actor(imagined_trajectories)
update_critic(imagined_trajectories)
# Update world model
update_world_model(replay_buffer)
""")
# 4. MULTI-AGENT RL
print("\n=== Multi-Agent RL ===")
print("""
Setting:
- Multiple agents in shared environment
- Cooperation, competition, or mixed
Challenges:
- Non-stationarity (others learning)
- Credit assignment
- Scalability
- Coordination
Methods:
- MADDPG (Multi-Agent DDPG)
- MAPPO (Multi-Agent PPO)
- QMIX (value factorization)
- VDN (Value Decomposition)
- COMA (Counterfactual)
Applications:
- Team games (Dota 2, StarCraft)
- Traffic control
- Swarm robotics
- Economic simulation
Future:
- Emergent communication
- General-sum games
- Open-ended learning
""")
# 5. SAFE RL
print("=== Safe RL ===")
print("""
Goal:
Maximize reward while satisfying safety constraints
Formulation:
maximize E[return]
subject to E[cost] ≤ limit
Approaches:
- Constrained MDPs
- Lagrangian methods
- CPO (Constrained Policy Optimization)
- PPO-Lagrangian
- Safe exploration
Applications:
- Healthcare
- Autonomous driving
- Industrial control
- Finance
Future:
- Provable safety
- Runtime verification
- Human oversight
- Regulatory compliance
""")
# 6. FOUNDATION MODELS FOR RL
print("=== Foundation Models for RL ===")
print("""
Combining LLMs with RL:
Decision Transformer:
- Sequence modeling approach
- No value function needed
- Trained on offline data
Gato (DeepMind):
- Multi-task, multi-modal
- Text, images, actions
- One model for everything
RT-2 (Google):
- Vision-Language-Action
- Robot control via LLM
Foundation policies:
- Pretrained on diverse tasks
- Fine-tune for specific tasks
- Generalization to new tasks
Future:
- Unified RL + LLM
- Language-conditioned policies
- World models from internet data
- Embodied AI
""")
# 7. KEY CHALLENGES
print("=== Key Challenges ===")
print("""
1. SAMPLE EFFICIENCY
- Real-world data is expensive
- Sim-to-real transfer
- Offline RL methods
2. SCALABILITY
- Large state/action spaces
- Long horizons
- Multi-agent systems
3. STABILITY
- Hyperparameter sensitivity
- Non-stationarity
- Catastrophic forgetting
4. SAFETY
- Constrained optimization
- Provable guarantees
- Human oversight
5. GENERALIZATION
- New tasks
- New environments
- Distribution shift
6. INTERPRETABILITY
- Understanding policies
- Explaining decisions
- Building trust
7. EVALUATION
- Off-policy evaluation
- Safety metrics
- Real-world impact
""")
# 8. FUTURE VISION
print("=== Future Vision ===")
print("""
Near-term (1-2 years):
- Better RLHF methods
- Multi-modal RL
- Efficient offline RL
- Foundation policies
Mid-term (3-5 years):
- AGI-level agents
- Sim-to-real mastery
- Multi-agent coordination
- Real-world deployment
Long-term (5-10 years):
- General-purpose agents
- Autonomous systems
- Scientific discovery
- Human-AI collaboration
Grand challenges:
- Open-ended learning
- Lifelong learning
- Human-level reasoning
- Safe superintelligence
""")
# 9. RECOMMENDATIONS
print("=== Recommendations ===")
print("""
For practitioners:
1. Start with PPO
2. Use GAE for advantages
3. Normalize everything
4. Vectorized environments
5. Monitor KL, entropy
6. Tune hyperparameters
7. Test on multiple tasks
For researchers:
1. Focus on sample efficiency
2. Safety and robustness
3. Offline and hybrid methods
4. Foundation models
5. Multi-agent systems
6. Real-world validation
For LLM alignment:
1. RLHF or DPO
2. Careful reward modeling
3. KL regularization
4. Red teaming
5. Continuous monitoring
For real-world deployment:
1. Safety first
2. Offline pretraining
3. Online fine-tuning
4. Human oversight
5. Monitoring and rollback
""")
# 10. FINAL THOUGHTS
print("\n=== Final Thoughts ===")
print("""
Policy gradient methods have evolved from:
Simple REINFORCE → Modern PPO/RLHF
They've enabled:
- Game-playing champions
- Robotic dexterity
- Aligned LLMs
- Business optimization
They will continue to evolve:
- More sample efficient
- More robust
- More general
- More aligned
Key insight:
Policy gradients bridge supervised learning
and interactive intelligence.
They're the foundation for:
- AGI
- Real-world AI
- Human-AI collaboration
"Reinforcement learning is the first
computational approach to learning
from interaction that is both
powerful and general."
— Richard Sutton
The future is interactive.
The future is RL.
The future is now.
""")