Policy Gradient Methods

ℹ️ Definition Policy Gradient Methods are a class of reinforcement learning algorithms that directly optimize the policy function using gradient ascent on expected rewards, enabling learning of stochastic policies and handling continuous action spaces.

Learning Objectives

By the end of this lesson, you will:

• Understand the difference between value-based and policy-based RL
• Learn the policy gradient theorem and its derivation
• Implement the REINFORCE algorithm from scratch
• Apply variance reduction techniques (baselines)
• Handle continuous action spaces with Gaussian policies
• Debug common policy gradient training issues

Introduction

In Lessons 2-3, we learned value-based methods (Q-Learning, DQN) that learn Q-values and derive policies implicitly:

Learn Q(s,a) → Policy: π(s) = argmax_a Q(s,a)

Policy gradient methods take a different approach: directly learn the policy.

Learn π(a|s; θ) directly using gradient ascent

This seemingly simple change unlocks powerful capabilities.

Value-Based vs Policy-Based RL

Value-Based Methods (Q-Learning, DQN)

Approach:

1. Learn value function Q(s,a)
2. Derive policy: choose action with max Q-value
3. Policy is deterministic and implicit

Strengths:

• Sample efficient (reuse experiences)
• Well-understood convergence properties
• Works well for discrete actions

Weaknesses:

• Discrete actions only (argmax requires enumeration)
• Deterministic policies (always pick best action)
• Indirect optimization (learn Q, hope policy is good)

Policy-Based Methods (Policy Gradients)

Approach:

1. Parameterize policy π(a|s; θ)
2. Optimize θ directly to maximize expected reward
3. Policy is explicit and can be stochastic

Strengths:

• Handles continuous actions naturally
• Stochastic policies (important for exploration, game theory)
• Direct optimization of what we care about (policy)
• Guaranteed convergence to local optimum

Weaknesses:

• High variance (noisy gradients)
• Sample inefficient
• Can converge to local optima (not global)

When to Use Policy Gradients

Use policy gradient methods when:

One. Continuous Action Spaces

Example - Robot Control:

# Value-based: Need to discretize (bad approximation)
actions = [0.0, 0.1, 0.2, ..., 1.0]  # 10 discrete actions
Q(state, action) for each discrete action

# Policy-based: Output continuous action directly
action = π(state; θ)  # Can be any value in [0, 1]

2. Stochastic Policies Needed

Example - Rock-Paper-Scissors:

• Deterministic policy: Opponent learns and exploits
• Stochastic policy: π(Rock) = π(Paper) = π(Scissors) = 1/3 (unexploitable)

3. High-Dimensional Action Spaces

Example - Multi-Joint Robot:

• DQN: Discretize each joint -> 10^6 action combinations (intractable)
• Policy Gradient: Output vector of joint angles (scalable)

Policy Parameterization

Discrete Actions: Softmax Policy

For discrete action space with n actions:

π(a|s; θ) = softmax(f(s; θ))_a = exp(f_a(s; θ)) / Σ_a' exp(f_a'(s; θ))

Where:

• f(s; θ): Neural network that outputs logits for each action
• π(a|s; θ): Probability of taking action a in state s

Example:

logits = network(state)  # [2.0, 1.0, 0.5]
probs = softmax(logits)  # [0.59, 0.24, 0.16]
action = sample(probs)   # Sample from distribution

Continuous Actions: Gaussian Policy

For continuous action space:

π(a|s; θ) = N(μ(s; θ), σ²)

Where:

• μ(s; θ): Neural network outputs mean
• σ: Standard deviation (fixed or learned)
• N(μ, σ²): Gaussian distribution

Example:

mu = network(state)           # μ = 0.7
sigma = 0.5                   # Fixed σ
action = mu + sigma * randn() # Sample: a ~ N(0.7, 0.5²)

The Policy Gradient Theorem

Objective

Maximize expected return:

J(θ) = E_τ~π_θ [R(τ)]

Where:

• τ: Trajectory (s₀, a₀, r₀, s₁, a₁, r₁, ..., sₜ, aₜ, rₜ)
• R(τ): Total return of trajectory
• π_θ: Policy parameterized by θ

The Theorem

Policy Gradient Theorem:

∇_θ J(θ) = E_τ~π_θ [Σ_t ∇_θ log π_θ(a_t|s_t) * R(τ)]

