RouteRAG: Efficient Retrieval-Augmented Generation from Text and Graph via Reinforcement Learning
July 1, 2026 · View on GitHub
💡 Overview
This is the official code release of the following paper:
RouteRAG: Efficient Retrieval-Augmented Generation from Text and Graph via Reinforcement Learning
🌟 Introduction
RouteRAG is a reinforcement learning-based framework that enables LLMs to perform multi-turn and adaptive graph-text hybrid retrieval-augmented generation (RAG). Unlike existing RAG systems that rely on fixed or handcrafted retrieval pipelines, RouteRAG jointly optimizes the entire generation process via RL, empowering the model to dynamically decide:
- 🤔 When to reason;
- 🔍 What to retrieve (and from unstructured texts or structured knowledge graphs);
- ✅ When to produce the final answer.
✨ Highlights
- Adaptive Hybrid Retrieval: Seamlessly routes between text retrieval and graph retrieval based on the current reasoning state, leveraging the complementary strengths of both retrieval modes.
- Two-Stage RL Training: A two-stage RL training framework that first optimizes task accuracy (EM reward) and then incorporates retrieval efficiency reward, enabling the model to exploit hybrid evidence while avoiding unnecessary retrieval overhead.
- End-to-End RL Optimization: The full generation policy is trained end-to-end without any handcrafted rules or fixed pipeline constraints.
📦 1. Installation
Install all dependencies:
conda create -n routerag python=3.10
conda activate routerag
pip install -r requirements.txt
🚀 2. Model Training
⚙️ Step 1: Start Services
Set up the environment variables and start all required services:
# Set GPU device for IE model
export CUDA_VISIBLE_DEVICES=4
# Start the OpenIE model service for graph construction
vllm serve /path/to/OpenIE_model \
--max_model_len 4096 \
--gpu-memory-utilization 0.95 \
--port 8000
Parameter Description:
OpenIE_model: OpenIE model path (Llama-3.1-8B-Instruct/Llama-3.3-70B-Instruct in our experiments)--max_model_len: Maximum sequence length--gpu-memory-utilization: GPU memory utilization limit--port: Service port for the OpenIE model
In a new terminal, start the retrieval service:
# Set GPU device for retrieval service
export CUDA_VISIBLE_DEVICES=5
# Start retrieval service
python retrieval_api.py \
--llm_model_name /path/to/OpenIE_model \
--llm_base_url http://localhost:8000/v1 \
--embedding_model_name /path/to/embedding_model \
--dataset hotpotqa_10k \
--corpus_path train/RL_dataset/corpus.json \
--port 8001
Parameter Description:
--llm_base_url: vLLM service API address--embedding_model_name: Embedding model path (contriever/NV-Embed-v2 in our experiments)--dataset: Training corpus name--corpus_path: Training corpus path--port: Retrieval service port
🎯 Step 2: Two-Stage Reinforcement Learning
Navigate to the train directory:
cd train
Execute the two-stage training script:
bash train_grpo_grag_two_stages.sh
📝 Note: Please modify the paths in train_grpo_grag_two_stages.sh to match your local setup, including TRAIN_DATA_DIR, TEST_DATA_DIR, BASE_MODEL, SWITCH_STEP, SEED, retriever.url, and trainer.default_local_dir.
Two-Stage Training Process:
- Stage 1 (Steps 1-<SWITCH_STEP>): EM reward only reinforcement learning training
- Stage 2 (Steps <SWITCH_STEP+1>-40): Add efficiency reward to the training
🧪 3. Model Testing
Run the test script in the main directory (need to start the OpenIE model service first):
python test.py --dataset <dataset_name> \
--llm_base_url http://localhost:8000/v1 \
--llm_name /path/to/OpenIE_model \
--reader_name /path/to/trained/model \
--embedding_name /path/to/embedding_model \
--deepsearch
📁 Directory Structure
routerag/
├── routerag/ # Main routerag framework
│ ├── embedding_model/ # Embedding model implementations
│ ├── evaluation/ # Evaluation implementations
│ ├── information_extraction/ # Information extraction implementations
│ ├── llm/ # LLM model implementations
│ ├── prompts/ # Prompt templates and management
│ ├── utils/ # Utility functions
│ ├── routerag.py # Main routerag class
│ ├── embedding_store.py # Embedding storage management
│ ├── rerank.py # Reranking functionality
│ └── README.md # routerag module documentation
├── train/ # Training framework
│ ├── verl/ # VERL framework
│ ├── routerag/ # routerag training components
│ │ └── llm_agent/ # LLM agent implementations
│ ├── RL_dataset/ # Reinforcement learning datasets
│ ├── train_grpo_grag_two_stages.sh # Two-stage training script
│ ├── setup.py # Training package setup
│ ├── pyproject.toml # Python project configuration
│ ├── LICENSE # License file
│ └── Notice.txt # Notice file
├── dataset/ # Evaluation datasets
├── outputs/ # Output directory for results
├── retrieval_api.py # FastAPI retrieval service
├── test.py # Testing script
├── requirements.txt # Main dependencies
└── README.md # This file
🙏 Acknowledgement
We would like to thank the following projects for their foundational work:
- Search-R1: We built upon their reinforcement learning framework for model training.
- HippoRAG: Our graph retriever is implemented based on HippoRAG 2.
📚 Citation
If you find our work useful, please kindly cite:
@inproceedings{guo-etal-2026-routerag,
title = "{R}oute{RAG}: Efficient Retrieval-Augmented Generation from Text and Graph via Reinforcement Learning",
author = "Guo, Yucan and
Su, Miao and
Guan, Saiping and
Sun, Zihao and
Jin, Xiaolong and
Guo, Jiafeng and
Cheng, Xueqi",
editor = "Liakata, Maria and
Moreira, Viviane P. and
Zhang, Jiajun and
Jurgens, David",
booktitle = "Findings of the {A}ssociation for {C}omputational {L}inguistics: {ACL} 2026",
month = jul,
year = "2026",
address = "San Diego, California, United States",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2026.findings-acl.1502/",
pages = "30042--30059",
ISBN = "979-8-89176-395-1
}