InstructRAG: Leveraging Retrieval-Augmented Generation on Instruction Graphs for LLM-Based Task Planning
Zheng Wang, Shu Xian Teo, Jun Jie Chew, Wei Shi
TL;DR
This paper tackles the limitations of TAO-based LLM planning by introducing InstructRAG, a retrieval-augmented, graph-grounded framework that jointly optimizes enlargability and transferability. It embeds past instruction paths in an instruction graph and deploys two cooperative agents: an RL-Agent to expand graph coverage and a meta-learning ML-Agent to enhance generalization via in-context exemplars. The approach is trained end-to-end in a multi-task, meta-reinforcement learning setting and evaluated across four diverse task-planning datasets, showing up to 19.2% improvements over the best baselines and rapid adaptation to unseen tasks. The work advances practical LLM-based task planning by grounding generation in task-relevant retrieved plans and enabling efficient transfer to new domains and tasks.
Abstract
Recent advancements in large language models (LLMs) have enabled their use as agents for planning complex tasks. Existing methods typically rely on a thought-action-observation (TAO) process to enhance LLM performance, but these approaches are often constrained by the LLMs' limited knowledge of complex tasks. Retrieval-augmented generation (RAG) offers new opportunities by leveraging external databases to ground generation in retrieved information. In this paper, we identify two key challenges (enlargability and transferability) in applying RAG to task planning. We propose InstructRAG, a novel solution within a multi-agent meta-reinforcement learning framework, to address these challenges. InstructRAG includes a graph to organize past instruction paths (sequences of correct actions), an RL-Agent with Reinforcement Learning to expand graph coverage for enlargability, and an ML-Agent with Meta-Learning to improve task generalization for transferability. The two agents are trained end-to-end to optimize overall planning performance. Our experiments on four widely used task planning datasets demonstrate that InstructRAG significantly enhances performance and adapts efficiently to new tasks, achieving up to a 19.2% improvement over the best existing approach.
