Learning RL from the Math: From Bellman Equations to LLM Post-Training

type

status

date

slug

summary

Why this book is my main recommendation

The official repository provides the book PDF, chapter PDFs, slides, and links to English and Chinese lectures. The book assumes probability and linear algebra, but not prior RL knowledge. Its examples use a grid world, so the environment stays simple while the ideas become more demanding.

That last choice matters to me. When both the environment and the algorithm change, it is hard to tell what caused a new difficulty. A small, familiar environment lets me concentrate on the mathematical change.

The book’s ten-chapter structure moves from values and Bellman equations through sampling, stochastic approximation, temporal-difference learning, function approximation, policy gradients, and actor–critic methods.

My recommendation has a clear boundary. This is a foundation for understanding LLM post-training, not a complete guide to implementing a production PPO or GRPO system. Those connections are work I need to add after the relevant chapters.

I prefer that to pretending the foundations are already obvious.

Read it as a chain of questions

Here is how I would organize the book. The chapter coverage follows the published contents; the questions and checkpoints are my proposed study route.

Chapters	The question to answer	A concrete checkpoint
1–3: Basic concepts, state values, Bellman equations, optimal values	What exactly are we trying to predict or optimize?	Define a tiny MDP. Compute its returns and values. Separate evaluation of a fixed policy from finding an optimal one.
4: Value iteration and policy iteration	How do the equations become algorithms when the environment model is known?	Implement both methods on the same small grid. Check their answers against direct calculations.
5–7: Monte Carlo, stochastic approximation, temporal-difference methods	How do we learn when we observe samples instead of exact expectations?	Compare MC and TD value estimates. Then explain why SARSA and Q-learning can update differently on the same transition.
8: Value function methods	What changes when a table becomes a parameterized function?	Compare approximate values with a tabular reference. Identify which parts of the target are held fixed during an update.
9–10: Policy gradients and actor–critic	How can we improve a policy directly, and where does its learning signal come from?	Derive a policy gradient, prove the baseline identity, and check a tiny implementation against an exact gradient.

The route is not “collect ten algorithms.” It is “remove one assumption at a time and understand what must replace it.”

1. Bellman is a consistency equation, not a vague trade-off

One of my early descriptions was that Bellman balances current reward against future reward. That intuition points toward long-term consequences, but it misses the useful mathematical idea.

First, separate three objects. A reward is feedback from one step. A return adds rewards along a trajectory, possibly with discounting. A value is an expectation of that return under a specified policy.

Using for the reward received after action :

The superscript matters. A state does not have a single value independent of what the agent will do next.

Now split the return after its first reward:

Taking the appropriate conditional expectation gives the Bellman equation:

The future term is a value, not just another reward. It already summarizes the later rewards and the policy’s future decisions. This is the central relationship developed in the book’s chapters on state values and Bellman equations.

My better interpretation is: a value estimate should agree with the immediate reward plus the value of where the policy takes us next.

For a finite discounted MDP, the fixed-policy equation can also be written as:

Under the usual bounded-reward assumption and , its solution is unique. This makes Bellman a precise consistency condition, not merely advice to think about the future.

Evaluation and optimality ask different questions

Policy evaluation asks how well a particular policy performs. Optimality asks for the best achievable value. Writing for the expected return after taking action and then following , we have . At optimality, the policy-weighted average is replaced by a choice of the best action.

But “best action” does not mean the largest immediate reward.

Consider a toy choice. Stopping gives reward 1. Continuing costs 0.2 now and gives reward 2 one step later. With discount 0.9, continuing is worth . A greedy choice based only on immediate reward selects the wrong action.

My checkpoint: explain why the equation is true before trying to memorize a value-iteration update.

2. Stochastic approximation connects an equation to a learning algorithm

Chapter 6 deserves special attention. It connects the exact quantities in the earlier chapters to updates based on noisy observations.

A running average is the simplest place to start:

With , this reproduces the ordinary sample mean after initialization. The update moves the current estimate toward the next observation.

Why is that sensible? At an estimate , the expected correction is . It points toward the desired mean and becomes zero at the answer.

