Chapter 4: Policy Gradient Methods

4.1 Introduction to Policy Gradient

Unlike value-based methods that learn Q(s,a) and derive a policy, policy gradient methods directly parameterize and optimize the policy π_θ(a|s). This approach is particularly powerful for continuous action spaces and stochastic policies.

Why Policy Gradient?

• Continuous Actions: Natural for continuous action spaces
• Stochastic Policies: Can learn probabilistic strategies
• Theoretical Guarantees: Converges to local optima
• Better Convergence: Small changes in θ → small changes in policy
• Works with Function Approximation: Compatible with deep learning

4.2 Policy Gradient Theorem

The policy gradient theorem provides the foundation for all policy gradient algorithms:

∇_θ J(θ) = E_π[∇_θ log π_θ(a|s) Q^π(s,a)]

In words: The gradient of the objective equals the expected gradient of log probability weighted by Q-value.

Intuition: Increase probability of good actions (high Q), decrease probability of bad actions (low Q)

Derivation Intuition

J(θ) = E_π[R]  // Expected return

∇_θ J(θ) = ∇_θ E[Σ r_t]
         = E[Σ ∇_θ log π_θ(a_t|s_t) Σ_{t'≥t} r_t']
         = E[Σ ∇_θ log π_θ(a_t|s_t) Q^π(s_t,a_t)]

This is the REINFORCE gradient estimator.

4.3 REINFORCE Algorithm

REINFORCE is the simplest policy gradient algorithm, using Monte Carlo returns to estimate Q^π(s,a):

Algorithm: REINFORCE (Monte Carlo Policy Gradient)

Initialize policy parameters θ randomly

For each episode:
    Generate episode following π_θ: s_0, a_0, r_1, s_1, a_1, ..., s_T

    For t = 0 to T-1:
        G_t = Σ_{k=0}^{T-t-1} γ^k r_{t+k+1}  // Return from time t

        θ ← θ + α γ^t G_t ∇_θ log π_θ(a_t|s_t)  // Gradient ascent

REINFORCE Implementation

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=128):
        super(PolicyNetwork, self).__init__()
        self.network = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, action_dim),
            nn.Softmax(dim=-1)  # Output probabilities
        )

    def forward(self, state):
        return self.network(state)

def reinforce(env, episodes=1000, gamma=0.99, lr=1e-3):
    """
    REINFORCE algorithm for episodic tasks.
    """
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    policy = PolicyNetwork(state_dim, action_dim)
    optimizer = optim.Adam(policy.parameters(), lr=lr)

    for episode in range(episodes):
        # Generate episode
        states, actions, rewards = [], [], []
        state = env.reset()
        done = False

        while not done:
            state_tensor = torch.FloatTensor(state).unsqueeze(0)

            # Sample action from policy
            with torch.no_grad():
                action_probs = policy(state_tensor)
                action_dist = torch.distributions.Categorical(action_probs)
                action = action_dist.sample()

            next_state, reward, done, _ = env.step(action.item())

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

            state = next_state

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

        returns = torch.FloatTensor(returns)

        # Normalize returns (reduces variance)
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)

        # Compute policy gradient
        loss = 0
        for state, action, G in zip(states, actions, returns):
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            action_probs = policy(state_tensor)
            action_dist = torch.distributions.Categorical(action_probs)

            # -log π(a|s) * G  (negative because we minimize)
            loss += -action_dist.log_prob(action) * G

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

        if (episode + 1) % 100 == 0:
            print(f"Episode {episode+1}, Total Reward: {sum(rewards)}")

    return policy

4.4 Baseline Functions

Adding a baseline b(s) reduces variance without introducing bias:

∇_θ J(θ) = E[∇_θ log π_θ(a|s) (Q^π(s,a) - b(s))]

Common baselines:
1. b(s) = V^π(s)  // State value function
2. b(s) = average return
3. b(s) = moving average of returns

REINFORCE with Baseline

