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

1 Introduction

• Existing methods either train construction with retrieval fixed or pair online retrieval RL with offline construction imitation, limiting coordinated adaptation.Construction-only methods include Memory-R2 (Yan et al., 2026a) and Mem-α (Wang et al., 2025); Mem-T (Yue et al., 2026) uses asymmetric joint learning.Source: §1 · p. 1
• Joint rollouts can confound stage-level learning signals, while answer-time rewards provide sparse feedback across hundreds of operations and memory updates.The paper notes that updates can overwrite items produced by earlier decisions, complicating credit attribution.Source: §1 · p. 2
• Stage-Decoupled GRPO alternates construction and retrieval optimization while controlling variation from the other stage within each rollout group.Construction uses evidence-indexed scoring that bypasses learned retrieval; retrieval rollouts share an identical memory state.Source: §1 · p. 2
• Event-Wise Credit Assignment localizes feedback to responsible events and propagates event-level advantages to the tool calls within them.It aligns event positions across rollouts and tracks events along the same memory-update chain.Source: §1 · p. 2
• UNI-MEM outperforms the strongest RL-based memory manager on LoCoMo and generalizes to LongMemEval, supporting joint learning of construction and retrieval.The paper identifies Stage-Decoupled GRPO and Event-Wise Credit Assignment as the mechanisms for that joint learning.Source: §1 · p. 2

2 Related Work

• Prior memory systems explore persistent records and reflections, hierarchical storage, structured organization, and iterative evidence gathering.Examples include Park et al. (2023) and Zhong et al. (2024) for persistent records; Packer et al. (2023), Kang et al. (2025), and Fang et al. (2026) for hierarchical storage; and Yan et al. (2025) for agentic retrieval.Source: §2 · pp. 2–3
• RL-based memory work targets compressed working states, external-memory construction, retrieval decisions, and coordination of long-term memory with short-term context.The passage cites MEM1 (Zhou et al., 2026), MemAgent (Yu et al., 2026a), and Age-Mem (Yu et al., 2026b), among others.Source: §2 · p. 3
• Memory-R1 trains construction with a separate answer agent, while Memory-R2 extends long-horizon training with local and global group-relative credit.Mem-T (Yue et al., 2026) combines online retrieval RL with offline imitation of hindsight-selected construction actions.Source: §2 · p. 3
• UNI-MEM differs by using end-to-end online RL for a shared-parameter manager and addressing confounded group comparisons and credit assignment with Stage-Decoupled GRPO and event-wise attribution..Source: §2 · p. 3

3.1 Problem Formulation

• The manager processes an ordered stream of dialogue chunks and queries to maintain a persistent external memory bank and retrieve relevant memories for question answering.In the stated datasets, queries arrive after the dialogue, so retrieval uses the bank built from the dialogue chunks.Source: §3.1 · p. 3
• Construction events revise the memory bank in response to dialogue chunks, while retrieval events gather supporting evidence for queries.The objective is to jointly optimize construction and retrieval under a shared policy to improve downstream answer accuracy.Source: §3.1 · p. 3

3.2 Unified Memory Management Agent

• The memory manager is a tool-using agent governed by a shared policy, and each memory record stores a unique ID, source chunk index, and text content.The record links stored content to the dialogue chunk from which it was created.Source: §3.2 · p. 3
• Both event types support multiple rounds of tool calls; construction can search, create, update, and delete, while retrieval searches and submits selected evidence.The agent terminates each event with FINISH.Source: §3.2 · p. 3

4.1 Event-Wise Credit Assignment

• ECA links query evidence to responsible construction events and tracks record-update lineages so later updates or deletions do not erase event credit.Evidence-indexed scoring uses current versions of records associated with the query’s evidence chunks.Source: §4.1 · p. 4
• Construction events accumulate scores from related queries, while events unrelated to any query receive no return and are excluded from their position-aligned group.The stored-information score is evaluated using an evidence-indexed bank without the learned retrieval policy.Source: §4.1 · p. 4
• A failed retrieval is penalized only when the shared memory bank supports a correct answer; otherwise, that retrieval event is excluded from group statistics and policy optimization.This assigns the failure to construction when stored evidence is insufficient.Source: §4.1 · p. 4

4.2 Unified Manager Policy Optimization

• Stage-Decoupled GRPO alternates construction and retrieval optimization, evaluating construction without learned retrieval and holding a fixed memory bank across retrieval rollouts.Both phases update the same shared manager.Source: §4.2 · p. 5
• Corresponding events are grouped across rollouts by triggering dialogue chunk or query, and groups with fewer than two valid returns or zero variance are discarded.Construction events without returns and retrieval events marked ⊥ are excluded from group statistics and policy optimization.Source: §4.2 · p. 5
• All tool calls within a retained event share its advantage, and retained tool-call tokens are optimized with a clipped surrogate; training first cold-starts the shared policy with SFT.The token-level objective uses each token’s event advantage; no KL penalty is used.Source: §4.2 · p. 5

