Skip to main content
Chapter 3 built a policy that maps a token sequence to next-token probabilities. Chapter 4 built an environment that maps action text to observation text. A trainer needs neither text nor turns. It needs tensors: one token sequence per episode, the log-probability with which each policy token was sampled, and a mask that says which tokens the policy is responsible for. This chapter builds that bridge. The bridge has more failure modes than its size suggests. A chat template decides which tokens the model sees, and it may render the history differently from the way the model generated it. Decoding sampled tokens to text and encoding the text again can produce different tokens, so a trainer that works from text computes probabilities of tokens the model never sampled. And one off-by-one error in a mask silently trains the policy to imitate the environment. We trace each of these on the real Qwen3-0.6B tokenizer and model, then build TinyAgent’s tokenizer, the Trajectory record, a rollout that records exactly what was sampled, and a collate that turns a group of rollouts into batch tensors. Every training chapter from here on consumes the output of collate.
The saved code of this chapter calls functions from earlier chapters through arl. as well, so the import is part of the saved block. The remaining imports are for this chapter’s demonstrations only.

​
5.1 One sequence per episode

A multi-turn episode alternates between text written by different parties. In the clamp episode of Chapter 4, the issue comes from the task, the calls come from the policy, and the source listing and test results come from tools. To score the policy’s calls, the trainer concatenates everything into one token sequence and runs one causal forward pass over it, so that each policy token is conditioned on everything before it, including earlier observations. Every token in that sequence has an owner, the party whose output it is:
  • template tokens are inserted by the formatting code: role headers, delimiters, and the begin-of-sequence token.
  • prompt tokens are the task input: the system message, the issue text, and the tool descriptions.
  • policy tokens were sampled from the model during the rollout.
  • tool tokens are observations returned by the environment.
Only policy tokens are actions. The action mask is 1 on policy tokens and 0 elsewhere. It decides which log-probabilities enter the loss. Every other token still conditions the model through attention, but the model is not trained to produce it. Training on tool tokens would teach the policy to predict, and later to hallucinate, tool output. Training on template tokens is usually harmless but wastes gradient on text that the model never chooses. The mask must be aligned with the log-probabilities it selects. sequence_log_probs (Section 3.4) returns [B, T-1] values for the predicted positions ids[:, 1:], so the action mask also has shape [B, T-1], and column jj answers the question “was token j+1j+1 sampled by the policy?”. The same holds for the stored sampling log-probabilities. Section 5.5 builds that alignment into the data structure, so that no later chapter has to remember the shift.

​
5.2 A real chat template with tools