def reinforce_with_baseline(env, episodes=1000, gamma=0.99, lr=1e-3):
    """
    REINFORCE with learned value function baseline.
    """
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    policy = PolicyNetwork(state_dim, action_dim)
    value_net = ValueNetwork(state_dim)  # Baseline

    policy_optimizer = optim.Adam(policy.parameters(), lr=lr)
    value_optimizer = optim.Adam(value_net.parameters(), lr=lr)

    for episode in range(episodes):
        states, actions, rewards = generate_episode(env, policy)

        returns = compute_returns(rewards, gamma)
        returns = torch.FloatTensor(returns)

        # Update value network (baseline)
        states_tensor = torch.FloatTensor(states)
        values = value_net(states_tensor).squeeze()
        value_loss = nn.MSELoss()(values, returns)

        value_optimizer.zero_grad()
        value_loss.backward()
        value_optimizer.step()

        # Update policy with baseline
        advantages = returns - values.detach()  # A(s,a) = G - V(s)

        policy_loss = 0
        for state, action, advantage in zip(states, actions, advantages):
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            action_probs = policy(state_tensor)
            action_dist = torch.distributions.Categorical(action_probs)
            policy_loss += -action_dist.log_prob(action) * advantage

        policy_optimizer.zero_grad()
        policy_loss.backward()
        policy_optimizer.step()

    return policy

4.5 Trust Region Policy Optimization (TRPO)

TRPO ensures stable learning by constraining policy updates to a "trust region":

maximize E[π_θ_new / π_θ_old * A(s,a)]
subject to E[KL(π_θ_old || π_θ_new)] ≤ δ

where:
- KL divergence constrains how much the policy can change
- δ is the trust region size
- This prevents catastrophic policy updates

Natural Policy Gradient

TRPO uses the natural gradient, which accounts for the geometry of the policy space:

Natural gradient: ∇̃_θ J = F^{-1} ∇_θ J

where F is the Fisher Information Matrix:
F = E[∇_θ log π(a|s) ∇_θ log π(a|s)^T]

The natural gradient points in the direction of steepest ascent in probability space, not parameter space.

4.6 Proximal Policy Optimization (PPO)

PPO simplifies TRPO while maintaining its benefits through a clipped objective:

L^CLIP(θ) = E[min(r_t(θ) Â_t, clip(r_t(θ), 1-ε, 1+ε) Â_t)]

where:
r_t(θ) = π_θ(a_t|s_t) / π_θ_old(a_t|s_t)  // Probability ratio
Â_t = estimated advantage
ε = clipping parameter (typically 0.1 or 0.2)

PPO Algorithm

class PPO:
    def __init__(self, state_dim, action_dim, lr=3e-4, gamma=0.99,
                 epsilon=0.2, epochs=10, batch_size=64):
        self.policy = PolicyNetwork(state_dim, action_dim)
        self.old_policy = PolicyNetwork(state_dim, action_dim)
        self.old_policy.load_state_dict(self.policy.state_dict())

        self.value_net = ValueNetwork(state_dim)

        self.optimizer = optim.Adam([
            {'params': self.policy.parameters()},
            {'params': self.value_net.parameters()}
        ], lr=lr)

        self.gamma = gamma
        self.epsilon = epsilon
        self.epochs = epochs
        self.batch_size = batch_size

    def select_action(self, state):
        state = torch.FloatTensor(state).unsqueeze(0)
        with torch.no_grad():
            action_probs = self.policy(state)
            dist = torch.distributions.Categorical(action_probs)
            action = dist.sample()
        return action.item()

    def update(self, states, actions, rewards, next_states, dones):
        """
        Update policy using PPO.
        """
        states = torch.FloatTensor(states)
        actions = torch.LongTensor(actions)
        rewards = torch.FloatTensor(rewards)
        next_states = torch.FloatTensor(next_states)
        dones = torch.FloatTensor(dones)

        # Compute returns and advantages
        with torch.no_grad():
            values = self.value_net(states).squeeze()
            next_values = self.value_net(next_states).squeeze()

            # GAE (Generalized Advantage Estimation)
            advantages = rewards + self.gamma * next_values * (1 - dones) - values
            returns = advantages + values

        # Normalize advantages
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        # Get old action probabilities
        with torch.no_grad():
            old_probs = self.old_policy(states)
            old_dist = torch.distributions.Categorical(old_probs)
            old_log_probs = old_dist.log_prob(actions)

        # PPO update for multiple epochs
        for _ in range(self.epochs):
            # Current policy probabilities
            new_probs = self.policy(states)
            new_dist = torch.distributions.Categorical(new_probs)
            new_log_probs = new_dist.log_prob(actions)

            # Ratio of new and old probabilities
            ratio = torch.exp(new_log_probs - old_log_probs)

            # Clipped surrogate objective
            surr1 = ratio * advantages
            surr2 = torch.clamp(ratio, 1 - self.epsilon, 1 + self.epsilon) * advantages
            policy_loss = -torch.min(surr1, surr2).mean()

            # Value loss
            new_values = self.value_net(states).squeeze()
            value_loss = nn.MSELoss()(new_values, returns)

            # Entropy bonus (encourages exploration)
            entropy = new_dist.entropy().mean()

            # Total loss
            loss = policy_loss + 0.5 * value_loss - 0.01 * entropy

            # Update networks
            self.optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(self.policy.parameters(), 0.5)
            self.optimizer.step()

        # Update old policy
        self.old_policy.load_state_dict(self.policy.state_dict())

