All-in-One Memory: Process-Rewarded Reinforcement Learning for Memory Agents

Bingbing Wang; Jingjie Lin; Weiming Qiao; Jing Li; Ruifeng Xu

cs.AIOct 15, 2026

All-in-One Memory: Process-Rewarded Reinforcement Learning for Memory Agents

Authors: Bingbing Wang, Jingjie Lin, Weiming Qiao, Jing Li, Ruifeng Xu

Organizations: Harbin Institute of Technology, Shenzhen · The Hong Kong Polytechnic University · Fudan University · Peng Cheng Laboratory

Abstract

Large language models increasingly rely on external memory to answer queries grounded in long-horizon histories such as multi-session dialogues and lengthy documents. However, existing memory methods often optimize retrieval and memory construction separately, without a shared reward space for evaluating search, writing, and stopping decisions. To address this, we propose All-in-One Memory (AiO-Mem), which organizes these operations as transitions over a shared evidence state and directly evaluates the semantic quality of each action in a unified reward space. AiO-Mem compiles heterogeneous histories into source-traceable multi-view events, providing a shared substrate for retrieval and memory writing. A unified memory policy actor then rolls out complete interaction trajectories over this shared evidence state. A separate trajectory-conditioned process reward model is trained through evidence-grounded supervised fine-tuning and trajectory-level DPO to assign process scores to the collected trajectories. These scores are aggregated into an action-balanced trajectory reward and used by Process-Aligned DAPO (PA-DAPO) to optimize the memory policy. Across long-context memory benchmarks, AiO-Mem consistently improves answer accuracy over both static memory systems and RL-trained memory agents.

Figures & tables

Figure 1: Comparison of existing RL-driven memory methods and AiO-Mem.

Figure 2: Overview of our AiO-Mem method.

Reference
Method	Multi-hop		Temporal		Open-domain		Single-hop		Overall
Method	F1	B1	F1	B1	F1	B1	F1	B1	F1	B1
Long-Context	23.2	16.8	19.0	13.4	13.9	10.7	40.7	31.5	31.5	23.9
RAG	23.5	17.1	43.1	37.4	20.2	14.9	49.5	44.9	41.6	36.5
Agent Memory Harness
MemGPT	16.7	14.0	12.6	10.9	11.6	9.2	14.0	11.8	14.1	11.8

Table 1: Main results on LoCoMo .

Figure 3: Main results on MemoryAgentBench .

Memory Representation and Interaction Ablations
Method	Multi-hop		Temporal		Open-domain		Single-hop		Overall
Method	F1	B1	F1	B1	F1	B1	F1	B1	F1	B1
Base Model	39.3	30.9	65.0	58.4	25.9	21.9	61.3	56.1	55.8 (-5.2)	49.8 (-5.2)
w/o Event Comp.	43.2	34.6	66.9	61.0	27.2	22.6	64.3	58.7	58.6 (-2.4)	52.5 (-2.5)
w/o Multi-view Ret.	41.3	33.4	65.8	60.4	29.5	24.4	64.2	58.6	58.1 (-2.9)	52.1 (-2.9)
w/o WM Constr.	33.2	26.4	57.0	50.9	24.3	20.0	57.6	52.3	50.9 (-10.1)	45.2 (-9.8)

Table 2: Ablation on LoCoMo . Row abbreviations: Event Comp. = Event Compilation, Multi-view Ret. = Multi-view Retrieval, WM Constr. = Working-Memory Construction, Evid. SFT = Evidence-Grounded SFT, and Traj. DPO = Trajectory-Level DPO. Red values denote drops over AiO-Mem.

Figure 4: Parameter analysis on LoCoMo over top-

K

(

K\in\{5,10,15\}

) and maximum interaction steps

T_{\max}\in\{4,6,8,10\}

, reporting F1 across four question types (MH = Multi-hop, Tmp = Temporal, OD = Open-domain, SH = Single-hop) and overall. To facilitate comparison, all values are normalized for visualization on the chart.

Figure 5: Two cases between the MEM-T and our AiO-Mem method.

Appendix figures & tables14 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Reward-model trajectories by behavioral region in t-SNE space.

top- $K$	Max Step	Multi-hop		Temporal		Open-domain		Single-hop		Overall
top- $K$	Max Step	F1	B1	F1	B1	F1	B1	F1	B1	F1	B1
$K=5$	$T_{\max}=4$	44.6	35.7	66.2	59.6	30.3	25.6	62.7	57.3	58.0	51.8
	$T_{\max}=6$	43.5	35.4	62.4	53.3	34.9	30.9	65.9	61.5	59.2	53.2
	$T_{\max}=8$	44.0	36.6	63.9	55.7	34.4	30.3	65.4	60.9	59.2	53.5
	$T_{\max}=10$	43.0	35.1	61.9	53.1	33.5	29.3	65.8	61.1	58.8	52.8
