CAST: Critique-Aware Supervision for Training Reliable Long-Horizon Tool-Calling Agents

2026-08-31Computation and Language

Computation and Language
AI summary

The authors studied how to make large language model (LLM) agents more reliable when they have to perform many steps in complex tasks, where one wrong move can cause big problems. They created a new training method called CAST, which helps the model understand and explain when actions might be wrong by turning overall task results into detailed feedback for each step. By teaching the model to critique its own actions, they improved its performance on tasks like retail and telehealth compared to other strong models. Their work shows that adding these explanations during training helps LLM agents work better in complicated, realistic situations.

Large Language ModelsAgent ReliabilityCritique-aware TrainingPolicy OptimizationAction-level SupervisionPartial ObservabilityDynamic EnvironmentsTrajectory AnalysisTool-calling BenchmarksQwen3 Models
Authors
Amir Saeidi, Zehua Zhang, Rishitosh Singh, Naman Ahuja, Vivek Gupta, Ali Payani, Gaowen Liu, Jayanth Srinivasa, Chitta Baral
Abstract
Large language model (LLM) agents are increasingly deployed in long-horizon, interactive, and stateful environments. In these settings, a single wrong action, such as refunding the wrong purchase, can cause irreversible task failure and must be intercepted before execution. Such failures may not appear in every single run, but can emerge across repeated trials, making reliability across steps and trials critical. However, ensuring agentic reliability is challenging: even frontier LLMs struggle to explain why an action may be wrong, especially in long, intertwined trajectories governed by domain-specific policies. Much recent work relies on prompt-based critique agents, while optimization-based methods lack a systematic way to produce rich verification rationales for training. We address this gap with CAST, a critique-aware training framework that converts sparse task outcomes into action-level supervision for critique learning and policy optimization. CAST analyzes agent trajectories to synthesize structured rationales explaining action validity under partial observability. The resulting critique model is used to construct critique-aware training data for optimizing the policy model. Fine-tuning Qwen3-family models on dynamic tool-calling benchmarks, CAST improves reliability across domains, outperforming GPT-OSS-120B by over 10% pass^4 on Retail tasks and yielding an additional 9% improvement on Telehealth in an out-of-domain setting. These results demonstrate that critique-aware training improves the robustness of LLM agents in realistic dynamic environments.