Vol. XIX · Reinforcement Learning for Reasoning Plate 1 of 12reinforcement learning, policy, reward, return, episode, sampling, expected reward, language modelLE

Policies, Rewards and Returns.

A prompt is a state, a token is an action, a finished answer is an episode — and the reward comes at the end.

In reinforcement learning an agent in a state takes an action according to its policy, receives a reward, and moves to a new state; the sum of rewards over an episode is the return, and the goal is a policy with high expected return. A language model fits the frame directly. The state is the prompt plus the tokens written so far, each action is the next token, the policy is the model’s next-token distribution, and an episode ends with the end-of-turn token. For reasoning tasks the reward usually comes once, at the end: 1 if the answer is right, 0 if not. So the return of an episode is just its final reward, and the expected return of the policy on a prompt is the probability that a sampled answer is right.

The RL world of this volume is the post-training world’s chat model (plate XVII.1): fine-tuned on 30 examples, it answers about half the held-out sums correctly. A verifier checks each answer by doing the arithmetic, and every run here — sampling, rewards, advantages, updates — is real training of that model in a Web Worker. The workshop samples it 32 times on a chosen sum at a chosen temperature and shows the answers it gives and the reward each earns.

Below: the answers the policy gave in 32 samples for the chosen sum, green if the verifier accepts them; readouts give the expected reward.

What you will learn

Map states, actions, policies and rewards onto a language model.
Define the return of an episode with a terminal reward.
Estimate expected reward by sampling.
Explain what an RL update must do.

Workshop

The fine-tuned chat model of plate XVII.1, sampled 32 times on “U:a+b?|A:” at the chosen temperature, each answer ending at the end-of-turn mark or after four characters. Reward 1 if the answer equals a + b, else 0; expected reward estimated as the mean.

The model as a policy

For 5 + 6, which the fine-tuned model never saw in chat form, 32 samples at temperature 1 give 2 different answers and an expected reward of 0.56. For 7 + 8 the policy is confident and right (1.00); for 4 + 5 it is confident and wrong (0.00), and no amount of sampling would produce a reward to learn from. Temperature changes the policy that is sampled: at 0.3 the expected reward on 5 + 6 is 0.91, at 1.5 it is 0.53. Reinforcement learning works on the temperature-1 policy itself, raising the probability of the answers the verifier accepts.

Seen this way, the earlier volumes fit the same frame. Supervised fine-tuning imitates given actions; RLHF (Vol. XVIII) is RL with a learned reward. What is new here is a reward that can be checked — by arithmetic, unit tests or a proof checker — so it cannot be flattered and need not be learned. The policy explores by sampling, and its own successes become the training signal: the model is taught by its best attempts rather than by examples someone wrote.

With a single reward at the end, every token of the episode shares one return; deciding which tokens deserve the credit is the credit-assignment problem, and the subject of plate XIX.8.

Terms

Policy: The model’s distribution over actions — here, next tokens — given the state.
Reward: A number judging an action or an episode; here, 1 for a right answer.
Return: The total reward of an episode.
Episode: One complete response, from prompt to end-of-turn.

A language model is a policy over tokens; with a verifiable reward at the end, its expected return on a prompt is the probability its sampled answer is right.

Sources

Sutton and Barto, Reinforcement Learning: An Introduction, 2nd ed., 2018, chapter 3.
The course’s RL world: js/post.js.