Policy gradient and REINFORCE

Q-learning learns values and derives the policy from them. Policy gradient skips the middle step and optimises the policy directly.

The policy $\pi_\theta(a \mid s)$ is a neural network that gives probabilities over actions. We want to adjust $\theta$ so that good outcomes become more likely.

The idea in one sentence:

It went well — do more of what you did. It went badly — do less.

Mathematically: multiply the log probability of every action by how well the episode turned out, and walk up the gradient.

Why care, when Q-learning exists?

	Value-based	Policy gradient
Continuous actions	hard (a max over infinitely many)	natural
A stochastic policy	no	yes — needed in games with bluffing
Convergence	can oscillate with function approximation	more stable guarantees
Sample efficiency	better (off-policy, replay)	worse (on-policy)

The last row is the price: REINFORCE throws away all its data after every update.

The objective is the expected return over trajectories $\tau$:

$J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)] = \int p_\theta(\tau) R(\tau)\, d\tau$

The problem: the gradient lands on the probability distribution, which we cannot differentiate through directly.

The log-derivative trick solves it. Since $\nabla_\theta p_\theta = p_\theta \nabla_\theta \log p_\theta$:

$\nabla_\theta J = \int \nabla_\theta p_\theta(\tau) R(\tau)\, d\tau = \int p_\theta(\tau)\, \nabla_\theta \log p_\theta(\tau)\, R(\tau)\, d\tau = \mathbb{E}_\tau\!\left[\nabla_\theta \log p_\theta(\tau)\, R(\tau)\right]$

Now it is an expectation again — and expectations can be estimated from samples.

The next step: the probability of a trajectory is

$p_\theta(\tau) = \rho(s_0)\prod_t P(s_{t+1} \mid s_t, a_t)\, \pi_\theta(a_t \mid s_t)$

Take the logarithm and it becomes a sum, and everything that does not depend on $\theta$ disappears on differentiation — including the environment's dynamics $P$. That is the whole point: we need no model of the world.

$\nabla_\theta J = \mathbb{E}\left[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, R(\tau)\right]$

Two improvements that do not change the expectation but greatly reduce the variance:

1. Causality. An action cannot affect what has already happened. Swap $R(\tau)$ for the return forwards from step $t$: $G_t = \sum_{k \geq t}\gamma^{k-t} r_k$.
2. A baseline. For any function $b(s)$ we have $\mathbb{E}[\nabla_\theta \log \pi_\theta(a\mid s)\, b(s)] = 0$, since $\sum_a \nabla_\theta \pi_\theta(a \mid s) = \nabla_\theta 1 = 0$. So we can subtract $b(s)$ for free.

With $b(s) = V(s)$, $G_t - V(s_t)$ becomes the advantage $A_t$ — and then we have arrived at actor–critic.

$\boxed{\nabla_\theta J = \mathbb{E}\left[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t)\, A_t\right]}$

Why the variance is the problem: the estimate is built on the random outcomes of whole episodes. Without a baseline the gradient can point in different directions between two runs with the same policy. The baseline removes the part of the signal that is common to every action in a state — and that part is nothing but noise.

import torch, torch.nn as nn

class Policy(nn.Module):
    def __init__(self, obs, actions, hidden=128):
        super().__init__()
        self.f = nn.Sequential(nn.Linear(obs, hidden), nn.Tanh(), nn.Linear(hidden, actions))

    def forward(self, s):
        return torch.distributions.Categorical(logits=self.f(s))

def returns(rewards, gamma=0.99):
    out, G = [], 0.0
    for r in reversed(rewards):
        G = r + gamma * G
        out.append(G)
    return list(reversed(out))

def reinforce(env, policy, opt, episodes=1000, gamma=0.99):
    for _ in range(episodes):
        s, done = env.reset()[0], False
        logp, rews = [], []
        while not done:
            d = policy(torch.as_tensor(s, dtype=torch.float32))
            a = d.sample()
            logp.append(d.log_prob(a))
            s, r, term, trunc, _ = env.step(int(a))
            rews.append(r); done = term or trunc

        G = torch.tensor(returns(rews, gamma))
        G = (G - G.mean()) / (G.std() + 1e-8)          # a baseline: greatly reduces the variance
        loss = -(torch.stack(logp) * G).sum()          # the minus sign: we are maximising J
        opt.zero_grad(); loss.backward()
        nn.utils.clip_grad_norm_(policy.parameters(), 1.0)
        opt.step()

Three things that usually go wrong:

Symptom	Cause
The policy quickly becomes deterministic and stops improving	the entropy has collapsed — add an entropy bonus
No learning at all	the minus sign is missing, or the normalisation removed all the signal
Very shaky learning	normalising over a single short episode; batch several

The normalisation on line 24 is standard but not innocent: do it over a single episode and you are only comparing actions within that episode, and the information that the whole episode was bad is lost.