Source-linked AI summary
Agent-R: Training Language Model Agents to Reflect via Iterative Self-Training
Siyu Yuan, Zehui Chen, Zhiheng Xi, Junjie Ye, Zhengyin Du, Jiecao Chen
TL;DR
Interactive language agents often cannot recover from errors, while step-level reflection data is costly to collect. Agent-R uses MCTS and model-guided trajectory revision for iterative self-training, and experiments across three environments show improved error recovery and loop avoidance over baselines.
Problem
Existing expert-trajectory approaches struggle with dynamic error detection and revision, while step-level reflection data is difficult and expensive to collect.
Method
Agent-R uses MCTS and model-guided reflection to construct revision trajectories that recover correct behavior from erroneous trajectories through iterative self-training.
Results
Experiments across three interactive environments show that Agent-R significantly outperforms baselines while improving error correction and avoiding loops.
Takeaways & Limitations
Agent-R equips language agents to reflect and correct actions in real time without relying on human or expert-model supervision.
Takeaways & Limitations
Contrastive learning methods such as ETO can improve performance but do not equip agents with self-reflection capabilities.
Abstract
from arXiv · showhide
Large Language Models (LLMs) agents are increasingly pivotal for addressing complex tasks in interactive environments. Existing work mainly focuses on enhancing performance through behavior cloning from stronger experts, yet such approaches often falter in real-world applications, mainly due to the inability to recover from errors. However, step-level critique data is difficult and expensive to collect. Automating and dynamically constructing self-critique datasets is thus crucial to empowering models with intelligent agent capabilities. In this work, we propose an iterative self-training framework, Agent-R, that enables language Agent to Reflect on the fly. Unlike traditional methods that reward or penalize actions based on correctness, Agent-R leverages MCTS to construct training data that recover correct trajectories from erroneous ones. A key challenge of agent reflection lies in the necessity for timely revision rather than waiting until the end of a rollout. To address this, we introduce a model-guided critique construction mechanism: the actor model identifies the first error step (within its current capability) in a failed trajectory. Starting from it, we splice it with the adjacent correct path, which shares the same parent node in the tree. This strategy enables the model to learn reflection based on its current policy, therefore yielding better learning efficiency. To further explore the scalability of this self-improvement paradigm, we investigate iterative refinement of both error correction capabilities and dataset construction. Our findings demonstrate that Agent-R continuously improves the model's ability to recover from errors and enables timely error correction. Experiments on three interactive environments show that Agent-R effectively equips agents to correct erroneous actions while avoiding loops, achieving superior performance compared to baseline methods (+5.59%).
1 Introduction
Existing expert-trajectory and reward-based approaches struggle to support timely, step-level error correction in interactive environments. Agent-R addresses this gap with MCTS-based, model-guided revision trajectories and improves self-correction across three environments.
- Expert-trajectory behavior cloning struggles to detect and revise errors dynamically, causing cascading failures and suboptimal performance.
- Step-level reflection data is costly to annotate, while end-of-rollout correction delays learning and can leave early catastrophic errors unaddressed.
- Agent-R uses MCTS and model-guided reflection to identify a suitable revision step and connect erroneous trajectories to subsequent correct paths.
- Across three interactive environments, Agent-R outperforms baselines and improves error correction while avoiding loops.
- Agent-R iteratively constructs and revises trajectories without relying on human or expert-model supervision.
2 Preliminary
The paper formulates language-agent tasks as partially observable interactive decision processes in which agents generate actions from instructions and trajectory history. It introduces MCTS as a tree-search procedure that evaluates actions through expansion, simulation, and value updates.
- Interactive tasks are modeled as POMDPs with instruction, state, action, observation, transition, and reward spaces.
- A trajectory records successive actions and observations, with each observation obtained after executing its corresponding action.
- At each step, the actor model generates an action from the instruction and historical trajectory, and the environment provides a terminal reward.
- MCTS builds a decision tree and estimates action values through selection, expansion, simulation, and backpropagation.
3 Method
Agent-R uses MCTS and model-guided reflection to construct revision trajectories, then trains agents on them iteratively. Its method identifies timely transition points between erroneous and corrected behavior while combining revision and good trajectories during training.
- Framework: Agent-R has two phases: MCTS-based revision-trajectory generation followed by agent training on the collected trajectories, repeated iteratively.Phase I constructs revision data; Phase II optimizes the agent using that data.
- Reflection Trajectories: Revision trajectories splice corrected bad-trajectory prefixes with corresponding good-trajectory segments at a transition marked by a revision signal.The signal briefly communicates reflection on the agent’s error before continuing along the corrected path.
- Trajectory Collection: MCTS explores action paths through selection, expansion, simulation, and backpropagation, using multiple rollouts to estimate rewards and collect diverse trajectories.UCT balances exploration and exploitation through state rewards and visit counts.
- Transition Point Determination: The actor model identifies an incorrect action and sets the transition point there, truncating the bad trajectory before combining it with the good trajectory.This model-guided transition avoids waiting until rollout termination, when early catastrophic errors may already have propagated.
- Iterative Self-Training: Training mixes revision trajectories with good trajectories while progressively increasing the quality threshold for good trajectories across self-training iterations.This addresses the cold-start problem and improves both error reflection and correct-trajectory identification.
4 Experiment
Agent-R is evaluated across three interactive environments using iterative training on revision trajectories, with comparisons against diverse models and training strategies. The results indicate stronger performance, earlier error correction, fewer loops, and benefits from iterative and multi-task training.
- Experimental Environments: Experiments cover WebShop, ScienceWorld, and TextCraft, representing web shopping, scientific reasoning, and text-based crafting environments.The evaluation uses average final reward for ScienceWorld and WebShop, and success rate for TextCraft.
- Main Results: Agent-R significantly improves performance over advanced closed-source models and agents trained on expert trajectories.The comparison includes GPT-4o, AgentLM, and Agent-Flan.
- Main Results: Early error detection enables faster recovery and more stable learning than direct-revision methods.Agent-R identifies a suitable reflection step during trajectory construction rather than waiting until the end.
- Findings with Analysis: Iterative self-training gradually enhances model capabilities, while mixing revision and optimal trajectories improves performance; adding only good trajectories decreases performance.Revision trajectories provide recovery information that offsets noise in otherwise imperfect good trajectories.
- Findings with Analysis: Agent-R improves recognition of errors within trajectories, and multi-task training is more effective than single-task training.Shorter revision lengths indicate more efficient path adjustment across iterations.
- Findings with Analysis: Revision trajectories reduce repeated actions and dead loops compared with training on optimal trajectories, encouraging exploration of new actions.MCTS-derived optimal trajectories can contain noisy middle actions that trap agents in local loops.
5 Related Work
Prior work improves agents through prompting, experience summaries, and self-correction data, but interactive environments remain difficult because they lack explicit intermediate error signals and involve longer trajectories.
- Prompt-based strategies summarize experiences into transferable skills or hints that are stored in model memory.
- Self-correction methods often perform ineffectively, while prompt engineering may fail to produce intrinsic correction or degrade performance.
- Existing revision-data approaches target code repair, tool use, and mathematics, where explicit error signals or critical-model guidance are available.
- Interactive environments require self-correction over longer trajectories without explicit error signals, and prior expert-based correction does not resolve this setting.
6 Conclusion
Agent-R addresses error correction for language agents by using iterative self-training, MCTS, and model-guided reflection to revise erroneous trajectories. Across three interactive environments, it improves error recovery, prevents loops, and benefits from self-generated revision trajectories.
- Agent-R uses MCTS and model-guided reflection construction to enable timely revision of erroneous trajectories.
- Experiments across three interactive environments show improved error correction and reduced looping and suboptimal behavior compared with baseline methods.
- Self-generated revision trajectories improve overall agent performance and support more self-reflective agents.
B Trajectory Definition
Agent-R distinguishes initial, bad, good, and revision trajectories. A revision trajectory keeps a bad trajectory until a transition point, then continues with a good trajectory while marking the revision process.
- The framework defines initial, bad, good, and revision trajectories as central trajectory types.
- An initial trajectory is the sequence of actions and observations generated from the user instruction.
- A bad trajectory extends the initial trajectory with suboptimal actions and observations until terminal step T_b.
- A good trajectory extends the initial trajectory with optimal or high-reward actions and observations until terminal step T_g.
- A revision trajectory follows the bad trajectory through transition point t′, then switches to the good trajectory and includes a revision signal.The revision signal marks the revision process without altering the environment.
C.1 Training and Evaluation Details
The experiments use separate environment services and iterative supervised fine-tuning, with reduced later-iteration training to mitigate overfitting. Evaluation uses temperature 0 and either few-shot or zero-shot prompting depending on fine-tuning.
- Experiments use eight A100-80GB GPUs, with different environment services deployed on separate ports of one server.
- Table 4 provides prompt templates for determining the transition point in Agent-R.
- Iterative SFT uses 3 epochs initially and 1 epoch later to mitigate overfitting, with learning rate 2e−5 and a cosine scheduler.
- Evaluation temperature is 0; non-fine-tuned models use few-shot evaluation, whereas models fine-tuned on expert or revision trajectories use zero-shot evaluation.
C.2 Multi-task training v.s. Single-task Training
Multi-task training with Agent-R trajectories outperforms single-task training, with cross-task learning especially improving performance in later iterations.
- Multi-task training outperforms single-task training when using trajectories generated by Agent-R.
- Cross-task learning enhances performance across datasets, particularly in later iterations.
- Agent-R trajectories appear particularly beneficial for multi-task training.
C.3 Case Study
Agent-R-trained agents can recover from incorrect trajectories, recognize erroneous actions at adaptive transition points, and avoid actions that are invalid, mismatched with observations, or irrelevant to task goals.
- Error Correction in Trajectory Generation: Agent-R-trained agents can recover from incorrect trajectories and identify the correct action.
- Adaptive Transition Point Evaluation: Iterative SFT helps agents recognize when they are on bad trajectories with erroneous actions.
- Adaptive Transition Point Evaluation: Erroneous actions are categorized as invalid, mismatched with observations, or irrelevant to the task goal.
- Adaptive Transition Point Evaluation: Agent-R-trained agents are better equipped to identify and avoid erroneous actions.
- Adaptive Transition Point Evaluation: The resulting decisions better distinguish permissible, aligned, and task-relevant actions in complex environments.
- Error Correction in Trajectory Generation: Table 7 documents an error-correction example from a Llama-3.1-8B agent trained on Agent-R trajectories during the third iteration.
- Adaptive Transition Point Evaluation: Table 8 reports adaptive-transition-point judgments from Llama-3.1-8B trained on Agent-R trajectories during the third iteration.