Uni-Mem: Learning a Long-Horizon Agentic Memory Manager with Unified Online Policy Optimization

Zhiqian Xia; Jielong Tang; Ling-An Zeng; Shixiong Kai; Rong Yao; Hao Kong; Xiao Dong; Wei Wang

cs.CLOct 21, 2026

Uni-Mem: Learning a Long-Horizon Agentic Memory Manager with Unified Online Policy Optimization

Authors: Zhiqian Xia, Jielong Tang, Ling-An Zeng, Shixiong Kai, Rong Yao, Hao Kong, Xiao Dong, Wei Wang,

Organizations: Southeast University · Huawei Technologies, Noah’s Ark Lab · Sun Yat-sen University · Shandong Normal University

Abstract

Long-term agent memory is shifting from predefined workflows toward learnable policies for memory construction and retrieval. Existing learning-based methods either optimize construction alone or combine retrieval-side online reinforcement learning (RL) with offline construction supervision, limiting co-adaptation between the two stages. In this paper, we explore end-to-end online RL for a unified memory manager (Uni-Mem) that learns both capabilities under a shared policy. However, joint optimization with existing RL algorithms such as GRPO faces confounded group comparisons and sparse feedback over long memory-management trajectories. To this end, we propose Stage-Decoupled GRPO, which alternates construction and retrieval optimization while controlling variation from the other stage within each rollout group, yielding stage-specific learning signals. We further introduce Event-Wise Credit Assignment (ECA), which localizes feedback to responsible events and propagates event-level advantages to their constituent operations. Experiments on LoCoMo and LongMemEval show that Uni-Mem outperforms the strongest RL-based memory, demonstrating the effectiveness of jointly learning memory construction and retrieval via Stage-Decoupled GRPO and ECA.

Figures & tables

Figure 1: Comparison of training paradigms in different methods. Construction-only RL keeps retrieval fixed, while asymmetric joint learning combines offline construction imitation with online retrieval RL. Uni-Mem jointly optimizes both stages under a shared policy through alternating online RL and event-wise credit.

Figure 2: Overview of Uni-Mem . (a) Uni-Mem unifies memory construction and retrieval through a shared tool-using agent, continually processing an ordered stream of events. (b) Event-Wise Credit Assignment attributes sparse outcome rewards to relevant events. (c) Event-level feedback is propagated to the corresponding memory operations (tool calls). (d) Stage-Decoupled GRPO alternates between construction and retrieval optimization.

Training-free Methods
Model	Multi	Open	Single	Temp.	Overall
RAG	15.06	19.28	12.79	5.43	12.17
A-MEM	40.78	31.32	46.90	28.68	40.78
Mem0	61.92	48.19	65.20	38.76	53.30
MemoryOS	48.12	38.55	63.14	24.81	51.26
Training-based methods

Table 1: Performance comparison on LoCoMo (LLM-as-a-Judge, %). Bold and underlined results indicate the best and second-best, respectively. † marks our re-evaluation of the officially released checkpoints under the same protocol as Uni-Mem . Other baseline results are taken from Yan et al. (2026a) under their original protocol, as reference only (Appendix E.1 ).

Figure 3: OOD accuracy (

J

, %) on LongMemEval-Oracle and LongMemEval-S. SS: single-session; asst.: assistant; pref.: preference. Both panels use a 0–100% scale. Mem-T results are shown only for Oracle.

Qwen3-4B-Instruct
Model	Multi	Open	Single	Temp.	Overall
$M_{0}$	78.61	70.97	84.04	57.81	76.58
$M_{1}$	81.09	69.35	85.57	54.01	76.95
$M_{2}$	82.59	72.58	86.59	55.27	78.24
$M_{3}$	84.08	66.13	87.10	65.82	80.72
$M_{4}$	81.09	75.81	88.79	62.87	80.99