5.1 Experiment Setup

• The evaluation uses LoCoMo for training, validation, and testing with a 2:1:7 split, and LongMemEval for out-of-distribution testing.Scores use the official Mem0 LLM-as-a-Judge prompt.Source: §5.1 · pp. 5–6

5.2 Main Results

• 88.80% overall J, 10.93 points above Memory-R2 (OSS) and 12.31 above Mem-T; multi-hop and temporal gains are 12.93 and 9.29 points.UNI-MEM also scores 80.99% without SFT, 3.12 points above the strongest baseline.Source: Table 1 · p. 6, §5.2 · p. 6
• On LongMemEval-Oracle, UNI-MEM leads on single-session-preference and multi-session questions; on LongMemEval-S, it beats Memory-R2 variants across all six categories.On Oracle, it also gains in knowledge update and temporal reasoning over Mem-T.Source: Figure 3 · p. 6, §5.2 · p. 6

5.3 Ablation Study

• ECA improved overall J by 1.29 and 3.03 points without and with SFT, respectively; Stage-Decoupled GRPO added 2.48 points without SFT.Update-robust attribution further improved performance by tracing credit along memory update lineages.Source: §5.3 · p. 7

5.4 Further Analysis

• With base retrieval fixed, trained memories add 8.18 points; trained retrieval adds 4.50 and 4.04 points on base and trained memory banks, respectively.Combining both trained capabilities yields the highest overall score, 88.80%.Source: Table 3 · p. 7, §5.4 · p. 7
• A 1:1 schedule outperformed 32:32, 1:8, and 8:1 schedules; SFT followed by RL raised the 4B manager’s J score from 76.58% to 88.80%.UNI-MEM also outperformed Qwen3.5-27B (87.05%) and Qwen3.5-397B-A17B (85.86%).Source: §5.4 · pp. 7–8, Figure 4 · p. 8

5.5 Case Study

• In the construction case, UNI-MEM corrects a conflicting entry with SEARCH and UPDATE, while the base manager uses CREATE and leaves contradictory claims unresolved.In retrieval, UNI-MEM refines a later query using a ceramic-bowl reference found in its initial search.Source: Figure 5 · p. 8, §5.5 · p. 8

• UNI-MEM jointly learns agentic memory construction and retrieval through end-to-end online reinforcement learning under a shared policy.Source: Conclusion · p. 8
• Stage-Decoupled GRPO and Event-Wise Credit Assignment provide stage-specific signals and assign credit to relevant memory operations.Source: Conclusion · p. 8
• Experiments report that UNI-MEM outperforms strong RL-based memory baselines on LoCoMo and generalizes to LongMemEval.Source: Conclusion · p. 9

Limitations

• Reward computation requires downstream QA and judge-model evaluations, making training resource-intensive and adding cost and latency.The paper notes that concurrent judge evaluations or asynchronous RL could alleviate latency, while future work will explore process-level supervision to reduce reliance on outcome-based evaluations.Source: Limitations · p. 9
• Future work will explore process-level supervision to reduce reliance on outcome-based evaluations.Concurrent judge evaluations or asynchronous RL are also proposed as ways to alleviate latency.Source: Limitations · p. 9

Appendix Overview

• The appendix organizes supplementary material on implementation, training, baselines, ablations, detailed results, and prompt templates to support reproducibility.It also covers memory-agent positioning and experimental analyses.Source: Appendix Overview · p. 11

A Positioning UNI-MEM with Related Methods

• Table 4 compares construction and retrieval by interaction pattern and training, while UNI-MEM interacts with memory over multiple turns and observes tool results.Most prior managers act in one generation from fixed context; Mem-T uses multi-turn query-time search but a fixed construction pipeline.Source: A Positioning UNI-MEM with Related… · p. 11, Table 4 · p. 11
• UNI-MEM’s trained retrieval policy improves overall J score by 4.50 and 4.04 points on the same memory banks.The paper links this comparison to agentic retrieval, which lets the manager reformulate queries and follow intermediate results.Source: A Positioning UNI-MEM with Related… · pp. 11–12

A.1 Differences from Mem-T

• UNI-MEM trains construction and retrieval with on-policy RL, whereas Mem-T uses online retrieval RL and offline imitation for construction.Mem-T’s hindsight scores depend on retrieval outcomes and may become outdated as shared parameters change; UNI-MEM’s evidence-indexed construction scoring does not depend on its learned retriever.Source: §A.1 · p. 12
• Longer alternation blocks lower overall J score from 88.80% to 83.75%, consistent with the paper’s concern that one stage can lag behind the other.The comparison is reported in Figure 4(c).Source: §A.1 · p. 12

B.1 Memory Representation and Provenance Tracking

