Policy Gradient Methods

Why Does This Exist?

Q-learning selects actions via $\arg\max_a Q(s, a)$. When the action space is continuous — robot joint angles, throttle levels, portfolio weights, drug dosages — this argmax is intractable. Policy gradient methods sidestep this entirely by directly parameterizing the policy $\pi_\theta(a \mid s)$ and optimizing it via gradient ascent.

A second motivation: stochastic policies. Value-based methods produce deterministic greedy policies. But some tasks require mixed strategies — a deterministic policy in rock-paper-scissors is exploitable. Policy gradient naturally represents and learns stochastic policies.

Think of It Like This

How It Actually Works

The policy gradient theorem

We want to maximize expected return $J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[G(\tau)]$. The policy gradient theorem gives:

$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\left[\sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t \mid s_t) \cdot G_t\right]$

Where $G_t = \sum_{k=t}^{T} \gamma^{k-t} r_k$ is the return from timestep $t$.

The key term $\nabla_\theta \log \pi_\theta(a_t \mid s_t)$ is the score function: it points in the direction that increases the probability of action $a_t$ in state $s_t$. Multiplying by $G_t$ means: if $G_t > 0$ (good outcome), push toward $a_t$; if $G_t < 0$ (bad outcome), push away from it.

REINFORCE (Monte Carlo policy gradient)

The simplest policy gradient algorithm:

1. Sample trajectory $\tau = (s_0, a_0, r_0, \ldots, s_T)$ under current $\pi_\theta$
2. Compute returns $G_t$ for each timestep $t$
3. Update: $\theta \leftarrow \theta + \alpha \sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t) \cdot G_t$

Variance reduction with a baseline

The REINFORCE gradient estimator has high variance — returns vary dramatically across episodes. Subtracting a baseline $b(s_t)$ that doesn't depend on $a_t$ leaves the gradient unbiased while reducing variance:

$\nabla_\theta J \approx \mathbb{E}\left[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t) \cdot \underbrace{(G_t - b(s_t))}_{\text{advantage } A_t}\right]$

The baseline is typically $V(s_t)$, estimated by a separate value network. The advantage $A_t = G_t - V(s_t)$ measures how much better action $a_t$ was compared to what was expected on average in state $s_t$.

Policy parameterization

• Discrete actions: Softmax over learned logits — $\pi_\theta(a \mid s) = \text{softmax}(\text{MLP}(s))_a$
• Continuous actions: Gaussian policy — $a \sim \mathcal{N}(\mu_\theta(s), \sigma_\theta(s))$, where the network outputs mean and log-std

Code

1import gymnasium as gym2import torch3import torch.nn as nn4import torch.optim as optim5
6# ── Policy network (discrete actions) ────────────────────────────────────────7class PolicyNet(nn.Module):8    def __init__(self, obs_dim, n_actions, hidden=128):9        super().__init__()10        self.net = nn.Sequential(11            nn.Linear(obs_dim, hidden), nn.ReLU(),12            nn.Linear(hidden, hidden), nn.ReLU(),13            nn.Linear(hidden, n_actions),14        )15    16    def forward(self, x):17        return torch.softmax(self.net(x), dim=-1)18
19# ── REINFORCE with baseline ───────────────────────────────────────────────────20env = gym.make("CartPole-v1")21obs_dim   = env.observation_space.shape[0]22n_actions = env.action_space.n23
24policy    = PolicyNet(obs_dim, n_actions)25baseline  = nn.Sequential(nn.Linear(obs_dim, 128), nn.ReLU(), nn.Linear(128, 1))26opt_p = optim.Adam(policy.parameters(), lr=3e-4)27opt_b = optim.Adam(baseline.parameters(), lr=1e-3)28
29gamma = 0.9930
31for episode in range(800):32    states, actions, rewards = [], [], []33    state, _ = env.reset()34    done = False35
36    while not done:37        s_tensor = torch.tensor(state, dtype=torch.float32)38        probs  = policy(s_tensor)39        action = torch.multinomial(probs, 1).item()40        next_state, reward, terminated, truncated, _ = env.step(action)41        states.append(state); actions.append(action); rewards.append(reward)42        state = next_state43        done  = terminated or truncated44
45    # Compute discounted returns46    G, returns = 0.0, []47    for r in reversed(rewards):48        G = r + gamma * G49        returns.insert(0, G)50    returns = torch.tensor(returns, dtype=torch.float32)51
52    states_t  = torch.tensor(states, dtype=torch.float32)53    actions_t = torch.tensor(actions, dtype=torch.long)54
55    # Baseline update (value regression)56    values = baseline(states_t).squeeze()57    baseline_loss = nn.MSELoss()(values, returns)58    opt_b.zero_grad(); baseline_loss.backward(); opt_b.step()59
60    # Policy update (REINFORCE with advantage)61    with torch.no_grad():62        advantages = returns - baseline(states_t).squeeze()63
64    # Normalize advantages for stability65    advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)66
67    log_probs = torch.log(policy(states_t).gather(1, actions_t.unsqueeze(1)).squeeze() + 1e-8)68    policy_loss = -(log_probs * advantages).mean()69    opt_p.zero_grad(); policy_loss.backward(); opt_p.step()70
71    if episode % 100 == 0:72        ep_return = sum(rewards)73        print(f"Episode {episode:>4}: return={ep_return:.0f}, steps={len(rewards)}")74
75env.close()

Watch Out For

The Quick Version

• Policy gradients directly optimize $\pi_\theta(a \mid s)$ via $\nabla_\theta J \propto \nabla_\theta \log \pi_\theta \cdot G_t$.
• REINFORCE: collect Monte Carlo trajectories, compute returns, update policy — simple but high variance.
• Baselines reduce variance: advantage $A_t = G_t - V(s_t)$ keeps the gradient unbiased while reducing noise.
• Works with continuous actions and naturally represents stochastic policies (unlike Q-learning).
• PPO and TRPO extend policy gradient with update-size constraints for stability.