Lecture 21 · Part 6

RL with Sequence Models & LLMs

Three lectures in one arc. First, what breaks when observations replace states — which RL algorithms survive partial observability and which quietly become wrong. Then the reverse direction, RL as a training signal for language models: the single-step RLHF formulation behind ChatGPT, with reward models learned from human preferences. Finally the two combined — multi-turn dialogue and tool use as a POMDP solved with offline value-based RL over history states.

▶ Lecture videos (3 parts) Slides (current revision, PDF) ↗

POMDPs are weird

The lecture opens where the course began: what happens beyond fully observed MDPs. The trouble with observations is that they “don’t obey the Markov property” — $o_t$ may not suffice to infer the state, so previous observations can genuinely help. With a true state that never happens: $s_t$ d-separates the future from the past. In the cheetah-and-gazelle example, a car driving in front of the cheetah changes the observation without changing the state; and in most real problems — a car in your blind spot, a robot inferring a human’s mental state, a dialogue where “the history of the interaction really matters” — partial observability is a matter of degree, sometimes ignorable (Atari) and sometimes central.

“Partially observed MDPs can be really, really weird.” Two things happen in POMDPs that never happen in MDPs. First, information-gathering actions become optimal: solving a distribution of mazes from first-person view, it pays to peek over the wall — an action that gets you no closer to the exit but reveals the unobserved maze configuration. Second, stochastic policies can beat every deterministic policy:

So “handle partial observability” needs a definition: an algorithm handles it if, feeding $o_t$ to the policy in place of $s_t$ , it finds the best policy in the class of memoryless reactive policies. Memory does better — that comes later — but first, which of the course’s algorithm families pass even this test?

Which algorithms survive swapping s for o

Policy gradients: yes, with care. The policy gradient derivation “never actually assumed the Markov property” — only the chain rule of probability, which always holds. So the grad-log-pi estimator with a Monte Carlo return is fine on $\pi_\theta(a_t \mid o_t)$ , and the causality trick (summing rewards from $t$ onward) is also fine — it only used the fact that the future doesn’t influence the past. What is not fine is a bootstrapped advantage $r_t + \hat{V}(o_{t+1}) - \hat{V}(o_t)$ : training $\hat{V}(o_t)$ silently assumes the value is the same every time you see $o_t$ regardless of how you got there — false without the Markov property. Curiously, $\hat{V}(o_t)$ as a baseline is still legal: all baselines are unbiased, whatever they are; this one just might not reduce variance as much as you’d like.

Value-based methods: no. Q-learning relies on exactly the property that broke the bootstrapped critic — the value upon visiting $s$ is history-independent. Substitute $o$ for $s$ and the update rule is simply invalid. A quick sanity check: greedy argmax extraction always yields a deterministic policy, and we just saw a POMDP where no deterministic policy is optimal.

