PaperlayerView exact revision on arXiv

Source-linked AI summary

Watch Every Step! LLM Agent Learning via Iterative Step-Level Process Refinement

Weimin Xiong, Yifan Song, Xiutian Zhao, Wenhao Wu, Xun Wang, Ke Wang, Cheng Li, Wei Peng, Sujian Li

arXiv:2406.11176v2cs.CLcs.AIcs.LG

TL;DR

Existing agent training often treats a trajectory as a whole and prioritizes outcome rewards, leaving step-level process information underused. IPR estimates step rewards with Monte Carlo sampling and iteratively compares new actions with expert steps to form contrastive training pairs. Across WebShop, InterCodeSQL, and ALFWorld, it exceeds the current leading method by 5.8%, 7.2%, and 3.2%, respectively.

  • Problem

    Prior methods prioritize trajectory-level outcome rewards, while agent environments often provide only final or sparse subgoal feedback, leaving step-level process information underused.

    Source: §1 · p. 2

  • Method

    IPR uses Monte Carlo sampling to estimate step rewards, then iteratively compares new actions with corresponding expert steps to create contrastive training pairs.

    Source: §1 · p. 2

  • Results

    IPR surpasses the current leading method by 5.8% on WebShop, 7.2% on InterCodeSQL, and 3.2% on ALFWorld.

    Source: §1 · p. 2

  • Takeaways & Limitations

    IPR provides a way to optimize agent decision-making trajectories and enhance performance at the action level.

    Source: §7 · p. 9

    With limited training data, iterative preference learning on self-generated samples can lead to overfitting.

    Source: Limitations · p. 9

Abstract

from arXiv · show

Large language model agents have exhibited exceptional performance across a range of complex interactive tasks. Recent approaches have utilized tuning with expert trajectories to enhance agent performance, yet they primarily concentrate on outcome rewards, which may lead to errors or suboptimal actions due to the absence of process supervision signals. In this paper, we introduce the Iterative step-level Process Refinement (IPR) framework, which provides detailed step-by-step guidance to enhance agent training. Specifically, we adopt the Monte Carlo method to estimate step-level rewards. During each iteration, the agent explores along the expert trajectory and generates new actions. These actions are then evaluated against the corresponding step of expert trajectory using step-level rewards. Such comparison helps identify discrepancies, yielding contrastive action pairs that serve as training data for the agent. Our experiments on three complex agent tasks demonstrate that our framework outperforms a variety of strong baselines. Moreover, our analytical findings highlight the effectiveness of IPR in augmenting action efficiency and its applicability to diverse models.

1 Introduction

Prior trajectory-tuning methods prioritize final rewards and can overlook useful interaction-process information; step-level supervision is difficult when environments provide only outcome feedback or sparse sub-goal feedback. IPR estimates step rewards and iteratively trains agents with contrastive actions, outperforming leading methods on three interactive tasks.

1 Introduction

  • Prior trajectory-tuning methods treat each trajectory as a whole and prioritize its final reward, overlooking potentially useful information from the interaction process.The cited work includes expert-success-trajectory tuning and ETO, which uses failure trajectories.Source: §1 · pp. 1–2
  • Step-level supervision can provide granular guidance, but most agent environments offer only final-outcome feedback and available sub-goal feedback may be sparse.How to use step rewards effectively for long trajectories and complex action spaces remains unexplored in the cited discussion.Source: §1 · p. 2
  • IPR estimates step rewards with Monte Carlo sampling, compares newly generated actions with expert steps, and trains on contrastive pairs using outcome-level DPO, step-level DPO, and SFT losses.The iterative process is intended to identify errors and refine actions along the expert trajectory.Source: §1 · p. 2
  • IPR surpasses the leading method by 5.8% on WebShop, 7.2% on InterCodeSQL, and 3.2% on ALFWorld.The evaluation covers online shopping, interactive SQL, and textual embodied tasks.Source: §1 · p. 2
  • IPR is presented as the first integration of step-level process supervision into LLM-agent training, enabling fine-grained adjustment of task completion.The paper also reports that IPR increases reward per step and that automatically constructing a step reward model can reduce MC-method training costs.Source: §1 · p. 2

2 Task Formulation

The task is framed as an LLM agent interacting with an environment, receiving observations and execution feedback until completion or a maximum step limit. The formulation specifies the instruction, state, action, observation, transition, and reward spaces.

2 Task Formulation

  • The task is formulated as a partially observable Markov decision process with instruction, state, action, observation, transition, and reward elements.Instructions, actions, and observations are subsets of natural-language space.Source: §2 · p. 2
  • At each step, the agent acts on its interaction history, receives an environment observation, and continues until task completion or the maximum number of steps.The final trajectory receives an outcome reward, and the formulation defines the remaining trajectory from step t as e_t:n.Source: §2 · p. 2
  • The IPR iteration has the agent explore new actions along the expert trajectory, scores steps to form contrastive action data, and optimizes the agent with a mixed loss.This is the architecture described for a single iteration.Source: Figure 2 · p. 3

