[Paper Review] AGILE: A Novel Reinforcement Learning Framework of LLM Agents

Peiyuan Feng, Yichen He+5 more|arXiv (Cornell University)|May 23, 2024

Multi-Agent Systems and Negotiation4 citations

TL;DR

AGILE is a novel reinforcement learning framework for LLM agents that integrates memory, tools, expert consultation, and reflection into a unified system, trained via PPO to optimize end-to-end performance. It achieves state-of-the-art results on ProductQA and MedMCQA, outperforming GPT-4 and GPT-3.5, with ablation studies confirming the critical role of all components.

ABSTRACT

We introduce a novel reinforcement learning framework of LLM agents named AGILE (AGent that Interacts and Learns from Environments) designed to perform complex conversational tasks with users, leveraging LLMs, memory, tools, and interactions with experts. The agent possesses capabilities beyond conversation, including reflection, tool usage, and expert consultation. We formulate the construction of such an LLM agent as a reinforcement learning (RL) problem, in which the LLM serves as the policy model. We fine-tune the LLM using labeled data of actions and the PPO algorithm. We focus on question answering and release a dataset for agents called ProductQA, comprising challenging questions in online shopping. Our extensive experiments on ProductQA, MedMCQA and HotPotQA show that AGILE agents based on 7B and 13B LLMs trained with PPO can outperform GPT-4 agents. Our ablation study highlights the indispensability of memory, tools, consultation, reflection, and reinforcement learning in achieving the agent's strong performance. Datasets and code are available at https://github.com/bytarnish/AGILE.

Motivation & Objective

• To unify multiple LLM agent capabilities—reasoning, planning, reflection, tool-use, memory, and expert consultation—into a single, end-to-end trainable framework.
• To address the lack of benchmarks that comprehensively evaluate the integration of multiple agent components in complex question-answering tasks.
• To develop a new benchmark, ProductQA, with 88,229 real-world product queries across 26 categories to evaluate agent performance in diverse, realistic scenarios.
• To demonstrate that reinforcement learning, specifically PPO fine-tuning, significantly improves agent performance beyond supervised fine-tuning alone.
• To validate the necessity and effectiveness of expert consultation as a mechanism for improving accuracy and adaptability in complex, out-of-distribution tasks.

Proposed method

• The AGILE framework structures the agent as a system with four core modules: LLM (policy model), memory, tools, and executor, where the executor orchestrates actions and context updates.
• The LLM generates actions (tokens) that are interpreted and executed by the executor, which manages memory retrieval, tool invocation, and expert consultation.
• Expert consultation is modeled as a proactive mechanism where the agent seeks human feedback when uncertain, then reflects on and stores the feedback for future use.
• The agent is trained via a two-stage process: first, supervised fine-tuning (SFT) on labeled action sequences; second, policy optimization using Proximal Policy Optimization (PPO) to maximize task rewards.
• The framework supports end-to-end training of reasoning, planning, reflection, and consultation abilities within a single RL objective.
• A new benchmark, ProductQA, is constructed with real Amazon customer queries, including fact-based, reasoning, and recommendation questions, with disjoint training and test sets to evaluate zero-shot generalization.

Experimental results

Research questions

• RQ1 Can a unified LLM agent framework effectively integrate memory, tool-use, reflection, and expert consultation in a single, end-to-end trainable system?
• RQ2 Does reinforcement learning with PPO significantly improve agent performance compared to supervised fine-tuning alone in complex question-answering tasks?
• RQ3 How critical are individual components—memory, tools, expert consultation, and reflection—to overall agent performance?
• RQ4 Can the proposed agent framework generalize to new product categories and outperform strong baselines like GPT-4 in zero-shot settings?
• RQ5 To what extent does expert consultation improve accuracy and adaptability in handling complex, ambiguous, or out-of-distribution queries?

Key findings

• The AGILE agent fine-tuned with PPO (agile-vic13b-ppo) achieved a 9.2% relative improvement in total score over GPT-4 on the ProductQA benchmark.
• The same agent outperformed GPT-3.5 by 90.8% in relative total score, demonstrating the effectiveness of the PPO training regime.
• Ablation studies showed that removing memory or tools increased expert consultation requests by 17.4% and 25.9%, respectively, indicating their role in reducing reliance on human input.
• Disabling the expert consultation function led to a 10.7% drop in accuracy, confirming its critical role in handling complex queries.
• On MedMCQA, the agile-mek7b-ppo agent achieved 85.2% accuracy, surpassing GPT-4-MedPrompt’s SOTA of 79.1% and the base Meerkat-7b’s 53.4%.
• The PPO fine-tuning stage improved performance by 2.3% relative gain over the SFT baseline, confirming the value of reinforcement learning in policy optimization.