• Each record tracks its source-chunk provenance across creation and updates, and that metadata supports evidence-indexed scoring and training credit assignment.The manager cannot see the provenance metadata; the implementation uses it to identify current surviving record versions associated with evidence chunks.Source: §B.1 · p. 12

B.2 Tool Specifications and Execution Protocol

• The tool interface specifies inputs, outputs, and execution semantics; calls run sequentially, with results or errors observed by the manager on its next decision.FINISH must be issued alone, and the RL configuration caps events at eight searches and 24 manager responses.Source: §B.2 · p. 12, Tool B.5: FINISH · p. 13

Tool B.1: SEARCH

• SEARCH queries the current memory bank without modifying it and returns matching record IDs and text, or an empty result.The query is a string and the requested result limit is optional; construction and retrieval have different default limits.Source: Tool B.1: SEARCH · p. 12

Tool B.2: CREATE

• CREATE inserts the supplied agent-written content verbatim as a new record, preserving names, dates, and quantities.The tool returns the created record’s ID, memory, and ADD event; new facts are stored separately.Source: Tool B.2: CREATE · p. 12

Tool B.3: UPDATE

• UPDATE replaces a record’s text and embedding while preserving its ID, for correction or supersession of the same stored fact.New facts should be stored through CREATE.Source: Tool B.3: UPDATE · pp. 12–13

Tool B.4: DELETE

• DELETE removes a specified record from the memory bank when it is wrong and no longer needed.Source: Tool B.4: DELETE · p. 13

Tool B.5: FINISH

• FINISH closes an event; retrieval passes accumulated evidence to the answer model, while construction advances to the next input.Retrieval permits only SEARCH and FINISH, and construction permits interleaved searches and writes after a successful search response.Source: Tool B.5: FINISH · p. 13

C.1 Supervised Cold-Start Training

• SFT demonstrations use four source tasks spanning 162 to 1,054 turns, with 117 sessions, 2,290 turns, and 773 question–answer pairs.GPT-5.4 generates interactive demonstrations for memory construction and tool-based retrieval, conditioning decisions on prior tool observations.Source: §C.1 · p. 13, Table 5 · p. 14
• SFT fine-tunes Qwen3-4B-Instruct for three epochs, with loss applied only to assistant tokens and checkpoint evaluation every 431 optimizer updates.Across 4,257 updates, mean loss decreases from 1.579 over the first 100 updates to 0.213 over the last 100; the 95% tool-format gate measures format compliance, not execution success.Source: §C.1 · p. 14, Figure 6 · p. 14
• Step 862 is selected for RL initialization, balancing task performance, tool-use behavior, and policy diversity.Compared with step 431, pass@1 improves from 81.93% to 83.35%, finish-bundled turns decrease from 182 to 54, while entropy remains at 0.4058 nats/token and the gap at 11.06 percentage points; step 1724 has the highest pass@1 and F1, but lower entropy and sampling gap than step 862.Source: §C.1 · p. 14, Table 6 · p. 15

C.2 Event-Wise Credit Computation

• Event-wise credit expands recipients to later updates or deletions of evidence records while keeping memory selection tied to the original evidence chunks.Events that create, update, or delete the same evidence record can share its score; deletion alone does not imply a zero score, and lineage-based assignment needs no extra answer-model calls.Source: §C.2 · pp. 14–15

C.3 RL Training Hyperparameters and Computational Resources

• RL alternates construction and retrieval iterations at a 1:1 ratio, using frozen Qwen3.8-27B answer and judge models during training.Evaluation instead uses GPT-OSS-120B as the answer model and GPT-4o-mini as judge; the paper says the reported gains suggest transfer across answer models.Source: §C.3 · p. 15

D Details of Baselines.

• The documented baselines include RL-trained managers such as Memory-R1, Memory-R2, and Mem-T, alongside RAG, Mem0, and MEM1.Their described approaches range from fixed similarity retrieval and extraction pipelines to internal-state memory and offline construction imitation.Source: D Details of Baselines · pp. 15–16

E.1 Memory-R2 Re-evaluation and Judge-Prompt Sensitivity

• 9.73 and 9.23 overall points separate reported from re-evaluated Memory-R2 and Memory-R2 (OSS) results under the original versus Mem0-aligned judge prompts.Table 8 presents reported and re-evaluated Memory-R2 results on LoCoMo using the Mem0-aligned prompt for re-evaluation.Source: Table 8 · p. 17, §E.1 · p. 17
• The authors treat originally reported baseline scores as reference points, relying on re-evaluated baselines that share the test split, judge prompt, and judge model with Uni-Mem.The reported and re-evaluated Memory-R2 results differ because the evaluation prompts differ.Source: §E.1 · pp. 16–17

E.2 Single-Stage GRPO and Ablation Configurations