This is a useful entry point to the book’s stochastic approximation chapter: we may not know an expectation exactly, but we can use noisy observations to approach a root of an expected update.

The familiar Robbins–Monro step-size conditions are:

They express two competing requirements: keep moving enough to reach the solution, while making accumulated noise manageable. They are not sufficient on their own. Convergence also depends on assumptions about the underlying function, noise, and stability. They certainly do not prove that an arbitrary deep-RL training loop converges.

Monte Carlo and TD use different targets

For policy evaluation, a Monte Carlo update can move a value estimate toward an observed return:

A one-step TD update instead uses:

Monte Carlo waits for the relevant sampled return. TD uses an estimate of what happens after the next state. That use of another estimate is called bootstrapping. Both methods use samples; sampling is not what distinguishes them. These targets are developed in chapters 5–7.

The distinction helped clarify a later note I wrote about language models. With only a terminal reward and no discounting, every token’s Monte Carlo return equals that final reward. Its TD target can still differ, because TD includes the next prefix’s estimated value.

Same objective. Different learning signal.

My study rule here is to label every quantity as an observation, an expectation, or an estimate. A surprising amount of confusion disappears once those labels are explicit.

3. On-policy is not the same as online

I asked why Q-learning was off-policy. Then I asked why SARSA could not be off-policy too.

The key is to distinguish the behavior policy, which generates data, from the target policy, whose behavior we are evaluating or improving.

In ordinary SARSA, the target uses the next action actually sampled by the policy being followed:

Q-learning uses the greedy next-action value instead:

These are the contrasting targets in the book’s temporal-difference chapter.

For a concrete example, suppose the next state has action values 2 and −1. Exploration selects the second action. If the current reward is 1 and the discount is 0.9, the SARSA target is 0.1. The Q-learning target is 2.8.

Both algorithms can receive that transition immediately after interacting with the environment. Both can therefore operate online. But Q-learning’s target does not evaluate the exploratory next action that was actually selected.

A replay buffer is not the definition of off-policy learning. Live data collection is not proof that an update is on-policy.

My follow-up about SARSA was also reasonable. Off-policy SARSA variants exist; the standard update is on-policy, not the name of an algebraic expression that can never be adapted. Estimating another policy’s returns requires appropriate correction or an appropriate expectation, along with data coverage. Simply relabeling the target does not solve the distribution mismatch.

My checkpoint: point to the exact place where the behavior policy and target policy enter an update.

4. Policy gradient differentiates an expectation, not a sampled token

This was one of the most important connections to LLMs.

A reward model can be differentiable with respect to its own parameters or continuous inputs. That does not make the hard token-sampling operation between the policy and the reward model differentiable in the usual way.

Start with a simpler problem. Fix a prompt , let be a complete response, and assume the reward function does not itself depend on the policy parameters. For clarity, imagine a finite response space.

The reward attached to a particular response may be fixed. The probability of producing that response is not.

Differentiate the probabilities:

This is the score-function identity behind the relevant policy-gradient estimator. The book develops policy gradients in chapter 9.

The derivation does not require differentiating a token ID. It requires knowing the probability the policy assigned to the sampled response.

A small Bernoulli example makes the distinction concrete. Let action 1 earn reward 1 and action 0 earn reward 0. If the probability of action 1 is , then:

A particular hard sample does not smoothly change from 0 to 1 when I nudge the parameter. Yet the expected reward changes smoothly because the probability changes.

Backpropagation has not disappeared

For fresh on-policy samples, one way to implement this estimator is to minimize a surrogate loss:

Backpropagation differentiates the policy log-probability. The sampled feedback is held fixed for that actor update. This distinction between an estimator and a differentiable surrogate is also explicit in the original PPO paper, Section 2.1.

This is not the full PPO loss, and repeatedly optimizing it on unchanged old samples needs further care.

It also answers another question I had: an instruction to increase a useful action’s probability is not an instruction to increase every parameter. The gradient determines a direction in parameter space. Its components can have different signs.

My checkpoint: identify the sampling distribution, the reward, and the differentiation path separately. Do not let one expectation symbol hide all three.

5. A baseline changes the estimate, not the goal

