20.2 Multi-Turn RL Formulation

22.1 Overview used the flight-booking example to show the fundamental difference between Agentic RL and single-turn RL. This section turns those differences into precise mathematical objects — adopting the POMDP formulation from the AppWorld paper, which explicitly separates "tokens generated by the model" from "tokens returned by the environment." That separation is the foundation for everything that follows: action masks, step-level advantage, credit assignment.

The Simplified View of Single-Turn RL

The GRPO covered in earlier chapters is fundamentally a degenerate MDP. The model receives a prompt, autoregressively generates a sequence of tokens, and at the end a reward model or verifier hands back a single scalar reward.

• State $s$: the current token context (prompt + tokens generated so far)
• Action $a$: the next token
• Transition: deterministic append — the sampled token gets added to the context
• Reward $r$: given once, after the whole rollout ends

The action at each step is sampled from the LLM's next-token distribution — every token is an independent action. The optimization objective is to maximize the expected reward of the single-turn output:

$${\mathbb{E}}_{a\sim {\pi }_{\theta }}[r(a)]⟶\underset{\theta }{\max }$$

The key assumption behind this view is that every token is generated by the model, so every token participates in the gradient update. That assumption breaks down in multi-turn interaction.

The POMDP of Multi-Turn Interaction

Once the model stops generating in a vacuum and can call tools or change the environment's state at every step, the state space has to expand. Write a trajectory as $x$; the full state is then:

$$z_t=[s_0,c,x_{1:t}]$$

The three pieces mean:

• $s_0$: the hidden initial environment state — the database snapshot in AppWorld, the initial state of a Python REPL, the contents of a filesystem. The model cannot see it directly, only observe it indirectly through tool calls.
• $c$: the task context — the user request, the system prompt, the specifications of the available tools.
• $x_{1:t}$: the complete token history up to the current point — including both the thought/action tokens the model generated and the observation tokens the environment returned.

This is called partially observable (the PO in POMDP) because the model can only see the text history $(c,x_{1:t})$; the hidden environment state $s_0$ and how it evolves over time are invisible to the model. The model can call an API to check a calendar, but it cannot directly "read the entire world state."

Actions: Text Tokens and Tool Calls

At the token level, the model is still doing next-token prediction:

$$x_{t+1}\sim p_{\theta }(\cdot \mid c,x_{1:t})$$

But semantically, the token stream splits into two kinds:

• Plain text tokens (thoughts, pieces of code): these only update the context, and the transition is $[s_0,c,x_{1:t}]\to [s_0,c,x_{1:t+1}]$
• Structured tool-call tokens (like <tool_call>...</tool_call>): these trigger the environment to execute — the environment runs code or calls an API, and appends the result back into the context too: $[s_0,c,x_{1:t}]\to [s_0,c,x_{1:t+1+k}]$

The extra $k$ tokens here are environment observations, not policy actions. The JSON an API returns, for example, influences the model's next decision, but the model did not sample it itself.

Chain Decomposition of Trajectory Probability

To decompose the probability of a full trajectory, the key move is to multiply only over the positions where the model actually generated a token. Let $a(x)$ denote the set of token positions in trajectory $x$ generated by the LLM (the "action token" positions) — the positions where the environment returned an observation are not in this set. The trajectory distribution is:

$${\rho }_{\theta }(x\mid s_0,c)=\mathbb{I}(s_0,x)\prod _{t\in a(x)}{p}_{\theta }(x_t\mid c,x_{1:t-1})$$

Here $\mathbb{I}(s_0,x)$ simply folds the environment dynamics into the formula: given the initial database, the REPL state, and the API calls made so far, what observation the environment returns is determined by the environment (deterministically or stochastically) — it is not something the model freely generates. This formula is the starting point for the action-mask derivation that follows: only tokens in $a(x)$ have a gradient with respect to $\theta$.

The Optimization Objective

The objective is to maximize the expected return given the initial state and task context:

$$\underset{\theta }{\max }{\mathbb{E}}_{(s_0,c)\sim \mathcal{D}}\left[{\mathbb{E}}_{x\sim {\rho }_{\theta }(\cdot \mid s_0,c)}[R(s_0,c,x)]\right]$$

The outer expectation samples tasks from the training set $\mathcal{D}$; the inner expectation samples trajectories generated by the policy ${\rho }_{\theta }$ under a fixed task. The reward $R(s_0,c,x)$ evaluates whether the entire trajectory completed the task — this is the typical form of an ORM (outcome reward model).

Four Kinds of Reward

In practice, reward is never just "was the final answer right." A flight-booking agent has not succeeded just because it says "booked" — the database needs to actually contain a new order that satisfies the constraints, and no other field should have been changed by mistake. XiaoRed5's introductory material splits reward into four categories:

Type	Meaning	Example
Outcome	Whether the final answer is correct, or the final environment state satisfies the task	QA answer matches, AppWorld unit tests pass
Format	Whether the action can be parsed and executed by the environment	JSON arguments are complete, tool name is spelled correctly
Cost	Trajectory length, number of tool calls, API spend	Cap rollouts at 20 steps, penalize repeated searches
Process	Whether an intermediate step actually advances the task	Did the search find valid evidence, did the code pass an intermediate test

Getting started usually means implementing just Outcome + Format, so training can run end to end. Process reward is the topic of the credit-assignment chapter later — it turns a sparse outcome signal dense.