4.7 Continuous Action Spaces

Policy gradient methods excel in continuous control. We parameterize the policy as a Gaussian:

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

where:
- μ_θ(s): mean action (neural network output)
- σ_θ(s): standard deviation (can be learned or fixed)

Continuous Action Policy Network

class ContinuousPolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=256):
        super(ContinuousPolicyNetwork, self).__init__()

        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )

        # Mean of action distribution
        self.mean = nn.Linear(hidden_dim, action_dim)

        # Log std of action distribution
        self.log_std = nn.Parameter(torch.zeros(action_dim))

    def forward(self, state):
        features = self.shared(state)
        mean = self.mean(features)
        std = torch.exp(self.log_std)
        return mean, std

    def sample_action(self, state):
        mean, std = self.forward(state)
        dist = torch.distributions.Normal(mean, std)
        action = dist.sample()
        log_prob = dist.log_prob(action).sum(dim=-1)
        return action, log_prob

# Usage for continuous control (e.g., MuJoCo)
policy = ContinuousPolicyNetwork(state_dim=17, action_dim=6)
state = torch.FloatTensor(state)
action, log_prob = policy.sample_action(state)

4.8 Generalized Advantage Estimation (GAE)

GAE balances bias and variance in advantage estimation:

A^GAE(λ)_t = Σ_{l=0}^∞ (γλ)^l δ_{t+l}

where:
δ_t = r_t + γV(s_{t+1}) - V(s_t)  // TD error

λ ∈ [0, 1] controls the bias-variance tradeoff:
- λ = 0: A = δ_t (high bias, low variance, like TD)
- λ = 1: A = Σ r_t - V(s_t) (low bias, high variance, like MC)

GAE Implementation

def compute_gae(rewards, values, next_values, dones, gamma=0.99, lambda_=0.95):
    """
    Compute Generalized Advantage Estimation.

    Args:
        rewards: list of rewards
        values: list of value estimates V(s_t)
        next_values: list of value estimates V(s_{t+1})
        dones: list of done flags
        gamma: discount factor
        lambda_: GAE parameter
    """
    advantages = []
    gae = 0

    for t in reversed(range(len(rewards))):
        # TD error
        delta = rewards[t] + gamma * next_values[t] * (1 - dones[t]) - values[t]

        # GAE
        gae = delta + gamma * lambda_ * (1 - dones[t]) * gae
        advantages.insert(0, gae)

    return advantages

4.9 Practical Considerations

Hyperparameters

Tips for Success

• Normalize observations: Use running mean/std
• Normalize returns: Reduces variance
• Clip gradients: Prevents explosion
• Use orthogonal initialization: Better for deep networks
• Tune entropy bonus: Balances exploration
• Monitor KL divergence: Detect unstable updates

4.10 Comparison: Policy Gradient vs Value-Based