$\mathbf{K=10}$	$T_{\max}=4$	43.6	35.6	63.4	54.6	33.7	30.2	65.9	60.9	59.3	53.3

Appendix

Table 3: Parameter analysis on LoCoMo with different top-

K

values and maximum interaction steps

T_{\max}

Panel of Figure 7 — Figure 7: Complementary evidence coverage across retrieval views. Panel 7(a) summarizes the major view combinations, panel 7(b) expands the remaining combinations, and panel 7(c) shows the sources of evidence detected by exactly one view across question categories.

Figure 8: Sensitivity to retrieval modes and RRF parameters. Rows and columns denote the summary and raw-turn retrieval modes, respectively, while each panel uses a different aggregation RRF parameter. Red and blue indicate higher and lower evidence Recall.

Figure 9: Trade-off between retrieval Recall and context cost.

Panel of Figure 10 — Figure 10: Sample-level Reward scores versus evidence metrics.

Panel of Figure 11 — Figure 11: Training dynamics.

Panel of Figure 12 — Figure 12: Effects of policy training on evidence flow through the memory pipeline. The three panels report Search Recall, Final working memory Recall, and Answer F1, respectively.

System	(I) Segment & Retrieval	(II) Construction	(III) Inference
OMS	$O(N)\cdot\mathcal{C}_{\mathrm{seg}}$ $+\,O(N\ell)\cdot\mathcal{C}_{\mathrm{LLM}}$ $+\,O(NM_{1})\cdot\mathcal{C}_{\mathrm{emb}}$	$O(NM_{1})\cdot\mathcal{C}_{\mathrm{sim}}$ $+\,O(NM_{1}R_{1})\cdot(\mathcal{C}_{\mathrm{LLM}}+\mathcal{C}_{\mathrm{emb}})$	$O(N\ell+NM_{1}R_{1})\cdot\mathcal{C}_{\mathrm{LLM}}$ $+\,O(NM_{1}(1+R_{1}))\cdot\mathcal{C}_{\mathrm{emb}}$ $+\,O(NM_{1})\cdot\mathcal{C}_{\mathrm{sim}}$
AiO-Mem	$O(T_{\max}M)\cdot\mathcal{C}_{\mathrm{sim}}$ $+\,O(N\ell)\cdot\mathcal{C}_{\mathrm{LLM}}$ $+\,O(M+T_{\max})\cdot\mathcal{C}_{\mathrm{emb}}$	$O(T_{\max}K_{\mathrm{ret}}\ell_{e})\cdot\mathcal{C}_{\mathrm{LLM}}$	$O(N\ell+T_{\max}K_{\mathrm{ret}}\ell_{e})\cdot\mathcal{C}_{\mathrm{LLM}}$ $+\,O(M+T_{\max})\cdot\mathcal{C}_{\mathrm{emb}}$ $+\,O(T_{\max}M)\cdot\mathcal{C}_{\mathrm{sim}}$

Appendix

Table 6: Complexity comparison between AiO-Mem and OMS.

Panel of Figure 13 — Figure 13: Efficiency comparison between AiO-Mem and Baseline on LoCoMo.

Symbol	Definition
$\mathcal{H}=(h_{1},\ldots,h_{N})$	Long-horizon history consisting of $N$ heterogeneous source units.
$N$	Number of source units in $\mathcal{H}$ .
$\ell$	Average number of tokens in each source unit $h_{n}$ .
$\mathcal{B}=\mathcal{C}(\mathcal{H})$	Traceable event bank constructed from $\mathcal{H}$ by the structure-aware compiler $\mathcal{C}$ .
$M$	Number of traceable events in the event bank $\mathcal{B}$ .
$M_{1}$	Number of persistent memory entries produced per source unit by OMS.

Appendix

Table 7: Notation used in the complexity analysis.

Stage	Tool call and function
DECIDE	search_memory_units(query) : Retrieves event-level memory units concerning information missing from the current notes.
DECIDE	answer_from_notes() (no arguments): Ends memory retrieval and proceeds to the answering stage using the existing working memory.
EXTRACT	add_note(event_id, datetime, content, evidence_turns) : Saves a clue, its event time, and the original dialogue supporting it. An EXTRACT step may invoke add_note multiple times, with each invocation writing one note.
EXTRACT	skip_extraction(reason) : Skips extraction when the current search results contain no relevant evidence to write.

Appendix

Table 8: Tool calls available to the memory agent.

