Context-attended Adversarial Reinforcement Learning for Robust Multi-step Retrieval Augmented Generation

🌟 Introduction

Official implementation of ACL 2026 findings paper Context-attended Adversarial Reinforcement Learning for Robust Multi-step Retrieval Augmented Generation.

Key Features

• 🔄 Adversarial Reforcement Training: Alternates between RL training and dynamic attack optimization
• ⚡ TextGrad-based Attack Optimization: Uses gradient-based text optimization to craft stronger adversarial documents
• 🛡️ Enhanced Robustness: Significantly improves model resilience to retrieval poisoning attacks

---

🔧 Installation

Clone this repository, then create a conda environment and install the required packages.

# Clone repository
git clone https://github.com/yingtaoren/CARE.git
cd CARE

# Create conda environment
conda create -n care python=3.10
conda activate care

# Install packages
pip install -r requirements.txt

💡 Note: If you encounter any issues when installing the Python packages, we recommend following the official installation instructions provided by FlashRAG#Installation.

If you are using Docker, run the command below:

# pull vllm docker image
docker pull vllm/vllm-openai:v0.10.2

# And install Additional Package
pip install textgrad
pip install falshrag

---

📊 Dataset

We conduct our training and evaluation across multiple multi-hop QA benchmarks:

Type	Dataset	Train Num	Test Num
Multi-hop QA	HotpotQA	90447	7405
Multi-hop QA	2WikiMultiHopQA	113284	7405
Multi-hop QA	MuSiQue	19938	2417
Multi-hop QA	Bamboogle	Null	125

Data Collection

For initial training data collection (multi-hop QA with retrieved documents), please refer to the DRAG repository which provides the data collection pipeline.

The dataset using in our experiment are provided here.

---

📚 Document Corpus & Index

We use the wiki18_100w dataset as the document corpus. Both the document corpus and the index can be downloaded from:

• ModelScope Dataset (look in retrieval_corpus folder)

---

🤖 Models

Retriever Model

We use e5-base-v2 as the default retriever. Download from: 🤗 intfloat/e5-base-v2

Generator Models

This project supports all LLMs compatible with HuggingFace and vLLM. Recommended models:

Model	Link
Qwen3-4B-Instruct	🔗 Link
Qwen3-30B-A3B-Instruct	🔗 Link
Llama-3.1-8B-Instruct	🔗 Link

Note on Model Configuration: We sincerely apologize for a typo in the published version of our paper. The Implementation Details section incorrectly states that we use Qwen-2.5-3B-Instruct as the backbone model. The correct model used throughout our 4B experiments is Qwen3-4B-Instruct.

This is a textual error and does not affect the reported experimental results. We apologize for any confusion. If you have any questions about our implementation or experiments, please feel free to contact us.

---

🏋️ Training Pipeline

CARE follows an iterative curriculum-based adversarial training pipeline:

Step 1: Data Collection

Collect multi-hop QA training data with retrieved documents using the DRAG framework:

# Follow DRAG repository for data collection
# https://github.com/Huenao/Debate-Augmented-RAG

Step 2: Difficulty Measurement

Measure the difficulty of each training sample using sampling-based evaluation:

python src/difficulty_measurement.py \
    --model_path /path/to/base_model \
    --input_file data/train_data/pure_data/pure_hotpotqa_train.json \
    --output_file data/train_data/pure_data/pure_hotpotqa_train_with_difficulty.json \
    --num_samples 20 \
    --batch_size 8

Step 3: Initial Attack Generation

Generate initial adversarial documents with wrong answer injection:

python src/gen_ini_attack.py \
    --model_path /path/to/attack_generator_model \
    --input_file data/train_data/pure_data/pure_hotpotqa_train_with_difficulty.json \
    --output_file data/train_data/attack_data/attack_round1_train_rl_prompt.json

Step 4-6: Iterative RL Training

We use OpenRLHF for reinforcement learning. Run the training script:

# Round 1 Training
bash train.sh  # Configure with attack_round1_train_rl_prompt.json

After each round, optimize attacks using TextGrad:

# Dynamic Attack Optimization (after round N, before round N+1)
python src/dynamic_attack.py \
    --model_path /path/to/trained_model_roundN \
    --input_file data/train_data/pure_data/pure_data_for_round.json \
    --output_file data/train_data/attack_data/attack_roundN+1_train_rl_prompt.json

Then continue training with the new adversarial data:

# Round N+1 Training
bash train.sh  # Configure with attack_roundN+1_train_rl_prompt.json

📈 Evaluation

Running Evaluation

Evaluate CARE-trained models on poisoned/counterfactual test sets:

python src/answer_vllm_pipeline.py \
    --model_path /path/to/care_trained_model \
    --test_file data/test_data/Poisoned/poisoned_hotpotqa_test.jsonl \
    --output_file results/poisoned_hotpotqa_results.json \
    --tensor_parallel_size 2

Evaluation Metrics

• Exact Match (EM): Strict matching after normalization
• F1 Score: Token-level F1 between prediction and ground truth

---

📁 Project Structure

CARE/
├── data/
│   ├── train_data/
│   │   ├── pure_data/              # Clean training data
│   │   └── attack_data/            # Adversarial training data
│   └── test_data/
│       ├── Poisoned/               # Poisoned test sets
│       └── counterfact/            # Counterfactual test sets
├── src/
│   ├── difficulty_measurement.py   # Sample difficulty evaluation
│   ├── gen_ini_attack.py           # Initial attack document generation
│   ├── dynamic_attack.py           # TextGrad-based attack optimization
│   └── answer_vllm_pipeline.py     # vLLM-based evaluation pipeline
├── logs/                           # Training logs
├── output/                         # Model outputs
├── results/                        # Evaluation results
├── train.sh                        # RL training script
├── requirements.txt                # Python dependencies
└── README.md                       # This file

---

Citation

If you find our work useful, please consider citing our paper:

Paper: Context-attended Adversarial Reinforcement Learning for Robust Multi-step Retrieval Augmented Generation

@inproceedings{ren-etal-2026-context,
    title = "Context-attended Adversarial Reinforcement Learning for Robust Multi-step Retrieval Augmented Generation",
    author = "Ren, Yingtao  and
      Luo, Xiao  and
      Chang, Yu-Cheng  and
      Lin, Chin-teng",
    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.856/",
    doi = "10.18653/v1/2026.findings-acl.856",
    pages = "17305--17319",
    ISBN = "979-8-89176-395-1"
}

---

✨ Acknowledgments

We gratefully acknowledge the following projects:

• FlashRAG: A Python toolkit for the reproduction and development of Retrieval Augmented Generation (RAG) research.
• DRAG: Debate-Augmented RAG framework for training data collection.
• OpenRLHF: An Easy-to-use, Scalable and High-performance RLHF Framework.
• TextGrad: Automatic "Differentiation" via Text for gradient-based text optimization.

---