PPO for Language Models › The Token Action Space
Language AI HandbookAll chapters PPO for Language Models › The Token Action SpaceOne section · Sections in this chapter3 / 15 ↓
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 action space in language generation is the model's vocabulary. This discrete set of tokens defines every possible action the policy can take at each step. Unlike continuous control problems where actions might be forces or velocities, or even discrete games where actions represent button presses, the language model's action space consists of linguistic units. These might be complete words, word fragments, punctuation marks, or special tokens that signal the end of generation.
Think of the vocabulary as a massive library of building blocks, and generating a response as repeatedly choosing which block to place next. Each choice is constrained by what came before: you cannot place a closing parenthesis before an opening one without creating grammatical nonsense, and you cannot begin a response mid-sentence. The model has learned an enormous amount about which sequences of choices produce coherent, meaningful text, but at each step it must select from the entire library.
For modern language models, this vocabulary typically contains between 30,000 and 100,000 tokens. The exact size depends on the tokenization algorithm used during pretraining. Byte-pair encoding, the most common approach, creates vocabularies that balance coverage of common words with the ability to represent rare words through subword decomposition. The vocabulary must be large enough to efficiently represent common text patterns while remaining small enough for the softmax computation to be tractable. This tradeoff explains why vocabulary sizes cluster around 32,000 to 65,000 tokens for most production models.
This scale is much larger than traditional RL domains. In classic control, action spaces are small. For example, a robot arm might have six joints. In Atari games, agents choose among roughly 18 discrete actions representing joystick directions and button combinations. Language models must select from tens of thousands of possible tokens at every single step. This represents an increase of three to four orders of magnitude in the number of discrete choices.
Implications for Exploration and Credit Assignment
The enormous action space has several important consequences for how PPO behaves when applied to language models. Each of these implications changes how we design and tune the algorithm.
Exploration is implicit in language generation. With such a large action space, the model cannot systematically try all options. There is no way to enumerate the consequences of every possible token choice. Instead, exploration emerges from the inherent stochasticity of sampling from high-entropy distributions during generation. When the model is uncertain, it assigns probability mass to many tokens, and sampling naturally explores these alternatives. This is quite different from classical RL, where explicit exploration strategies like epsilon-greedy or optimistic initialization are necessary. In language generation, sampling temperature takes the place of explicit exploration: higher temperature spreads probability mass more evenly, encouraging exploration; lower temperature concentrates it on the most probable tokens, encouraging exploitation.
Credit assignment is diffuse across the token sequence. When a response receives a reward, determining which specific token choices contributed to that reward is challenging. Did the response succeed because of a precise word choice in the third sentence, or because of the overall argument structure established from the first token? The signal must somehow propagate back through dozens or hundreds of individual decisions. The advantage function addresses this by using a value function to assess the quality of each intermediate state, but the value function itself must learn from the same sparse reward signal. This creates a bootstrapping challenge: we need good value estimates to produce good advantages, but we need good advantages to produce good policy updates that lead to rewards that improve value estimates.
Despite these challenges, the discrete nature of the action space simplifies some aspects of PPO. We can compute exact action probabilities rather than approximating them, as we would need to do in continuous action spaces. We can directly enumerate the KL divergence between policies by summing over vocabulary positions. The softmax function gives us a proper probability distribution that sums to one, avoiding the density estimation challenges that arise with continuous distributions.
import torch
import torch.nn.functional as F
# Illustrate the action space concept
vocab_size = 32000 # Typical vocabulary size
hidden_dim = 768
# A simplified "policy head" that maps hidden states to action probabilities
class PolicyHead(torch.nn.Module):
def __init__(self, hidden_dim, vocab_size):
super().__init__()
self.lm_head = torch.nn.Linear(hidden_dim, vocab_size, bias=False)
def forward(self, hidden_state):
# hidden_state: (batch, seq_len, hidden_dim)
logits = self.lm_head(hidden_state) # (batch, seq_len, vocab_size)
return logits
def get_action_probs(self, hidden_state, temperature=1.0):
logits = self.forward(hidden_state)
return F.softmax(logits / temperature, dim=-1)
policy_head = PolicyHead(hidden_dim, vocab_size)
num_params = sum(p.numel() for p in policy_head.parameters())Action space size: 32,000 possible tokens Policy head parameters: 24,576,000
Each position in the generated sequence requires selecting from this massive action space. A 100-token response involves 100 sequential decisions, each choosing among 32,000 options. Even a short response spans a large search space of possible sequences. This is why language generation uses learned policies to navigate the space efficiently.