Learning to Program Retrievals: Composing Search Primitives via Reinforcement Learning
Organizations: KAIST AI · Meta Superintelligence Labs
Abstract
Retrieval-augmented language models increasingly use multi-step search to solve knowledge-intensive tasks, yet existing systems typically expose retrieval through fixed pipelines or collections of atomic tools. We introduce Retrieval Program Learning (RPL), a programmable search framework that allows a policy to compose search primitives into executable programs with local state and control flow. This enables adaptive search that can retrieve broadly, process intermediate results, and selectively construct model context without tying retrieval breadth to context consumption. However, the resulting flexibility creates a difficult exploration problem: useful programs may require several coordinated decisions, while reinforcement learning often supplies only sparse feedback. To address this challenge, we propose program mutation, a grammar-constrained and execution-verified procedure that generates local variants of suboptimal programs sampled during RL, converts them into training trajectories, and distills them back into the policy. We evaluate RPL on open-domain question answering and corpus-level reasoning tasks, comparing it with a range of retrieval and reasoning baselines. RPL consistently improves open-domain QA performance across model sizes and in- and out-of-domain benchmarks, with particularly large gains on multi-hop tasks. Notably, a 4B model trained with RPL outperforms all 8B model baselines. On corpus-level reasoning, RPL achieves the highest answer F1 and document-coverage F1 for both model sizes while substantially reducing context consumption and retrieval turns. Our results show that programmable search expands the retrieval strategies that agents can learn, while program mutation provides an effective mechanism for exploring useful program structures beyond direct on-policy sampling.
Figures & tables
| Primitive | Description |
| sparse_retrieval(q, top_k) | Retrieve a ranked list of documents using lexical matching |
| dense_retrieval(q, top_k) | Retrieve a ranked list of documents using semantic similarity |
| fuse(*lists) | Fuse ranked lists using reciprocal-rank fusion |
| dedup(docs) | Remove duplicates while preserving their first occurrence |
| rerank(q, docs, top_k) | Rerank documents with respect to and retain the top- |
| Methods | Single-Hop QA | Multi-Hop QA | Avg | |||||
| NQ † | TriviaQA ⋆ | PopQA ⋆ | HotpotQA † | 2Wiki ⋆ | MuSiQue ⋆ | Bamb ⋆ | ||
| Qwen3-4B | ||||||||
| Direct | 0.142 | 0.362 | 0.138 | 0.170 | 0.261 | 0.026 | 0.072 | 0.167 |
| CoT | 0.161 | 0.417 | 0.156 | 0.186 | 0.242 | 0.036 | 0.304 | 0.215 |
| RAG | 0.319 | 0.589 | 0.389 | 0.243 | 0.140 | 0.065 | 0.208 | 0.279 |
| IRCoT | 0.316 | 0.574 | 0.372 | 0.292 | 0.229 | 0.133 | 0.456 | 0.339 |
| Methods | TopK | Count | Sort | MinMax | Avg | |||||
| F1 | dF1 | F1 | dF1 | F1 | dF1 | F1 | dF1 | F1 | dF1 | |
| Qwen3-4B | ||||||||||
| RAG | 4.97 | 12.80 | 0.40 | 13.86 | 8.45 | 11.68 | 3.23 | 12.91 | 4.27 | 12.81 |
| IRCoT | 4.30 | 12.76 | 1.11 | 13.84 | 5.07 | 13.07 | 2.49 | 12.91 | 3.24 | 13.14 |
| Search-R1 | 15.82 | 20.39 | 4.19 | 20.38 | 24.47 | 20.21 | 11.53 | 19.35 | 14.00 | 20.08 |
| MARAG-R1 | 13.26 | 23.79 | 6.98 | 24.08 | 16.95 | 25.45 | 11.44 | 24.08 | 12.16 | 24.35 |
| Methods | Single-Hop QA | Multi-Hop QA | Avg | |||||
| NQ † | TrivQA ⋆ | PopQA ⋆ | HotpotQA † | 2Wiki ⋆ | MuSiQue ⋆ | Bamb ⋆ | ||
| Qwen3-4B | ||||||||
| RPL (Dense) | 0.504 | 0.656 | 0.476 | 0.481 | 0.475 | 0.204 | 0.512 | 0.473 |
| RPL (Dual) | 0.502 | 0.670 | 0.487 | 0.501 | 0.471 | 0.199 | 0.512 | 0.478 |
| RPL (SDK) | 0.516 | 0.684 | 0.510 | 0.517 | 0.463 | 0.230 | 0.488 | 0.487 |
| RPL (PM) | 0.536 | 0.701 | 0.517 | 0.539 | 0.491 | 0.258 | 0.560 | 0.515 |
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
| Methods | Single-Hop QA | Multi-Hop QA | Avg | |||||
| NQ † | TriviaQA ⋆ | PopQA ⋆ | HotpotQA † | 2Wiki ⋆ | MuSiQue ⋆ | Bamb ⋆ | ||
| Qwen2.5-3B | ||||||||
| ReSearch | 0.365 | 0.571 | 0.395 | 0.351 | 0.272 | 0.095 | 0.266 | 0.331 |
| MR-Search | 0.477 | 0.635 | 0.460 | 0.419 | 0.401 | 0.165 | 0.344 | 0.414 |
| GiGPO | 0.420 | 0.595 | 0.424 | 0.369 | 0.370 | 0.126 | 0.641 | 0.421 |
| IG-Search | 0.450 | 0.614 | 0.460 | 0.431 | 0.440 | 0.191 | 0.424 | 0.430 |
| Methods | Single-Hop QA | Multi-Hop QA | Avg | |||||
| NQ | TriviaQA | PopQA | HotpotQA | 2Wiki | MuSiQue | Bamboogle | ||
| Qwen3-4B | ||||||||
| CoT | 0.146 / 0.161 | 0.382 / 0.417 | 0.147 / 0.156 | 0.183 / 0.186 | 0.233 / 0.242 | 0.043 / 0.036 | 0.248 / 0.304 | 0.197 / 0.215 |
| RAG | 0.207 / 0.319 | 0.537 / 0.589 | 0.349 / 0.389 | 0.221 / 0.243 | 0.149 / 0.140 | 0.047 / 0.065 | 0.176 / 0.208 | 0.241 / 0.279 |
| IRCoT | 0.299 / 0.316 | 0.527 / 0.574 | 0.376 / 0.372 | 0.256 / 0.292 | 0.211 / 0.229 | 0.090 / 0.133 | 0.264 / 0.456 | 0.289 / 0.339 |
| Search-o1 | 0.335 / 0.285 | 0.554 / 0.547 | 0.382 / 0.354 | 0.270 / 0.283 | 0.197 / 0.226 | 0.084 / 0.122 | 0.232 / 0.400 | 0.293 / 0.317 |
| Hyperparameter | Value | |
| Training | Epochs | 1 |
| Batch size | 32 | |
| Optimization | Learning rate | 1e-6 |
| LR scheduler | constant | |
| Warm-up steps | none | |
| Weight decay | 0.01 |
| Hyperparameter | Value | |
| Algorithm | RL algorithm | GRPO |
| Epochs per rollout | 1 | |
| KL coefficient | 0.001 | |
| Discount factor | 1.0 | |
| Training steps (open-domain QA) | 1,000 | |
| Training steps (GlobalQA) | 500 |