Lesson 6: Policy Gradient — REINFORCE & Actor-Critic

Reinforcement Learning: From Basics to Advanced

Part 2: Deep Reinforcement Learning — Neural Networks meet RL

xdev.asia

Introduction

Policy Gradient methods directly optimize policy π(a|s;θ) — instead of learning value function and then deriving policy. Advantages: handle continuous actions, stochastic policies, and convergence guarantees.

---

1. Policy Gradient Theorem

$$\nabla_\theta J(\theta) = \mathbb{E}{\pi\theta}[\nabla_\theta \log \pi_\theta(a|s) \cdot G_t]$$

Intuition: Increasing the probability of actions leads to high return, decreasing the probability of actions leads to low return.

---

2. REINFORCE Algorithm

import torch
import torch.nn as nn
from torch.distributions import Categorical

class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, 128),
            nn.ReLU(),
            nn.Linear(128, action_dim),
            nn.Softmax(dim=-1),
        )
    
    def forward(self, x):
        return self.net(x)

def reinforce(env, policy_net, optimizer, num_episodes, gamma=0.99):
    for episode in range(num_episodes):
        log_probs = []
        rewards = []
        state, _ = env.reset()
        done = False
        
        while not done:
            probs = policy_net(torch.FloatTensor(state))
            dist = Categorical(probs)
            action = dist.sample()
            log_probs.append(dist.log_prob(action))
            
            state, reward, terminated, truncated, _ = env.step(action.item())
            rewards.append(reward)
            done = terminated or truncated
        
        # Compute discounted returns
        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)
        
        # Policy gradient loss
        loss = sum(-lp * G for lp, G in zip(log_probs, returns))
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

---

3. Baseline & Variance Reduction

Use baseline b(s) to reduce variance without changing the expected gradient:

$$\nabla_\theta J(\theta) = \mathbb{E}[\nabla_\theta \log \pi_\theta(a|s) \cdot (G_t - b(s))]$$

---

4. Actor-Critic

class ActorCritic(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.shared = nn.Sequential(nn.Linear(state_dim, 128), nn.ReLU())
        self.actor = nn.Sequential(nn.Linear(128, action_dim), nn.Softmax(dim=-1))
        self.critic = nn.Linear(128, 1)
    
    def forward(self, x):
        feat = self.shared(x)
        return self.actor(feat), self.critic(feat)

def actor_critic_update(model, optimizer, state, action, reward, next_state, done, gamma):
    probs, value = model(torch.FloatTensor(state))
    _, next_value = model(torch.FloatTensor(next_state))
    
    # Advantage = TD error
    advantage = reward + gamma * next_value * (1 - done) - value
    
    # Actor loss
    dist = Categorical(probs)
    actor_loss = -dist.log_prob(torch.tensor(action)) * advantage.detach()
    
    # Critic loss
    critic_loss = advantage.pow(2)
    
    loss = actor_loss + 0.5 * critic_loss
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

---

5. A2C — Advantage Actor-Critic

A2C = synchronous, multiple workers → reduce variance + faster training.

---

Summary

Method	Type	Variance	Bias	Online
REINFORCE	Policy	High	None	❌
REINFORCE + baseline	Policy	Medium	None	❌
Actor-Critic	Both	Low	Some	✅
A2C	Both	Lower	Some	✅