3 Method

The method first gives the agent basic capabilities through supervised fine-tuning, then estimates step rewards and iteratively refines actions using contrastive trajectories and mixture optimization.

3.1 Supervised Fine-tuning

  • Supervised fine-tuning on ReAct-style expert trajectories gives the agent basic task capabilities.The expert trajectories follow the ReAct pattern (Yao et al., 2022b).Source: §3.1 · p. 3

3.2 Step-level Reward Acquisition

  • Step rewards estimate anticipated outcome rewards through subsequent exploration, with Monte Carlo sampling and an SFT-trained agent serving as the scorer.The method uses sampled trajectories to estimate the expectation; at the final step, it uses the observed trajectory’s outcome reward.Source: §3.2 · pp. 3–4

3.3 Iterative Agent Optimization

  • Exploring from each expert-trajectory prefix yields candidate actions; an action is marked mistaken when its step reward falls below the expert action by more than τ and its trajectory outcome is worse.The resulting higher- and lower-reward continuations form contrastive pairs.Source: §3.3 · p. 4
  • Optimization combines outcome-DPO, step-DPO, and SFT losses, then reuses the updated agent to collect more pairs until the iteration limit.The added SFT loss directly increases the likelihood of successful trajectories.Source: §3.3 · pp. 4–5

4 Experiments

IPR is evaluated on WebShop, InterCodeSQL, and ALFWorld against prompt-based, outcome-refinement, and process-refinement baselines. It outperforms all baselines across the three tasks, with robustness also reported on ALFWorld unseen cases.

4.1 Experiment Settings

  • The evaluation covers WebShop, InterCodeSQL, and ALFWorld, using average reward across all tasks.ALFWorld includes 140 seen and 134 unseen test cases.Source: §4.1 · p. 5
  • GPT-4 generates expert trajectories in the ReAct pattern, which are filtered to retain correct trajectories.The experiments use Llama-2-7B as the base model, with training capped at four IPR iterations.Source: §4.1 · p. 5

4.2 Results

  • IPR outperforms ETO by 5.8%, 7.2%, 2.5%, and 3.2% on WebShop, InterCodeSQL, ALFWorld seen, and ALFWorld unseen, respectively, averaging 4.5%.IPR also outperforms all baselines across the three tasks and shows robustness on ALFWorld unseen cases.Source: §4.2 · pp. 5–6, Table 2 · p. 6
  • Prompt-based models perform significantly worse than outcome- and process-refinement methods, while Step-PPO is unstable beyond InterCodeSQL.Step-PPO performs well on InterCodeSQL but poorly on the other two tasks.Source: §4.2 · pp. 5–6, §4.1 · p. 5

5 Analysis

The analyses find consistent IPR gains across base models and tasks, while ablations highlight the roles of training losses and iteration limits. Step rewards reach up to 82% accuracy, average step rewards improve across tasks, and a learned reward model generalizes but remains below Monte Carlo estimation.

5.1 Different Base Models

  • IPR outperforms ETO and SFT across all tested models and datasets, with stronger initial agents tending to achieve more pronounced post-IPR performance.The tested models are Mistral-7B, Llama-2-13B, and Llama-3-8B, on WebShop and InterCodeSQL.Source: §5.1 · p. 6, Table 3 · p. 6

5.2 Ablation Study

  • Removing each training module lowers performance; removing SFT loss causes the largest drop, and removing step-DPO hurts more than removing outcome-DPO.The ALFWorld ablation evaluates the unseen test set.Source: §5.2 · p. 6, Table 4 · p. 7
  • Further iterations initially improve performance by learning from incorrect actions, but excessive exploration may reduce it through overfitting.The authors describe overfitting as a possible explanation for the decline.Source: §5.2 · p. 6, Table 4 · p. 7

5.3 Step Reward Estimation Quality

  • Monte Carlo step-reward estimation reaches up to 82% accuracy despite scorer noise, with accuracy assessed by whether contrastive-pair ordering agrees with the scoring function.The analysis varies scorer agents and Monte Carlo sampling times.Source: §5.3 · p. 7, Figure 3 · p. 7

5.4 Average Reward Per Step

  • IPR improves average reward per step over SFT and ETO across three tasks, with more pronounced improvement on datasets where SFT has higher average step reward.The reported step rewards are averaged over actions within trajectories and then over the test set.Source: §5.4 · p. 7, Figure 4 · p. 7