Raw reward is not always the most useful way to judge a sampled action. An outcome can be good in absolute terms but disappointing relative to what the policy normally achieves in that state.

That motivates a baseline. For a baseline that does not depend on the sampled action, conditional on state :

Subtracting this baseline therefore preserves the expected score-function gradient. In an actor loss, the baseline is treated as fixed feedback rather than differentiated through as part of the weight.

Choosing a useful baseline can reduce variance. Choosing an arbitrary baseline does not guarantee that. The advantage function makes the comparison explicit:

An advantage is not an extra environmental reward. It measures an action against the policy’s usual value at that state. A critic estimates values to help construct the actor’s learning signal; it is not the same object as a model that scores completed answers. The book develops these connections in chapter 10.

A small check beats a persuasive explanation

The following standard-library Python example enumerates both actions. It checks the expected policy gradient against a finite difference and compares two baselines. No training framework is needed.

The output is:

Same expected gradient. Different variance.

There is another useful detail. With baseline 0.3, sampling the zero-reward action gives a negative advantage. Its negative score is multiplied by that negative advantage, producing a positive update to the logit of the rewarding action. Learning from a disappointing action is not mysterious once the signs are written out.

This tiny example is not evidence that a particular baseline always helps. It is a way to verify exactly which claim is being made.

6. Why RLHF can look like a bandit and still have token-level states

My notes and conversations kept returning to a modeling question: is the action a token or a whole response?

Both can be useful descriptions of a single-response task.

In a completion-level contextual bandit, the context is the prompt, the action is a full response, and feedback arrives after that response. Each episode has one decision. That does not mean every prompt is the same state.

In a token-level MDP, the state is the prompt plus the generated prefix. The action is the next token. Appending the token changes the state, even when that transition is deterministic. Generation ends at EOS or another stopping condition.

For the same autoregressive policy, the two descriptions connect through:

In a finite episode with no intermediate rewards and , each token’s Monte Carlo return is the final response reward. With discounting, the return depends on how far the token is from that reward. With intermediate costs or rewards, the equality also changes.

This is the precise version of the distinction I eventually wrote down in my own notes: a sampled return and a bootstrapped value target are not interchangeable, even when the only task reward arrives at the end.

Bellman relationships remain valid for the token-level MDP. A particular training method may simply not solve a Bellman equation or train a critic explicitly. REINFORCE can use sampled returns directly. Absence of a Bellman update is not absence of a sequential model. For the LLM-specific formulation, see Nathan Lambert’s RLHF and Post-Training Book.

For a tool-using agent, I need to be even more explicit. A tool call can change an external system, reveal new information, or fail. The state must include the relevant interaction history or another sufficient representation. “The environment hardly changes during text generation” is not a safe assumption to carry into every agent problem.

7. The bridge to PPO and GRPO should come after the foundation

I do not want to turn this book into a claim that classical RL and LLM post-training are identical. The foundation gives me a way to ask better questions about the differences.

GAE connects to TD errors. Once I understand the one-step error , I can study how generalized advantage estimation combines errors across multiple steps. The GAE paper is a natural extension, rather than another formula to memorize without context.

PPO connects to learning from a fixed rollout batch. Its probability ratio compares the changing policy with the policy that generated the data:

The clipped policy objective uses the smaller of the unclipped and clipped advantage-weighted terms:

The PPO paper motivates this surrogate for multiple updates on collected data. Clipping is not a hard constraint that guarantees every probability ratio or global KL stays within a bound. In a sequential problem, an action ratio at recorded states is also not a complete correction for every change in the state distribution.

This helps answer a recurring source of confusion: collecting rollouts online does not keep the policy equal to the behavior policy throughout later updates. At the start they can match; after an update, they need not.

GRPO connects to how we construct relative feedback. In the original DeepSeekMath paper, responses to the same prompt are compared within a group, avoiding a separate learned value model. Outcome-supervised GRPO uses group-centered, standardized rewards as token advantages. That is an extension beyond the book, not one of its chapters.

The baseline proof above must not be applied carelessly. A group mean that includes a response’s own reward is not an action-independent baseline for that sample. Even before standardization and clipping, this changes the exact estimator. With independent samples, subtracting the self-inclusive group mean scales the expected plain score-function gradient by . Standardization introduces further effects. The details deserve their own derivation, not the slogan “baselines never change anything.”