Model-based RL: emphatically no. Levine’s illustration is two doors, one locked at random each episode, with the lock state unobserved. Fit a Markovian model $p(o' \mid o, a)$ and it learns that opening a door succeeds with probability $0.5$ — suggesting the absurd strategy of trying the same locked door until it “unlocks.” In reality, “no matter how many times you tried, it’ll remain locked.” The model’s structure — next observation depends only on current observation and action — cannot represent the environment’s structure.

Memory: state-space models and history states

Memoryless is an artificial restriction — if you tried the door, you should remember. One principled fix is a state-space model: train a sequence VAE (from the variational inference lecture) with Markovian latent dynamics $p(z_{t+1} \mid z_t, a_t)$ , an observation model $p(o_t \mid z_t)$ , and an encoder mapping observation histories to $z_t$ . Because the $z$ ‘s are trained to have Markovian dynamics, substituting $z$ for $s$ in Q-learning is valid. The catch: fully predicting observations — every pixel of a MuJoCo scene — may be “more difficult than recovering the optimal policy.” The shortcut: the encoder infers $z_t$ from $o_1, \ldots, o_t$ anyway, so skip the latent and use the history itself.

Architecturally that means an RNN, LSTM, or transformer. Concatenating a fixed short window (four frames, DQN-style) is “not as terrible an idea as it seems” but fails when you must remember the whole maze for the whole episode. Naively storing histories in a replay buffer costs $O(T^2)$ memory — each of $T$ steps carries up to $T$ observations — so a practical trick is to store RNN hidden states in the buffer and resume from them, treating $h_t$ as an approximate state (approximate because hidden states go stale as the RNN trains; see “Recurrent Experience Replay in Distributed Reinforcement Learning,” which works very well on Atari). To Levine’s knowledge no one has figured out the analogous trick for transformers, which lack a single Markovian hidden state.

Language models as one-step RL

The second part reverses the direction: RL to train better sequence models. A language model is a transformer predicting next tokens, modeling $p(x_t \mid x_1, \ldots, x_{t-1})$ — decidedly non-Markovian. Supervised training matches the data distribution; but to maximize a reward — satisfy human preferences, use tools, achieve dialogue goals — is RL, and we must pick the MDP. In the basic formulation the prompt is the state $\mathbf{s}$ and the whole completion is the action $\mathbf{a}$ , with the policy factored over language time steps: for a four-token prompt and two-token completion,

$\pi_\theta(\mathbf{a} \mid \mathbf{s}) = p(x_5 \mid x_{1:4})\, p(x_6 \mid x_{1:4}, x_5)$

There are now “two notions of time step — super confusing”: many language time steps, but exactly one RL time step. It is a bandit problem. The policy gradient applies directly — $\nabla_\theta \log \pi_\theta(\mathbf{a} \mid \mathbf{s})$ is the sum of per-token log-probability gradients, the same gradients as a cross-entropy backward pass. Estimate it REINFORCE-style with on-policy samples, or importance-weighted with samples from an earlier policy $\bar{\pi}$ :

$\nabla_\theta J \approx \frac{1}{N} \sum_i \frac{\pi_\theta(\mathbf{a}_i \mid \mathbf{s}_i)}{\bar{\pi}(\mathbf{a}_i \mid \mathbf{s}_i)}\, \nabla_\theta \log \pi_\theta(\mathbf{a}_i \mid \mathbf{s}_i)\, r(\mathbf{s}_i, \mathbf{a}_i)$

The importance-weighted (PPO-style) version dominates in practice because sampling from an LLM — and evaluating rewards — is slow, so you want many gradient steps per batch of generations:

Code walkthrough — Importance-sampled policy gradient loop for LLMs

repeat:
    B = sample_completions(pi_theta, prompts)   # slow: LLM generation
    R = [reward(s, a) for (s, a) in B]          # slow: reward evaluation
    pi_bar = snapshot(pi_theta)
    for k in range(K):                          # fast: many reused steps
        M = sample_minibatch(B)
        theta += alpha * grad_hat(theta, pi_bar, M)   # IS-weighted estimator

Batch $\mathcal{B}$ might be ~1,000 completions, minibatches of 64, $K$ inner steps — the classic importance-sampled / PPO loop, and “a very popular way to train language models with RL.”

Rewards from human preferences

The loop above needs $r(\mathbf{s}, \mathbf{a})$ for any completion — not just Paris ( $+1$ ) but “it’s a city called Paris” ( $0.9$ ), “I don’t know” ( $-0.1$ ), “London” ( $-1$ ), and “why are you asking such a stupid question” ( $-10$ ). That demands a learned, very powerful reward model $r_\psi$ . Humans struggle to emit absolute numbers when utility is subjective, but comparing two answers is easy. So model the preference probability, in a form “derived from the same mathematical foundations as maximum-entropy inverse RL”:

The full recipe — the foundation of InstructGPT and ChatGPT — is: (1) supervised fine-tuning gives the initial $\pi_\theta$ ; (2) sample $K$ answers per prompt; (3) humans label preferences; (4) fit $r_\psi$ ; (5) run many steps of policy optimization against $r_\psi$ ; repeat, sparingly. Human labels take days or weeks, so most preference data comes from the initial supervised model, and the poor man’s version runs the outer loop exactly once — making this, in effect, an offline model-based RL algorithm: no dynamics model (one step), but a reward model trained rarely and optimized against heavily.

That inherits model-based RL’s disease: distributional shift, here called over-optimization — the policy exploits the reward model, improves for a while, then gets worse. The standard fix folds a KL penalty against the supervised model into the reward, $\tilde{r}(\mathbf{s}, \mathbf{a}) = r_\psi(\mathbf{s}, \mathbf{a}) + \beta\left(\log \pi_{\text{sup}}(\mathbf{a} \mid \mathbf{s}) - \log \pi_\theta(\mathbf{a} \mid \mathbf{s})\right)$ , and makes the reward model a large pretrained transformer, “powerful enough that it can resist all that optimization pressure from reinforcement learning.”

Multi-step: dialogue as a POMDP

Part three puts the two halves together. In visual dialogue, a questioner bot asks a sequence of questions to guess which image the answerer has in mind: actions are the bot’s utterances, observations are the answerer’s replies, the state is the history — exactly the history states of the first section — and the reward arrives only at the end. Genuinely different from RLHF: there the episode was one answer, a bandit; here information-gathering questions, multiple turns, and partial observability all matter. The same structure covers assistant chatbots, tool use (text into a database, terminal, or calculator), and text games.

Policy gradients work — they handle histories — but on-policy rollouts against a human per gradient step are brutally expensive (less so against a tool). Value-based methods are appealing precisely because offline RL applies: train directly on logged human–human conversations or past deployments. The delicate design choice is the time step:

Per utterance: each turn is one RL step. Horizons are short (~10) and the action/observation alternation is natural, but the action space is all possible sentences — exponential in length. The critic encodes history and candidate utterance with sequence models into $Q(\mathbf{s}, \mathbf{a})$ , and the max over actions must be approximated: beam search, a learned actor, or sampling candidates from a supervised model and taking the highest-Q one.
Per token: every token is an action. The action space is just the vocabulary — enumerable, so the Bellman backup looks like a language-model head outputting Q-values instead of probabilities, maxing over next-token values (or taking the dataset token when the environment chooses). The price: horizons of thousands of steps even for short dialogues.

Both appear in the literature with no established winner. The usual value-based machinery applies — target networks, replay buffers, double-Q — plus offline regularization: policy constraints, IQL-style backups, or a CQL-style penalty, which for per-token discrete actions “conveniently works out to be the same as a cross-entropy loss.” Concrete instantiations: Jaques et al. (per-utterance actor-critic with a KL policy constraint, sentiment rewards), the CHAI chatbot (CQL penalty, negotiation revenue as reward, approximate max via GPT-2 samples), and Snell et al.’s Implicit Language Q-Learning (IQL backup plus CQL penalty, per-token, evaluated on visual dialogue). This is “very much at the bleeding edge of current research as of 2023.”

Back of the chalk tray — what to remember

POMDPs permit behaviors MDPs never require: information-gathering actions and strictly-stochastic optimal policies.
Policy gradients with Monte Carlo returns (and the causality trick, and any baseline) survive swapping $s \to o$ ; bootstrapped critics, Q-learning, and Markovian models do not — they bake the Markov property into their function class.
History states $(o_1, \ldots, o_t)$ provably obey the Markov property, so any RL algorithm becomes valid again once policies, values, and models are sequence models; RNN-state replay makes it efficient.
RLHF is a one-step bandit: prompt = state, completion = action, policy factored over tokens, trained with importance-weighted policy gradients so one batch of expensive generations funds many gradient steps.
Reward models are trained from pairwise human preferences via an exponential (max-ent-IRL-style) preference model; the whole pipeline is offline model-based RL, so over-optimization is fought with a KL penalty toward the supervised model and a very large reward model.
Multi-turn dialogue and tool use form a true POMDP over history states; offline value-based RL applies with a per-utterance (huge actions, short horizon) or per-token (small actions, long horizon) time step — no established standard yet.

Check yourself

Conceptual and derivation exercises in the lecture’s spirit — work them before reading the solutions.

Problem 21.1 memoryless policy value

In the three-state POMDP ( $A$ , $B$ , $C$ in a line; reward at $B$ ; start 50/50 at $A$ or $C$ ; observations uninformative), consider the memoryless stochastic policy that moves right with probability $p$ and left with probability $1-p$ . Starting from $A$ (where only “right” makes progress toward $B$ and “left” bumps the wall), the expected number of steps to reach $B$ is $1/p$ ; from $C$ it is $1/(1-p)$ by symmetry. Which $p$ minimizes the expected time to reward, and what does this say about deterministic policies?

Show solution

The expected hitting time averaged over starts is $T(p) = \tfrac{1}{2}\left(\tfrac{1}{p} + \tfrac{1}{1-p}\right)$ . Differentiating: $T'(p) = \tfrac{1}{2}\left(-\tfrac{1}{p^2} + \tfrac{1}{(1-p)^2}\right) = 0$ gives $p = 1/2$ , and $T(1/2) = 2$ . As $p \to 0$ or $p \to 1$ , $T(p) \to \infty$ : a deterministic policy ( $p \in \{0, 1\}$ in the limit) never reaches the reward from one of the two starts, so its expected time — and its failure probability on half the episodes — is unbounded. The uniformly random policy is the optimal memoryless policy, confirming the lecture’s claim that in POMDPs only stochastic policies may be optimal, which also rules out greedy extraction from any Q-function over observations.

Problem 21.2 Markovian model pathology

In the two-door environment, exactly one door (left or right, 50/50 per episode) is unlocked, and the lock state is unobserved. You fit a Markovian model $p(o_{\text{pass}} \mid o_{\text{left}}, \text{open})$ by maximum likelihood on data from many episodes. (a) What probability does the model learn? (b) Under this model, what is the predicted probability of passing within $k$ tries of the same door, and what is the true probability? (c) State precisely which structural assumption fails.

Show solution

(a) Half of all episodes have the left door unlocked, so the empirical frequency — and the maximum-likelihood estimate — is $0.5$ . (b) The model treats tries as independent: predicted success within $k$ tries is $1 - 0.5^k \to 1$ . In truth the outcome is determined once per episode: the probability is $0.5$ whether you try once or a thousand times — “no matter how many times you tried, it’ll remain locked.” (c) The model asserts $p(o' \mid o_t, a_t)$ is independent of $o_{1:t-1}, a_{1:t-1}$ given $(o_t, a_t)$ . But conditioning on a past failed attempt changes the true success probability from $0.5$ to $0$ . The environment’s conditional distribution depends on history; the Markovian model family cannot represent any such distribution, so no amount of data fixes it — the bias is structural.

Problem 21.3 preference model mechanics

A reward model assigns $r_\psi(\mathbf{s}, \mathbf{a}_1) = 2.0$ to “Paris” and $r_\psi(\mathbf{s}, \mathbf{a}_2) = 0.0$ to “I don’t know.” (a) Compute the modeled probability that a human prefers $\mathbf{a}_1$ . (b) Show that the preference probability depends only on the reward difference, and explain why this means the learned reward is only identified up to an additive constant per state. (c) Does that ambiguity change the RLHF policy optimization in the bandit setting?

Show solution

(a) $p(\mathbf{a}_1 \succ \mathbf{a}_2) = \frac{e^{2}}{e^{2} + e^{0}} = \frac{e^2}{e^2 + 1} \approx 0.88$ . (b) Dividing numerator and denominator by $\exp(r_\psi(\mathbf{s}, \mathbf{a}_2))$ gives $p(\mathbf{a}_1 \succ \mathbf{a}_2) = \sigma\!\left(r_\psi(\mathbf{s}, \mathbf{a}_1) - r_\psi(\mathbf{s}, \mathbf{a}_2)\right)$ , a logistic function of the difference. Adding any $c(\mathbf{s})$ to both rewards leaves every preference probability — hence the training likelihood — unchanged, so the data cannot pin down the per-state offset. (c) No: in the one-step problem a state-dependent constant added to the reward is exactly a baseline, and baselines leave the policy gradient unbiased (and the optimal policy unchanged) — the expectation of $\nabla_\theta \log \pi_\theta \cdot c(\mathbf{s})$ is zero. The ambiguity is harmless for optimization, though it means raw reward values are not comparable across states.

Problem 21.4 time-step design tradeoff

A dialogue task has $H = 10$ bot turns, each utterance at most $L = 20$ tokens from a vocabulary of size $V = 30{,}000$ . Compare the per-utterance and per-token formulations: give the RL horizon and the size of the per-step action space for each, and explain which term of the Q-learning procedure becomes the bottleneck in each case.

Show solution

Per utterance: horizon $10$ ; the action space is all token sequences up to length $L$ , on the order of $V^{20} = 30{,}000^{20}$ — astronomically large. The bottleneck is the $\max_{\mathbf{a}'} Q(\mathbf{s}', \mathbf{a}')$ in the Bellman target: it cannot be enumerated and must be approximated by beam search, a learned actor, or sampling candidates from a supervised model and taking the highest-Q one — an approximate max that introduces its own bias. Per token: the horizon becomes every generated token, up to $10 \times 20 = 200$ RL steps (the lecture notes horizons can reach thousands), but each step’s action space is just the $30{,}000$ tokens, so the max is a single argmax over the output head, exactly like a language-model softmax layer reinterpreted as Q-values. The bottleneck shifts to long-horizon credit assignment and value propagation over hundreds of backups, with the terminal reward arriving only at the end. Neither dominates: the literature has explored both with no established standard.