5.5 Exploration of Step Reward Modeling

  • A learned step-reward model improves Llama-3-8B over ETO despite its actions being absent from training data, but remains below Monte Carlo estimation.The model is trained on 70,000 actions and uses mean squared error loss.Source: §5.5 · p. 8, Table 5 · p. 8

6 Related Work

Prior LLM-agent tuning often learns from successful trajectories or contrasts success and failure but focuses on final outcomes. Related process-supervision and self-improvement work motivates finer-grained guidance, which this paper applies to agent learning.

6.1 LLM as Agents

  • Recent fine-tuning methods for open-source LLM agents learn from successful trajectories or use contrastive information from failed trajectories, yet only leverage final outcome reward, with no studies to date integrating process information to improve agent performance.Source: §6.1 · p. 8

6.2 Step-level Process Supervision

  • Process-supervision methods use step-level guidance, while pseudo-label approaches avoid manual labels but PPO-based optimization can have efficiency and stability issues.The paper describes its mixture trajectory optimization as enhancing agent performance.Source: §6.2 · p. 8

6.3 Self-Improvement

  • Self-improvement research spans alignment and agent planning, tool use, and communication; this work aims to add fine-grained guidance to that process.The cited agent work explores self-improvement in diverse domains.Source: §6.3 · p. 8

7 Conclusion

IPR combines process-level supervision with contrastive action pairs and Monte Carlo step rewards to refine LLM agents. Experiments and analyses report improved benchmark performance and action efficiency.

7 Conclusion

  • IPR integrates process-level supervision so agents learn from contrastive action pairs.Monte Carlo estimates step rewards in environments that provide only outcome rewards.Source: §7 · pp. 8–9
  • IPR consistently outperforms existing baselines across three benchmarks.The analyses also validate each framework component and its action efficiency.Source: §7 · p. 9
  • Iterative agent optimization provides a way to optimize agent decision-making trajectories.Source: §7 · p. 9

Limitations

The authors identify three limitations: overfitting risk with limited training data, incomplete use of step-reward values, and a step reward model trained on only one task.

Limitations

  • Limited training data can lead iterative preference learning on self-generated samples to overfit.Source: Limitations · p. 9
  • IPR identifies error actions for contrastive datasets but does not fully exploit step-reward values to prioritize error severity.The authors suggest correcting more severe errors first as a possible curriculum-learning direction.Source: Limitations · p. 9
  • The step reward model is trained on a single agent task, affecting its generalizability across tasks.The authors propose developing a general agent step reward model for different tasks.Source: Limitations · p. 9

A Dataset Details

The appendix describes three environments with final-outcome feedback and explains how WebShop page states receive step scores based on progress toward the target state.

A Dataset Details

  • WebShop simulates ecommerce over 1.8 million products; agents search or click to match product instructions, with a final reward based on attribute and price matching.Source: A Dataset Details · p. 12
  • InterCodeSQL has agents retrieve database information and submit SQL queries, with final reward measured by IoU against the gold execution output.Its database derives from Spider and the environment was modified for this evaluation framework.Source: A Dataset Details · p. 12
  • ALFWorld tasks require room exploration and commonsense actions, with success measured by completion within the allowed steps.Seen tasks assess in-distribution generalization; unseen instances assess out-of-distribution generalization.Source: A Dataset Details · p. 12

B Details of the Scoring Function

  • WebShop step rewards score search, product-description, and order-confirmation pages by product match and progress toward the target state.The search page uses its highest product score, the description page the highest score across options, and the confirmation page the selected product's score.Source: B Details of the Scoring Function · p. 12

C Training Efficiency Analysis

On WebShop, IPR's three-iteration training takes longer than SFT or ETO, while the authors report nearly 6% higher performance without added inference-time costs.

C Training Efficiency Analysis

D Case Study

The case studies contrast IPR with ETO on WebShop and ALFWorld, while the appendices describe expert-trajectory collection and evaluation prompts. The examples show IPR correcting specific action errors and reducing unnecessary exploration.

D Case Study

  • In WebShop and ALFWorld, IPR corrected ETO’s product-selection and potato-handling errors, with more direct actions in the ALFWorld example.ETO bought a dress shirt instead of the requested tank top and cooled the potato without picking it up; IPR chose the tank top and navigated directly to the potato.Source: D Case Study · pp. 12–13, Figure 5 · p. 14, Figure 6 · p. 16

E Expert Trajectories Collection

  • Expert trajectories came from Song et al. (2024), task datasets, and GPT-4 annotations; WebShop retains rewards above 0.7, InterCodeSQL retains 1.0, and GPT-4 supplies missing action thoughts.ALFWorld provides human-annotated trajectories; InterCodeSQL trajectories were annotated by the authors.Source: E Expert Trajectories Collection · p. 17

F Prompt for Evaluation

Loading 2406.11176v2…