Production models do not see raw turns. They see the output of a chat template, a program (usually Jinja) shipped with the tokenizer that renders a list of messages, and optionally a list of tool schemas, into a single string with special tokens. To know which tokens a model sees and which it generates, we have to run the actual template. We render a two-turn clamp episode with Qwen3-0.6B’s template. The tool schemas are the MCP-style schemas from Section 4.7, converted to the function-calling format that the template expects. The assistant messages are hand-authored: no model has generated them yet.
Output:
This is the complete conversation part of the rendered string. The system part before it, which lists the tool schemas, is shown in the next block. Several template decisions are visible. Tool results are wrapped in <tool_response> tags and placed in a turn with the user role, so the model sees observations as user input. Each assistant turn renders its reasoning inside <think> tags, followed by the call as JSON inside <tool_call> tags, and ends with <|im_end|>. The rendering ends with an open assistant header, the generation prompt, after which the model writes its next turn.
Output:
The four tool schemas cost a fixed preamble of a few hundred tokens in every episode, before the agent has done anything. That is about two thirds of this short episode. Every delimiter the template uses is a single token, so a policy can learn to emit <tool_call> as one decision, and a parser can find tool calls by token id instead of by string matching. The template’s output has to be split by owner. The tokenizer reports, for each token, the character span it covers (offset_mapping), so we label characters first, then give each token the label of its first character. User content is prompt, tool-response content is tool output, and each assistant turn from just after its header through its <|im_end|> is policy output, because in deployment the model generates all of it, including <think> and the closing <|im_end|>.
The function skips user turns that are really tool results: the tool span inside them is labeled tool, and the wrapper tags stay template. We count the tokens per owner and decode the policy tokens of the first assistant turn.
Output:
Most of this rendering is template: in this labeling, the system preamble counts as template because the template generates it from the schemas. The issue is 37 prompt tokens, the two assistant turns are 55 policy tokens, and the two tool results contribute only 23. In real SWE episodes the proportions reverse: file listings and test logs dominate the context, and policy tokens are often a small minority. Either way, the action mask selects a small and irregular subset of positions. Figure 5.1: The Qwen3-0.6B rendering of the two-turn clamp episode, token by token, outlined by owner. The system preamble with the tool schemas is collapsed into one box. Figure 5.1 shows the same tokens in order. Follow the outlines: dotted prompt cells for the issue, blue policy turns that each begin with <think> and end with <|im_end|>, solid-outlined tool spans enclosed in unboxed template wrappers, and the open assistant header at the end. The loss uses only the blue cells. Attention uses all of them. The gym’s tools are an abstraction (Section 4.7). A production agent makes the same repair with Read, Edit, and Bash, writes the replacement text itself, and ends by answering without a tool call. Rendering that episode through the same template shows what the abstraction removes from the token budget. The assistant turns are again hand-authored.
Output:
Compare these counts with the gym episode above. Six general tool schemas make the fixed preamble longer: 585 template tokens instead of 350, paid in every episode. The policy writes 107 tokens instead of 55, even without any <think> reasoning, and the single Edit call is 45 of them. Most of those 45 tokens are not free choices. The old_string must reproduce the source exactly, character for character, or the tool returns an error, while the gym’s edit call names an op in a handful of tokens. That is the trade the BugFix gym makes for TinyAgent. A pretrained model that already copies and writes code accurately can explore the production surface, so its RL signal is about which edit to write. A model trained from scratch would first have to learn to copy text, and the reward for a fix would be spread across dozens of tokens that must all be right.

​
5.3 What the template does to history

A template renders the whole history every time. That creates a problem. The text the model sees for an earlier turn at step kk may differ from the text it generated at that turn. The Qwen3 model card states that “the historical model output should only include the final output part”, meaning that earlier reasoning should not be carried forward, and its template implements this (Qwen3-0.6B model card, inspected 2026-10-13). We check exactly when. The test is prefix stability. Render the conversation up to a turn with the generation prompt, which is exactly what the model saw when it generated that turn. Then render a longer conversation and ask whether the first rendering is a prefix of the second. If it is, a single forward pass over the final rendering reproduces the context of every earlier turn.
Output:
Within one user request, the history is stable: reasoning from earlier tool-calling turns is kept, and every earlier rendering is a prefix of the later one. As soon as a new user message arrives, the template drops the reasoning from every assistant turn before it, and earlier renderings stop being prefixes. A trainer that re-renders such a conversation and scores the earlier assistant turns computes their log-probabilities in a context the model never saw while sampling them, without their reasoning. The template also has a switch for non-thinking mode, which changes the generation prompt itself:
Output:
In non-thinking mode, the template itself writes an empty <think>\n\n</think>\n\n block after the assistant header. Those tokens are template-owned, so they belong in the context but not in the action mask, even though in thinking mode the same tokens would be policy-owned. Ownership is a property of who produced a token in this particular rollout, not of the token’s text.

​
5.4 Retokenization drift

