Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Policy Gradient Methods › The REINFORCE Algorithm

Michael BrenndoerferDecember 26, 20255 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 277Policy Gradient Methods › The REINFORCE Algorithm

Reading section 6 of 14 · The REINFORCE Algorithm · 5 minSection 6 of 14 · 5 min

Section 6 of 14 · 5 min read

Policy Gradient Methods › The REINFORCE AlgorithmOne section · Sections in this chapter6 / 14 ↓

Sections in this chapter14

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

The policy gradient theorem gives us a theoretical result: the gradient of expected reward equals the expected product of the score function and the reward. REINFORCE (also called the Monte Carlo policy gradient) is the algorithm that turns this theory into practice through Monte Carlo estimation. The name comes from the idea that we reinforce behaviors that lead to good outcomes: when a sequence receives high reward, we reinforce the probability of generating that exact sequence.

Think of REINFORCE as a simple loop: try something, evaluate whether it worked, update your approach based on the result. A human learning to write better essays follows essentially the same process. They write an essay, receive feedback (reward), and adjust their writing style to emphasize the patterns that received positive feedback and reduce the patterns that received criticism. REINFORCE formalizes this intuition mathematically.

The algorithm proceeds in four steps for each training iteration:

Sample a complete sequence from the current policy
Compute the reward for that sequence using the reward model
For each token in the sequence, compute the gradient of its log-probability scaled by the reward
Sum these token-level gradients and update parameters via gradient ascent

In practice, we implement gradient ascent by minimizing the negative objective (gradient descent on the negation):

In[7]:

Code

def reinforce_loss(log_probs, reward):
    """
    Compute REINFORCE loss for a single sequence.

    Args:
        log_probs: Log probabilities of each selected token, shape (seq_len,)
        reward: Scalar reward for the complete sequence

    Returns:
        Loss to minimize (negative of policy gradient objective)
    """
    # Sum log probs to get log pi(y|x)
    log_prob_sequence = log_probs.sum()

    # REINFORCE: maximize E[log pi(y|x) * R]
    # For gradient descent, we minimize the negative
    loss = -log_prob_sequence * reward

    return loss

The reinforce_loss function captures the entire algorithm in two lines of logic. We sum the per-token log-probabilities to get the sequence log-probability, then multiply by the reward and negate (to turn maximization into minimization). When PyTorch computes gradients of this loss with respect to model parameters, it produces exactly the policy gradient we derived above.

To use this loss, we first need to generate a sequence and collect the log-probabilities of each chosen token:

In[8]:

Code

import torch
import torch.nn.functional as F


def generate_and_get_log_probs(
    model, prompt_ids, max_length=20, vocab_size=1000
):
    """Generate a sequence and collect log probabilities."""
    model.eval()

    current_ids = prompt_ids.clone()
    log_probs = []

    with torch.no_grad():
        for _ in range(max_length):
            # Get action probabilities
            logits = model(current_ids)
            probs = F.softmax(logits[:, -1, :], dim=-1)

            # Sample action (next token)
            action = torch.multinomial(probs, 1)

            # Store log probability of selected action
            log_prob = torch.log(probs.gather(1, action))
            log_probs.append(log_prob)

            # Update state
            current_ids = torch.cat([current_ids, action], dim=1)

    return current_ids, torch.cat(log_probs, dim=1).squeeze()

Notice the use of torch.no_grad() during generation. We generate the sequence without tracking gradients, then separately compute the loss and allow gradients to flow through that computation. This is a memory optimization: we do not need to store the full computation graph during generation, only during the loss computation.

A Worked Numerical ExampleLink Copied

Let's trace through a concrete numerical example to build intuition. Suppose our model generates a three-token response, and at each position the model assigns the following probabilities to the token it chose:

Token 1: probability 0.30 (for example, "The")
Token 2: probability 0.50 (for example, "answer")
Token 3: probability 0.20 (for example, "is")

The reward model gives this short response a reward of $R = +2.0$ .

Step 1: Compute per-token log-probabilities.

\begin{aligned} \log \pi(y_1|s_1) &= \log(0.30) \approx -1.204 \\ \log \pi(y_2|s_2) &= \log(0.50) \approx -0.693 \\ \log \pi(y_3|s_3) &= \log(0.20) \approx -1.609 \end{aligned}

Step 2: Sum to get the sequence log-probability.

\log \pi_\theta(y|x) = -1.204 + (-0.693) + (-1.609) = -3.506

This corresponds to a sequence probability of $e^{-3.506} \approx 0.030$ , which matches $0.30 \times 0.50 \times 0.20 = 0.030$ directly.

Step 3: Compute the REINFORCE loss.

\mathcal{L} = -\log \pi_\theta(y|x) \cdot R = -(-3.506) \times 2.0 = +7.012

Step 4: Interpret the gradient.

When we call loss.backward(), the gradient will be computed with respect to model parameters. This gradient points in the direction that would increase the loss, which means the optimizer (doing gradient descent to minimize the loss) will adjust parameters in the direction that decreases the loss. Decreasing the loss means increasing $\log \pi_\theta(y|x) \cdot R$ , which means increasing the log-probability of generating this sequence. Since $R = +2.0$ is positive, this is exactly what we want: after the update, the model is more likely to generate this sequence.

Negative reward case: Suppose instead $R = -1.5$ . Then the loss equals $-(-3.506) \times (-1.5) = -5.259$ . Minimizing (making more negative) this loss requires making $-\log \pi \cdot R$ decrease. Since $R = -1.5 < 0$ , decreasing $-\log \pi \cdot (-1.5)$ means decreasing $\log \pi$ , making the sequence less likely. This is again the correct behavior: negative reward makes the sequence less probable.

In[9]:

Code

import torch

# Demonstration: REINFORCE gradient direction
log_probs = torch.log(torch.tensor([0.3, 0.5, 0.2]))
log_prob_total = log_probs.sum()

# Positive reward: gradient increases sequence probability
reward_positive = 2.0
loss_positive = -log_prob_total * reward_positive

# Negative reward: gradient decreases sequence probability
reward_negative = -1.5
loss_negative = -log_prob_total * reward_negative

Out[10]:

Console

Log probability of sequence: -3.507
Sequence probability (exp): 0.0300

With positive reward (+2.0): loss = 7.013
  -> Minimizing this loss increases sequence probability

With negative reward (-1.5): loss = -5.260
  -> Minimizing this loss decreases sequence probability

The key insight is that REINFORCE increases the probability of sequences that receive high rewards and decreases the probability of sequences that receive low or negative rewards. Over many updates, the policy gradually shifts to favor the kinds of responses that the reward model prefers. This is the essence of learning from rewards: try things, evaluate what works, do more of what works.