Verified Critical Step Optimization for LLM Agents

AI summary

Overview

Research area: Post-training of large language model (LLM) agents — specifically preference-based learning for long-horizon, tool-using agents.

Technical level: Intermediate. Readers benefit from familiarity with supervised fine-tuning (SFT), reinforcement learning (RL), Direct Preference Optimization (DPO), and process reward models (PRMs).

Scope: The paper proposes Critical Step Optimization (CSO), a post-training method that applies preference learning only to outcome-verified "critical steps" in failed agent trajectories, and evaluates it on two agent benchmarks using an 8B-parameter policy model.

What This Paper Is About

LLM agents that solve long-horizon tasks through repeated tool use are hard to post-train: outcome-only rewards cannot say which intermediate step was responsible for failure, step-level reward estimates are noisy, and Monte Carlo sampling of step rewards is prohibitively expensive. CSO addresses this by finding decision points where swapping in a different action demonstrably flips a task from failure to success, and training preference pairs only at those verified points. The goal is precise credit assignment with supervision on a small fraction of trajectory steps.

Key Contributions

1. Diagnosis and framing. The authors identify limitations of existing post-training paradigms — trajectory-level optimization (coarse credit assignment), dense step-level optimization (noisy estimated rewards), and hybrid Monte Carlo approaches such as IPR (prohibitive cost) — and propose focusing on verified critical steps instead.

2. The CSO method. Starting from failed policy trajectories rather than expert demonstrations, CSO uses a process reward model (PRM) to nominate candidate critical steps, has an expert model propose alternative actions, then continues execution with the policy model itself to verify which alternatives actually produce task success. Only verified alternatives become DPO training pairs.

3. Empirical results. CSO achieves 37% and 26% relative improvement over the SFT baseline on GAIA-Text-103 and XBench-DeepSearch respectively, enabling the open-source 8B CK-Pro-8B model to match GPT-4.1, while requiring supervision at only 16% of trajectory steps.

4. Evidence that PRM identification plus outcome verification pinpoints critical decision steps, supported by ablations on data source, PRM selection, verification, number of branch candidates, PRM quality/usage, computational cost, and iterative refinement.

Main Findings

• Overall accuracy gains. CSO reaches 49.5% overall accuracy on GAIA-Text-103 versus 35.9% for the CK-Pro-8B SFT baseline (a 37% relative improvement), and 29.0 on XBench-DeepSearch2505 versus 23.0 for SFT (a 26% relative improvement). CSO outperforms all post-training baselines by at least +5.0 points on both benchmarks.

• Matches a proprietary model. With CSO, the open-source 8B model CK-Pro-8B matches GPT-4.1, which scores 45.6 overall on GAIA-Text-103. Claude-3.7-Sonnet scores 62.1 on GAIA-Text-103 and 41.0 on XBench-DeepSearch2505.

• Baseline comparison on GAIA-Text-103 (overall). ETO 38.9, RFT 34.9, Step-DPO 38.9, IPR 44.6, CSO 49.5, SFT 35.9. On XBench-DeepSearch2505: ETO 22.0, RFT 20.0, Step-DPO 25.0, IPR 24.0, CSO 29.0, SFT 23.0, Qwen3-8B 7.0.

• Trajectory-level methods are limited. RFT shows no significant change (34.9%, versus 35.9% for SFT) training on the policy's own successful trajectories. ETO improves to 38.9% (+3.0 points) by contrasting expert successes against policy failures, but both apply outcome-based rewards uniformly across all steps.

• Step-level methods degrade on hard tasks. Step-DPO achieves +7.1 points on simple L1 tasks but shows no improvement on harder L2/L3 tasks, which the authors attribute to PRM accuracy degradation. IPR improves through outcome verification to 44.6% overall but still propagates trajectory-level outcome signals to all steps. CSO exceeds IPR by +5.0 points.

• Best preference-pair configuration. Contrasting expert successes with policy failures scores 49.5 on GAIA-Text, beating expert success + expert failure (46.6) and policy success + policy failure (42.7).

• Verification is what matters; PRM selection adds efficiency. PRM + Verification: 49.5% with 671 preference pairs. Without PRM selection: 48.5% but with 1,967 pairs (nearly 3x more). Without outcome verification: 43.6% with 4,126 pairs. Verification-based supervision covers only 16% of trajectory steps.

• Branch candidate count. k=3 gives 46.6 on GAIA-Text and 26.0 on XBench; k=5 gives 49.6 and 29.0; k=7 gives 49.6 and 28.0. Increasing beyond 5 yields no further gain but raises verification cost, so k=5 is chosen as the balance point.

• PRM quality and usage. On GAIA-Text-103-L1 with k=5, Claude-3.7-Sonnet as PRM yields CSO 61.5 versus step-level Best-of-N 56.2; GPT-4.1 as PRM yields CSO 53.3 versus BoN 48.7. Stronger foundation models produce more reliable process rewards, and CSO outperforms PRM-guided search for the same PRM.

