Training AI agents to use tools through multi-turn interactions presents a deceptively tricky problem. A single decision in a long sequence can make or break the entire task, yet when training uses trajectory-level rewards, it is unclear which decisions actually deserve gradient updates. Worse, reward variation across rollouts can be misleading: a model might produce the same action in two different runs and get different outcomes, but that difference might reflect randomness in later interactions rather than the quality of the action itself. You cannot tell whether a call would benefit from training just by observing whether the overall trajectory succeeded or failed.
Critical-State RL: Diagnosing Trainable States for Multi-Turn Tool Use, by Zixiang Chen and colleagues at Salesforce AI Research, addresses this diagnostic problem directly. The paper introduces a method that identifies which specific model call in a multi-turn interaction is worth training, and which is not, before any training takes place. On the Berkeley Function Calling Leaderboard (BFCL) v4, the method achieves a 14 percentage point improvement on missing-function tasks by training precisely the right response, while training alternative responses leaves performance flat or even degrades it.
Why reward variation alone is insufficient
In multi-turn tool use, a trajectory-level reward aggregates the outcome of every step into a single signal. When a model fails, the failure could stem from any turn: a premature tool call, a missed argument, or a wrong recovery strategy after a tool becomes available. The standard approach of training on all turns with trajectory-level reward conflates these distinct failure modes, applying gradient updates uniformly even when only one turn is actually responsible.
The deeper problem is statistical. Consider a scenario where a required tool is unavailable at one point in the conversation but becomes available later. The model must decide whether to issue a read-only query or wait, then recover once the tool reappears. The reward depends on both of these decisions, but they are entangled with downstream randomness: the same decision before tool availability might be followed by success or failure depending on what happens after. The variance in observed rewards mixes the action's true effect on success with the noise introduced by subsequent interactions.
This means a mixed-reward group, where some rollouts succeed and others fail, does not automatically indicate that a given call is trainable. The paper demonstrates this with a concrete example: in a pilot study on no-think Gemma, a mixed-reward prompt group showed reward swings of plus or minus 1 with a within-group standard deviation of 0.477, yet every surviving rollout produced a clean refusal with no tool call. The reward variation did not distinguish between different decision behaviors; it was noise.
The Critical-State RL diagnostic
Critical-State RL introduces a training-free diagnostic that separates action-dependent reward variation from continuation noise before training begins. The method operates on a simple insight: if changing the current action changes the expected reward, then the action is trainable; if the reward variation is dominated by what happens after the action, training that call will not help.
The diagnostic works through nested within-prefix sampling. Given a set of task-defined candidate calls, the method samples many policy prefixes, draws alternative actions at each fixed prefix, and then holds each action fixed while resampling only the reward-only continuation. This separation is crucial: by freezing the prefix and action while resampling what comes after, the method isolates the action's own contribution to reward variance from downstream randomness.
The diagnostic evaluates each candidate call against three conditions. The first, action-sufficiency, requires that the reward mediates the relationship between the action and the benchmark return, which the task structure supplies. The second, headroom, requires that some action's label mean exceeds the reference policy's mean, indicating room for improvement. The third, trainability, requires that action-conditioned label means actually vary under the base policy, confirming that the action's effect is not deterministic.
Under a small KL budget, the method provides a theoretical guarantee: action-dependent variance governs the best unrestricted local-label improvement to leading order at a fixed prefix and continuation kernel. This means the diagnostic's ranking of candidates by action-dependent variance is not merely heuristic; it is grounded in the fundamental limits of what local training can achieve.
The BFCL four-cell study
The paper tests the diagnostic in a controlled four-cell study on BFCL v4 multi_turn using no-think Gemma-4-26B-A4B. BFCL v4 provides four task categories: Base, Miss Func (required tool is withheld until a specific turn), Miss Param (required argument is missing), and Long Context. The study focuses on Miss Func and Miss Param, which represent the most common critical-state scenarios.
For Miss Func, the tool is unavailable until turn k, while the substantive user request appears at turn k minus 1. The model must defer or issue a read-only query at turn k minus 1, then emit the held call once the bridge announces tool availability. The diagnostic identifies the recovery turn after tool availability as the trainable state. For Miss Param, the function is present but a required argument is missing until the following turn; the decision-turn failure is a premature state-changing call, so the diagnostic selects the decision turn before the missing argument arrives.
The results are decisive. Training the diagnostic-selected recovery turn improves Miss Func accuracy from 0.14 to 0.283 across four seeds, a gain of 14.3 percentage points. Training the alternative decision turn in the same category leaves accuracy at 0.095, actually below the base rate of 0.14. For Miss Param, the pattern reverses: training the selected decision turn improves accuracy from 0.435 to 0.473, while training the alternative recovery turn gives a negligible 1 percentage point gain.
The four-cell design is the key experimental insight: the same occurrence-local RL interface is applied to all four cells, with the diagnostic selecting the right turn in both categories while fixed rules (always-decision or always-recovery) improve one category but leave the other flat or worse. This confirms that the diagnostic is not merely selecting an easy training target but is identifying the causally relevant decision.
Applications beyond the benchmark
The paper demonstrates that the diagnostic recipe generalizes across models and tasks. In a logged repeat-call avoidance application, the method identifies and trains on occurrences where the model repeats a failed tool call, reducing redundant calls. For Nemotron, the missing-function application trains the decision before tool availability rather than Gemma's recovery response after availability, illustrating that the training location is determined by the specific task mechanics rather than a fixed rule.
Memory management applications on xLAM and Gemma use a different label construction: an additive reward for correct operation order and preservation of answer-critical information, with hard gates for invalid actions. A keyword extraction proxy grades storage text against required terms, including meaningful negations such as "not married." The diagnostic selects the storage response as the trainable state, and training improves the full storage-to-answer chain.
The paper also demonstrates transfer to novel conditions: replacing the canonical bridge wording at evaluation with an unseen phrase that names no tools or functions, accuracy on Miss Func rises from 0.095 to 0.225 after recovery-turn RL, matching the canonical bridge's gain. This confirms that the improvement is not memorized to specific bridge wording but reflects genuine learning of the underlying decision.
Limitations and practical considerations
The method's reliance on nested sampling means the diagnostic phase requires many model rollouts without parameter updates, which can be computationally expensive. The finite-sample analysis in the appendix addresses selection error bounds, but the method's efficiency depends on the number of candidate occurrences and the variance of the noise. The paper also notes that different configurations can select different turns, and the diagnostic is task-specific rather than universal.
The occurrence-local training approach restricts gradient flow to a single call per training unit, which means the policy updates are local and may not capture long-range dependencies across the full trajectory. The paper provides a theoretical condition under which small local updates improve full-rollout return, but this condition requires that drift outside the selected occurrence is controlled.
What this means in practice
For teams building tool-using agents, Critical-State RL provides a systematic answer to a question that previously relied on intuition or trial and error: where in a multi-turn interaction should training effort go? The 14 percentage point improvement on Miss Func is significant for a benchmark where base rates are already low, and the diagnostic approach ensures that training data and compute are spent on the calls that actually matter.
The broader principle, that reward variation must be decomposed before it can guide training, has implications beyond tool use. Any multi-step interaction where downstream randomness confounds local decision quality could benefit from the nested sampling diagnostic. The paper's theoretical grounding in action-dependent variance as a proxy for local improvement potential provides a framework that other practitioners can adapt to their own domains.
As multi-turn agentic systems become more complex, the cost of training on the wrong states grows. Critical-State RL offers a principled way to avoid that waste, turning the question from "should I train this trajectory?" to "which call in this trajectory is worth training, and why?"