Intuition:

• ∇_θ log π_θ(a_t|s_t): Direction to increase probability of action a_t
• R(τ): Weight by how good the trajectory was
• Effect: Increase probability of actions in good trajectories

Simplified Form (Episodic)

For episodic tasks:

∇_θ J(θ) = E_τ [Σ_t ∇_θ log π_θ(a_t|s_t) * G_t]

Where:

• G_t: Return from timestep t (G_t = r_t + γr_t+1 + γ²r_t+2 + ...)

Interpretation: Push up probability of actions that led to high returns.

REINFORCE Algorithm

REINFORCE (Monte Carlo Policy Gradient) is the simplest policy gradient algorithm.

Algorithm

Initialize policy network π(a|s; θ) with random θ
Set learning rate α

for episode in range(num_episodes):
    # Generate trajectory
    τ = []
    s = env.reset()
    done = False

    while not done:
        # Sample action from policy
        a ~ π(·|s; θ)
        s', r, done = env.step(a)
        τ.append((s, a, r))
        s = s'

    # Compute returns
    G = 0
    returns = []
    for (s, a, r) in reversed(τ):
        G = r + γ * G
        returns.insert(0, G)

    # Policy gradient update
    for t, (s_t, a_t, r_t) in enumerate(τ):
        G_t = returns[t]

        # Compute gradient
        ∇J = ∇_θ log π_θ(a_t|s_t) * G_t

        # Gradient ascent (maximize reward)
        θ ← θ + α * ∇J

PyTorch Implementation

import torch
import torch.nn as nn
import torch.optim as optim
import torch.distributions as distributions