Retokenization drift is the most consequential hazard in multi-turn RL infrastructure, because it fails silently. A rollout samples token ids. The environment needs text, so the ids are decoded. A common and wrong design keeps only the text, and when the trainer needs token ids later, it re-encodes the text, often after re-rendering it through the template. Decoding and then re-encoding does not reproduce the sampled ids in general. BPE tokenizers have many token sequences for the same string. The sampler can produce any of them, but encoding always returns one canonical segmentation. When the two differ, the trainer computes log-probabilities, ratios, and gradients for a sequence the model never sampled. The first way this happens is through the template. A model may write a tool call with compact JSON. A system that parses the call into a message and renders it again through the template reformats it, as the following hand-authored example shows.
Output:
The template serializes the arguments with a space after every colon, so the 19 sampled tokens become 21 different tokens. Every log-probability from the first differing token onward refers to a sequence the model did not produce. The second way needs no template at all. We sample eight 128-token continuations from the real Qwen3-0.6B at temperature 1, with no truncation, decode each to text, and encode the text again.
Output:
In this run, one of the eight samples drifts. The model emitted - and hi as two tokens. The decoded text contains -hi, and the tokenizer’s canonical segmentation of that text is -h followed by i. The decoded strings are identical, and the token sequences are not. Drift is rare per token, but an RL batch contains millions of tokens, and long agent trajectories make every episode a candidate. Worse, drift concentrates on unusual segmentations, which are exactly the low-probability samples whose log-probabilities matter most to the gradient. The fix is a discipline rather than an algorithm, and we call it token-in/token-out. The rollout keeps the token ids that were sampled, appends the environment’s observation as token ids, and never re-encodes text that the policy produced. Text exists only at the boundary, as the decoded action given to the environment’s parser and the observation string it returns. The trainer receives ids, masks, and sampling log-probabilities, never text. Everything we build from here follows that rule.
In production. Systems that serve the model through an OpenAI-style chat API receive text back and are exposed to both kinds of drift unless the engine also returns token ids. This is a general engineering observation; we have not audited how specific frameworks handle it. The book’s recommendation is to treat any path that re-encodes generated text as a bug, and to test for it by comparing the trainer’s recomputed log-probabilities with the sampler’s at step 0 (Section 3.7).

​
5.5 TinyAgent’s tokenizer

TinyAgent needs a tokenizer whose vocabulary covers every string the gym can produce, and whose sequences stay short. A word-level tokenizer fits. All gym text is already whitespace-separated words (the read tool returns the compact projection of Section 4.4), so a token is a word, a newline is the special word <nl>, and decoding joins words with spaces. A word-level tokenizer has exactly one segmentation for each string, so it cannot drift. That is a convenience of the toy, and the real-model lesson of Section 5.4 still governs the design.
The vocabulary is built by enumerating gym text, not by guessing. For every task we collect the issue, the read and test observations of the buggy source, and the same observations after each edit op, plus every well-formed call, every fixed message, and the integers from -20 to 150, which cover the values that tests and wrong edits produce. Pass all tasks, train and held-out, so that held-out function names have their own ids. They stay untrained until a held-out task is seen, which is the word-level analog of a rare token.
Output:
The vocabulary has 316 entries, so TinyAgent’s output layer is about 316 by 64. The read observation of clamp is 23 tokens long, including one <nl>, and decodes back exactly. A call with an argument is four tokens, and the special tokens have fixed ids 0 to 7 (<call> is 5 and </call> is 6), which the rollout relies on. Words outside the vocabulary map to <unk>. The next block checks that the gym never produces one, by running every scripted policy of Chapter 4 on every task and encoding every observation.
Output:
All observation tokens of all four policies on all 30 tasks are in the vocabulary. The integer range is the one assumption a long episode of unusual edits could break, for example by applying fix_range_end repeatedly to a sum. Such a value would encode as <unk>, which is harmless for training because observations are never actions.

​
5.6 Trajectories and masks

