Module 8 — Policy gradient methods

DQN learns Q, then extracts a policy from it by argmax. Policy gradient methods invert that logic: parameterize the policy directly as π_θ(a | s) and improve θ by gradient ascent on expected return. Three things get easier in the bargain — stochastic policies, continuous actions, and the objective the algorithm optimizes.

Parameterizing the policy

For discrete actions, the network outputs a logit per action and a softmax gives probabilities:

import torch
import torch.nn as nn

class PolicyNet(nn.Module):
    def __init__(self, n_obs, n_actions, hidden=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(n_obs, hidden), nn.Tanh(),
            nn.Linear(hidden, hidden), nn.Tanh(),
            nn.Linear(hidden, n_actions),      # raw logits
        )
    def forward(self, x):
        return self.net(x)

    def sample(self, x):
        logits = self.forward(x)
        dist = torch.distributions.Categorical(logits=logits)
        action = dist.sample()
        return action, dist.log_prob(action)

The last layer produces raw logits, and sampling is done through Categorical. Tanh activations are common in policy gradients because they bound the intermediate features and make gradients smoother — a small detail that matters more than for DQN.

The policy gradient theorem

We want θ* = argmax_θ J(θ) where J(θ) = E_π_θ[Σ γ^t r_t]. The theorem states:

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

Two things to notice. First, the gradient is an expectation, so it can be estimated by sampling episodes. Second, G_t is the return from step t — the same one Monte Carlo used in module 4. That is why REINFORCE is sometimes called "Monte Carlo policy gradient": it uses full-episode returns.

The intuition matters more than the derivation. Increase the log-probability of actions that led to high returns; decrease it for those that led to low returns. Nothing about which action was "correct" is required — only the return afterward.

REINFORCE in code

import gymnasium as gym

env = gym.make("CartPole-v1")
policy = PolicyNet(env.observation_space.shape[0], env.action_space.n)
optimizer = torch.optim.Adam(policy.parameters(), lr=1e-3)
gamma = 0.99

for episode in range(1000):
    state, _ = env.reset()
    log_probs, rewards, done = [], [], False

    # 1. Roll out one full episode with the current policy.
    while not done:
        state_t = torch.tensor(state, dtype=torch.float32)
        action, log_prob = policy.sample(state_t)
        state, reward, term, trunc, _ = env.step(int(action))
        log_probs.append(log_prob)
        rewards.append(reward)
        done = term or trunc

    # 2. Compute returns from step t to end.
    returns, G = [], 0.0
    for r in reversed(rewards):
        G = r + gamma * G
        returns.insert(0, G)
    returns = torch.tensor(returns, dtype=torch.float32)

    # 3. Baseline: subtract the mean, divide by std for scale stability.
    returns = (returns - returns.mean()) / (returns.std() + 1e-8)

    # 4. Policy gradient loss: sum of -log_prob(a_t) * return_t.
    loss = -torch.stack([lp * G for lp, G in zip(log_probs, returns)]).sum()

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Two conventions in this loop that are worth naming. The loss is negated because PyTorch minimizes, but we want to maximize. And returns are computed backward in one pass — the standard efficient form of the definition.

The variance problem

REINFORCE works. It is also famously noisy. Two full runs on CartPole from different seeds can produce reward curves that look nothing alike, and getting to the 500-reward plateau reliably takes several thousand episodes. The reason: G_t is a sum of tens or hundreds of stochastic rewards, and a single episode gives you one sample of that sum. The gradient estimate inherits that variance.

Two mitigations are standard.

Baseline subtraction. Rewrite the gradient with a baseline b(s) that does not depend on the action:

$$\nabla_\theta J = \mathbb{E}\!\left[\sum_t \nabla_\theta \log \pi_\theta(a_t | s_t) \cdot (G_t - b(s_t))\right]$$

The expected value is unchanged because E[∇log π · b(s)] = 0, but the variance drops dramatically when b(s) tracks the average return. The simplest baseline is the mean return of the current batch; the more effective baseline is a learned value function V_φ(s), which is exactly what actor-critic methods do in module 9.

Return normalization. The (returns - mean) / std trick used in the code above is not a proper baseline, but it stabilizes gradient magnitudes across episodes at essentially zero cost. Every serious REINFORCE implementation uses it.

Continuous actions

For discrete actions the policy outputs logits; for continuous actions it outputs the parameters of a distribution, typically a Gaussian:

class GaussianPolicy(nn.Module):
    def __init__(self, n_obs, n_actions, hidden=128):
        super().__init__()
        self.trunk = nn.Sequential(nn.Linear(n_obs, hidden), nn.Tanh(),
                                    nn.Linear(hidden, hidden), nn.Tanh())
        self.mu = nn.Linear(hidden, n_actions)
        self.log_std = nn.Parameter(torch.zeros(n_actions))
    def sample(self, x):
        h = self.trunk(x)
        dist = torch.distributions.Normal(self.mu(h), self.log_std.exp())
        action = dist.sample()
        return action, dist.log_prob(action).sum(-1)

Note that log_std is a state-independent parameter, learned separately from the trunk. This is the convention that most policy-gradient references use, and departing from it — making log_std state-dependent — introduces subtle instabilities. Continuous control is where policy gradient truly beats DQN: for a robot arm with three joints, the discrete action space blows up combinatorially, while a Gaussian policy scales linearly with the number of joints.

REINFORCE is on-policy, so replay does not help

Unlike DQN, REINFORCE must use data collected under the current policy — the gradient theorem's expectation is with respect to π_θ. Reusing old episodes biases the estimate, and the algorithm silently drifts. Off-policy policy gradient exists (importance sampling), and it comes with its own instabilities that PPO in module 9 explicitly addresses.

Summary

• Policy gradient parameterizes the policy directly and improves it by gradient ascent on expected return, without ever fitting Q.
• REINFORCE uses full-episode returns as its gradient weights; unbiased but high variance, and slow to converge without tricks.
• Baseline subtraction and return normalization are the two low-cost variance reducers; a learned value baseline is actor-critic.
• Continuous actions are the natural home of policy gradient: parameterize a Gaussian and sample from it, one line at inference.

Next module: combine policy and value into actor-critic, then A2C and PPO — and benchmark a homemade PPO against Stable-Baselines3 on LunarLander.