Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Policy Gradient Methods › Language Models as Policies

Michael BrenndoerferDecember 26, 20256 min read

Language AI HandbookAll chapters

Part 37 of 67Alignment and RLHF

Chapter 277Policy Gradient Methods › Language Models as Policies

Reading section 2 of 14 · Language Models as Policies · 6 minSection 2 of 14 · 6 min

Section 2 of 14 · 6 min read

Policy Gradient Methods › Language Models as PoliciesOne section · Sections in this chapter2 / 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

In reinforcement learning terminology, a policy is a function that maps states to actions. For language models, this mapping has a natural interpretation that becomes clear once we examine how text generation works. When you ask a language model to complete a sentence, it does not produce the entire response instantaneously. Instead, it generates one token at a time, each choice depending on everything that came before. This sequential decision-making process is precisely what reinforcement learning was designed to handle.

Think of a language model as an author writing a letter one word at a time. At each moment, the author has a specific context in mind (the prompt and everything written so far), and from that context, must choose the next word. Some words are clearly more appropriate given the context; others would seem out of place. The author's internal sense of which words fit a context is exactly what a language model encodes as a probability distribution. That probability distribution is the policy.

Policy

A policy $\pi_\theta(a|s)$ is a probability distribution over actions $a$ given state $s$ , parameterized by $\theta$ . In language models, the policy is the model itself: given a context (state), it produces a probability distribution over the next token (action).

To understand why this framing is so powerful, consider what happens during text generation. The model receives a prompt, which establishes the initial context. Based on this context, the model must decide which token to produce first. This decision changes the context, as the newly generated token becomes part of the history. The model then faces a new decision: given the original prompt plus the first generated token, what should the second token be? This process continues, with each token choice altering the state and presenting a fresh decision problem.

The reinforcement learning vocabulary for this process is precise and useful. At each timestep $t$ , the system occupies a state $s_t$ , the agent takes an action $a_t$ , the state transitions to $s_{t+1}$ , and eventually the agent receives a reward $R$ . For language generation, this maps directly:

State $s_t$ : The current context, consisting of the prompt plus all tokens generated so far
Action $a_t$ : The next token to generate
Policy $\pi_\theta(a_t|s_t)$ : The probability the model assigns to token $a_t$ given context $s_t$

This framing reveals why language model alignment is fundamentally a reinforcement learning problem. At each timestep, the model takes an action (selects a token), the state changes (the context grows), and eventually the complete sequence receives a reward from our reward model. The challenge is learning which tokens led to that reward. Unlike a game where individual moves might receive immediate feedback, language generation only reveals whether the response was good or bad after the entire sequence is complete. This delayed reward structure is exactly the scenario where reinforcement learning excels.

One important property of this MDP (Markov Decision Process) formulation is that the state transition is deterministic given the action. Once we pick token $y_t$ , the new state is simply the old state with $y_t$ appended. There is no external randomness in the transition. All the stochasticity lives in the policy: the model samples tokens from its probability distribution. This simplification makes the mathematics more tractable than general RL settings where transitions are random.

In[3]:

Code

import torch.nn as nn
import torch.nn.functional as F


# A simple language model acting as a policy
class LanguageModelPolicy(nn.Module):
    def __init__(self, vocab_size, hidden_size=256, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_size)
        self.lstm = nn.LSTM(
            hidden_size, hidden_size, num_layers, batch_first=True
        )
        self.output = nn.Linear(hidden_size, vocab_size)

    def forward(self, input_ids):
        # State: all tokens seen so far
        embeds = self.embedding(input_ids)
        lstm_out, _ = self.lstm(embeds)
        # Action distribution: probability over next token
        logits = self.output(lstm_out)
        return logits

    def get_action_probs(self, state):
        """Get probability distribution over actions (next tokens)"""
        logits = self.forward(state)
        # Policy: pi_theta(a|s) - probability of each action given state
        return F.softmax(logits[:, -1, :], dim=-1)

Out[4]:

Visualization

Diagram showing states expanding as tokens are added, with actions connecting consecutive states and reward at the end. — Sequential token generation modeled as a Markov Decision Process. Each state $s_t$ encompasses the full context (prompt and generated history), while actions $a_t$ correspond to next-token selections. The visualization highlights the temporal gap between individual actions and the final sparse reward $R$, illustrating the credit assignment challenge that makes language model alignment a reinforcement learning problem.

The model's parameters $\theta$ define the policy. Our goal is to adjust these parameters so that the model generates responses that receive high rewards. This means we need to find parameter values that make high-reward sequences more probable and low-reward sequences less probable. The question is: how do we know which direction to adjust the parameters? This is the central technical question that policy gradient theory answers.