A trajectory is the token-level record of one episode. It stores the ids of the whole sequence, and, aligned to the predicted positions ids[1:], the action mask and the sampling log-probabilities. It also stores the reward and the environment’s final info. Two extra fields, aligned to ids, record each token’s owner and the turn it belongs to. Chapters 7 and 12 use the turn index to assign credit per turn.
An Episode builds a trajectory incrementally, in the token-in/token-out style of Section 5.4. It starts with <bos> and the encoded issue. act appends the policy’s call ids as they were sampled, decodes them only to hand the text to env.step, and appends the encoded observation between <obs> and </obs> template tokens. It also enforces the two limits that the environment does not know about: the rollout’s turn limit, and the model’s context length max_len. If the next observation and another call would not fit, the episode ends as truncated, with reward 0. The observation after the final call is not appended, because no further action depends on it.
make_trajectory converts the lists into tensors. The alignment happens here, in one place: the action mask is “the owner of token j+1j+1 is the policy”, and the sampling log-probabilities drop the first entry, which belongs to <bos> and is never predicted.
scripted_rollout runs a Chapter 4 text policy through the same path. Its trajectories are valid demonstrations for supervised fine-tuning (Chapter 15). Their sample_logp is 0, which is the correct log-probability for a deterministic policy such as the oracle and a placeholder otherwise, so it must not be used in an importance ratio. We encode the oracle’s clamp episode and print the first sixteen predicted positions.
Output:
The issue occupies columns 0 to 5, so the first action token, <call>, is at column 6, and the first observation begins after </call>. Each column describes the token that the logits at that position predict, which is the convention sequence_log_probs uses. The oracle’s episode is 76 tokens long, of which 13 are actions: four calls, three of three tokens and the edit of four. Figure 5.2: Token ownership in the oracle's clamp trajectory, with the action mask below. Only policy tokens (blue) have mask 1. Figure 5.2 draws the whole trajectory. The blue policy spans are short and regular, and the solid-outlined tool spans vary from one token (ok) to about twenty (the source listing). The blue rule under the strip, the action mask, marks exactly the policy cells. That is the invariant every later loss relies on: gradients flow only through tokens the policy sampled. We check the invariant and the lengths across all four scripted policies.
Output:
Episodes are short: the oracle uses about 70 tokens, and even the random policy, which often reads the file several times, stays well below 200. Action tokens are a fifth to two fifths of each sequence. These lengths keep TinyAgent’s training fast, and they fit comfortably in its 256-position context.

​
5.7 Sampling rollouts

Now the policy writes the calls. generate_call samples one call token by token with the KV-free loop of Section 3.5, stopping at </call> or after MAX_CALL_TOKENS. It records each token’s log-probability under the sampling distribution. rollout repeats that until the episode ends. Both work with any model that maps ids [B, T] to logits [B, T, V].
What does a freshly initialized TinyAgent do in the gym? Chapter 3 found its next-token distribution is close to uniform over the vocabulary, so we expect it to write nonsense. We roll out one episode and print it.
Output:
The untrained policy emits six random words per turn, never closes a call, and receives error : malformed call six times until the budget ends the episode. We measure the same behavior over 64 episodes, and check that the stored sampling log-probabilities equal what a trainer recomputes from the finished sequences.
Output:
All 384 calls of the 64 episodes are malformed, and the reward is 0 everywhere. A policy-gradient method learns from differences in reward between samples, and here there are none: every advantage is zero and the gradient is exactly zero. This is the motivation for the rest of the book in one number. RL needs a policy that already succeeds sometimes. Chapter 6 starts from a TinyAgent that first imitates well-formed scripted calls, the toy counterpart of starting RL from a pretrained model, and Chapter 15 treats supervised warm starts in depth. The last line confirms token-in/token-out: on action tokens, the trainer’s recomputation matches the sampler to fp32 rounding. The single-episode rollout runs one forward pass per generated token. RL samples many episodes per step, so we also save a batched rollout that advances many episodes in lockstep. Each step pads the active sequences on the right, runs one forward pass, and reads each row’s logits at its own last position. Right padding is safe under a causal mask, because padding positions come after the position we read. An episode that finishes a call steps its environment, and an episode that ends leaves the batch.
The batched rollout samples from the same policy, but it consumes the random stream in a different order, so its episodes differ from the sequential ones for the same seed. We compare speed and check the log-probabilities again.
Output:
On the check host, 64 sequential rollouts of the untrained model took about 1.5 seconds and the batched version about 0.2 seconds while this chapter was written. Those timings vary by machine and are not checked by the build. The batched rollout’s log-probabilities agree with recomputation as well. The tiny residual comes from running the same weights on differently shaped batches, the toy version of the batch-size effects that Chapter 13 measures in inference engines. The last step turns a list of trajectories into batch tensors. collate right-pads every field to the longest episode: ids with the pad id, masks with 0, log-probabilities with 0, and turn indices with -1. attention_mask marks real tokens. Under a causal model, right padding never affects real positions, so TinyAgent ignores it, but losses and metrics need it. Every [B, T-1] field is aligned with sequence_log_probs.
A GRPO step (Chapter 9) samples a group of GG episodes on the same task and compares their rewards. We collate a group of four oracle demonstrations of different tasks with four random-model rollouts, so the lengths differ, and verify that the masked log-probabilities match a batched recomputation.
Output:
The batch has the shapes the training chapters expect: [8, 76] for ids and attention, and [8, 75] for the three predicted-position fields. The oracle rows have reward 1 and the random rows 0. The oracle rows’ sample_logp is the scripted placeholder, so the log-probability check uses only the sampled rows, where recomputation over the padded batch still matches. No padding position has action mask 1.

