18. REINFORCE Policy-Gradient Loss

Problem

The REINFORCE algorithm is the simplest way to train a stochastic policy directly from reward: increase the log-probability of actions that led to high return, decrease it for actions that led to low return. No value function, no Bellman equation — just the score-function gradient estimator, E[∇ log π(a|s) · G_t], applied episode by episode. It's also the conceptual starting point for every policy-gradient method DeepMind has published on top of it, from AlphaGo's policy network to AlphaStar's league training.

Given a batch of trajectories' per-timestep rewards and the log-probabilities the policy assigned to the actions it actually sampled, this problem chains two things together: turn the raw rewards into discounted returns-to-go via a backward recursion, then combine those returns with the log-probabilities (and an optional baseline, to reduce variance) into the scalar loss whose gradient is the REINFORCE update.

Source: src/18_reinforce_policy_gradient.py

def reinforce_policy_gradient_loss(
    rewards: torch.Tensor, log_probs: torch.Tensor, gamma: float,
    baseline: torch.Tensor | None = None,
) -> torch.Tensor: ...

>>> rewards = torch.tensor([[1.0, 0.0, 1.0]])
>>> log_probs = torch.tensor([[-0.2, -0.5, -0.1]])
>>> reinforce_policy_gradient_loss(rewards, log_probs, gamma=0.99).shape
torch.Size([])

Step-by-Step Approach

1. Recognize the two separate computations chained together: a discounted return-to-go over time, then a policy-gradient loss over those returns.
2. Compute the return-to-go with a backward loop over timesteps: G_t = r_t + gamma * G_{t+1}, starting from G_{T} = 0 past the last timestep.
3. If a baseline is given, subtract it from the returns to form the advantage; if not, use the raw returns directly — both paths should share the same loss line.
4. Detach the baseline (defensively, even if it's already a plain tensor) so no gradient accidentally flows into the loss through a would-be value-function baseline.
5. Combine into the loss: -(log_probs * advantage).mean() — the leading minus sign turns gradient ascent on expected return into gradient descent on the loss, which is what .backward() and an optimizer expect.
6. Average over every (batch, timestep) entry, not just over the batch dimension, so trajectory length doesn't skew the effective learning rate.

The entire method is two small, separable pieces: a purely mechanical discounting recursion, and a one-line loss formula — the only real way to get this wrong is mixing up the sign (ascent vs. descent) or forgetting that the baseline must never carry gradient.

Reference solution

import torch


def reinforce_policy_gradient_loss(
    rewards: torch.Tensor, log_probs: torch.Tensor, gamma: float,
    baseline: torch.Tensor | None = None,
) -> torch.Tensor:
    batch, T = rewards.shape

    # backward recursion for discounted return-to-go: G_t = r_t + gamma * G_{t+1}
    returns = torch.zeros_like(rewards)
    running = torch.zeros(batch)
    for t in reversed(range(T)):
        running = rewards[:, t] + gamma * running
        returns[:, t] = running

    # baseline is a variance-reduction trick and must not carry gradient into
    # the policy loss -- detach defensively even though callers usually pass
    # a plain tensor already
    advantage = returns if baseline is None else returns - baseline.detach()

    # score-function estimator: loss = -E[log pi(a|s) * advantage], mean over
    # every (batch, timestep) sample so batch size doesn't skew the gradient
    return -(log_probs * advantage).mean()

Key Functions & Tricks

• torch.zeros_like(rewards) — allocates the returns tensor with matching shape and dtype before filling it in the backward loop
• reversed(range(T)) — walks timesteps back-to-front so each G_t can reuse the already-computed G_{t+1}
• baseline.detach() — cuts the baseline out of the autograd graph — a variance-reduction term must never receive its own gradient from this loss
• tensor.mean() — reduces the elementwise log_probs * advantage product to the single scalar a training loop calls .backward() on

How to Recognize This Pattern