AgentRM: Enhancing Agent Generalization with Reward Modeling
This work addresses the challenge of agent generalization for AI systems, offering a robust method that enhances performance on both unseen and held-in tasks, though it is incremental as it builds on existing reward modeling techniques.
The paper tackles the problem of poor generalization of LLM-based agents to unseen tasks by proposing AgentRM, a reward model that guides policy models during test-time search, achieving an average improvement of 8.8 points on nine agent tasks and surpassing the top general agent by 4.0 points.
Existing LLM-based agents have achieved strong performance on held-in tasks, but their generalizability to unseen tasks remains poor. Hence, some recent work focus on fine-tuning the policy model with more diverse tasks to improve the generalizability. In this work, we find that finetuning a reward model to guide the policy model is more robust than directly finetuning the policy model. Based on this finding, we propose AgentRM, a generalizable reward model, to guide the policy model for effective test-time search. We comprehensively investigate three approaches to construct the reward model, including explicit reward modeling, implicit reward modeling and LLM-as-a-judge. We then use AgentRM to guide the answer generation with Best-of-N sampling and step-level beam search. On four types of nine agent tasks, AgentRM enhances the base policy model by $8.8$ points on average, surpassing the top general agent by $4.0$. Moreover, it demonstrates weak-to-strong generalization, yielding greater improvement of $12.6$ on LLaMA-3-70B policy model. As for the specializability, AgentRM can also boost a finetuned policy model and outperform the top specialized agent by $11.4$ on three held-in tasks. Further analysis verifies its effectiveness in test-time scaling. Codes will be released to facilitate the research in this area.