​
5.8 Design problems in practice

Problem: context growth across turns. Symptom: step time and memory grow over training while the number of episodes is unchanged, and late turns of long episodes are slow to sample. Cause: every turn appends the full observation to the context. The sequence length is the sum of all turns, attention cost grows quadratically with it, and without a KV cache sampling cost grows quadratically too (Section 3.5). Agent RL tends to lengthen episodes as policies learn to gather more information. Options: (a) cap observations (Section 4.8), which loses information. (b) Cap turns, which changes the task. (c) Summarize or drop old observations, which changes the policy’s input distribution (Chapter 12). (d) Pay for it with token-budget batching (Chapter 14). Default: log the length per turn from the start, and set the context limit from measured lengths, not guesses. The next block shows how the context grows turn by turn in the oracle’s episode.
Output:
Each turn adds its call and its observation, so the context grows from a prompt of 7 tokens to 76. Without a cache, scoring each new token reprocesses everything before it.
Problem: truncation policy shapes the training distribution. Symptom: a growing fraction of episodes ends with stop='truncated', and the policy learns to be terse or to submit early. Cause: a truncated episode has no outcome, and the rollout assigns it reward 0. If long episodes are systematically truncated, RL learns that length itself is bad, whatever the episode was doing. Options: (a) reward 0, which is simple and biased toward short episodes. (b) Drop truncated episodes from the batch, which is unbiased for the episodes kept but filters the data toward short ones. (c) Bootstrap a value estimate at the cut (Chapter 7), which needs a critic. Default: track the truncation rate as a first-class metric, keep it near zero by sizing the context, and treat a rising rate as a bug. The next block forces truncation with a model whose context is only 48 positions.
Output:
Every episode stops as truncated before the turn budget runs out, and none exceeds the context.
Problem: template changes between SFT and RL. Symptom: reward collapses in the first RL steps after a successful SFT stage, or the step-0 importance ratio is far from 1. Cause: the SFT data was rendered with one template (or one thinking mode, Section 5.3) and RL rollouts with another, so the RL policy sees prefixes that differ from those it was trained on. Template details that look cosmetic, such as JSON spacing or an empty <think> block, are distribution shift at the token level. Options: (a) pin the tokenizer and template versions in the run configuration, which is cheap. (b) Re-render SFT data with the RL template, which costs a data pass. (c) Add a check that renders a fixed conversation with both pipelines and compares token ids. Default: (a) and (c). A one-line diff of rendered ids catches these errors before any training.
Problem: stripping reasoning from history. Symptom: multi-turn runs where the step-0 ratio is fine for the last turn and not for earlier turns, or where learning on long conversations is unstable. Cause: templates that remove earlier reasoning (Section 5.3) make the final rendering differ from the context each turn was sampled in. Scoring all turns in one forward pass over the final rendering evaluates earlier turns in the wrong context. Options: (a) store each turn’s tokens and context as sampled and score each turn against its own context, which costs one forward pass per turn. (b) Keep the reasoning in history during training, which makes the training context differ from deployment. (c) Train only on the last turn, which discards data. Default: score turns against the context they were sampled in. Within one tool-calling request, Qwen3’s template is prefix-stable, so a single pass is correct there.
Problem: mask bugs that silently train on tool output. Symptom: the policy starts to write text that looks like observations, such as test results or file contents, after its calls. The loss looks normal. Cause: the action mask is misaligned with the log-probabilities, usually because one of them was shifted and the other was not. Options: (a) build masks from per-token ownership in one place (as make_trajectory does), which is cheap. (b) Assert in the trainer that every masked token was sampled, using the ownership tensor. Default: both. The next block shows what an unshifted mask selects.
Output:
The correctly shifted mask selects the 13 policy tokens and nothing else. The unshifted mask selects 12 positions, and 3 of them are wrong: the first token of every call (<call>) drops out of the loss, the final </call> falls off the end, and for each of the first three calls the slot after </call> now trains the <obs> template token. Nothing crashes, and the policy is trained to predict that an observation begins after its own call. With longer observations and a mask built from tool spans, the same off-by-one trains on tool output.
Problem: stop-token handling. Symptom: calls that never end, episodes that always hit the call-length limit, or a model that keeps generating after its tool call. Cause: the stop token (</call> here, <|im_end|> for Qwen3) is generated by the policy and must be in the action mask. If the rollout stops on it but does not record it, or records it with mask 0, the policy never receives a gradient for stopping. Options: (a) always include the sampled stop token in the trajectory with its log-probability, as generate_call does. (b) Inject the stop token as a template token when the call-length limit is hit, which ends the call cleanly but adds tokens the policy did not sample. Default: (a). If a limit cuts a call off, keep it unterminated and let the parser reject it, as the BugFix gym does, so the policy learns that unterminated calls fail.

