RAPO: Teaching LLM Agents When to Think for Themselves and When to Peek at Notes

This post introduces RAPO (Retrieval-Augmented Policy Optimization), a method for expanding the exploration capabilities of LLM agents during reinforcement learning. The core analogy: a student allowed to peek at an "expert notes" book during an exam — but who must learn *when* consulting the notes actually helps.

The Problem: On-Policy RL Only Reinforces Existing Habits

LLM agents that perform multi-step tool use and reasoning are typically trained with reinforcement learning. But pure on-policy exploration means the agent only learns from its own generated reasoning trajectories — like a student who only drills their own mistake notebook. Recent empirical work confirms that on-policy RL largely amplifies existing behaviors of the base model rather than discovering new reasoning strategies.

Existing off-policy methods (e.g., LUFFY) inject complete trajectories from stronger models, but only at the trajectory level. For multi-step tasks that get stuck at step 3 and run smoothly at step 5, such coarse guidance offers limited value.

RAPO's Core Components

1. Step-Level Trajectory Cache

• A stronger off-policy model (AEPO-Qwen3-14B) generates many reasoning trajectories on the training set.
• Only trajectories with correct final answers are kept.
• Each trajectory is decomposed into step-level Key-Value pairs:
• Key: the reasoning context before the step (the current "predicament state")
• Value: the reasoning action taken at that context

2. Hybrid-Policy Sampling

At training time, the agent produces: 1. Purely on-policy trajectories (fully self-generated reasoning) 2. Hybrid-policy trajectories where retrieval is triggered at each step with 50% probability. The current reasoning history queries the buffer; the retrieved off-policy snippet is inserted into the context and the agent continues reasoning from there ("Retrieval-then-Reasoning").

3. Retrieval Reward: Learning *When* to Consult

External information can help or mislead. RAPO scores retrieval using entropy:
• Retrieval quality: the entropy drop after retrieval (passed through tanh). Lower entropy = more confident model = helpful retrieval.
• Retrieval timing: higher entropy before retrieval means the model was confused, so retrieval is more likely valuable.

The product forms the retrieval reward, encouraging the agent to seek help exactly when it is most confused and the help is genuinely useful — and to ignore retrieval otherwise.

Key Experimental Findings

Experiments cover 14 datasets, 3 task types, 3 base models, and 13 baselines. Highlights:

• 100% noisy retrieval: even when every retrieval returns pure random noise, RAPO still beats baselines. The model learns via the entropy-based reward to simply ignore unhelpful retrievals — a transferable meta-capability.
• Weak teachers can help strong students: a buffer built from Qwen2.5-3B (a relatively weak model) still yields performance gains, showing the value lies in diverse step-level perspectives rather than copying full answers.
• Step-level vs. trajectory-level ablation: a trajectory-level retrieval variant ("w/ TO") performs worst, confirming that fine-grained, state-matched retrieval is the key design.

Theoretical View: Information Bottleneck

The paper provides analysis showing RAPO's retrieval-aware optimization implicitly forms an information bottleneck: the model learns to extract only the useful parts of retrieved information and discard the rest, improving generalization rather than overfitting to specific retrieval results.

Takeaway

RAPO teaches agents a form of metacognition: not just *how* to use external experience, but *when* to use it. The authors suggest future work could build richer step-level knowledge bases from multiple models with different strengths (math, search, coding), letting agents truly "stand on the shoulders of giants."

Reference: RAPO: Expanding Exploration for LLM Agents via Retrieval-Augmented Policy Optimization. arXiv:2603.03078, 2025. https://arxiv.org/abs/2603.03078