PPO for Language Models › The Language Model as a Policy

In reinforcement learning, a policy maps states to action probabilities. Given a state $s$, the policy $\pi(a|s)$ tells us the probability of taking action $a$. This definition describes decision-making: the agent observes its state, checks its policy, and chooses an action. Language models do the same thing: given a context, the model outputs a probability distribution over the next token. The parallel is exact, and recognizing it is the conceptual foundation for everything that follows.

Think of the language model's forward pass as the policy evaluation step in RL. Every time you call the model with a sequence of tokens, it is computing the policy, reading the current state and determining how to act. The transformer layers process the input, build rich contextual representations through self-attention, and project the final hidden state to a distribution over vocabulary tokens. This distribution is the policy output. Sampling from it is the action selection step. Appending the sampled token to the context creates the next state, and the cycle repeats until the model generates a stop token.

To see this clearly, consider what happens when you type a prompt. The model processes your input, builds internal representations through its transformer layers, and produces a probability distribution over its vocabulary. This distribution assigns higher probabilities to tokens that would naturally continue the text and lower probabilities to tokens that would seem out of place. When the model generates a response, it samples from this distribution (or selects greedily), appends the chosen token to the context, and repeats the process. Each step involves observing a state and selecting an action according to a probability distribution, which is exactly what a policy does.

The key insight is that autoregressive generation is policy rollout in disguise. Starting from the initial state consisting of the prompt, we sample actions according to the policy, each action extends the state, and we continue until reaching a terminal condition. This is a precise mathematical correspondence that allows us to apply policy gradient theorems and optimization techniques to language generation.

Let $x = (x_1, x_2, \ldots, x_n)$ denote the input prompt and $y = (y_1, y_2, \ldots, y_T)$ the generated response. At each generation step $t$, the language model computes the conditional probability of the next token given everything that came before:

$$\pi_\theta(y_t | x, y_{<t})$$

where:

• $\pi_\theta$: the policy defined by model parameters $\theta$
• $y_t$: the token generated at the current step $t$
• $x$: the input prompt sequence
• $y_{<t}$: the sequence of tokens generated prior to step $t$, that is $(y_1, \ldots, y_{t-1})$

This formula represents the probability of generating token $y_t$ given the full context of the prompt and previous tokens. The notation emphasizes that the policy depends on everything that came before: the original prompt establishes the task, and each previously generated token shapes what should come next. The model's parameters $\theta$ encode the learned patterns that determine how context maps to token probabilities.

The probability of a complete response $y$ given prompt $x$ is the product of these conditional probabilities across all timesteps. This product structure means that the probability of a response decomposes multiplicatively into the probability of each individual token choice. Understanding this decomposition is needed for understanding how policy gradient methods can work here: updating the parameters $\theta$ changes the probability of individual tokens at each step, and those changes compound across the response to shift the overall probability of different response types.

Comparing language models to standard RL shows a direct correspondence between the two frameworks. The state at each step is the concatenation of the prompt and all tokens generated so far: $s_t = (x, y_1, \ldots, y_{t-1})$. The action is the next token to generate: $a_t = y_t$. The policy is the LLM's softmax output distribution over the vocabulary: $\pi_\theta(a_t|s_t)$. The trajectory is the complete prompt-response pair: $\tau = (x, y)$.

Each of these mappings changes the interpretation for how we set up the optimization. The state includes the entire generation history, so the policy is not Markovian in the traditional sense. The model cannot simply look at one step and decide what to do; it must consider the full context. This is why transformer architectures, which can attend to arbitrarily distant tokens, are well-suited for this role. Older architectures like RNNs, which compress history into a fixed-size vector, are at a disadvantage here.

The autoregressive generation process that produces a response is exactly a policy rollout. Starting from the initial state $s_1 = x$, we sample actions according to the policy, each action extends the state, and we continue until generating a stop token. This precise mathematical correspondence allows us to apply policy gradient methods to language generation. Policy optimization theorems and techniques apply to steering language models toward desired behaviors.