AgenticRag-R1: stack memory plus process rewards lift 3B HotpotQA F1 by 6.76

AgenticRag-R1: Agentic Reinforcement Learning with Stack Memory for Multi-Step Reasoning, Retrieval and Memorizing

Xinke Jiang, Yue Fang, Zhibang Yang, Jiaran Gao, Zhixin Zhang, Tao Feng, Rihong Qiu, Wentao Zhang, Hongxin Ding, Ruizhe Zhang, Yongxin Xu, Yuheng Huang, Xu Chu, Junfeng Zhao, Yasha Wang

cs.MA, cs.AI

2026-08-30

Peking University trains agentic RAG with stack memory, action-level rewards, and information-aware rejection. On Qwen2.5-3B, HotpotQA F1 reaches 44.00, 6.76 above Search-R1.

What problem this solves

Retrieval-augmented generation (RAG) helps with facts, but multi-hop questions need the model to retrieve while it reasons and to revise intermediate claims. Recent RL-trained agentic RAG systems usually expose only coarse <think> and <search> actions, append memory, and score the final answer. Noisy documents and bad inferences then sit in context for good. Policy learning also prefers short, easily rewarded templates.

Peking University splits that into two failures: the action space is too coarse, and the training trajectories are too easy. AgenticRag-R1 adds a stack memory, action-level process rewards, and an information-aware filter so the policy has to practice long-horizon retrieval and repair.

Method

Working memory is a last-in, first-out stack. Each entry is one interpretable action. Six actions:

Reward is hierarchical. An outcome term (Exact Match or F1 on the final answer) is spread over unmasked generated tokens; retrieved observations are masked so they do not send gradients. Process rewards land on action tokens: format (1.0 for conclusion, 0.5 for other valid tags), retrieval relevance from a general reward model in [0, 1], and whether a memory edit helped, scored by the same model on the current stack. Process mass is capped and scaled, so the agent cannot farm points by spamming tools.

At each step the policy rolls out K trajectories per question. Across the mini-batch each rollout is scored as variance squared times reward plus the per-question mean. Only the top fraction is kept. High variance means the question supports distinct paths; high reward means that path is actually good.

Results

Training uses multi-hop subsets of HotpotQA and 2Wiki, dropping items that need no retrieval or only one trivial search. Evaluation uses a Wikipedia dump dated 1 Nov 2023 on Qwen2.5 at 1.5B, 3B, and 7B. The headline metric is F1.

BackboneMethod2WikiHotpotQAAvg
3BSearch-R129.9037.2430.16
3BAgenticRag-R132.9244.0033.55
7BSearch-R135.0338.89
7BAgenticRag-R138.3445.1536.60

On 3B, HotpotQA rises 6.76 absolute over the strongest baseline and the seven-task average rises 3.39. On 7B, Bamboogle reaches 49.21, about 7.2 above the best agentic baseline; Natural Questions falls to 23.60 versus Search-R1 at 29.59, so single-hop open-domain items are not along for the ride.

Dropping memory actions cuts the 3B average from 33.55 to 27.37; dropping rejection cuts it to 25.43; dropping both retrieval and memory process rewards leaves 19.94. Raising the step cap from 10 to 30 lifts TriviaQA from 37.53 to 55.28, while Search-R1 slides from 34.21 to 32.28. Cold-start RL averages 33.55; SFT then RL reaches 35.20; SFT alone is 22.40.

Why it matters

If you train a retrieval agent, this paper treats reversible memory as a first-class action rather than a log pasted into the prompt. Budgeted process rewards fight "search more, score more." The filter admits that most multi-hop items are shallow, and unfiltered RL will sit on short traces.

The cost is real. The stack and dynamic retrieval are heavier. The rejection threshold needs tuning; F1 looks better in the 0.6-0.8 band. The 7B drop on NQ is a reminder that a long-horizon design is not a free upgrade for single-hop RAG. Reuse the action design and the filter, not an anonymous checkpoint.

Limitations

The authors flag compute overhead and sensitivity of rejection hyperparameters. The Bamboogle long-horizon curve is not monotone: 10.03 at 10 steps, 14.29 at 20, 35.47 at 30, which sits poorly next to a "more steps, more stable" story. The 7B NQ regression is not diagnosed. Process rewards depend on an external general reward model whose biases are not reported. Training filtered out easy items while the eval mix still contains single-hop questions, so the two distributions are offset.

Terms

Source

What people are saying

Related papers

All paper explainers