Critical-State RL: Diagnosing Trainable States for Multi-Turn Tool Use
Zixiang Chen, Huan Wang and colleagues at Salesforce AI Research introduce Critical-State RL, a method that decides which model call in a multi-turn tool-use trajectory should be trained before any training happens.
Ask this paper
The problem. When reward depends on later turns, its variation can come from downstream randomness rather than from the current action, so trajectory-level rewards do not say which call to train.
Diagnostic. For each candidate call, nested sampling separates action-dependent reward variance from continuation noise and checks whether there is headroom over a reference policy. Training then uses contextual-bandit updates at the selected calls only.
BFCL v4 result. On the missing-function category, training the diagnostic-selected turn (the response after the tool becomes available) improves accuracy by about 14 points (+14.3 pp), while training the alternative turn leaves performance flat or lower. For missing-argument tasks the diagnostic picks the turn before the argument arrives (+3.8 pp).
Selection is model-dependent. The best training location shifts across models and reasoning settings, which is why the diagnostic runs per configuration.
Beyond one model. The same recipe applies to logged repeat-call avoidance and to memory management with xLAM, and gives a +4.4 point paired missing-function gain over a Nemotron starting checkpoint.
Abstract
Multi-turn tool-use failures can hinge on a single model call, yet reward variation alone does not reveal which call would benefit from training. When rewards depend on later interactions, their variation can reflect downstream randomness rather than differences between the current actions. We introduce Critical-State RL to identify trainable states in multi-turn interactions. Given task-defined candidate calls and local rewards, the method assesses whether each reward captures the action's effect on task success and whether improvement over a reference policy is possible. It then uses nested sampling to separate action-dependent reward variation from continuation noise and optimizes the policy at the selected states using contextual-bandit training. Experiments on the Berkeley Function Calling Leaderboard (BFCL) v4 compare training at diagnostic-selected states with training at alternative states. For missing-function tasks, the diagnostic selects the response after the tool becomes available; for missing-argument tasks, it selects the response before the missing argument is supplied. Training the selected responses improves performance, including about 14 percentage points on the missing-function task, while training the alternatives leaves performance flat or worse. We further apply the recipe across models and tasks, including logged repeat-call avoidance and memory management.