Verified Critical Step Optimization for LLM Agents

TL;DR

The paper introduces Critical Step Optimization (CSO), a new method for improving large language model agents by focusing on verified critical decision points, leading to significant performance gains with less supervision.

Contribution

CSO is a novel approach that targets critical decision steps using verified supervision, starting from failed trajectories and leveraging expert models for high-quality alternatives.

Findings

CSO achieves 37% and 26% relative improvements on GAIA-Text-103 and XBench-DeepSearch.

It outperforms other post-training methods while supervising only 16% of trajectory steps.

The method enhances policy robustness by focusing on verifiable critical decisions.

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.