Search-R1 Example

Overview

This example implements Search R1 within Agent Lightning. It also serves as a demonstration of a framework-free agent training pipeline, showing how to run end-to-end RL training without relying on specialized frameworks. It's tested and compatible with Agent-lightning v0.2.x.

The example is designed to run on a single node with 8 GPUs, each having at least 40 GB of memory.

Included Files

Prepare Data and Environment

Run the following script once to prepare data and the retriever environment:

bash data_process.sh

This script performs the following steps:

• Creates a new conda environment named retriever.
• Downloads the Wikipedia data used to build the retrieval database.
• Downloads the training and testing datasets.
• Stores all data under the newly created data/ directory.

The environment setup and data-processing logic are adapted from PeterGriffinJin/Search-R1.

Prepare Retrieval Server

To start the retrieval server, run:

bash retrieval_launch.sh

This script activates the previously created retriever environment and starts a retrieval server at http://127.0.0.1:8000 using the downloaded Wikipedia data. The server receives user queries and returns a ranked list of retrieved text passages.

The retrieval server implementation is based on search_r1/search/retrieval_server.py](https://github.com/PeterGriffinJin/Search-R1/blob/main/search_r1/search/retrieval_server.py).

⚠️ Note: Keep the retrieval server running during training (for example, in a separate tmux session or terminal window).

Run RL Training (GRPO) with Llama-3.2-3B-Instruct

1. Start Ray

   bashCopy

   bash ../../scripts/restart_ray.sh
   

   If you plan to use WandB for experiment tracking, set the environment variable WANDB_API_KEY before starting Ray.

2. Start the Training Server In another terminal, run:

   bashCopy

   python train_search_r1_agent.py llama
   

   This script starts the RL training. Each agent follows the Search-R1 workflow, retrieving information from the database and generating answers accordingly.

Benchmark Results

We evaluated Search-R1 across seven diverse question-answering benchmarks, covering both General QA (NQ, TriviaQA, PopQA) and complex multi-hop reasoning tasks (HotpotQA, 2WikiMultiHopQA, Musique, and Bamboogle).

The following tables compare the performance of the original Search-R1 implementation and the Agent-Lightning version across various base models.