class PolicyNetwork(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 = torch.relu(self.fc1(state))
        x = torch.relu(self.fc2(x))
        logits = self.fc3(x)
        return logits

    def select_action(self, state):
        logits = self.forward(state)
        probs = torch.softmax(logits, dim=-1)
        dist = distributions.Categorical(probs)
        action = dist.sample()
        log_prob = dist.log_prob(action)
        return action.item(), log_prob

# Training loop
policy = PolicyNetwork(state_dim=4, action_dim=2)
optimizer = optim.Adam(policy.parameters(), lr=0.01)

for episode in range(1000):
    states, actions, rewards, log_probs = [], [], [], []

    # Collect trajectory
    state = env.reset()
    done = False
    while not done:
        state_tensor = torch.FloatTensor(state)
        action, log_prob = policy.select_action(state_tensor)
        next_state, reward, done, _ = env.step(action)

        states.append(state)
        actions.append(action)
        rewards.append(reward)
        log_probs.append(log_prob)

        state = next_state

    # Compute returns (discounted cumulative rewards)
    returns = []
    G = 0
    for r in reversed(rewards):
        G = r + gamma * G
        returns.insert(0, G)

    returns = torch.FloatTensor(returns)
    returns = (returns - returns.mean()) / (returns.std() + 1e-8)  # Normalize

    # Compute policy gradient loss
    log_probs = torch.stack(log_probs)
    loss = -(log_probs * returns).sum()  # Negative for gradient ascent

    # Update policy
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

The Log Probability Trick

Why Log Probabilities?

Mathematical convenience:

∇_θ π_θ(a|s) / π_θ(a|s) = ∇_θ log π_θ(a|s)

This transforms a difficult derivative into a simple gradient of log probability.

Computing Log Probabilities

Discrete actions (Categorical):

logits = policy_network(state)
log_probs = F.log_softmax(logits, dim=-1)
log_prob_action = log_probs[action]

Continuous actions (Gaussian):

mu = policy_network(state)
dist = Normal(mu, sigma)
log_prob = dist.log_prob(action)

Variance Reduction Techniques

Problem: High Variance

REINFORCE has very high variance in gradient estimates:

• Different trajectories give wildly different returns
• Gradients are noisy -> slow, unstable learning

Solution 1: Baseline Subtraction

Subtract a baseline b(s) from returns to reduce variance:

∇_θ J(θ) ≈ Σ_t ∇_θ log π_θ(a_t|s_t) * [G_t - b(s_t)]

Common baselines:

1. Constant baseline: b(s) = average return
2. State-value baseline: b(s) = V(s) (learned value function)

Why it works:

• Reduces variance without introducing bias
• Centers returns around zero
• Actions better than baseline -> positive gradient
• Actions worse than baseline -> negative gradient

Implementation:

returns = torch.FloatTensor(returns)
baseline = returns.mean()
advantages = returns - baseline
loss = -(log_probs * advantages).sum()

Solution 2: Normalization

Normalize returns to have mean 0 and std 1:

returns = (returns - returns.mean()) / (returns.std() + 1e-8)

Benefits:

• Stabilizes learning across different reward scales
• Reduces sensitivity to hyperparameters

Solution 3: Advantage Function

Use advantage A(s,a) = Q(s,a) - V(s):

∇_θ J(θ) = E [Σ_t ∇_θ log π_θ(a_t|s_t) * A(s_t, a_t)]

This leads to Actor-Critic methods (next lesson).

Continuous Action Spaces

Gaussian Policy

For continuous actions, use a Gaussian policy:

class GaussianPolicy(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.mean = nn.Linear(128, action_dim)
        self.log_std = nn.Parameter(torch.zeros(action_dim))  # Learnable std

    def forward(self, state):
        x = torch.relu(self.fc1(state))
        x = torch.relu(self.fc2(x))
        mean = self.mean(x)
        std = torch.exp(self.log_std)
        return mean, std

    def select_action(self, state):
        mean, std = self.forward(state)
        dist = torch.distributions.Normal(mean, std)
        action = dist.sample()
        log_prob = dist.log_prob(action).sum()  # Sum over action dimensions
        return action, log_prob

Action Clipping

Clip actions to valid range:

action = action.clamp(env.action_space.low, env.action_space.high)

Hyperparameters

Common Issues and Debugging

Issue 1: Policy Collapse

Symptom: Policy becomes deterministic too quickly, stops exploring.

Cause: Exploitation too early.

Fix: Add entropy bonus to encourage exploration:

entropy = -(probs * torch.log(probs)).sum()
loss = -(log_probs * returns).sum() - entropy_weight * entropy

Issue 2: Slow Learning

Symptom: Policy improves very slowly.

Causes:

• High variance (noisy gradients)
• Learning rate too low
• Poor reward scaling

Fixes:

• Use baseline and return normalization
• Increase learning rate
• Scale/normalize rewards

Issue 3: Instability

Symptom: Performance oscillates wildly.

Causes:

• Learning rate too high
• No gradient clipping
• Extreme returns

Fixes:

• Reduce learning rate
• Clip gradients: torch.nn.utils.clip_grad_norm_(policy.parameters(), max_norm=1.0)
• Clip returns to reasonable range

Advantages of Policy Gradients

1. Continuous actions: Natural handling of continuous control
2. Stochastic policies: Can learn probabilistic strategies
3. Convergence guarantees: Guaranteed to converge to local optimum
4. Simplicity: Conceptually simple (maximize reward directly)
5. Effective in high dimensions: Works well for robotics

Limitations of Policy Gradients

1. High variance: Noisy gradient estimates
2. Sample inefficient: Needs many trajectories
3. Slow convergence: Compared to value-based methods
4. Local optima: Only guaranteed local convergence
5. Sensitive to hyperparameters: Learning rate, baseline critical

Key Takeaways

1. Policy gradients directly optimize the policy to maximize expected reward
2. REINFORCE uses Monte Carlo sampling to estimate policy gradients
3. Log probability trick enables tractable gradient computation
4. Variance reduction is critical (baselines, normalization)
5. Continuous actions handled naturally with Gaussian policies
6. Trade-offs: Flexibility vs sample efficiency

Looking Ahead

Policy gradients opened new possibilities, but their high variance limits performance. In the next lesson:

• Lesson 5: Actor-Critic methods combine value and policy learning
• Lesson 6: PPO improves stability and sample efficiency
• Lesson 8: Applications to robotics and continuous control

Next lesson: Learn how to reduce variance by combining policy gradients with value functions!

Summary

• Policy gradient methods directly optimize policies using gradient ascent
• REINFORCE algorithm samples trajectories and updates policy to increase probability of high-reward actions
• Log probability trick simplifies gradient computation
• Variance reduction techniques (baselines, normalization) are essential for stable learning
• Continuous action spaces handled with Gaussian policies
• Trade-off: Handles continuous actions and stochastic policies, but has high variance and sample inefficiency