Verified Critical Step Optimization for LLM Agents

2026-02-03

AI Summary

This work addresses the challenges of sparse rewards, high noise, and expensive post-training costs faced by large language model agents in complex, long-horizon tasks. The authors propose Critical Step Optimization (CSO), a novel method that begins with failed trajectories and leverages a Process Reward Model (PRM) to identify critical decision points. At these points, an expert model generates alternative actions, and only those leading to successful outcomes—verified by the agent’s ability to complete the task correctly—are retained as preference data for supervised learning. By focusing supervision exclusively on verifiable, pivotal steps, CSO circumvents the inefficiencies of trajectory-level coarse labeling and step-level noise, all without requiring expert demonstrations. Evaluated on GAIA-Text-103 and XBench-DeepSearch, CSO improves over SFT baselines by 37% and 26%, respectively, using supervision signals from merely 16% of trajectory steps.

Technology Category

• Search and Optimization: Learning to Search
• Planning, Routing, and Scheduling: Planning with Language Models
• Multiagent Systems: Adversarial Agents

Application Category

• Semantics and Knowledge: Data modeling to support human-machine intelligence, including LLMs agents, intelligent system behavior, explanations, and user-friendly interactions
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

As large language model agents tackle increasingly complex long-horizon tasks, effective post-training becomes critical. Prior work faces fundamental challenges: outcome-only rewards fail to precisely attribute credit to intermediate steps, estimated step-level rewards introduce systematic noise, and Monte Carlo sampling approaches for step reward estimation incur prohibitive computational cost. Inspired by findings that only a small fraction of high-entropy tokens drive effective RL for reasoning, we propose Critical Step Optimization (CSO), which focuses preference learning on verified critical steps, decision points where alternate actions demonstrably flip task outcomes from failure to success. Crucially, our method starts from failed policy trajectories rather than expert demonstrations, directly targeting the policy model's weaknesses. We use a process reward model (PRM) to identify candidate critical steps, leverage expert models to propose high-quality alternatives, then continue execution from these alternatives using the policy model itself until task completion. Only alternatives that the policy successfully executes to correct outcomes are verified and used as DPO training data, ensuring both quality and policy reachability. This yields fine-grained, verifiable supervision at critical decisions while avoiding trajectory-level coarseness and step-level noise. Experiments on GAIA-Text-103 and XBench-DeepSearch show that CSO achieves 37% and 26% relative improvement over the SFT baseline and substantially outperforms other post-training methods, while requiring supervision at only 16% of trajectory steps. This demonstrates the effectiveness of selective verification-based learning for agent post-training.

Problem

• LLM agents
• post-training
• step-level reward
• long-horizon tasks
• credit assignment

Innovation

• Critical Step Optimization
• Verified Supervision
• Preference Learning
• Process Reward Model
• DPO