​
Summary

  • A multi-turn episode becomes one token sequence whose tokens are owned by the template, the prompt, the policy, or tools. The action mask selects policy tokens. It is aligned with the predicted positions ids[:, 1:], as are the stored sampling log-probabilities.
  • Qwen3-0.6B’s template renders tool schemas as a system preamble of a few hundred tokens, observations as user turns in <tool_response> tags, and calls as JSON in <tool_call> tags. It keeps reasoning within a tool-calling request and drops it once a new user message arrives, which breaks prefix stability.
  • Decoding and re-encoding generated tokens can change them. Measured here: one of eight 128-token Qwen3 samples drifted, and re-rendering a compact JSON call changed 19 tokens into 21. Rollouts must be token-in/token-out.
  • TinyAgent’s word-level tokenizer has 316 entries and covers all gym text. Episodes are about 70 tokens for the oracle and stay under 200 for every scripted policy.
  • rollout and rollout_batch record sampled ids, sampling log-probabilities, and masks. collate pads them into ids, attention_mask [B, T], action_mask, sample_logp, turn_ids [B, T-1], and rewards [B]. An untrained TinyAgent emits only malformed calls and earns reward 0, so RL needs a starting policy that sometimes succeeds.

​
Exercises

  1. Render the episode of Section 5.2 with enable_thinking=False and without reasoning_content. Count tokens per owner. Which template tokens would wrongly enter the action mask if you labeled everything after the assistant header as policy?
  2. Repeat the drift experiment of Section 5.4 with 32 samples and temperatures 0.6 and 1.0. Does lowering the temperature reduce drift? Explain in terms of which tokens drift.
  3. Change an assumption: let the rollout append the observation after the final call as well. Which fields of the trajectory change, and does any loss in Part II change its value? What changes in the cost?
  4. Write assert_token_in_token_out(traj, tok) that fails if any action span, decoded and re-encoded, differs from the stored ids. Why can it never fail for TinyAgent’s tokenizer, and what would you check for Qwen3 instead?
  5. Modify collate to pad on the left instead of the right. What else must change for TinyAgent to compute the same log-probabilities (consider positions and the causal mask)? Why do inference engines often prefer left padding for generation?
  6. Design review: your inference service returns completions as text through a chat API, and the trainer re-tokenizes them. A colleague argues that drift is too rare to matter. Design a measurement that bounds its effect on the gradient, and propose the minimum change to the serving interface that removes the problem.