Action Masks: Model Output and Environment Output Must Be Kept Separate

Translating the trajectory probability formula ${\rho }_{\theta }$ above into a loss gives the mathematical basis for the action mask. The policy gradient should only update tokens the model actually generated:

$$\nabla J(\theta )={\mathbb{E}}_{x\sim {\rho }_{\theta }}\left[R(x)\sum _{t\in a(x)}{\nabla }_{\theta }\log p_{\theta }(x_t\mid c,x_{1:t-1})\right]$$

If observation tokens returned by the environment were also allowed to carry gradient, that would be equivalent to making the model "learn to predict the web page content the environment returns" — this contaminates the policy gradient and destabilizes training.

In implementation, the action mask is a 0/1 vector the same length as the trajectory:

# The token sequence of one rollout and its corresponding action mask
# 1 = a token the model generated (participates in the gradient)
# 0 = prompt / tool output / padding (does not participate in the gradient)

# <prompt>       <think>...search</think>  <search>query</search>  <information>...web page content...</information>  <answer>...</answer>
#  0 0 0 0 0      1 1 1 1 1 1 1 1 1      1 1 1 1 1 1 1 1 1      0 0 0 0 0 0 0 0 0 0 0 0 0                  1 1 1 1 1 1 1

Search-R1 is a minimal working example that makes this very clear: it segments the token stream with four kinds of tags — <think> is model reasoning (trained), <search> is a model action (trained), <information> is an observation returned by the retriever (masked, not trained), and <answer> is the final answer (trained). The state_masking=true config flag is exactly this.

Agent-R1 goes further and finds that excluding non-agent tokens entirely is not optimal — you can apply an SFT loss to environment tokens (learning to predict environment behavior), which amounts to learning a policy and a world model at the same time. This line of work has been extended further by Echo, PaW, and other follow-ups.

Step-Level Trajectory Structure

In theory, the trajectory $\tau =(s_0,a_0,s_1,a_1,\dots ,a_T)$ is already sufficient. But in industrial implementations, how the trajectory is stored directly affects training stability and engineering efficiency.

The Problem with Flat Token Sequences

The simplest storage scheme is to flatten the whole trajectory into a single token sequence. This has two problems:

1. Implicit step boundaries: which tokens belong to "the model's output at turn 3" and which belong to "the tool output at turn 3" is determined entirely by special-token splitting, and error handling is easy to get wrong.
2. Retokenization drift: during rollout the model generates in token space, but storage often parses this into a message list, and training re-tokenizes the messages again. Tokenization is not a reversible operation — the same text can map to different token sequences, which leaves the training data inconsistent with what happened during rollout.

Step-Level Records (Agent-R1 Style)

Store the trajectory as a structured step-level record, explicitly saving, for each step:

@dataclass
class Step:
    state_before: str          # context at the start of this step
    action_tokens: List[int]   # raw token ids generated by the model (never re-tokenized)
    observation: str           # tool output (if any)
    reward: float              # reward for this step (nonzero under process reward)
    is_terminal: bool          # whether this is the final step

This gives three benefits: precise step boundaries, no retokenization drift, and flexible context-management strategies (append-only, sliding-window, LLM summarization, selective retention). Agent-R1's experiments show that sliding-window outperforms append-only on GSM8K — "less is more": the model does not need to see the entire history to make good decisions.

Comparison with Single-Turn RL

	Single-Turn RL (GRPO)	Multi-Turn Agentic RL
State	prompt + tokens generated so far	$[s_0,c,x_{1:t}]$ — hidden environment state + task context + token history
Action	plain text token	text token + structured tool call (both end up as tokens, but different semantics)
Transition	deterministic append	text tokens append deterministically; tool calls trigger environment dynamics (possibly stochastic)
Observation	not distinguished (all model-generated)	must explicitly distinguish observation tokens from action tokens
Reward	single-step scalar $r(a)$	four categories: Outcome / Format / Cost / Process
Objective	${\mathbb{E}}_{a\sim {\pi }_{\theta }}[r(a)]$	${\mathbb{E}}_{(s_0,c)}[{\mathbb{E}}_x[R(s_0,c,x)]]$
Rollout latency	hundreds of milliseconds	seconds to minutes (dominated by environment latency)
Training representation	token sequence	step-level structured record (Agent-R1)

Section Summary

This section built the formal skeleton of Agentic RL. Adopting the POMDP formulation from the AppWorld paper, the state splits into three parts, $[s_0,c,x_{1:t}]$, which is the mathematical foundation for everything that follows. The chain decomposition of trajectory probability, ${\rho }_{\theta }(x|s_0,c)=\mathbb{I}(s_0,x){\prod }_{t\in a(x)}{p}_{\theta }(x_t|\cdot )$, directly implies the action mask — only action tokens the model generated participate in the policy gradient. The step-level trajectory structure (Agent-R1) solves the retokenization drift and context-management problems that come up in industrial implementations.

The central question that comes next: the trajectory probability formula gives $R(x)$ as a trajectory-level scalar — how does that get broken back down into a per-step advantage? The formulation tells us that gradients should only be computed on action tokens, but it does not tell us how much advantage each action token should be multiplied by. That is credit assignment — 22.3 Trajectory Credit Assignment.