Policy Gradient Methods: REINFORCE and the Policy Gradient Theorem

Value-based methods learn $Q$ and act greedily. That works for discrete actions, but struggles with continuous actions (the max over a continuum is itself an optimisation problem) and cannot naturally represent stochastic policies, which are sometimes optimal (rock–paper–scissors, partially observable problems). Policy gradient methods parameterise the policy directly, $\pi_\theta(a \mid s)$, and adjust $\theta$ by gradient ascent on expected return. They are the foundation of actor–critic methods, PPO — and RLHF for language models.

The objective#

Maximise the expected return of trajectories $\tau = (s_0, a_0, r_1, s_1, \dots)$ generated by the policy:

J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}\big[G(\tau)\big]

The difficulty: $J$ depends on $\theta$ through the distribution of trajectories, and the environment's dynamics are unknown and not differentiable.

The log-derivative trick#

For any distribution $p_\theta(x)$:

\nabla_\theta\mathbb{E}_{p_\theta}[f(x)] = \mathbb{E}_{p_\theta}\big[f(x)\,\nabla_\theta\log p_\theta(x)\big]

since $\nabla p = p\nabla\log p$. The trajectory probability is $p_\theta(\tau) = p(s_0)\prod_t\pi_\theta(a_t \mid s_t)\,p(s_{t+1} \mid s_t, a_t)$; in its logarithm, the dynamics terms do not depend on $\theta$, so they vanish from the gradient:

\nabla_\theta\log p_\theta(\tau) = \sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)

We never need to know the environment's model.

The policy gradient theorem#

Using causality (an action cannot affect rewards received before it), we obtain

\nabla_\theta J(\theta) = \mathbb{E}_{\pi_\theta}\left[\sum_{t}\nabla_\theta\log\pi_\theta(a_t \mid s_t)\,G_t\right]

where $G_t$ is the return from time $t$ ("reward-to-go"). More generally, the policy gradient theorem (Sutton et al., 2000) states $\nabla J(\theta) \propto \mathbb{E}\big[Q^{\pi}(s, a)\nabla_\theta\log\pi_\theta(a \mid s)\big]$.

Intuition: increase the log-probability of actions in proportion to how good the outcome was. Actions followed by high returns become more likely; actions followed by low returns become less likely.

REINFORCE#

Williams's (1992) algorithm estimates the gradient from sampled episodes:

Run the policy for an episode; record states, actions and rewards.
Compute returns $G_t$ for each step.
Update $\theta \leftarrow \theta + \alpha\sum_t G_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)$.

Reducing variance with a baseline#

REINFORCE is unbiased but very high variance: returns fluctuate enormously between episodes. Subtracting a baseline $b(s_t)$ that does not depend on the action keeps the estimator unbiased (because $\mathbb{E}_a[\nabla_\theta\log\pi_\theta(a \mid s)] = 0$) while reducing variance:

\nabla_\theta J \approx \sum_t\nabla_\theta\log\pi_\theta(a_t \mid s_t)\,\big(G_t - b(s_t)\big)

A natural baseline is a learned state-value estimate $\hat{V}(s_t)$; then $G_t - \hat{V}(s_t)$ estimates the advantage — how much better the action was than average. This leads directly to actor–critic methods.

Implementation#

python

import numpy as np
import torch, torch.nn as nn
import gymnasium as gym

env = gym.make("CartPole-v1")
policy = nn.Sequential(nn.Linear(4, 128), nn.Tanh(), nn.Linear(128, 2))       # logits
value = nn.Sequential(nn.Linear(4, 128), nn.Tanh(), nn.Linear(128, 1))        # baseline
opt = torch.optim.Adam(list(policy.parameters()) + list(value.parameters()), lr=3e-3)
gamma = 0.99

for episode in range(600):
    s, _ = env.reset(seed=episode); done = False
    logps, rewards, states = [], [], []
    while not done:
        st = torch.as_tensor(s, dtype=torch.float32)
        dist = torch.distributions.Categorical(logits=policy(st))
        a = dist.sample()
        s, r, term, trunc, _ = env.step(a.item()); done = term or trunc
        logps.append(dist.log_prob(a)); rewards.append(r); states.append(st)
    G, returns = 0.0, []
    for r in reversed(rewards):                                # reward-to-go
        G = r + gamma * G; returns.insert(0, G)
    returns = torch.tensor(returns)
    V = value(torch.stack(states)).squeeze(-1)
    advantage = returns - V.detach()                            # baseline subtraction
    advantage = (advantage - advantage.mean()) / (advantage.std() + 1e-8)
    policy_loss = -(torch.stack(logps) * advantage).sum()
    value_loss = ((returns - V) ** 2).mean()
    opt.zero_grad(); (policy_loss + 0.5 * value_loss).backward(); opt.step()
    if episode % 100 == 0:
        print(f"episode {episode}: return {sum(rewards):.0f}")

Strengths and weaknesses#

Strengths:

Directly optimises the objective; handles continuous actions (e.g. Gaussian policies $\mathcal{N}(\mu_\theta(s), \sigma_\theta(s)^2)$).
Learns stochastic policies; smooth changes in behaviour as parameters change.
Good convergence properties (to a local optimum) under standard conditions.

Weaknesses:

High variance, hence sample inefficiency.
On-policy: samples must come from the current policy; old data is discarded after each update.
Sensitive to step size: a single large update can collapse performance — motivating trust-region methods (TRPO, PPO).

Connection to language models#

In RLHF and RL for reasoning, the "policy" is the language model, an "action" is a token (or a whole response), and the reward comes from a reward model or a correctness check. The gradient estimator is exactly REINFORCE-style — with baselines (value networks in PPO, or group averages in GRPO) and KL penalties for stability.

#policy gradient#reinforce#baseline#variance reduction#stochastic policies

Written by

Janin A Apurba

B.Sc. in CSE, AUST · Advanced ICT Officer, CNRS-UNHCR. Teaching AI, ML and Deep Learning to the next generation of engineers and researchers.