Stage

Prompt template

DECIDE

You are a memory assistant. Given the current question and saved working memory, choose the next action: search memory, or stop retrieval and proceed to the answering stage. Rules: (1) Judge whether the question can be answered completely using only the working memory, including the relationships and all necessary parts requested by the question. If the evidence is sufficient, call answer_from_notes . (2) Otherwise, generate a query targeting one specific fact or clue that is still missing, and call search_memory_units . The query should identify the relevant person or object, relationship, and information needed. Avoid repeating previous searches that produced no useful notes. (3) At the final decision step, call answer_from_notes . Input: Question {question}, working memory {working_notes}, previous queries {previous_searches}, and step {step}/{max_steps}. Output: A brief explanation of the decision and one <tool_call> .

EXTRACT

You are a memory system analyst. Given the question, the latest search results, and the existing working memory, identify new evidence useful for answering the question and write it into the working memory. Search results are temporary; useful information not written into the notes will not be available to subsequent decisions. Rules: (1) Read the visible original dialogue in retrieved events. Highlighted turns are only hints; relevant evidence may also appear in other turns. (2) For each new useful clue, call add_note to write a concise, faithful note, preserving the exact event ID, recorded time, and original dialogue supporting its content. Each note should capture one natural clue and should be compared with existing notes to avoid duplication. (3) If the search results contain no relevant evidence, call skip_extraction . Input: Question {question}, search observation {search_observation}, and working memory {working_notes}. The search observation contains event identifiers, recorded times, retrieved fields, summaries, and visible original dialogue. Output: A brief assessment of the evidence and the required <tool_call> ; distinct clues may be written as separate notes.

Appendix

Table 9: Prompt templates for the LoCoMo memory agent.

Prompt template

Task and input You are a reward model evaluator of memory trajectories. Given a complete scoreable memory trajectory, evaluate each step according to whether the required evidence was found, whether useful evidence was written into working memory, and whether retrieval stopped when the evidence was sufficient. You may use trajectory context to understand the current action, but each score must evaluate only the current step; do not evaluate the quality of the final answer model’s response. Input: {question}, {reference_answer}, {gold_evidence}, {trajectory_information}, and {scoreable_steps}. Objectives by action type decide_search : Does the query target information missing from the current notes and help locate evidence needed to answer the question? A reasonable query should not receive the lowest score solely because it failed to retrieve a match. extract_add_note : Does the action faithfully write useful, previously unsaved evidence from the latest search into the notes, and do the cited turns support the note’s content? extract_skip : Does the latest search truly contain nothing worth writing? Skipping clearly critical evidence is a severe problem. decide_answer_from_notes : Are the existing notes sufficient to support an answer when retrieval stops? Evidence assessment When judging whether the notes are sufficient, consider both their content and the original dialogue they cite. Annotated evidence indicates information needed for the question, but it is not a mechanical hit-rate scoring formula: unannotated dialogue turns may also support the answer semantically, and matching annotated turns does not automatically warrant a high score. Error severity Category A, high severity: Failure of the current stage’s core objective. Category B, high severity: Factual or attribution errors in a note, or citations that do not support key content. Category C, medium severity: Problems in assessing the current state or respecting stage boundaries. Category D, medium severity: An incomplete contribution or substantial noise in the action. Category E, low severity: Wording or formatting issues. Score ranges 90–100: The action fulfills the stage’s core objective very well, with at most minor Category E issues. 75–89: The action is clearly useful and largely fulfills the core objective; it has no Category A or B issues, though minor Category C, D, or E issues are allowed. 60–74: The action makes a partial contribution but falls short of the core objective; it has no Category A or B issues, though Category C or D issues are more pronounced. 40–59: The action has a severe problem but retains some useful contribution; it typically contains one Category A or B issue. 1–39: The action provides almost no useful contribution or takes a clearly wrong direction; severe problems cause the current step to largely fail. Output requirements Output a strict JSON object containing step_scores in input order. Each step must contain step_index , step_type , quality_level , and an integer score from 1 to 100. Do not omit or add steps, and do not output scoring explanations. quality_level roughly corresponds to the score as follows: bad (1–30), weak (31–55), partial (56–70), good (71–85), and excellent (86–100).

Appendix

Table 10: Prompt template and scoring rules for the stepwise reward model.

Persistent Memory for Language Models

Persistent memory for language models studies how models retain, update, and use information across interactions and beyond their context windows.

See all 267 papers

Jul 23, 2026

AttriMem: Attribution-Guided Process Feedback for Agent Memory Construction

