19.2 Multi-Turn Reinforcement Learning

19.1 Overview used the example of booking a flight to illustrate the fundamental differences between Agentic RL and single-turn RL. This section formalizes these differences into precise mathematical objects—using the POMDP formulation from the AppWorld paper, which explicitly distinguishes between "model-generated tokens" and "environment-return tokens," forming the basis for subsequent discussions on action masking, step-level advantage, and credit assignment.

Simplified Perspective of Single-Turn RL

The previous chapters' GRPO is essentially a degenerate MDP. The model receives a prompt, autoregressively generates a token sequence, and finally receives a scalar reward from a reward model or verifier.

• State $s$: The current token context (prompt + generated tokens)
• Action $a$: The next token
• Transition: Deterministic append—adding the selected token to the context
• Reward $r$: Given once after the entire rollout

Each action is sampled from the LLM's next-token distribution—each 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 of this perspective is: All tokens are generated by the model, so they all participate in gradient updates. This assumption no longer holds in multi-turn interactions.

POMDP with Multi-turn Interaction

When the model is no longer confined to generating within a closed environment, but instead can invoke tools and alter the environment state at each step, the state space must be expanded. Let a trajectory be denoted as $x$, and the full state is written as:

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

The three components correspond to:

• $s_0$: Hidden initial environment state — a snapshot of the database in AppWorld, the initial state of the Python REPL, or the contents of the file system. The model cannot see this directly and can only observe it indirectly through tool invocations.
• $c$: Task context — the user request, the system prompt, and the specifications of the available tools.
• $x_{1:t}$: Complete token history up to the current step — this includes both the model's generated thought/action tokens and the observation tokens returned by the environment.

The term partially observable (PO in POMDP) arises because the model can only observe the text history $(c,x_{1:t})$; the hidden environment state $s_0$ and its evolution over time are completely invisible to the model. The model can invoke an API to check the calendar, but it cannot directly "read the entire world state."

Actions: Text Tokens and Tool Invocations

At the token level, the model still performs next-token prediction:

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

Semantically, the token stream is divided into two categories:

• Ordinary text tokens (thoughts, parts of code): These only update the context, and the transition is written as $[s_0,c,x_{1:t}]\to [s_0,c,x_{1:t+1}]$
• Structured tool invocation tokens (e.g., <tool_call>...</tool_call>): These trigger the environment to execute, run code or an API, and append the returned results into the context: $[s_0,c,x_{1:t}]\to [s_0,c,x_{1:t+1+k}]$

The additional $k$ tokens are environment observations, not policy actions. For example, the JSON returned by an API influences the model's next decision, but it is not sampled by the model itself.

Chain Rule Decomposition of Trajectory Probability

Breaking down the probability of a complete trajectory involves multiplying only at the positions where the LLM generates tokens. Let $a(x)$ denote the set of token positions in trajectory $x$ generated by the LLM (i.e., "action tokens"), while the observation positions returned by the environment are not included in this set. The trajectory distribution can be written as:

$${\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)$ incorporates the environment dynamics into the formula: given the initial database, REPL state, and previous API calls, the observation returned by the environment is determined by the environment (deterministic or stochastic), and is not freely generated by the model. This formula serves as the starting point for deriving action masks—only the tokens in $a(x)$ contribute to the gradient with respect to $\theta$.

Optimization Objective

Maximize the expected return given an 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}$, while the inner expectation samples trajectories generated by the policy ${\rho }_{\theta }$ for a fixed task. The reward $R(s_0,c,x)$ evaluates whether the entire trajectory completes the task—this is a typical form of outcome reward (ORM).

Four Types of Reward

In practical engineering, the reward is not merely "whether the final answer is correct." For a ticket booking agent, success is not simply saying "booked," but rather that the database actually contains a new order that meets the constraints, and no other fields have been mistakenly modified. XiaoRed5's introductory materials categorize rewards into four types:

Type	Meaning	Example
Outcome	Whether the final answer is correct, or the final environment state satisfies the task	QA task answer matches, AppWorld unit test passes
Format	Whether the action can be parsed and executed by the environment	JSON parameters are complete, tool name is spelled correctly
Cost	Trajectory length, number of tool calls, API cost	Limit rollout to no more than 20 steps, penalize repeated searches
Process	Whether intermediate steps genuinely advance the task	Whether search finds valid evidence, whether code passes intermediate tests

