Uncertainty aware method improves multi step decision training
UOPD: Uncertainty-Aware Intervention for On-Policy Distillation of Multi-Turn Agents
Machine LearningArtificial Intelligence
Summary
Training AI agents that make several decisions in a row can be tricky because one wrong choice can ruin everything that follows. The authors found a way to spot when the AI is unsure about what to do and step in to correct it using advice from a more knowledgeable teacher. This approach, called UOPD, only intervenes when needed instead of all the time, helping the AI learn better decision-making. Tests show UOPD leads to better performance on tasks like online shopping and web searches.
What this means in practice
- •For conversational ai developers: Improve training of multi-turn dialogue agents by selectively correcting uncertain decisions during learning.
- •For automated web search teams: Enhance search agents’ performance by using uncertainty guided interventions to refine decision sequences.
Authors
Wenbo Zhang, Pengcheng Xu, Weizhi Du, Jing Zhang, Hengrui Cai
Abstract
On-policy distillation (OPD) trains a student on its own rollouts using dense supervision from a teacher. In multi-turn environments, a mistake at a critical decision step can redirect the subsequent rollout toward poor outcomes. We use low teacher confidence on student actions to select high-uncertainty steps for correction. In a controlled ALFWorld study, a single teacher correction at a low-confidence step improves subsequent student behavior and task success, motivating selective intervention during distillation. We propose UOPD, an uncertainty-aware intervention method for on-policy distillation. At low-uncertainty turns, UOPD executes student actions and applies the standard OPD loss. At high-uncertainty turns, it samples and executes teacher actions and trains the student to imitate them through supervised fine-tuning, which minimizes forward Kullback-Leibler divergence in expectation. UOPD utilizes adaptive uncertainty thresholds to target a scheduled intervention rate. Empirically, we evaluate UOPD across a broad range of agentic tasks, including ALFWorld, WebShop, and Search, demonstrating its superior performance over OPD methods and their variants. UOPD improves WebShop score by up to $15.8\%$ relative to standard OPD.