Effective memory is crucial for LLM agents, yet constructing it effectively remains challenging. A memory-construction policy decides what information to extract, store, update, compress, or discard as interactions accumulate. Heuristic memory methods rely on subjective, task-specific rules, which can misalign with downstream objectives and limit cross-task adaptability. RL-based methods, by contrast, learn from task feedback but mainly use outcome- or module-level rewards. These coarse signals indicate task success but cannot identify which intermediate memory contents support the final answer, creating a fine-grained credit-assignment bottleneck. However, constructing such process feedback is prohibitively difficult because intermediate memory decisions lack unique ground-truth targets, while the appropriate credit varies with the agent's uncertain reasoning trajectory and therefore cannot be specified in advance. We propose AttriMem, an attribution-guided process-feedback framework for learning memory-construction policies with RL. AttriMem augments the global outcome reward with local rewards derived from token-level contributions to the final answer. Experiments on long-horizon dialogue question answering show that AttriMem outperforms retrieval-based, heuristic, and RL-based baselines, generalizes across benchmarks and answer models, stabilizes RL optimization.

Qinfeng Li, Yuntai Bao, Xinyan Yu +7

1Zhejiang University · 2Washington University in St. Louis · University of Science and Technology of China

LLM Agent Memory Policy Learning Agent Memory Management RL for LLM Agents

Ask PDF

Feb 13, 2026

Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning

Long-context LLMs and Retrieval-Augmented Generation defer state tracking and evidence consolidation to query time, which is brittle when facts evolve and answers depend on latent states. We introduce Unified Memory Agent (UMA) for a one-to-many setting: query-agnostic external memory is constructed once from a stream and reused across multiple future QA sessions. A single policy maintains a structured Memory Bank through CRUD operations and answers using both the Memory Bank and raw context. Task-Stratified GRPO uses the mean reward of QA trajectories branching from each sampled memory state to supervise memory maintenance, while normalizing memory and per-question QA groups separately. We also introduce Ledger-QA, a diagnostic benchmark for long-horizon state tracking over accumulated updates. At the 16k budget, UMA-Generalist achieves the highest average score among compared methods across the test-time-learning and accurate-retrieval benchmarks and transfers to Ledger-QA without task-specific training; UMA-Specialist further improves long-horizon tracking after task adaptation. These results support learned proactive memory management for long-context reasoning.

Kehao Zhang, Shangtong Gui, Sheng Yang +2

1Key Laboratory of Intelligent Information Processing, Institute of Computing Technology, Chinese Academy of Sciences (ICT/CAS) · University of Chinese Academy of Sciences, Beijing, China · 4Li Auto Inc. +1

Persistent Memory for Language Models Long-Context Inference Agent Memory Management RL for LLM Agents Memory-Augmented Language Model Agents

Ask PDF

Jul 27, 2026

MemChain: Learning Interpretable Memory Traces for Memory-Augmented LLM Agents

Memory-augmented LLM agents typically answer queries by retrieving relevant memories and feeding them directly to an answer model. This retrieval-as-evidence paradigm assumes retrieved memories are already suitable for reasoning, leaving the answer model to resolve redundancy, conflicts, and weak relevance while incurring substantial context overhead in long-term memory tasks. We propose MemChain, a trainable post-retrieval memory policy that transforms retrieved candidates into answer-facing active memory, represented as a compact and grounded evidence context. Given a user query and retrieved candidates, MemChain first generates a question-conditioned evidence plan, then constructs an ordered grounded evidence trace that organizes retrieved memories according to their semantic roles and dependencies, and finally executes explicit memory actions to produce a concise evidence context for answer generation. To train the mediator, we introduce a two-stage learning framework. Supervised trace learning first teaches the policy to generate structurally valid plans, traces, actions, and evidence contexts. We then propose Trace-Guided Memory Policy Optimization (TMPO), a reinforcement learning objective that optimizes the memory policy using downstream answer quality while jointly encouraging trace grounding, evidence support, structural validity, and answer stability across multiple rollouts. Experiments on LoCoMo and LongMemEval-S demonstrate that MemChain consistently achieves state-of-the-art performance across both closed-source and open-weight frozen answer models while substantially reducing the memory context passed to the answer model.

Yiwen Ma, Songjun Tu, Qichao Zhang +3

Institute of Automation, Chinese Academy of Sciences · School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences +1

LLM Agent Memory Persistent Memory for Language Models Long-Term Conversational Memory RL for LLM Agents Memory-Augmented Language Model Agents

Ask PDF

Abstract

Persistent Memory for Language Models

Explore similar work

AttriMem: Attribution-Guided Process Feedback for Agent Memory Construction

Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning

MemChain: Learning Interpretable Memory Traces for Memory-Augmented LLM Agents