Critical-State RL improves training for multi-step tool use tasks
Critical-State RL: Diagnosing Trainable States for Multi-Turn Tool Use
Machine LearningComputation and Language
Summary
Sometimes, when using AI to complete tasks that need many steps, a single choice can cause failure. Just looking at the task's success score doesn’t always show which step needs better training. The authors introduce Critical-State RL, a method to find which steps in a multi-step process can be improved by training. They test it on coding tasks where their method helps pinpoint the right steps to train and shows noticeable improvements, unlike training other steps.
What this means in practice
- •For ai software developers: Identify and train the most impactful steps in multi-turn AI tool interactions to improve task success rates in function calling systems.
- •For cloud service engineers: Optimize resource usage by training models selectively on critical decision points in multi-turn operations like memory management or repeat-call avoidance.
Authors
Zixiang Chen, Wenting Zhao, Zhepeng Cen, Akshara Prabhakar, Jielin Qiu, Jianguo Zhang, Zhiwei Liu, Tulika Manoj Awalgaonkar, Liangwei Yang, Shelby Heinecke, Silvio Savarese, Huan Wang
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.