Source-linked AI summary
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
TL;DR
Existing RAG and agentic RL systems struggle with complex multi-step reasoning because they use coarse actions, irreversible memory updates, and weak trajectory-level reward assignment. AGENTICRAG-R1 integrates fine-grained actions with stack memory, hierarchical rewards, and information-aware rollout rejection, and consistently outperforms strong baselines across benchmark settings. The paper also identifies computational overhead and broader-domain evaluation as remaining boundaries.
Problem
Existing RAG systems struggle with adaptive retrieval and continuous intermediate-context revision, while agentic RL methods commonly use coarse actions and final-outcome rewards.
Method
AGENTICRAG-R1 combines fine-grained reasoning, retrieval, and memory actions with stack-based memory, action-aware hierarchical rewards, and information-aware trajectory rejection.
Results
AGENTICRAG-R1 consistently outperforms strong RAG and agentic baselines across in-domain and out-of-domain benchmarks and backbone scales; it improves TriviaQA by about +7.2% with the 1.5B backbone.
Takeaways & Limitations
The framework supports more robust, interpretable, memory-aware, and long-horizon reasoning behaviors across multi-hop, open-domain, and agentic reasoning settings.
Takeaways & Limitations
The memory stack and dynamic retrieval introduce relatively high computational overhead, creating processing-speed and efficiency challenges for large-scale or real-time applications.
Abstract
from arXiv · showhide
Retrieval-Augmented Generation (RAG) improves the factuality of large language models (LLMs), yet existing RAG systems often struggle with complex, multi-step reasoning that requires adaptive retrieval and continuous revision of intermediate contexts. Recent reinforcement learning (RL)-based agentic RAG methods partially alleviate this issue, but typically rely on coarse-grained action spaces and trajectory-level rewards, resulting in weak reward assignment and a bias toward short-horizon, stereotyped reasoning template. To address, we propose AgenticRag-R1, a RL framework that deeply integrates reasoning, retrieval, and memory via a memory stack and fine-grained action space, supported by hierarchical action-aware rewards and an information-aware trajectory rejection strategy to enable effective long-horizon learning. Experiments across a diverse set of multi-hop, open-domain, and agentic reasoning benchmarks, spanning multiple backbone model sizes, demonstrate that AgenticRag-R1 consistently outperforms strong baselines. Moreover, AgenticRag-R1 learns more robust, interpretable, and memory-aware reasoning behaviors, highlighting the effect of fine-grained action modeling and information-aware optimization for long-horizon reasoning. Our code is anonymous available at https://github.com/jiangxinke/Harness-RL/tree/AgenticRAG-R1-Whitebox.
1 Introduction
AGENTICRAG-R1 addresses the difficulty of integrating reasoning, retrieval, and memory in complex multi-step tasks. It combines fine-grained actions, reversible stack memory, action-aware rewards, and information-aware trajectory selection to support longer-horizon reasoning.
- Complex tasks require models to decide dynamically when and what to retrieve while integrating retrieved information into evolving internal reasoning states.
- Existing agentic RL methods often use coarse, irreversible reasoning-and-retrieval actions and final-outcome rewards, allowing noisy intermediate conclusions to accumulate.
- AGENTICRAG-R1 uses a structured multi-action space and stack-based memory that lets agents push, revise, and pop intermediate reasoning units.
- Information-aware trajectory rejection filters low-information rollouts and prioritizes challenging, long-horizon samples with diverse behaviors during RL training.
- Action-aware process rewards provide explicit supervision for heterogeneous cognitive actions and help mitigate memory error accumulation.
2 AGENTICRAG-R1
AGENTICRAG-R1 combines multi-action reasoning, retrieval, and memory operations with hierarchical rewards and information-aware rollout rejection. Its stack-based memory supports reversible reasoning, while training prioritizes high-quality and diverse trajectories for long-horizon exploration.
- Action Modeling with Hierarchical Rewards: The framework executes reasoning, retrieval, and memory operations through multi-turn, multi-action rollouts with hierarchical action-aware rewards.Outcome rewards cover final-answer quality, while process rewards are assigned to corresponding action tokens.
- Action Modeling with Hierarchical Rewards: A LIFO memory stack stores discrete reasoning units and supports push, pop, revision, and rollback operations for reversible decision tracking.Plan, Think, and Search push units; Backtrack pops elements before pushing revised content.
- Action Modeling with Hierarchical Rewards: Process rewards supervise action format, retrieval relevance, and the usefulness of memory operations rather than merely rewarding their occurrence.Search and memory operations are evaluated using relevance or reasonableness assessments tied to the current task context.
- Action Modeling with Hierarchical Rewards: A fixed per-trajectory process-reward budget proportionally scales excess process rewards to discourage redundant action invocation.
- Information-Aware Trajectories Rejection: Information-aware rejection scores rollouts using both within-input reward diversity and rollout quality, then retains only the top fraction globally.The variance term highlights inputs with diverse solution paths, while reward distinguishes stronger trajectories within each input.
3 Experiments
Experiments evaluate AGENTICRAG-R1 across multi-hop and out-of-domain QA, component and reward ablations, rejection-threshold sensitivity, and long-horizon step budgets. The method achieves strong benchmark performance, benefits from memory-aware actions and information-aware rejection, and improves as longer reasoning horizons are permitted.
- Main Results: AGENTICRAG-R1 achieves the best or near-best performance across backbone sizes, with gains of +3.02% on 2Wiki and +6.76% on HotpotQA using Qwen2.5-3B.Out-of-domain gains include about +7.2% on TriviaQA with the 1.5B backbone and +7.1% on Bamboogle with the 7B backbone over strongest Agentic RAG baselines.
- Learning Dynamics: Training increasingly favors reasoning, memory-management, and conclusion actions, while search remains comparatively stable or decreases for the smaller model.Rewards rise, action entropy generally falls, and response length and reasoning turns first increase then decline during training.
- Component Analysis: Removing memory actions lowers 2Wiki from 32.92% to 27.54% and HotpotQA from 44.00% to 28.77%, while removing rejection reduces average performance from 33.55% to 25.43%.Removing both components further lowers the average score to 24.24%, supporting complementary contributions.
- Reward Analysis: Removing RAG rewards lowers 2Wiki from 32.92% to 23.70%, removing memory rewards lowers MusiQue from 16.48% to 7.63%, and removing both yields a 19.94% average.These ablations examine the distinct effects of retrieval- and memory-based process rewards.
- Threshold Sensitivity: F1-Score reaches a minimum near rejection threshold 0.2 and peaks at 0.6 and 0.8, while rejected samples decline from over 20,000 to the fewest.The analysis varies the rejection threshold on Qwen2.5-3B.
- Long-Horizon Experiments: Performance consistently improves as the maximum step budget increases from 10 to 30 across evaluated datasets.AGENTICRAG-R1 leverages extended horizons for deeper reasoning and retrieval, unlike baselines with unstable, plateauing, or declining trends.
- Action Ablation: Removing <Plan>, <Summary>, and <Backtrack> decreases average F1 by 1.95%, 3.05%, and 3.85%, respectively, while removing both memory actions causes a 6.18% degradation.The leave-one-action-out analysis indicates complementary effects among memory operations.
- Initialization Analysis: Average F1 is 22.40±0.50% for SFT only, 33.55 ± 0.42% for cold-start RL, and 35.20 ± 0.40% for RL with SFT initialization.SFT is not required, while stronger initialization improves RL performance.
4 Related Works
Recent reinforcement-learning approaches aim to improve LLM reasoning and tool use, but efficiency remains a concern for on-policy methods.
- Reinforcement learning has improved LLM reasoning and tool-use capabilities, with PPO widely used but computationally expensive due to repeated on-policy updates.
5 Conclusion and Future Work
AGENTICRAG-R1 integrates fine-grained actions, structured memory control, and action-aware rewards for agentic retrieval-augmented reasoning. Across multiple reasoning benchmarks, it outperforms strong baselines and supports long-horizon reasoning, while several coordination and evaluation challenges remain.
- AGENTICRAG-R1 integrates fine-grained agentic actions, structured memory control, and action-aware rewards.
- AGENTICRAG-R1 consistently outperforms strong RAG and agentic baselines across multiple reasoning benchmarks.
- The framework supports more effective long-horizon reasoning.
- Future Work: Jointly optimizing high-level actions with their arguments and resolving conflicts among agentic actions remain difficult.
- Future Work: Future evaluations will target broader vertical domains and open-ended scenarios such as deep research and long-horizon tasks.
Limitations
The framework introduces computational and tuning costs that may constrain efficiency, scalability, and transfer to new tasks.
- The memory stack and dynamic retrieval processes increase computational costs, challenging processing speed and efficiency in large-scale or real-time applications.Speculative sampling is proposed to alleviate some overhead, but further optimization is needed for scalability.
- Rejection-based sampling requires careful hyperparameter tuning, which may complicate deployment and affect transfer to new datasets or tasks.
Ethical considerations
The evaluation uses publicly accessible benchmark datasets and reports no personally identifiable information or human or animal subjects. The cited methodological passages describe the rollout engine and its notation rather than additional ethical considerations.
- Experiments use publicly accessible benchmark datasets, including GAIA, HotpotQA, and MusiQue, under their licenses and usage policies.
- The study uses no personally identifiable information and involves no human or animal subjects.
- The rollout engine initializes a trajectory and memory stack, samples action segments, and updates memory through retrieval, backtracking, summarization, or conclusion.
- The framework's notation is documented in a dedicated table.
D.1 Training Dataset.
The training setup targets long-horizon, multi-step retrieval because common multihop datasets can often be solved with shallow reasoning. It constructs a curated dataset and evaluates across diverse question-answering benchmarks.
- Common HotpotQA and 2WikiMultiHopQA instances often require only shallow reasoning or zero-to-one retrieval steps, limiting supervision for complex search and planning.
- The curated training dataset filters candidate questions using ten independent TC-RAG trajectories, excluding universally incorrect and high-correctness instances.
- The experiments cover multi-hop, open-domain, and agent-evaluation datasets, including HotpotQA, 2WikiMultiHopQA, Bamboogle, FRAMES, MuSiQue, NQ, and TriviaQA.
E Baseline Implementation Details
The baselines span non-retrieval, heuristic retrieval, manually prompted agentic RAG, and reinforcement-learning-based agentic RAG methods. This comparison covers distinct approaches to reasoning, retrieval, memory, and tool interaction.
- The evaluation compares against baselines covering non-retrieval, naive RAG, agentic RAG, and RL-based agentic RAG paradigms.
- No-RAG and CoT test parametric knowledge and prompted intermediate reasoning without retrieval, respectively.
- FS-RAG and FL-RAG use fixed retrieval strategies based on sentences or token intervals rather than explicit reasoning–retrieval coordination.
- ReAct, IRCOT, and TC-RAG represent manually designed or memory-based reasoning–retrieval interaction strategies.
- RL-based baselines include ReSearch, Search-R1, AEPO, ARPO, and Mem1, covering learned search, entropy-aware optimization, and memory-constrained reasoning.
F Implementation Details
AgenticRag-R1 is implemented with GRPO, structured action tags, iterative search, and evaluator prompts for retrieval relevance and memory operations. Training uses fixed rollout, context, and optimization settings.
- The method uses a GRPO-based reinforcement-learning framework with batch size 16, learning rate 1 × 10^-6, 300 optimization steps, and four stochastic generations per prompt.
- Generations use temperature 0.7, diverse sampling, a diversity penalty of 1.0, up to 2048 new tokens, and aggregated context up to 4096 tokens.
- The action format supports Think, Search, Summary, Backtrack, and Conclusion operations, repeated as needed before the final answer.
- A relevance evaluator scores retrieved documents from 0.0 to 1.0 according to their semantic support for the question and task context.
- A memory evaluator scores whether each memory operation is reasonable and beneficial given the question, reasoning context, and stack before and after the operation.
H Reward Analysis
The reward analysis tests whether semantic process supervision is necessary by replacing LLM-judged retrieval and memory rewards with simpler heuristics. Performance decreases when these semantic rewards are weakened or removed.
- The process rewards rrag and rmem are scored by a general-purpose LLM judge through a fixed prompt without additional reward-model training.
- Replacing semantic rewards with ROUGE-L retrieval relevance and a positive score for valid memory operations reduces average F1 across seven benchmarks from 33.55 to 29.40.
- Removing both rrag and rmem further reduces average F1 across seven benchmarks to 19.94, supporting semantic process supervision.
I Additional Experiments
Additional experiments show that AGENTICRAG-R1 changes inference behavior, benefits from action-specific optimization, and supports memory-based error correction across reasoning settings.
- Inference Behavior and Tool Usage: 14.4% lower average inference latency accompanies 12.2% more actions, 5.1% longer responses, and a 5.8-point higher tool-call rate than TC-RAG.Latency decreases from 11.86 to 10.15 seconds, while tool-call rate rises from 91.7% to 97.5%.
- Action-Specific Importance Analysis: Action-specific neuron profiles broadly overlap but show clear divergence between Search and Think, and between Backtrack and Think.Summary and Backtrack exhibit relatively small pairwise differences, suggesting overlapping neural resources.
- Action-Specific Optimization: 72.2% MedQA accuracy under Self-Search with decoupled optimization exceeds 67.5% with Self-Search alone and 65.5% under External Search.Decoupled optimization also reaches 68.5% under External Search, indicating gains across both search paradigms.
- Cross-Domain Rollout Reward Statistics: Meaningful rollout-reward variance persists across domains, indicating that the trajectory rejection strategy does not collapse rewards into indistinguishable ranges.The analysis reports preserved reward variance across domains and model scales.
- Case Study: AGENTICRAG-R1 uses backtracking to revise noisy retrieval keys and progressively narrow searches before committing information to memory.The case study reaches the correct Battleship identification through repeated retrieval, inspection, and memory revision.
- Outcome-Only Reward Bias: For erroneous steps, generation probability increases whenever the final reward is positive, regardless of the step’s negative contribution.The analysis identifies this as monotonic amplification of erroneous reasoning steps without backtracking or summarizing.
- Pop-Based Stack and Context Masking: Pop-based attention masking removes erroneous actions from active context, evaluating gradients on a counterfactual trajectory where those actions never occurred.The mechanism structurally isolates erroneous steps rather than numerically redistributing rewards.
- Medical Question Answering: On MedQA, AGENTICRAG-R1 reaches 92% format accuracy and 87% answer accuracy, outperforming both the base model and ReAct on both metrics.ReAct improves format accuracy but lowers answer accuracy relative to the base model, whereas AGENTICRAG-R1 improves both.
M.3 Interactive Decision-Making Environments.
AGENTICRAG-R1 generalizes to interactive, long-horizon environments that require state maintenance, feedback integration, and strategy adaptation.
- Interactive Decision-Making Environments: These environments test whether agents can maintain state, incorporate feedback, and adapt strategies over extended trajectories.The evaluation extends beyond static question answering to sequential interactive environments.
- Interactive Decision-Making Environments: AGENTICRAG-R1 improves ALFWorld success rate from 36.5% to 51.1%.ALFWorld reports task-specific and overall success rates.
- Interactive Decision-Making Environments: AGENTICRAG-R1 raises WebShop task score from 14.7 to 27.8 and success rate from 0.8% to 2.3%.WebShop reports the official task score and success rate.