AgentRefine: Enhancing Agent Generalization through Refinement Tuning

TL;DR

This work tackles the generalization gap of open-source LLM-based agents by introducing AgentRefine, a refinement-tuning framework where models learn to correct mistakes via environment feedback. It combines a synthesis-based data pipeline with a trajectory refinement loss that masks erroneous steps, promoting self-refinement and exploration across diverse environments. Across five agent tasks, AgentRefine substantially outperforms prior agent-tuning methods on held-out sets, improves robustness to perturbations, and demonstrates reasoning generalization on HotpotQA. The approach also shows viability with open-source data and models, indicating practical impact for building more generalizable, robust LLM agents in real-world settings.

Abstract

Large Language Model (LLM) based agents have proved their ability to perform complex tasks like humans. However, there is still a large gap between open-sourced LLMs and commercial models like the GPT series. In this paper, we focus on improving the agent generalization capabilities of LLMs via instruction tuning. We first observe that the existing agent training corpus exhibits satisfactory results on held-in evaluation sets but fails to generalize to held-out sets. These agent-tuning works face severe formatting errors and are frequently stuck in the same mistake for a long while. We analyze that the poor generalization ability comes from overfitting to several manual agent environments and a lack of adaptation to new situations. They struggle with the wrong action steps and can not learn from the experience but just memorize existing observation-action relations. Inspired by the insight, we propose a novel AgentRefine framework for agent-tuning. The core idea is to enable the model to learn to correct its mistakes via observation in the trajectory. Specifically, we propose an agent synthesis framework to encompass a diverse array of environments and tasks and prompt a strong LLM to refine its error action according to the environment feedback. AgentRefine significantly outperforms state-of-the-art agent-tuning work in terms of generalization ability on diverse agent tasks. It also has better robustness facing perturbation and can generate diversified thought in inference. Our findings establish the correlation between agent generalization and self-refinement and provide a new paradigm for future research.

Figures (13)

• Figure 1: Overall progress score among 5 tasks. Agent-FLAN has been trained on Held-in task.
• Figure 2: Example of parameter memorization in Agent-FLAN.
• Figure 3: The success rate variation via perturbation
• Figure 4: The pipeline of AgentRefine data generation and refinement tuning.
• Figure 5: The model's performance as the AgentRefine train data scales up.