Finally, old policy and reference policy have different jobs. The old policy records how rollout data was sampled. A reference policy supplies an anchor for a KL regularizer. They may initially share weights without being conceptually interchangeable.

This is where the foundation becomes useful in code. A reward tensor with shape (B, G) must be grouped by prompt. Old log-probabilities must remain the probabilities of the sampling policy. Token masks, next-token alignment, and sequence-versus-token averaging define what is optimized. They are not cosmetic implementation details.

How I would study this without getting stuck in preparation

My reading plan puts a useful condition on “done”: reproduce the important derivations without looking, and connect them to the post-training methods I want to understand.

I would make that condition concrete in four passes.

Pass 1: Build an exact reference

Work through chapters 1–4 using one tiny environment. Write down its states, actions, transition probabilities, rewards, and termination rules. Compute a fixed policy’s value directly, then reproduce it with iterative updates. Implement policy iteration and value iteration.

Do not move on because the plot looks plausible. Move on when the small case has an answer you can check independently.

Pass 2: Replace exact expectations with samples

Use chapters 5–7 to estimate the same values with MC and TD. Compare estimates across several random seeds. Vary the step size. Inspect both error and variability.

Then add exploratory actions and compare SARSA with Q-learning. Before running the code, predict a transition on which their targets differ.

This makes sampling, bootstrapping, exploration, and policy mismatch separate experiments instead of one confusing training curve.

Pass 3: Add approximation and direct policy optimization

For chapter 8, replace the table with a simple function approximator while keeping the small environment as a reference. More parameters should not remove the obligation to check the answer.

For chapters 9–10, begin with an enumerable policy-gradient example such as the one above. Then implement REINFORCE and a small actor–critic. Compare gradient estimates and learning behavior, not just the final return from one seed.

These are proposed exercises and completion criteria. A checked reading-list item is not evidence that every implementation has already been completed.

Pass 4: Translate the math into an LLM update

Only now add GAE, PPO, and GRPO. For each method, write down the true objective, the distribution that supplied the samples, the estimator or surrogate used for training, and the quantities held fixed during differentiation.

Before running a large model, use a tiny synthetic batch to test masks, group reductions, probability ratios, and gradient flow. Change one assumption at a time.

The official repository provides grid-world environment code, not official implementations of every algorithm. It also links third-party implementations that the author says he has not verified. I would use those for comparison after an attempt, not as a substitute for making the update work myself.

How I want to use ChatGPT while learning

My conversations were most useful when a question exposed an actual inconsistency in my understanding. “Why does RLHF have no Bellman equation?” was more productive than asking for another overview of RL.

The tutoring loop I want is simple: explain my current model first, ask for one precise error, test it with the smallest counterexample, and then reconstruct the corrected idea without the explanation in front of me.

A prompt I would use is:

Here is my understanding of this update. Find the first incorrect step or missing assumption. Do not replace it with a full lecture yet. Give me a two-state or two-action example that distinguishes my explanation from the correct one. After I revise it, change one assumption and test me again.

The assistant’s agreement is not the test. A derivation, an independently computed answer, or a failing unit test is much stronger evidence.

The failure mode I want to avoid is familiar: an explanation sounds clear, so I mistake that feeling for an ability to reproduce it.

The distinction I most want to remember

For me, the hardest step is separating the objective, the estimator, and the differentiable loss.

The objective says what outcome I want on average. The estimator says how sampled data provides information about that objective or its gradient. The loss is what the training code differentiates to produce an update.

Those objects are connected. They are not interchangeable.

That distinction transfers beyond RL. It helps me reason about sampling bias, importance weighting, stopped gradients, noisy labels, and whether an implementation still matches its stated objective.

I strongly recommend Shiyu Zhao’s book and accompanying course as the main path through this part of becoming an LLM engineer. The English lectures and the author’s Chinese lecture page offer another way into the same material.

I do not need to master every theorem before building. But I do want to be able to look at a training update and explain what it is estimating, why that estimate makes sense, and which assumption would make it fail.