• Standard GRPO aggregates QA rewards into trajectory-level returns, normalizes them across sibling rollouts, and broadcasts each advantage to construction and retrieval decisions.The ablation rollouts construct a memory bank and then retrieve evidence from it; both event types contribute to the same policy update.Source: §E.2 · p. 17
• With ECA, query evidence localizes related construction events; an unsuccessful evidence-indexed re-answer attributes failure to construction and excludes the retrieval event.A successful evidence-indexed re-answer credits construction and assigns zero reward to the failed retrieval; ⊥ excludes an event from group statistics and optimization, unlike zero reward.Source: §E.2 · p. 17

F.1 Gains from the Harness versus the RL Algorithm

• Under the same harness, Uni-Mem RL adds 4.41 and 5.24 points across the two initializations, while standard GRPO yields only marginal gains.Table 9 reports overall J (%) on LoCoMo; the untrained manager performs on par with re-evaluated baselines, and Uni-Mem RL without SFT surpasses them.Source: Table 9 · p. 17, §F.1 · pp. 17–18

E.2 Single-Stage GRPO and Ablation Configurations

• Construction-event returns sum credits from related queries, exclude events without related queries, and are normalized across corresponding events before advantages are broadcast within each event.Source: §E.2 · p. 17

F.1 Gains from the Harness versus the RL Algorithm

• Figure 7’s light curve shows logged training-time judge scores, while its dark curve shows a trailing 20-step moving average.Source: Figure 7 · p. 18
• No off-the-shelf manager, including Qwen3.5-27B and Qwen3.5-397B-A17B, reaches the trained 4B manager; SFT and Uni-Mem RL are complementary.Uni-Mem RL adds 5.24 points on top of SFT, despite the smaller remaining headroom.Source: §F.1 · p. 18

F.2 Training Curves

• Mean J rises from 75.76% in the first 20 steps to 81.41% in the final 20 steps, a 5.65-percentage-point improvement over 200 logged steps.Individual scores fluctuate, but the moving average trends upward, with smaller changes later in training.Source: §F.2 · p. 18

F.3 Detailed Results on Training and Manager Scale

• The manager-scale evaluation compares Qwen3-4B with four larger Qwen3.5 models as off-the-shelf LoCoMo memory managers, while UNI-MEM uses supervised-fine-tuned Qwen3-4B optimized with Stage-Decoupled GRPO.The settings follow the main text’s evaluation protocol; Table 10 reports category-level J results, but the supplied passages do not give cell values.Source: §F.3 · p. 18, Table 10 · p. 18

F.4 Detailed Analysis of Alternating Training Schedules

• The alternating-schedule analysis varies construction-to-retrieval update allocation and, separately, the number of consecutive updates per stage while keeping the two stages equally allocated.Each C:R schedule performs C consecutive construction updates followed by R consecutive retrieval updates; Table 11 gives category-level J results, but supplied passages contain no values.Source: §F.4 · p. 18, Table 11 · p. 18

F.5 Retrieval Budget and Token Efficiency

• Increasing maximum retrieval steps from 1 to 3 raises J from 78.06% to 83.67%; increasing from 3 to 8 reaches 85.67% but raises token consumption from 10,889 to 58,642 (5.39×).The additional 2.00-point gain from k = 3 to k = 8 is accompanied by diminishing returns at a larger retrieval budget.Source: §F.5 · p. 19, Figure 8 · p. 19

G.1 Memory Manager Prompt

• A shared tool interface and system prompt govern memory construction and query-time retrieval, with the manager maintaining long-term memory and retrieving evidence for a separate answer model.The prompt calls for focused, self-contained memories and requires relevant retrieval evidence before memory writes.Source: §G.1 · pp. 19–20

G.2 LLM-as-a-Judge (J) Prompt

• J is the percentage of evaluated answers that the judge labels CORRECT; the judge compares a generated answer with the question and reference answer.The evaluation template is supplied as a user message and returns CORRECT or WRONG.Source: §G.2 · p. 20
• For time questions, the judge counts answers as correct when they refer to the same date or period as the gold answer, despite wording or format differences.The prompt asks for generous grading of equivalent relative or differently formatted time references.Source: §G.2 · p. 21
• The judge prompt requires a one-sentence reasoning explanation followed by a CORRECT or WRONG label in JSON under the key "label".It warns that including both labels can break the evaluation script.Source: §G.2 · p. 21

G.3 Answer Prompt

• The frozen answer model receives a question and retrieved memories, listed chronologically with date prefixes; braced fields are runtime inputs.The system message directs it to verify entities, combine related facts, handle dates relative to the reference date, and prefer specific evidence.Source: §G.3 · p. 21
• The answer model must not invent names, titles, places, or dates absent from memory, should use the configured fallback when no memory is remotely related, and outputs only “ANSWER: <final answer>”.Source: §G.3 · p. 21