• Cost overhead is moderate. Under identical conditions (same base policy, pool of 123k training tasks, average trajectory length of 1,381 tokens), CSO adds roughly 168M extra tokens — 1.19x the cost of Step-DPO (~141M extra tokens, normalized to 1.00x) and smaller than ETO's ~212M extra tokens (1.50x).

• Iterative refinement behavior. Across four online rounds on GAIA-Text-103-L1, ETO initially improves then degrades below the SFT baseline by Round 3; IPR reaches 56.4% at Rounds 2–3; CSO reaches 61.5% by Round 2 and maintains it through Round 3.

• What the critical steps are. Manual categorization of PRM-identified critical steps: Tool Invocation errors 26.1%, Reasoning Errors 25.1%, Other Errors 24.1%, Task Understanding 13.0%, Information Extraction 11.7%.

Methodology in Plain English

The agent is modeled as a ReAct-style loop: at each state it picks an action (reasoning plus tool call), the environment returns an observation, and the trajectory ends with a success or failure label.

CSO runs in six stages:

1. Collect failures. Deploy the current policy on training queries and keep only failed trajectories, so the training signal targets the policy's own weaknesses rather than expert behavior.

2. Propose alternatives. At each step of a failed trajectory, sample k alternative actions from a stronger expert model (Claude-3.7-Sonnet).

3. Score with a PRM. Use a rubric-prompted process reward model to score both the policy's original action and each expert alternative on a 0–1 scale. A step becomes a candidate critical step only if the policy action scores below a low threshold (γ_low = 0.45) and at least one expert alternative scores above a high threshold (γ_high = 0.65).

4. Verify by branch rollout. Replace the policy's action at the candidate step with a high-scoring expert alternative, then let the policy model itself continue the remaining steps to completion. This keeps successful branches within the policy's own reachable behavior.

5. Build preference pairs. Only branches that produce a verified successful outcome become training data, pairing the successful alternative action (positive) against the original failed action (negative) at the same state.

6. Train with DPO. Optimize the policy on these pairs with a KL penalty coefficient β = 0.5 against a reference policy. The whole loop can be repeated for up to 2 rounds, with each round collecting fresh failures from the updated policy and setting the reference policy to the previous iteration.

Setup details: the policy model is CK-Pro-8B (8B parameters, SFT-ed from Qwen3-8B), all interactions run through the open-source Cognitive Kernel Pro multi-agent framework (extended to support in-step PRM scoring and rollout continuation from intermediate states), and all post-training uses LlamaFactory. Preference data is built starting from CK-Pro-8B's SFT data of 47K task-trajectory pairs. Evaluation uses GAIA-Text-103 (103 questions across L1, L2, L3) and XBench-DeepSearch 2505 (100 complex tasks), with an LLM judging correctness against gold answers. Results are reported from three independent runs.

Why This Matters

Impact on research. The paper argues that agent trajectories, like token sequences in RLVR, contain a sparse subset of decision points that determine success — an analogy to the finding that only a small fraction of high-entropy tokens drive effective RL for reasoning. It offers an offline alternative to expensive online RL and to noisy dense step-level supervision, and it shows that PRM identification and outcome verification are complementary rather than competing.

Real-world applications:

• Deep research assistants that must plan multi-step web searches and synthesize evidence across sources.
• Web navigation agents that choose which pages to visit and how to formulate queries.
• Tool- and file-handling agents in the Cognitive Kernel Pro framework, where incorrect tool selection or suboptimal query formulation is the largest identified error category (26.1%).
• Post-training of small open-source models so they can approach proprietary-model accuracy in cost-sensitive deployments.

Industry relevance. The method enables an 8B open-source model to match GPT-4.1 on GAIA-Text-103 while adding only about 19% token overhead over Step-DPO for data construction, and it constructs preference data entirely automatically — no human annotation labor. That combination is attractive for teams that need agent post-training within bounded compute budgets.

Future Directions

• Reducing verification cost. Outcome verification requires executing trajectories to completion, which is time-consuming on complex tasks; the authors suggest early stopping heuristics or parallelized execution.
• Adapting CSO to online RL. The offline design avoids full RL overhead, but the verification bottleneck must be addressed before CSO can be used in online training.
• Open-source PRMs that can be jointly trained. The current implementation relies on closed-source models as the PRM, which cannot be optimized alongside the policy; jointly training the PRM as open models improve is left to future work.
• Following critical steps as they shift. Iterative refinement assumes critical steps move toward harder decision points as the policy improves; characterizing this shift is an open question raised by the iterative training design.

Target Audience

Researchers and engineers working on LLM agent post-training, preference optimization (DPO), process reward models, and credit assignment for long-horizon tasks. It is also relevant to practitioners who want to improve small open-source agent models without large-scale online RL infrastructure, and to readers interested in how entropy-based "critical token" intuitions from reasoning RL transfer to agent decision steps.

Authors’ 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.