Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

This paper introduces Verified Critical Step Optimization (CSO), a post-training method for LLM agents that targets verified critical decisions where alternative actions flip task outcomes, rather than applying uniform trajectory- or step-level supervision. CSO uses a Process Reward Model (PRM) to identify candidate critical steps, computes high-quality expert alternatives, and verifies successful branches by continuing execution with the policy to ensure feasibility. Preference data are constructed only from verified, outcome-correct branches and trained via Direct Preference Optimization, with iterative online refinements to progressively improve the policy. Experiments on GAIA-Text-103 and XBench-DeepSearch show that CSO achieves substantial gains over SFT (45+% relative) and matches or exceeds proprietary baselines while requiring supervision on only a fraction of trajectory steps, demonstrating effective, efficient credit assignment and practical post-training impact.

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.

Verified Critical Step Optimization for LLM Agents

TL;DR

This paper introduces Verified Critical Step Optimization (CSO), a post-training method for LLM agents that targets verified critical decisions where alternative actions flip task outcomes, rather than applying uniform trajectory- or step-level supervision. CSO uses a Process Reward Model (PRM) to identify candidate critical steps, computes high-quality expert alternatives, and verifies successful branches by continuing execution with the policy to ensure feasibility. Preference data are constructed only from verified, outcome-correct branches and trained via Direct Preference Optimization, with iterative online refinements to progressively improve the policy. Experiments on GAIA-Text-103 and XBench-DeepSearch show that CSO achieves substantial gains over SFT (45+% relative) and matches or exceeds proprietary baselines while requiring supervision on only a fraction of trajectory steps, demonstrating effective, efficient credit assignment and practical post-training impact.

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.
Paper Structure (38 sections, 4 equations, 5 figures, 5 tables)

This paper contains 38 sections, 4 equations, 5 figures, 5 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preference-based Post-training
  4. Entropy-enhanced Algorithms for RL
  5. Methodology
  6. Preliminary
  7. Problem Formulation
  8. Supervised Fine-Tuning
  9. Verified Critical Step Collection
  10. Collecting Failed Policy Trajectories
  11. Selecting Candidate Critical Steps with PRM Scoring
  12. Branch Rollout and Verification
  13. Preference Dataset Construction
  14. CSO Training
  15. CSO Training Objective
  16. ...and 23 more sections

Figures (5)

  • Figure 1: Comparing post-training paradigms for LLM agents. (a) Trajectory-level optimization assigns coarse outcome rewards to entire trajectories, while (b) step-level process optimization relies on intermediate rewards which could be inaccurate. (c) Verified Critical Step Optimization (CSO) focuses learning on verified critical steps where alternative actions demonstrably flip task outcomes, providing precise credit assignment with minimal supervision.
  • Figure 2: Overview of Critical Step Optimization. (1) The policy model generates trajectories on training queries and we collect failed trajectories; (2) an expert teacher model generates alternative actions at each step; (3) a Process Reward Model identifies candidate critical steps where policy actions have low quality but expert alternatives have high quality; (4) we perform branch rollouts with the policy model and verify task outcomes; (5) we construct preference pairs from verified critical steps and the original failed policy steps; (6) we train the policy via DPO on the verified critical step preference data.
  • Figure 3: Performance across online training iterations. All methods trained on CK-Pro-8B, evaluated on GAIA-Text-103-L1.
  • Figure 4: Distribution of critical step error types identified by CSO. Tool invocation and reasoning errors constitute the majority of critical decision points.
  • Figure 5: Process Reward Model prompt for evaluating the quality of agent actions (code) during task execution. The PRM critically assesses code correctness, task relevance, logical progression, information utilization, and thought quality, with special attention to data reliability issues.