Wiki-R1

Wiki-R1 is a reinforcement learning training framework for Knowledge-Intensive Visual Question Answering (KIVQA), built on top of verl. It features curriculum learning with label propagation to progressively train VLMs on open-domain visual QA tasks.

📚 Quick Links:Dataset Documentation | HuggingFace Dataset

Project Structure

Wiki_R1/
├── exps/                        # Training launch scripts
│   └── train_wiki_r1.sh         # Main training entry
├── recipe/dapo/src/             # Core training code
│   ├── main_kivqa_wiki_r1.py    # Hydra entry point
│   ├── kivqa_dapo_ray_trainer.py # Curriculum trainer with label propagation
│   ├── data_structure.py        # Ray actor state (SampleState, etc.)
│   └── config/
│       └── curriculum_trainer_infoseek_evqa.yaml  # Hydra config
├── verl/                        # verl framework (training infra)
├── data/                        # Training data directory (to be prepared, see below)
│   ├── annotation/              # JSON annotation files from HuggingFace
│   └── source/                  # Image files (OVEN/Infoseek/EVQA)
├── requirements.txt
├── DATASET_CARD.md              # Detailed dataset documentation
└── README.md

Installation

# 1. Create conda environment
conda create -n wiki_r1 python=3.10 -y
conda activate wiki_r1

# 2. Install dependencies
pip install -r requirements.txt

# 3. Install verl (editable mode)
pip install -e .

Key Dependencies

• Python >= 3.10
• PyTorch >= 2.1
• CUDA >= 12.1
• vLLM >= 0.8.2
• Ray >= 2.10
• Hydra / OmegaConf

Data Preparation

All annotation files are available on HuggingFace: https://huggingface.co/datasets/Artanic30/Wiki_R1_Train

For detailed descriptions of each file, see DATASET_CARD.md.

Step 1: Download annotations from HuggingFace

# Install huggingface_hub if needed
pip install huggingface_hub

# Download all JSON annotation files to data/annotation
huggingface-cli download Artanic30/Wiki_R1_Train --repo-type dataset --local-dir ./data/annotation

This will download the following files to data/annotation/:

• merge_infoseek_train_filtered_balance_sample20k_top5_and_evqa_train_sample20k_top5_w_I2T.json - Main training set (40k samples)
• final_related_KB_reflectiVA_v2.json - Knowledge base entity information
• oven_id2path.json - OVEN image ID to file path mapping
• final_data_v2_kb_sim.json - KB similarity matrix for label propagation

Step 2: Download image data (required separately)

The image files are NOT included in the HuggingFace dataset. You need to download them separately:

1. OVEN Dataset: Download from OVEN official source

2. EVQA Dataset: Download images from EVQA official source

Place all images under data/source/ directory. The exact structure should follow the paths specified in oven_id2path.json.

Step 3: Verify data paths

After downloading, your directory structure should look like:

data/
├── annotation/
│   ├── merge_infoseek_train_filtered_balance_sample20k_top5_and_evqa_train_sample20k_top5_w_I2T.json
│   ├── final_related_KB_reflectiVA_v2.json
│   ├── oven_id2path.json
│   ├── final_data_v2_kb_sim.json
└── source/                                             # Image files (organize according to oven_id2path.json)
    ├── EVQA_kb_img/                                    # EVQA knowledge base images
    ├── inaturalist/                                    # iNaturalist dataset images
    ├── landmarks/                                      # Landmarks dataset images
    └── oven_images/                                    # OVEN dataset images

Verify key files exist:

ls data/annotation/merge_infoseek_train_filtered_balance_sample20k_top5_and_evqa_train_sample20k_top5_w_I2T.json
ls data/annotation/final_related_KB_reflectiVA_v2.json
ls data/annotation/oven_id2path.json
ls data/annotation/final_data_v2_kb_sim.json

Training

Quick Start

# From project root
bash exps/train_wiki_r1.sh

Custom Configuration

You can override key parameters via environment variables:

# Use a different model
MODEL_PATH=/path/to/Qwen2.5-VL-3B-Instruct \
NUMGPU=8 \
gpu=0,1,2,3,4,5,6,7 \
EXP_NAME=my_experiment \
CKPTS_DIR=output/ckpts/my_exp \
bash exps/train_wiki_r1.sh

Key Hyperparameters

Parameter	Default	Description
NUMGPU	4	Number of GPUs
MODEL_PATH	Qwen/Qwen2.5-VL-3B-Instruct	Pretrained model path
total_epoch	200	Total training epochs
train_prompt_bsz	4	Training prompt batch size
n_resp_per_prompt	4	Number of responses per prompt (for GRPO)
label_prop_enable	True	Enable label propagation
ent_weight	0.5	Label propagation entity weight
prop_topk	100	Top-k for label propagation
diff_thr	0.552	Difficulty threshold for curriculum
max_prompt_length	8000	Max prompt token length
max_response_length	256	Max response token length

Resume Training

Training automatically resumes from the latest checkpoint if trainer.resume_mode=auto (default). Checkpoints are saved to CKPTS_DIR.

Monitoring

TensorBoard logs are saved to ${CKPTS_DIR}/tf_log:

tensorboard --logdir output/ckpts/Wiki_R1/default/tf_log

Testing

Testing scripts and evaluation pipeline will be added in the next update.

Acknowledgement

This project is built on verl (Volcano Engine Reinforcement Learning for LLMs).