Table 2: Ablation study on LoCoMo (

J

, %) with raw or SFT-initialized Qwen3-4B-Instruct as start point (

M_{0}

Memory	Search	Multi	Open	Single	Temp.	Overall
$P_{0}$	$Q_{0}$	78.61	70.97	84.04	57.81	76.58
$P_{0}$	$Q_{+}$	86.07	67.74	88.29	62.45	81.08
$P_{+}$	$Q_{0}$	77.61	69.35	89.98	81.86	84.76
$P_{+}$	$Q_{+}$	87.06	79.03	91.51	86.08	88.80

Table 3: Evaluation of memory construction and retrieval capabilities on LoCoMo (

J

, %) by cross-combining memory banks and retrieval policies.

P_{0}

and

P_{+}

denote memory banks constructed by the base Qwen3-4B manager and the trained Uni-Mem manager, respectively;

Q_{0}

and

Q_{+}

denote their corresponding retrieval policies.

Figure 4: Effects of manager training, scale, and alternating schedules on LoCoMo. (a) Training and manager scale. The arrow shows the gain from SFT followed by RL. (b) Construction–retrieval update ratio. (c) Alternation granularity at a fixed

1:1

ratio.

C

and

R

denote consecutive construction and retrieval updates per cycle.

Figure 5: Case studies comparing Uni-Mem and the base Qwen3-4B manager in (a) memory construction and (b) memory retrieval.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

	Interaction		Training
Method	Con.	Ret.	Con.	Ret.
MEM1	S	–	RL	–
MemAgent	S	–	RL	–
Memory-R1	S	Fixed	RL	Fixed
Mem- $\alpha$	S	Fixed	RL	Fixed
Memory-R2	S	Fixed	RL	Fixed

Appendix

Table 4: Interaction pattern and training of construction (Con.) and retrieval (Ret.). S: single-turn; P: fixed multi-step pipeline; M: multi-turn with tool observations; Fixed: similarity retrieval without learning; –: no explicit stage. Offline: imitation of hindsight-filtered actions.

					QA category
Length	Task	Sessions	Turns	QA pairs	Single-hop	Multi-hop	Temporal	Open-domain
Short	conv_05	10	162	114	66	21	16	11
Medium	conv_07	20	371	128	68	23	25	12
Long	conv_12	35	703	116	57	13	34	12
Extra-long	conv_19	52	1,054	415	198	74	113	30
Total	4 tasks	117	2,290	773	389	131	188	65

Appendix

Table 5: Statistics of the four source tasks for SFT demonstration generation. A turn is one nonempty message; session dates and speaker metadata are excluded. QA counts include all source questions, with no adversarial questions.

Figure 6: SFT training loss, measured as the mean negative log-likelihood over supervised assistant tokens (nats/token). The horizontal axis counts optimizer updates. The light curve shows the logged loss; the dark curve shows a trailing 100-step moving average.

Checkpoint	Tool format (%)	Pass@1 (%)	Pass@3 (%)	Gap (pp)	F1 (%)	Entropy	Finish-bundled
Base (pre-SFT)	99.32	79.57	90.45	10.89	24.80	0.1706	1,610
Step 431	99.60	81.93	94.94	13.01	26.52	0.4482	182
Step 862 (selected)	99.34	83.35	94.42	11.06	26.96	0.4058	54
Step 1293	99.70	83.59	94.09	10.50	26.98	0.3717	67
Step 1724	99.74	85.19	94.94	9.74	27.66	0.3776	21
Step 2155	99.81	85.06	94.55	9.48	27.24	0.3632	21

Appendix

Table 6: SFT checkpoint selection for RL initialization. Bold marks the selected checkpoint, not per-column maxima. All checkpoints pass the 95% tool-format gate. Pass@1, pass@3, and F1 use judged, non-adversarial QA records; gap is pass@3 minus pass@1 in percentage points (pp), computed before rounding. Entropy is a partial-sum estimate in nats/token on held-out SFT prompts. Tool format and finish-bundled are measured from the first evaluation run; finish-bundled is a raw count.

Hyperparameter	Value
Actor learning rate	$2\times 10^{-6}$
Task batch size	2
Rollouts per task ( $N$ )	8
PPO mini-batch size	1
Sampling temperature / top- $p$	0.7 / 1.0
Training epochs	200

Appendix

Table 7: Main RL training hyperparameters. Task batch size is specified before sampling the eight rollouts per task.

Memory-R2
Setting	Multi	Open	Single	Temp.	Overall
Reported	80.93	67.53	86.80	69.90	80.99
Re-evaluated	69.15	61.29	78.61	57.38	71.26
Memory-R2 (OSS)
Reported	85.07	74.14	90.58	83.33	87.10
Re-evaluated	74.13	70.97	81.15	74.68	77.87

Appendix

Table 8: Reported and re-evaluated Memory-R2 results on LoCoMo (

J

, %). Reported results are from Yan et al. (2026a) . Re-evaluated results are obtained from the officially released checkpoints using our Mem0-aligned judge prompt.

Re-evaluated baselines
Setting	Overall $J$	$\Delta$ vs. $M_{0}$
Memory-R2 (OSS)	77.87	–
Mem-T	76.49	–
Our harness, off-the-shelf manager
Qwen3.5-27B (strongest)	87.05	–
Qwen3.5-397B-A17B (largest)	85.86	–

Appendix

Table 9: Overall

J

(%) on LoCoMo under the same harness.

\Delta

is measured against the start point

M_{0}

of each block. Our RL combines ECA, Stage-Decoupled GRPO, and update-robust attribution; see Table 2 for per-component results and Table 10 for all off-the-shelf managers.

Figure 7: Training-time LLM-as-a-Judge scores (

J

, %) of Uni-Mem , with Qwen3.8-27B used for both answer generation and judging. The light curve shows the logged scores; the dark curve shows a trailing 20-step moving average.

Off-the-shelf managers
Manager	Multi	Open	Single	Temp.	Overall
Qwen3-4B	78.61	70.97	84.04	57.81	76.58
Qwen3.5-27B	85.07	74.19	89.13	86.92	87.05
Qwen3.5-35B-A3B	82.09	74.19	89.30	85.23	86.23
Qwen3.5-122B-A10B	84.08	67.74	88.46	85.65	85.86
Qwen3.5-397B-A17B	84.08	75.81	86.93	87.34	85.86

Appendix

Table 10: Detailed results on training and manager scale on LoCoMo (

J

, %). Multi, Open, Single, and Temp. denote multi-hop, open-domain, single-hop, and temporal questions, respectively. Best scores in each column are bold.

Construction–retrieval update ratio
Schedule	Multi	Open	Single	Temp.	Overall
$1{:}8$	81.09	66.13	89.47	79.75	84.48
$1{:}4$	86.07	64.52	90.83	81.01	86.32
$1{:}2$	87.56	70.97	91.85	83.97	88.15
$1{:}1$	87.06	79.03	91.51	86.08	88.80
$2{:}1$	89.55	74.19	89.81	82.28	87.24

Appendix

Table 11: Detailed results for alternating training schedules on LoCoMo (

J

, %).

C{:}R

specifies consecutive construction and retrieval updates per cycle. Best scores in each column within each block are bold.

Figure 8: Accuracy (

J

) versus token consumption under different maximum retrieval steps

k

. Highlighted points and annotations compare

k=3

and

k=8

; lines connect the measured settings.

Memory Management for LLM Agents

See all 88 papers

Learning How and What to Memorize: Cognition-Inspired Two-Stage Optimization for Evolving Memory
May 1, 2026Derong Xu, Shuochen Liu, Pengfei Luo +8Memorization in Language Models Memorization in Generative Models
University of Science and Technology of China & State Key Laboratory of Cognitive Intelligence · City University of Hong Kong · Dalian University of Technology +1
Memory-R2: Fair Credit Assignment for Long-Horizon Memory-Augmented LLM Agents
May 20, 2026Sikuan Yan, Ahmed Bahloul, Ercong Nie +4Group Relative Policy Optimization Credit Assignment in RL
Ludwig Maximilian University of Munich · Munich Center for Machine Learning · Huawei Heisenberg Research Center (Munich) +1
Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning
Feb 13, 2026Kehao Zhang, Shangtong Gui, Sheng Yang +2Persistent Memory for Language Models Long-Context Inference
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

May 1, 2026J/K move · Enter open · S save

Learning How and What to Memorize: Cognition-Inspired Two-Stage Optimization for Evolving Memory

Derong Xu, Shuochen Liu, Pengfei Luo, Pengyue Jia+7

University of Science and Technology of China & State Key Laboratory of Cognitive Intelligence · City University of Hong Kong · Dalian University of Technology+1

Memorization in Language Models Memorization in Generative Models Long-Term Conversational Memory Agent Memory Management RL for LLM Agents

Open paper AskPDF arXiv

Large language model (LLM) agents require long-term user memory for consistent personalization, but limited context windows hinder tracking evolving preferences over long interactions. Existing memory systems mainly rely on static, hand-crafted update rules; although reinforcement learning (RL)-based agents learn memory updates, sparse outcome rewards provide weak supervision, resulting in unstable long-horizon optimization. Drawing on memory schema theory and the functional division between prefrontal regions and hippocampus regions, we introduce MemCoE, a cognition-inspired two-stage optimization framework that learns how memory should be organized and what information to update. In the first stage, we propose Memory Guideline Induction to optimize a global guideline via contrastive feedback interpreted as textual gradients; in the second stage, Guideline-Aligned Memory Policy Optimization uses the induced guideline to define structured process rewards and performs multi-turn RL to learn a guideline-following memory evolution policy. We evaluate on three personalization memory benchmarks, covering explicit/implicit preference and different sizes and noise, and observe consistent improvements over strong baselines with favorable robustness, transferability, and efficiency.

Abstract

Memory Management for LLM Agents

Explore similar work

Learning How and What to Memorize: Cognition-Inspired Two-Stage Optimization for Evolving Memory

Memory-R2: Fair Credit Assignment for Long-Horizon Memory-Augmented LLM Agents

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