When starting out, it is common to first implement the Outcome and Format types, ensuring that training can proceed. The Process reward is a topic covered in later chapters on credit assignment, as it transforms sparse outcome signals into dense ones.

Action Mask: Model Generation and Environment Return Must Be Distinguished

Translating the trajectory probability formula ${\rho }_{\theta }$ into a loss function yields the action mask's mathematical foundation. The policy gradient should only update the tokens generated by the model:

$$\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 the observation tokens returned by the environment are also included in the gradient computation, it is equivalent to letting the model "learn to predict the environment's returned web content"—this contaminates the policy gradient and leads to unstable training.

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

# A sequence of tokens from a single rollout and the corresponding action mask
# 1 = token generated by the model (participates in gradient computation)
# 0 = prompt / tool return / padding (does not participate in gradient computation)

# <prompt>        ... thinking ... search
#  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

The minimal runnable example Search-R1 makes this point very clearly: it segments the token stream into four categories using four types of labels—<think> is the model's reasoning (participates in training), <search> is the model's action (participates in training), <information> is the returned observation from the retriever (masked out of training), and <answer> is the final answer (participates in training). The configuration option state_masking=true implements this.

Agent-R1 further discovers that completely excluding non-agent tokens is not optimal—it is possible to apply SFT loss on environment tokens (learning to predict environment behavior), which is equivalent to learning the policy and the world model simultaneously. This approach is further developed by subsequent works such as Echo and PaW.

Step-Level Trajectory Structure

The theoretical trajectory $\tau =(s_0,a_0,s_1,a_1,\dots ,a_T)$ is sufficient for theory, but how trajectories are stored in practice directly affects training stability and engineering efficiency.

Issues with Flat Token Sequence

The simplest way to store a trajectory is to flatten it into a single token sequence. However, this approach has two major issues:

1. Implicit Step Boundaries: It is unclear which tokens belong to "the model's output in the third turn" and which belong to "the tool's response in the third turn"—this is typically split using special tokens, which can lead to bugs in error handling.
2. Retokenization Drift: During rollout, the model generates in the token space, but when stored, it is often parsed into a message list. During training, these messages are then re-tokenized. Tokenization is not a reversible operation—different token sequences can correspond to the same text, leading to inconsistencies between training data and rollout.

Step-Level Recording (Agent-R1 Style)

Trajectories can be stored as structured step-level records, explicitly saving each step:

@dataclass
class Step:
    state_before: str          # The context at the beginning of the step
    action_tokens: List[int]   # The original token IDs generated by the model (without re-tokenizing)
    observation: str           # The observation returned by the tool (if any)
    reward: float              # The reward for this step (non-zero during reward processing)
    is_terminal: bool          # Whether this is the last step

This approach offers three benefits: precise step boundaries, no retokenization drift, and flexible context management strategies (append-only, sliding-window, LLM summarization, selective retention). Experiments with Agent-R1 show that sliding-window performs better on GSM8K than append-only—"less is more," as the model does not need to see all history to make good decisions.

Comparison with Single-turn RL

	Single-turn RL (GRPO)	Multi-turn Agentic RL
State	prompt + generated tokens	$[s_0,c,x_{1:t}]$ — hidden environment state + task context + token history
Action	plain text token	Text token + structured tool calls (ultimately all tokens, but with different semantics)
Transition	Deterministic append	Deterministic append of text tokens; 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
Optimization Objective	${\mathbb{E}}_{a\sim {\pi }_{\theta }}[r(a)]$	${\mathbb{E}}_{(s_0,c)}[{\mathbb{E}}_x[R(s_0,c,x)]]$
Rollout Cycle	Hundreds of milliseconds	Seconds to minutes (dominated by environment latency)
Training Representation	Token sequence	Step-level structured records (Agent-R1)

Summary of This Section

This section establishes the formal framework for Agentic RL. Drawing on the POMDP formalization from the AppWorld paper, the state is decomposed into three parts: $[s_0,c,x_{1:t}]$, which serves as the mathematical foundation for subsequent discussions. The chain rule 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 leads to the action mask—only the action tokens generated by the model participate in the policy gradient. The step-level trajectory structure (Agent-R1) addresses practical issues in industrial implementations, such as retokenization drift and context management.

The next core question is: How does the trajectory-level scalar $R(x)$ in the trajectory probability formula get decomposed into advantages at each step? The formalization tells us "gradients should only be computed on action tokens," but it does not specify "how much advantage each action token should be multiplied by." This is the problem of credit assignment—see 19.3 Trajectory Credit Assignment.