Source-linked AI summary
Verified Critical Step Optimization for LLM Agents
Mukai Li, Qingcheng Zeng, Tianqing Fang, Zhenwen Liang, Linfeng Song, Qi Liu, Haitao Mi, Dong Yu
TL;DR
Long-horizon agent post-training must improve credit assignment without relying on coarse trajectory rewards, noisy step estimates, or expensive rollout sampling. CSO learns from verified critical alternatives found in failed policy trajectories, and reports 37% and 26% relative gains over SFT on two benchmarks with supervision at only 16% of trajectory steps. Its main practical limitation is the cost of completing branched trajectories for outcome verification.
Problem
Long-horizon agent training faces off-policy shift, expensive full rollouts, and sparse or delayed rewards that hinder precise credit assignment.
Method
CSO identifies candidate critical steps in failed policy trajectories, verifies expert alternatives through policy branch rollouts, and trains with DPO on successful alternatives.
Results
37% and 26% relative improvements over SFT were reported on GAIA-Text-103 and XBench-DeepSearch, with supervision at only 16% of trajectory steps.
Takeaways & Limitations
Selective outcome-verified supervision at critical decisions provides an efficient framework for agent post-training and outperforms trajectory-level, dense step-level, and hybrid methods.
Takeaways & Limitations
Outcome verification requires completing trajectories, creating a time cost for complex tasks and an efficiency bottleneck for online reinforcement learning.
Abstract
from arXiv · showhide
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.
1 Introduction
CSO addresses coarse trajectory rewards, noisy step-level estimates, and costly rollout-based credit assignment by training on verified critical decisions. It identifies alternatives that flip failed outcomes and achieves strong benchmark gains with sparse supervision.
- Motivation: Trajectory-level rewards assign the same coarse outcome signal across a trajectory, while step-level rewards may inaccurately credit intermediate actions.These weaknesses make it difficult to distinguish reasonable actions from suboptimal decisions within agent trajectories.
- Core idea: CSO focuses preference learning on critical steps where alternative actions demonstrably change task outcomes from failure to success.This provides step-specific supervision without relying solely on estimated intermediate rewards.
- Method: CSO starts from failed policy trajectories, uses a PRM and expert alternatives to identify candidates, and verifies branches by continuing execution with the policy.Only alternatives that lead to successful completion are retained for preference learning.
- Results: 37% and 26% relative gains over the SFT baseline were achieved on GAIA-Text-103 and XBench-DeepSearch, respectively.The experiments used CK-Pro-8B and report that the model matched GPT-4.1 while outperforming baseline methods.
- Results: CSO requires supervision at only 16% of trajectory steps while substantially outperforming trajectory-level, dense step-level, and hybrid methods.The reported results support selective supervision at pivotal branches rather than uniform trajectory coverage.
2 Related Work
Related work spans preference learning, entropy-based reinforcement learning, and structured online optimization for agent training. CSO takes an orthogonal offline approach centered on outcome-verified critical steps and sparse supervision.
- Preference-based post-training: Preference-pair methods have been applied to report generation, debugging traces, and step-wise reasoning data for targeted agent post-training.CSO extends this preference-learning direction to agent domains by identifying critical execution steps.
- Entropy-based optimization: Entropy-based reinforcement-learning studies suggest that a small fraction of high-entropy tokens or entropy-weighted advantages can drive improved reasoning performance.These findings motivate searching for sparse pivotal decisions in long-horizon agent trajectories.
- Online optimization: Online reinforcement-learning methods address credit assignment through structured exploration, advantage modulation, and dynamic rollout allocation.These approaches refine reward distribution through online optimization mechanisms.
- Positioning CSO: CSO offers an orthogonal offline approach that focuses exclusively on outcome-verified critical steps with minimal step-level supervision.Its design targets precise credit assignment without requiring dense estimated rewards.
3 Methodology
CSO collects preference data from failed policy trajectories by identifying candidate critical steps, verifying expert alternatives through policy continuations, and training with DPO. Iterative refinement repeats this process on fresh policy failures.
- Preliminary formulation: The framework models agent execution as ReAct trajectories containing states, policy actions, environment observations, and a terminal success outcome.At each step, actions may include reasoning and tool invocation, and the environment returns the resulting observation.
- Preliminary formulation: The initial policy is obtained by supervised fine-tuning on successful agent trajectories but may still fail systematically at critical decision points.CSO subsequently targets these failures through preference learning.
- Candidate selection: CSO retains failed policy trajectories, generates expert alternatives at each step, and uses PRM scores to identify low-quality policy actions paired with high-quality alternatives.The procedure focuses training data on the policy’s reachable distribution and its specific weaknesses.
- Candidate selection: PRM-based candidate selection avoids exhaustive Monte Carlo rollouts, but candidates are still verified by checking whether alternatives produce successful task completion.Verification addresses noise from potentially inaccurate PRM estimates.
- Branch rollout and verification: For each candidate, the original action is replaced by an expert alternative and subsequent steps are executed by the policy before the branched outcome is evaluated.A successful branch defines a verified critical step and keeps the resulting target within the policy’s execution capability.
- Preference construction: Verified critical steps produce preference pairs contrasting the successful alternative with the original failed action in their shared state context.These pairs form the preference dataset used for targeted optimization.
- CSO training: DPO trains the policy on verified critical-step pairs, concentrating learning on validated decisions rather than applying signals uniformly across all steps.The objective uses a reference policy, a KL penalty coefficient, and a logistic function.
- Iterative refinement: Iterative refinement collects fresh failed trajectories after each policy update and repeats critical-step selection and training.Using the previous policy as reference combines offline DPO stability with semi-on-policy adaptation while avoiding full RL overhead.
4 Experiments
Experiments evaluate CSO against proprietary, open-source, and post-training baselines on GAIA-Text-103 and XBench-DeepSearch, showing strong performance from verified critical-step supervision. The method uses PRM-based candidate identification, expert alternatives, and iterative DPO training from failed policy trajectories.
- Experimental setup: CSO constructs preference data from failed policy trajectories by generating expert alternatives and continuing policy rollouts from candidate branching points.The implementation uses CK-Pro-8B, Claude-3.7-Sonnet for alternatives and PRM scoring, and iterative training for up to two rounds.
- Experimental setup: CSO is evaluated on GAIA-Text-103 and XBench-DeepSearch against proprietary models, open-source models, and multiple post-training baselines.GAIA-Text-103 uses 103 text-only questions across three difficulty levels, while outputs are evaluated against gold answers.
- Main results: 49.5% overall accuracy on GAIA-Text-103 represents a 37% relative improvement over the SFT baseline.CSO also enables CK-Pro-8B to match GPT-4.1 and outperforms post-training baselines by at least +5.0 points on both benchmarks.
- Main results: Trajectory-level methods provide limited gains because outcome-based rewards apply uniformly across steps, producing coarse credit assignment.RFT reaches 34.9%, while ETO reaches 38.9% and improves by +3.0 points over its comparison point.
- Main results: Step-level and hybrid methods face accuracy or reward-noise issues, whereas CSO gains +5.0 points over IPR through verified critical-step supervision.Step-DPO improves on L1 tasks but not harder L2/L3 tasks, while IPR reaches 44.6% overall and remains affected by outcome reward contamination.
5 Analysis
The analysis shows that CSO's strongest gains come from pairing expert successes with policy failures, combining PRM selection with outcome verification, and concentrating supervision on sparse critical steps. These design choices improve performance and efficiency while keeping verification and data-construction costs bounded.
- Preference data sources: Expert successes paired with policy failures achieve the best performance among the compared preference-data configurations.The configurations share failed policy trajectories and critical steps, differing only in the sources of positive and negative actions.
- PRM selection and verification: Outcome verification substantially outperforms PRM-only selection while requiring supervision at only 16% of trajectory steps.Skipping PRM selection achieves comparable performance but requires nearly 3× more preference pairs.
- PRM quality and usage: CSO with outcome verification outperforms PRM-guided search alone, and Claude-3.7-Sonnet provides a stronger PRM than GPT-4.1.The comparison fixes the number of candidates at k = 5 per step.
- Computational cost: CSO adds 19% token overhead relative to Step-DPO versus 50% for ETO while producing the largest quality gain and supervising about 16% of trajectory steps.All methods share full policy-trajectory sampling as the dominant cost.
- Critical-step error types: Tool invocation errors account for 26.1% of identified critical steps and reasoning errors for 25.1%, making them the two largest categories.Other errors account for 24.1%, followed by task-understanding errors at 13.0% and information-extraction errors at 11.7%.
6 Conclusion
The conclusion presents CSO as selective preference learning over verified critical steps, combining PRM selection with outcome verification to avoid coarse trajectory supervision and noisy step-level estimates. Across two benchmarks, it reports substantial gains with supervision applied to only a small fraction of steps.
- Contribution: CSO focuses preference learning on verified critical steps where alternative actions demonstrably flip task outcomes.The method is framed as a post-training approach for targeted preference optimization.
- Results: 37% and 26% relative improvements over SFT on GAIA-Text-103 and XBench-DeepSearch require supervision at only 16% of steps.The paper reports that CSO enables an open-source 8B model to match GPT-4.1 and outperforms trajectory-level, dense step-level, and hybrid methods.
Limitations
CSO's main limitations are the cost of completion-based outcome verification and its reliance on closed-source PRMs that cannot be jointly optimized with the policy.
- Verification cost: Outcome verification requires completing trajectories, creating a time cost that constrains CSO's direct application to online RL.The paper suggests early stopping or parallelized execution as possible future directions.
- PRM dependence: The current implementation relies on closed-source PRMs, preventing joint optimization of the PRM and policy.The paper identifies improved open-source models as a potential basis for future joint training.
Ethics Statement
The ethics statement reports open-source research practices, licensed data and models, automated preference-data construction, and limited use of a paid Google Search API.
- Research practices: The experiments use open-source data, models, and frameworks under their respective licenses, with Google Search API as the only paid service.The API is reported as being used according to its terms of service.
- Data annotation: Automated PRM scoring and outcome verification construct preference data without human labor.The paper links this design to avoiding annotator exploitation and human bias injection.
- Disclosure: ChatGPT was used for grammar checking in the manuscript.
A PRM Prompt
The PRM evaluates policy actions and expert alternatives on failed trajectories, producing scores used to identify candidate critical steps.
- The Process Reward Model evaluates the quality of both policy actions and expert alternatives at each step of failed trajectories.The implementation uses Claude 3.7 Sonnet with rubric-based evaluation prompts.
- PRM scores range from 0 to 1 and identify candidates when policy actions score below γlow while expert alternatives exceed γhigh.The thresholds compare low-quality policy actions with high-quality expert alternatives.
- Figure 5 presents the complete prompt used for PRM evaluation.
B Sensitivity to PRM Thresholds
CSO remains effective under a tighter PRM-threshold configuration, indicating that its gains are not tied to one threshold pair.
- The tighter threshold configuration selected slightly fewer candidate steps and caused a small drop in accuracy relative to the main setting.The main setting used (γlow, γhigh) = (0.45, 0.65).
- Both threshold settings substantially outperform all baselines reported in Table 1.
C Robustness to Expert Model Choice
CSO improves over the SFT base policy across proprietary and open-source expert teachers, although stronger experts produce larger gains.
- CSO yields improvements across both proprietary and open-source expert teachers.
- All three expert choices substantially improve over the 46.2% SFT base policy on GAIA-Text-103-L1.The comparison uses Claude-3.7-Sonnet, GPT-4.1, and Qwen3-235B-A22B as expert teachers.
- The open-source Qwen3-235B-A22B teacher recovers most of Claude-3.7-Sonnet’s gain.
- Claude-3.7-Sonnet produces higher-quality alternative actions and more reliable PRM scores than the other expert choices.