Table of Contents
Fetching ...

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.

InstructRAG: Leveraging Retrieval-Augmented Generation on Instruction Graphs for LLM-Based Task Planning

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.

Paper Structure

This paper contains 19 sections, 10 equations, 3 figures, 14 tables, 4 algorithms.

Figures (3)

  • Figure 1: Architecture of the proposed InstructRAG with multi-agent meta-reinforcement learning, illustrated on HotpotQA.
  • Figure 2: F1 scores and few-shot learning times wrt the number of unseen tasks or samples with DeepSeek-V2 on HotpotQA.
  • Figure 3: F1 scores and few-shot learning times wrt the number of unseen tasks or samples on HotpotQA, where (a)-(d) are for GLM-4, and (e)-(h) are for GPT-4o mini.