Source-linked AI summary
MemPO: Self-Memory Policy Optimization for Long-Horizon Agents
Ruoran Li, Xinghua Zhang, Haiyang Yu, Shitong Duan, Xiang Li, Wenxin Xiang, Chonghua Liao, Xudong Guo, Yongbin Li, Jinli Suo
TL;DR
Long-horizon agents face growing contexts that constrain interaction length, raise token costs, and can degrade performance, while external memory retrieval is not jointly optimized with task execution. MemPO lets the policy autonomously summarize and organize memory through a self-memory action and improves its credit assignment using memory-level rewards. Across five long-horizon benchmarks, it reports higher F1 and lower token usage than the base model and previous SOTA, while a stated limitation concerns bias from nonequivalent rollout states.
Problem
Growing interaction contexts constrain context capacity, increase token costs, and can degrade agent performance, while external memory retrieval remains passively optimized relative to task objectives.
Method
MemPO enables the policy to compress and reorganize historical information through a <mem> action, combining trajectory-level and memory-level advantages for task-aware credit assignment.
Results
25.98% and 7.1% absolute F1 gains over the base model and previous SOTA accompany 67.58% and 73.12% reductions in token usage across five long-horizon benchmarks.
Takeaways & Limitations
MemPO retains task-relevant information while reducing token consumption and jointly optimizes memory, reasoning, and tool invocation for long-horizon agents.
Takeaways & Limitations
Varying tool invocation makes rollout states nonequivalent, which may bias group-based advantages; more refined solutions and broader real-world generalization require investigation.
Abstract
from arXiv · showhide
Long-horizon agents face the challenge of growing context size during interaction with environment, which degrades the performance and stability. Existing methods typically introduce the external memory module and look up the relevant information from the stored memory, which prevents the model itself from proactively managing its memory content and aligning with the agent's overarching task objectives. To address these limitations, we propose the self-memory policy optimization algorithm (MemPO), which enables the agent (policy model) to autonomously summarize and manage their memory during interaction with environment. By improving the credit assignment mechanism based on memory effectiveness, the policy model can selectively retain crucial information, significantly reducing token consumption while preserving task performance. Extensive experiments and analyses confirm that MemPO achieves absolute F1 score gains of 25.98 over the base model and 7.1 over the previous SOTA baseline, while reducing token usage by 67.58% and 73.12%. The code is released at https://github.com/TheNewBeeKing/MemPO.
1 Introduction
Long-horizon agents suffer from expanding interaction contexts, which increase token costs and can degrade performance. MemPO makes memory management an intrinsic, task-aware policy capability and reports substantial performance gains with lower token usage.
- Motivation: Growing interaction histories limit context capacity, increase token costs, and can trigger “lost in the middle,” reducing agent performance.These constraints become more significant as tasks require more interaction rounds.
- Motivation: External memory systems retrieve relevant history through separate memory modules but lack joint optimization with the agent’s overarching task objectives.This leaves memory retrieval passive rather than proactively organized by the model.
- Approach: MemPO formalizes interaction around three actions—<mem>, <think>, and <tool_call>—so the agent compresses and reorganizes history before the next interaction.The inference process uses the previous-step interaction as input for the next step with a <mem> action.
- Approach: MemPO combines trajectory-level and memory-level advantages to improve credit assignment for memory actions, using answer likelihood conditioned on <mem> content as a dense memory reward.The method assigns additional memory-level advantages to <mem> tokens at each interaction step.
- Contributions: MemPO makes memory management intrinsic to the agent and jointly optimizes long-horizon memory, reasoning, and tool invocation.This distinguishes the method from external memory modules.
- Results: 25.98% and 7.1% absolute F1 gains over the base model and previous SOTA accompany 67.58% and 73.12% reductions in token usage.These results are reported across five long-horizon benchmarks.
2 Related Works
Prior memory systems commonly retrieve information from external stores using RAG, while reinforcement-learning approaches have only limitedly optimized agent memory. MemPO instead introduces dedicated memory-reward credit assignment to retain task-relevant information.
- External Memory: External memory systems such as MemGPT and Mem0 use structured memory management, extraction, consolidation, or retrieval to address LLM context limitations.These approaches expand or organize memory outside the policy’s intrinsic interaction process.
- Retrieval-Augmented Memory: RAG retrieves memory fragments through embedding similarity, but retrieval lacks flexibility and end-to-end joint optimization with the agent’s task.Similarity-based retrieval does not necessarily align retrieved information with the full task objective.
- RL for Memory: Existing reinforcement-learning work has rarely optimized agent memory, and approaches such as MEM1 do not explicitly design object-oriented memory credit assignment.The cited discussion identifies suboptimal memory representations as a limitation.
- RL for Memory: MemPO introduces dedicated credit assignment for memory rewards, encouraging retention of information most relevant to solving the target task.This directly targets the memory-representation limitation identified for prior approaches.
3 Preliminaries
The paper models each agent trajectory as a sequence of states containing memory, reasoning, tool calls, and tool responses. It then uses group-based reinforcement learning, supported by formatted trajectory data, to compare complete candidate trajectories.
- Task Formulation: A completed task trajectory is represented as T interaction states, τ = {s_1, s_2, . . . , s_T}, generated while the agent reasons and invokes tools.The formulation begins from a question-answer pair and multi-round environment interaction.
- Task Formulation: Each state decomposes into model-generated memory, reasoning, tool invocation, and returned tool information, followed by a predicted answer.Memory summarizes effective information from previous outputs and is enclosed by <mem> tags.
- Group Relative RL: Group-based RL generates N complete trajectories for one task and assigns each a scalar return measuring overall generation quality.The method compares candidate trajectories within a batch instead of modeling a per-trajectory value function.
- Group Relative RL: Group-based methods aggregate normalized or pairwise relative preferences, avoiding value-function estimation and reducing the memory overhead of extra networks.This makes them an efficient alternative for large-scale LLM reinforcement-learning training.
- Training Setup: Approximately 10K correctly answered, action-formatted trajectories are generated with GPT-4.1 and used to fine-tune the LLM before self-memory policy optimization.The trajectories are filtered from a publicly available training dataset.
4 Self-Memory Policy Optimization
MemPO supplements trajectory-level feedback with memory-level rewards that assess whether each generated memory preserves information useful for answering the task. It combines these signals to provide richer guidance for memory generation and uses compressed memory as the inference context.
- Motivation: Vanilla GRPO gives all tokens in a trajectory the same sparse reward, so final-answer correctness provides limited guidance for individual <mem> actions.The final answer cannot directly reflect the quality of each memory action during interaction.
- Memory-level advantage: MemPO evaluates the information content of each memory segment and computes an additional memory-level advantage alongside the trajectory-level advantage.This is designed to keep memories concise while preserving important information.
- Trajectory-level advantage: A trajectory receives reward 1 only when its predicted answer is correct and its output format is proper; otherwise, its trajectory-level reward is 0.The method evaluates trajectory format and final-answer accuracy as a coarse-grained assessment of overall trajectory quality.
- Memory-level advantage: Memory quality is estimated from conditional answer probabilities, with higher probability indicating that the retained memory preserves information relevant to generating the correct answer.The memory reward measures the effectiveness of information retained in the memory generated at each interaction step.
- Combination of advantages: For memory tokens, the final advantage sums trajectory-level and memory-level advantages; non-memory tokens receive only the trajectory-level advantage.This gives memory tokens richer feedback signals to guide memory generation.
- Inference: During inference, MemPO replaces the full prior context with the previous step’s effective memory, discarding earlier information while retaining its compressed content.The inference context is represented using the previous memory state rather than the entire interaction history.
5 Experiments
Experiments evaluate MemPO on increasingly difficult multi-objective tasks using accuracy and token-consumption metrics, ablations, and conditional-probability analyses. The method maintains strong performance with compact contexts, while its reward design improves memory effectiveness and token efficiency.
- Experimental Setup: Multi-objective experiments measure F1, EM, total tokens, and peak tokens across progressively harder tasks and compare MemPO with prompt-, RL-, and memory-based baselines.Tasks use 4, 6, 8, and 10 objectives; memory-based methods use truncated contexts, while other baselines use complete context.
- Main Results: MemPO achieves leading performance on harder-than-training tasks and remains competitive in a real-world web-search environment while minimizing token consumption.The reported results describe stable performance under distributional change and the highest performance with the least resource usage.
- Token Efficiency: At 10 objectives, MemPO uses approximately one-third as many tokens as ReSearch and has one-fifth its token peak, with the lowest consumption among all baselines.The comparison is narrower at simpler tasks, but the efficiency gap widens as task complexity increases.
- Memory Quality: MemPO explicitly guides retention toward context most relevant to solving the task, outperforming other memory-related baselines across datasets.The paper contrasts this task-linked memory selection with redundant or less relevant memory representations in other methods.
- Ablation and Context Retention: Performance remains stable with full context or one- and three-round retention, although more context helps short-horizon tasks and can weaken long-horizon performance.The authors associate the long-context degradation with attention dilution.
- Reward Analysis: MemPO’s reward design improves effective memory information: its conditional probabilities shift higher, and mean probabilities rise across the first ten steps while the baseline declines.After ten steps, MemPO’s probability decreases, consistent with the typical search-step count for a 10-objective task.
6 Conclusion
MemPO optimizes agent memory management through a reward design that retains relevant information while integrating memory, reasoning, and tool invocation. The method is reported to improve long-horizon performance and reduce token consumption, with future work targeting reward optimization and scalability.
- Conclusion: MemPO retains task-relevant information through a novel reward design, improving task performance while reducing computational costs.The method integrates memory, reasoning, and tool invocation through reinforcement learning.
Limitations
MemPO’s evaluation may be affected by unequal memory information across interaction steps, which can bias group-based advantage estimation.
- Varying tool invocation makes memory information differ across steps and rollout trajectories, so states are not completely equivalent for group-based advantage calculation.The authors introduce ϵ in Equation 4 to alleviate this bias, but more refined solutions may be needed in complex environments.
A.1 Single-Objective Tasks
MemPO is evaluated on seven single-objective question-answering benchmarks spanning diverse domains, including multi-hop, compositional, long-tail, and real-search-derived questions. It achieves strong performance across these benchmarks while substantially reducing token consumption, remaining competitive with models trained specifically for single-objective tasks.
- Datasets: Seven QA benchmarks cover diverse domains, including multi-hop reasoning, compositional questions, long-tail factual knowledge, and real anonymized search queries.The evaluation includes 2WikiMultiHopQA, HotpotQA, Bamboogle, Musique, Natural Questions, TriviaQA, and PopQA.
- Results: MemPO achieves strong performance across all evaluated single-objective benchmarks and reaches SOTA performance on several datasets, including TriviaQA.The results also report substantially reduced token consumption.
- Results: MemPO remains competitive on short-horizon tasks, matching or exceeding agent models trained specifically for single-objective settings.
A.2 Deep Research Tasks
MemPO is evaluated on deep-research benchmarks built from real web search results and therefore out-of-domain relative to training. It achieves accuracy comparable to DeepResearcher while using fewer tokens and performs relatively strongly on longer-horizon tasks such as GAIA.
- Datasets: Three deep-research benchmarks—GAIA, Frames, and WebWalkerQA—test generalization to complex out-of-domain tasks derived from real web search results.Frames and WebWalkerQA are evaluated using randomly sampled instances to reduce evaluation cost.
- Datasets: GAIA evaluates tool use and multi-step reasoning across 165 tasks spanning three difficulty levels.
- Results: MemPO achieves average accuracy comparable to DeepResearcher while significantly reducing token usage across the deep-research evaluation.DeepResearcher is trained specifically in real web-search environments.
- Results: MemPO delivers relatively strong performance on longer-horizon, higher-token-demand tasks such as GAIA.