Skip to content

MemPO: Self-Memory Policy Optimization for Long-Horizon Agents

MemoryAgentRL20+清华大学通义实验室阿里巴巴

Li R, Zhang X, Yu H, et al. MemPO: Self-Memory Policy Optimization for Long-Horizon Agents[J]. arXiv preprint arXiv:2603.00680v4, 2026.

https://github.com/TheNewBeeKing/MemPO


MemPO:面向长程智能体的自记忆策略优化

Abstract

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%.

长程智能体在与环境交互时面临上下文规模不断增长的挑战,这会降低性能和稳定性。 现有方法通常引入外部记忆模块,并从存储的记忆中查找相关信息,这使模型自身无法主动管理其记忆内容,也难以与智能体的整体任务目标对齐。 为解决这些限制,我们提出自记忆策略优化算法(MemPO),使智能体(策略模型)能够在与环境交互期间自主总结和管理自己的记忆。 通过改进基于记忆有效性的信用分配机制,策略模型能够选择性保留关键信息,在保持任务性能的同时显著降低 token 消耗。 大量实验和分析确认,MemPO 相比基础模型取得 25.98 的绝对 F1 分数提升,相比此前 SOTA 基线取得 7.1 的绝对 F1 分数提升,同时将 token 使用量分别降低 67.58% 和 73.12%。

1. Introduction

As large language models (LLMs) continue to evolve, LLM agents are becoming increasingly proficient in addressing more complex problems. In areas such as deep research, data analysis, and vibe coding, they have showcased remarkable performance. Long-horizon decision-making has always been one of the core capabilities for agents to solve complex user queries.

随着大语言模型(LLM)持续演进,LLM 智能体在处理更复杂问题方面正变得越来越熟练。 在深度研究、数据分析和 vibe coding 等领域,它们已经展现出显著性能。 长程决策一直是智能体解决复杂用户查询的核心能力之一。

Currently, the dominant method for the agent-environment interaction is ReAct paradigm. The feedback from the environment is attached to the previous interaction history and is used as a prompt, which then determines the next course of action. However, this approach causes the context to grow linearly with each round of interaction, resulting in longer contexts when tackling more complex problems, and presenting several challenges. Firstly, current LLMs have relatively limited context window sizes, which impose an explicit upper bound on the number of interactions. Secondly, long contexts lead to excessively high token costs, which impedes the widespread adoption of agent systems in practical scenarios. Furthermore, excessively long contexts can lead to the ``lost in the middle'' phenomenon, which degrades the model's ability, thereby reducing the overall performance of the agent.

目前,智能体-环境交互的主流方法是 ReAct 范式。 来自环境的反馈会被附加到先前的交互历史中并作为提示使用,随后决定下一步动作。 然而,这种方法会使上下文随着每一轮交互线性增长,在处理更复杂问题时产生更长上下文,并带来若干挑战。 首先,当前 LLM 的上下文窗口大小相对有限,这对交互次数施加了明确上界。 其次,长上下文会导致过高的 token 成本,阻碍智能体系统在实际场景中的广泛采用。 此外,过长上下文可能导致“lost in the middle”现象,这会削弱模型能力,从而降低智能体整体性能。

To address this challenge, a growing body of research is focusing on agent memory, with the aim of providing LLMs with historical interaction records to reduce the need for the entire context. Currently, the mainstream solution involves designing a memory module as an external knowledge database to maintain the agent's interaction history. When the memory module is accessed, relevant historical information is retrieved and integrated into the prompt based on the retrieval-augmented generation technique (RAG). However, the offline memory context compression method lacks the capacity for joint optimization oriented toward the agent task execution, making it difficult to effectively align with the agent's overarching task objectives. As a result, the model's memory retrieval remains passive, rather than leveraging its own capabilities to proactively select and organize information, and the latter would facilitate more effective task completion.

为应对这一挑战,越来越多研究开始关注智能体记忆,目标是为 LLM 提供历史交互记录,以减少对完整上下文的需求。 目前,主流方案是把记忆模块设计为外部知识数据库,用于维护智能体的交互历史。 当访问记忆模块时,会基于检索增强生成(RAG)技术检索相关历史信息,并将其整合到提示中。 然而,离线记忆上下文压缩方法缺乏面向智能体任务执行的联合优化能力,因此难以有效对齐智能体的整体任务目标。 因此,模型的记忆检索仍然是被动的,而不是利用自身能力主动选择和组织信息;后者会促进更有效的任务完成。

The self-memory inference process of MemPO
图1:The self-memory inference process of our method, which only uses the previous step interaction for next step input with <mem> action.

To this end, we formalize the agent interaction paradigm as autonomously refining and organizing historical information, while simultaneously reasoning and invoking tools with three actions <mem>, <think>, and <tool_call>. In this paradigm, the agent itself proactively compresses and reorganizes long-horizon historical information for the next step of interaction, making memory management an intrinsic part of its capabilities, as shown in Figure 1. To further enhance this ability, we propose self-memory policy optimization (MemPO), which incorporates the trajectory-level and memory-level information into advantage estimation to optimize the <mem> action for agent with task-objective awareness. Concretely, the tokens output by the agent are assigned trajectory-level advantages, and in each step of the interaction, the tokens of the <mem> action additionally take into account memory-level advantages, effectively alleviating the credit assignment problem in long-horizon, multi-turn interactions. In terms of dense rewards for <mem> action in each step, the conditional probability of the answer given <mem> content is designed to measure the quality of the <mem> action. Our contributions are as follows:

为此,我们将智能体交互范式形式化为自主精炼和组织历史信息,同时通过三个动作 <mem><think><tool_call> 进行推理并调用工具。 在这一范式中,智能体自身会主动压缩和重组用于下一步交互的长程历史信息,使记忆管理成为其能力的内在组成部分,如图1所示。 为进一步增强这一能力,我们提出自记忆策略优化(MemPO),它把轨迹级和记忆级信息纳入优势估计,以具有任务目标意识的方式优化智能体的 <mem> 动作。 具体而言,智能体输出的 token 会被分配轨迹级优势,而在每一步交互中,<mem> 动作的 token 还会额外考虑记忆级优势,从而有效缓解长程多轮交互中的信用分配问题。 就每一步 <mem> 动作的稠密奖励而言,我们设计了在给定 <mem> 内容时答案的条件概率来衡量 <mem> 动作的质量。 我们的贡献如下:

  • We render memory management an intrinsic part of the agent's own capabilities that differs from external memory modules, achieving joint optimization of long-horizon memory, reasoning, and tool invocation.
  • We propose MemPO, a self-memory policy optimization algorithm, which effectively addresses credit assignment and steers the <mem> action toward retaining the most relevant information for solving the task.
  • Extensive experiments on five long-horizon benchmarks confirm the efficacy of MemPO with 25.98% and 7.1% absolute F1 gains over the base model and previous SOTA, 67.58% and 73.12% reductions in token usage.
  • 我们使记忆管理成为智能体自身能力的内在组成部分,并区别于外部记忆模块,实现了长程记忆、推理和工具调用的联合优化。
  • 我们提出 MemPO,这是一种自记忆策略优化算法,能够有效处理信用分配,并引导 <mem> 动作保留对解决任务最相关的信息。
  • 在五个长程基准上的大量实验确认了 MemPO 的有效性:相比基础模型和此前 SOTA,绝对 F1 分别提升 25.98% 和 7.1%,token 使用量分别降低 67.58% 和 73.12%。

2.1 Memory for LLM agents

In recent years, researchers have introduced external memory and experience systems to address the limitations of LLM context windows. MemGPT proposes an operating-system-inspired virtual memory management framework that employs multiple memory hierarchies to manage contextual information. Mem0 enhances memory capacity through dynamic extraction, consolidation, and retrieval of conversational information. Despite their effectiveness in specific domains, most of these approaches rely on fixed workflows and limited optimization flexibility. They typically fail to support flexible cross-stage joint optimization, which constrains the adaptability and scalability of the overall system.

近年来,研究者引入外部记忆和经验系统,以解决 LLM 上下文窗口的限制。 MemGPT 提出了一种受操作系统启发的虚拟记忆管理框架,采用多个记忆层级来管理上下文信息。 Mem0 通过动态抽取、整合和检索对话信息来增强记忆容量。 尽管这些方法在特定领域中有效,但它们大多依赖固定工作流,优化灵活性有限。 它们通常无法支持灵活的跨阶段联合优化,这限制了整体系统的适应性和可扩展性。

2.2 RAG in Memory System

RAG has emerged as a powerful approach for enhancing LLM by incorporating external knowledge sources to improve model performance. In existing memory systems, the retrieval of relevant memory fragments is predominantly implemented based on RAG. While this approach can efficiently surface relevant information in certain scenarios, its major limitation lies in the lack of flexibility and end-to-end joint optimization. Specifically, retrieval relies solely on embedding similarity between the query and chunks, which does not necessarily yield information that is most useful for solving the target problem.

RAG 已成为增强 LLM 的有力方法,它通过引入外部知识源来提升模型性能。 在现有记忆系统中,相关记忆片段的检索主要基于 RAG 实现。 虽然这种方法在某些场景中能够高效浮现相关信息,但其主要局限在于缺乏灵活性和端到端联合优化。 具体而言,检索只依赖查询与片段之间的嵌入相似度,这不一定能产生对解决目标问题最有用的信息。

2.3 RL for LLM Agents

The recent success of reinforcement learning methods in LLMs has established RL as a central tool to enhance LLM-based agents to solve increasingly complex tasks. However, relatively few studies have explored applying RL to the optimization of agent memory. Existing approaches exhibit notable limitations. For example, MEM1 integrates memory into the reasoning process and applies RL optimization for the policy model. However, it does not explicitly design objectives for memory optimization, which can lead to suboptimal memory representations. In contrast, our method introduces a dedicated credit assignment mechanism for memory rewards, encouraging the model to retain information that is most relevant for solving the target task.

强化学习方法近期在 LLM 中的成功,已经使 RL 成为增强基于 LLM 的智能体并解决日益复杂任务的核心工具。 然而,探索将 RL 应用于智能体记忆优化的研究相对较少。 现有方法存在明显局限。 例如,MEM1 将记忆整合进推理过程,并对策略模型应用 RL 优化。 然而,它没有为记忆优化显式设计目标,这可能导致次优的记忆表示。 相比之下,我们的方法为记忆奖励引入专门的信用分配机制,鼓励模型保留与解决目标任务最相关的信息。

Overview of MemPO
图2:Overview of MemPO. At step t of any trajectory τi, the context is represented as {stmem,stthink,stcall,stresp}. The memory reward RM is calculated using conditional probabilities and contributes to the advantage AM. The final advantage is the sum of AM and the trajectory-level advantage AT. During inference, only the previous step’s content is used as context, discarding earlier information.

3. Preliminaries

3.1 Task Formulation

Given a question-answer pair (q,agt), when an LLM-based agent is tasked with solving the question q, it interacts with the external environment through multiple rounds of reasoning and tool invocation to acquire the information required for problem solving. If the agent completes the task after T steps, a full trajectory can be denoted as τ={s1,s2,,sT}.

给定问答对 (q,agt),当一个基于 LLM 的智能体被要求解决问题 q 时,它会通过多轮推理和工具调用与外部环境交互,以获取解决问题所需的信息。 如果智能体在 T 步后完成任务,则完整轨迹可以记为 τ={s1,s2,,sT}

Each state st is further decomposed into {stmem,stthink,stcall,stresp}. Specifically, stmem represents the model-generated summary of effective information from previous outputs s<t, which is enclosed by <mem> </mem>. stthink corresponds to the model's reasoning process and is wrapped by <think> </think>. stcall denotes the invocation of external tools by the model, which is represented as <tool_call> </tool_call>. stresp captures the information returned by the tool and is enclosed by <information> </information>. Once the agent has gathered sufficient information to answer the question q, it produces a predicted answer apred, wrapped by <answer> </answer>.

每个状态 st 进一步分解为 {stmem,stthink,stcall,stresp} 具体而言,stmem 表示模型从先前输出 s<t 中生成的有效信息摘要,并由 <mem> </mem> 包裹。 stthink 对应模型的推理过程,并由 <think> </think> 包裹。 stcall 表示模型对外部工具的调用,写作 <tool_call> </tool_call> stresp 捕获工具返回的信息,并由 <information> </information> 包裹。 一旦智能体收集到足够信息来回答问题 q,它会生成预测答案 apred,并由 <answer> </answer> 包裹。

3.2 Group Relative RL

In reinforcement learning for LLM, a class of group-based methods, exemplified by Group Relative Policy Optimization (GRPO), abandon per-trajectory value function modeling and instead performs relative comparison within a batch of candidate trajectories. Concretely, for a given task input q, the policy πθold generates N complete trajectories {τ1,τ2,,τN} in one shot, and each trajectory is assigned a scalar return R(τi) that measures the overall quality of the generated outcome. The algorithm then relies solely on statistics within this group to construct advantages, without explicitly learning a value network:

在面向 LLM 的强化学习中,以 Group Relative Policy Optimization(GRPO)为代表的一类基于组的方法,会放弃逐轨迹价值函数建模,转而在一批候选轨迹内进行相对比较。 具体而言,对于给定任务输入 q,策略 πθold 一次性生成 N 条完整轨迹 {τ1,τ2,,τN},并为每条轨迹分配一个标量回报 R(τi),用于衡量生成结果的整体质量。 随后,该算法只依赖这一组内的统计量来构造优势,而不显式学习价值网络:

A(τi)=GroupAgg({R(τj)}j=1N,i).

where GroupAgg() is the aggregation operator based on normalization, or pairwise comparison. The design bypasses the instabilities of value function estimation and reduces it to modeling relative preferences among a set of candidate answers. In large-scale LLM training, group-based methods can reduce the memory overhead of extra networks, making them an efficient alternative for RL training.

其中 GroupAgg() 是基于归一化或成对比较的聚合算子。 这种设计绕开了价值函数估计的不稳定性,并将其简化为建模一组候选答案之间的相对偏好。 在大规模 LLM 训练中,基于组的方法可以减少额外网络的记忆开销,使其成为 RL 训练的高效替代方案。

3.3 Behavior Cloning

To enable the model to better follow the action format, we first adopt GPT-4.1 to perform inference on the publicly available training dataset from the work of Tang et al., filter out trajectory with incorrect answers, and finally generate approximately 10K trajectories following the predefined action format in Section 3.1. Based on these trajectory data, we fine-tune the LLM and provide a promising starting point for self-memory policy optimization.

为了使模型更好地遵循动作格式,我们首先采用 GPT-4.1 在 Tang 等人的公开训练数据集上执行推理,过滤掉答案错误的轨迹,最终生成约 10K 条遵循第 3.1 节预定义动作格式的轨迹。 基于这些轨迹数据,我们对 LLM 进行微调,并为自记忆策略优化提供一个有希望的起点。

4. Self-Memory Policy Optimization

As mentioned above, memory is introduced to address long contexts of agents by removing irrelevant information and retaining key details. Vanilla GRPO computes rewards based on answer correctness and uses trajectory-level advantages, where tokens within the same trajectory share the same reward. It provides the sparse rewards and limited guidance for memory generation, as the correctness of the final answer can not directly reflect the quality of each <mem> action during the interaction.

如上所述,记忆被引入来处理智能体的长上下文问题,它通过移除无关信息并保留关键细节来发挥作用。 普通 GRPO 基于答案正确性计算奖励,并使用轨迹级优势,同一轨迹内的 token 共享相同奖励。 它为记忆生成提供的奖励稀疏且指导有限,因为最终答案的正确性无法直接反映交互期间每个 <mem> 动作的质量。

To address this, we propose MemPO, a self-memory policy optimization algorithm. We design a novel advantage computation method that, in addition to trajectory-level advantages, evaluates the information content of memory within <mem> </mem> at each step and computes an additional advantage, ensuring memory remains concise while preserving important information.

为解决这一问题,我们提出 MemPO,这是一种自记忆策略优化算法。 我们设计了一种新的优势计算方法:除了轨迹级优势外,它还评估每一步 <mem> </mem> 内记忆的信息含量并计算额外优势,确保记忆在保留重要信息的同时保持简洁。

4.1 Advantages of Global Trajectory

We first evaluate the trajectory format and the accuracy of the final answer to provide a coarse-grained assessment of the overall trajectory quality. Suppose that for a single training sample, we perform N rollouts and assign an overall score to each resulting trajectory, denoted as a group:

我们首先评估轨迹格式和最终答案准确性,以提供对整体轨迹质量的粗粒度评估。 假设对于单个训练样本,我们执行 N 次 rollout,并为每条生成轨迹分配一个整体分数,记为一个组:

GT={(τ1,RT(τ1)),(τ2,RT(τ2)),,(τN,RT(τN))}.

where τi denotes a trajectory, and RT(τi) represents the trajectory-level reward. In our method, the reward consists of evaluations of both the output format and the correctness of predicted answer. Concretely, the reward is set to 1 if and only if the predicted answer is correct and output format is proper; otherwise, it is set to 0.

其中 τi 表示一条轨迹,RT(τi) 表示轨迹级奖励。 在我们的方法中,奖励由输出格式和预测答案正确性两方面的评估组成。 具体而言,当且仅当预测答案正确且输出格式恰当时,奖励设为 1;否则设为 0。

To assess the global relative quality of each trajectory within the group, we adopt the advantage calculation strategy from GRPO, which normalizes the total reward using the mean and standard deviation computed over the group:

为了评估组内每条轨迹的全局相对质量,我们采用 GRPO 的优势计算策略,使用组内计算得到的均值和标准差对总奖励进行归一化:

AT(τi)=RT(τi)mean({RT(τj)}j=1N)std({RT(τj)}j=1N).

4.2 Advantages of Informative Memory

According to the probabilistic formulation of LLMs, the output of model is characterized as conditional probabilities given the preceding context, i.e., πθ(stq,s<t). As demonstrated and exploited in previous work, if the context s<t contains sufficient information to solve the problem, the probability that st is sampled as the answer-generation step sans will be relatively high. Similarly, this insight suggests that for an arbitrary context sany, a higher value of πθ(sansq,sany) indicates that sany contains more key information relevant to solving the question q, thereby increasing the model's confidence in generating the correct answer. Consequently, conditional probability can be used as a quantitative measure of the effective information content contained in a given context.

根据 LLM 的概率表述,模型输出被刻画为在给定前序上下文时的条件概率,即 πθ(stq,s<t) 如既有工作所展示和利用的,如果上下文 s<t 包含足以解决问题的信息,那么 st 被采样为答案生成步骤 sans 的概率会相对较高。 类似地,这一洞见表明,对于任意上下文 sanyπθ(sansq,sany) 的值越高,就说明 sany 包含越多与解决问题 q 相关的关键信息,从而提高模型生成正确答案的信心。 因此,条件概率可以用作给定上下文中有效信息含量的定量度量。

Based on this intuition, we design a step-level reward for the memory (<mem> action) generated at each interaction step, which reflects the quality of effective information retained in memory:

基于这一直觉,我们为每个交互步骤生成的记忆(<mem> 动作)设计了步级奖励,用以反映记忆中保留的有效信息质量:

RM(τi(stmem))=P[sansτi(stmem)]ϵ,1iN,1tT.

where τi(stmem) denotes the memory content within <mem> action in step t of trajectory τi, and sans={a1,a2,,aL} represents the correct answer string, where al denotes the l-th token of the answer. The ϵ represents P[sansτi(s<t)], which serves as a bias term. The term τi(s<t) corresponds to the trajectory of the first t-1 steps of τi. The operator P() denotes a posterior-probability-based measure. Specifically, P[sansτi(stmem)] can be represented as:

其中 τi(stmem) 表示轨迹 τi 中第 t<mem> 动作内的记忆内容,sans={a1,a2,,aL} 表示正确答案字符串,其中 al 表示答案的第 l 个 token。 ϵ 表示 P[sansτi(s<t)],它作为偏置项。 τi(s<t) 对应于轨迹 τit-1 步的轨迹。 算子 P() 表示一种基于后验概率的度量。 具体而言,P[sansτi(stmem)] 可表示为:

l=1Lπθ(alq,τi(stmem),a<l)L.

where πθ(alq,τi(stmem),a<l) denotes the probability of generating token al given the user query, the trajectory prefix up to step t1, and the preceding answer tokens a<l. Under this formulation, a larger value of RM(τi(stmem)) indicates that the memory generated at step t provides a more effective summary of the trajectory up to the first t-1 steps, and better preserves contextual information that is relevant to generating the correct answer.

其中 πθ(alq,τi(stmem),a<l) 表示在给定用户查询、截至第 t-1 步的轨迹前缀以及前序答案 token a<l 时生成 token al 的概率。 在这一表述下,RM(τi(stmem)) 的值越大,表示第 t 步生成的记忆越能有效总结前 t-1 步轨迹,并更好保留与生成正确答案相关的上下文信息。

Based on the reward formulation above, the resulting memory group can be expressed as:

基于上述奖励表述,得到的记忆组可以表示为:

GM={(τi(stmem),RM(τi(stmem)))1iN,1tT}.

We then normalize the rewards using the group-wise mean and standard deviation to obtain the corresponding advantages AM(τi(stmem)):

随后,我们使用组内均值和标准差对奖励进行归一化,得到对应的优势 AM(τi(stmem))

AM(τi(stmem))=RM(τi(stmem))M(τi(stmem))std({RM(τi(stmem))}).

where M(τi(stmem)) denotes the mean reward within the same group, defined as:

其中 M(τi(stmem)) 表示同一组内的平均奖励,定义为:

M(τi(stmem))=mean({RM(τi(stmem))(τi(stmem),RM(τi(stmem)))GM}).

The advantage AM(τi(stmem)) provides a quantitative assessment of memory quality, enabling finer-grained supervision over the model-generated memory content.

优势 AM(τi(stmem)) 提供了对记忆质量的定量评估,从而能够对模型生成的记忆内容进行更细粒度的监督。

4.3 Combination of Advantages

The final token-level advantage is obtained by combining the two types of advantages in Section 4.1 and Section 4.2. Let the k-th token of the i-th trajectory τi in a group be denoted as τi,k. The advantage assigned to this token Ai,k is defined as:

最终 token 级优势通过结合第 4.1 节和第 4.2 节中的两类优势得到。 令组中第 i 条轨迹 τi 的第 k 个 token 记为 τi,k 分配给该 token 的优势 Ai,k 定义为:

Ai,k={AT(τi)+AM(τi(stmem)),τi,kτi(stmem),AT(τi),otherwise.

That is, when τi,k corresponds to a token within the memory segment (<mem> action), its advantage is given by the sum of the trajectory-level advantage and the memory-level advantage; otherwise, only the trajectory-level advantage AT is used. In this way, tokens belonging to memory receive richer and more explicit feedback signals, which more effectively guide the rollout process for memory generation.

也就是说,当 τi,k 对应于记忆片段(<mem> 动作)中的 token 时,它的优势由轨迹级优势和记忆级优势之和给出;否则只使用轨迹级优势 AT 通过这种方式,属于记忆的 token 会接收更丰富且更明确的反馈信号,从而更有效地引导记忆生成的 rollout 过程。

4.4 Policy Optimization and Inference

Optimization. The policy optimization objective is to maximize J(θ), written as:

优化。 策略优化目标是最大化 J(θ),写作:

J(θ)=E[1Ni=1N1|τi|k=1|τi|min(γi,kAi,k,clip(γi,k,1ϵ,1+ϵ)Ai,k)βDKL(πθπref)].

where γi,k is the importance sampling ratio:

其中 γi,k 是重要性采样比率:

γi,k=πθ(τi,kq,τi,<k)πθold(τi,kq,τi,<k).

where qp(Q),{τi}i=1Nπθold. p(Q) denotes the distribution of queries in the training set, and β controls the weight of the KL-divergence regularization term.

其中 qp(Q),{τi}i=1Nπθold p(Q) 表示训练集中查询的分布,β 控制 KL 散度正则项的权重。

Inference. In vanilla ReAct framework, the t-th step inference is denoted as πθ(stq,s<t). In our method, since st1mem contains the effective information of s<t2, we replace s<t with st1mem as the inference context, represented as πθ(stq,st1mem).

推理。 在普通 ReAct 框架中,第 t 步推理记为 πθ(stq,s<t) 在我们的方法中,由于 st1mem 包含 s<t2 的有效信息,我们用 st1mem 替换 s<t 作为推理上下文,表示为 πθ(stq,st1mem)

表1:The accuracy and token consumption for multi-objective tasks of baselines. Text with bold means SOTA.
Local Wiki Search
Model2-objective4-objective6-objective8-objective10-objectiveAvg
F1EMF1EMF1EMF1EMF1EMF1 ↑EM ↑TT ↓PT ↓
Qwen2.5 (ReAct)33.7325.6010.597.005.374.005.924.252.631.9611.658.563.640.61
ReSearch47.4036.0024.1316.7020.8415.9310.807.705.083.5621.6515.983.290.71
DeepResearcher30.9424.7024.5218.1013.8810.739.126.855.073.5216.7112.784.290.77
A-MEM33.2425.2013.7110.109.807.076.865.105.563.6013.8310.212.620.38
MEM147.7437.1026.5118.9018.8114.0719.0413.5519.6113.3626.3419.401.380.20
GRPO (w/o <mem>)54.5742.9538.3128.6029.7822.6018.9713.6511.017.8430.5323.134.390.81
MemPO (Ours)56.4746.1542.7531.9034.3226.9330.4823.7024.1518.1637.6329.371.180.18
Online Web Search
Qwen2.5 (ReAct)45.7133.6016.8311.8012.129.079.666.807.254.6818.3113.193.140.34
ReSearch51.1739.2029.9221.1025.2118.7317.0912.5010.377.4426.7519.792.170.40
MEM150.5639.6030.4322.0021.6716.2019.4814.3018.0612.1228.0420.840.960.14
MemPO (Ours)57.4045.2041.4230.2037.9228.6034.3025.8022.9216.3238.7929.220.860.12

5. Experiments

5.1 Benchmarks

To evaluate the effectiveness of our approach, following the method of MEM1, we test on multi-objective tasks, where the number of interaction rounds required for the agent to solve a problem is significantly higher compared to single-objective tasks. This allows us to better assess the performance of our method in scenarios with long contexts. Additionally, we can observe the changes in agent performance by progressively increasing the number of objectives. We created a 2-objective task test set by combining queries from the validation sets of the HotpotQA and NQ QA datasets, and synthesized test sets with more objectives using the HotpotQA validation set. We conducted tests under both local wiki search engine and web search engine scenarios to enhance the credibility of the experiments.

为评估我们方法的有效性,我们遵循 MEM1 的方法,在多目标任务上进行测试;与单目标任务相比,智能体解决问题所需的交互轮数显著更多。 这使我们能够更好评估我们的方法在长上下文场景中的性能。 此外,我们可以通过逐步增加目标数量来观察智能体性能的变化。 我们通过组合 HotpotQA 和 NQ 问答数据集验证集中的查询,构造了一个 2-objective 任务测试集,并使用 HotpotQA 验证集合成了包含更多目标的测试集。 我们在本地 wiki 搜索引擎和 Web 搜索引擎两种场景下进行测试,以增强实验可信度。

Following previous work, we use the F1 score as a criterion for word-level matching, and Exact Match (EM) for exact matching. Furthermore, to evaluate the token consumption of the agent when solving a problem, we use the total number of tokens consumed to solve a question (TT), as well as the maximum number of tokens (peak tokens) consumed in a single step (PT).

遵循既有工作,我们使用 F1 分数作为词级匹配的标准,并使用 Exact Match(EM)进行精确匹配。 此外,为了评估智能体解决问题时的 token 消耗,我们使用解决一个问题所消耗的总 token 数(TT),以及单步消耗的最大 token 数(峰值 token,PT)。

5.2 Baselines

We compare our method with various baselines. For prompt-based baselines, we use ReAct. For agentic RL-based baselines, we adopt DeepResearcher and ReSearch. For agent memory-related baselines, we use RL-based method MEM1 and RAG-based method A-MEM. Additionally, we also trained a model without memory using GRPO in the exact same environment as a baseline. To ensure fairness, all methods use the 7B model from the Qwen2.5 series as the base model.

我们将我们的方法与多种基线进行比较。 对于基于提示的基线,我们使用 ReAct。 对于基于智能体 RL 的基线,我们采用 DeepResearcher 和 ReSearch。 对于智能体记忆相关基线,我们使用基于 RL 的方法 MEM1 和基于 RAG 的方法 A-MEM。 此外,我们还在完全相同环境下使用 GRPO 训练了一个不带记忆的模型作为基线。 为确保公平,所有方法都使用 Qwen2.5 系列的 7B 模型作为基础模型。

5.3 Implementation Details

We first performed inference using GPT-4.1 on the dataset from the work of Tang et al., and obtained approximately 10k trajectories containing memory. We then fine-tuned the model for one epoch using these data to enhance its ability to follow instructions related to the memory component. Simultaneously, we removed the memory component from the trajectories to use them for fine-tuning the baseline model, which was trained using GRPO and does not include memory, ensuring fairness in the comparison.

我们首先使用 GPT-4.1 在 Tang 等人的数据集上执行推理,并获得约 10k 条包含记忆的轨迹。 随后,我们使用这些数据对模型微调一个 epoch,以增强其遵循记忆组件相关指令的能力。 同时,我们从轨迹中移除了记忆组件,并使用这些轨迹微调基线模型;该基线模型使用 GRPO 训练且不包含记忆,以确保比较公平。

In the RL phase, we followed MEM1 and used the 2-objective task synthesized from HotpotQA and NQ as part of the training set. And we randomly sampled a subset from both datasets as another part of training set. The rollout group size N for group-based RL methods is set to 16, with a batch size of 128 and a learning rate of 1e-6. The maximum number of interaction rounds is set to 16. During training, we use the local wiki search engine as the search tool.

在 RL 阶段,我们遵循 MEM1,并使用由 HotpotQA 和 NQ 合成的 2-objective 任务作为训练集的一部分。 同时,我们从两个数据集中随机采样一个子集,作为训练集的另一部分。 对于基于组的 RL 方法,rollout 组大小 N 设为 16,batch size 为 128,学习率为 1e-6。 最大交互轮数设为 16。 训练期间,我们使用本地 wiki 搜索引擎作为搜索工具。

Ablation line resultAblation bar result
图3:The results of ablation study.
The result of step conditional probability analysis
图4:The result of step conditional probability analysis.

5.4 Experimental Results

Multi-objective task. The results of each baseline on the multi-objective task are shown in Table 1. We selected tasks with 4, 6, 8, and 10 objectives as progressively harder task groups and recorded F1 and EM for the answers from each baseline as precision metrics. Among the baselines, MEM1, A-MEM, and our method use truncated contexts, meaning the model only has access to the previous step of interactions, while the other baselines use the complete context. Additionally, we also recorded the total number of tokens required to solve a single problem (TT) and the peak token consumption per step (PT) during the model's execution. The more detailed results are presented in Appendix Table.

多目标任务。 各基线在多目标任务上的结果如表1所示。 我们选择包含 4、6、8 和 10 个目标的任务作为逐步更难的任务组,并记录各基线答案的 F1 和 EM 作为精度指标。 在这些基线中,MEM1、A-MEM 和我们的方法使用截断上下文,这意味着模型只能访问前一步交互,而其他基线使用完整上下文。 此外,我们还记录了模型执行期间解决单个问题所需的总 token 数(TT)和每步峰值 token 消耗(PT)。 更详细结果见附录表。

Conditional probability analysis. To investigate the impact of our reward design on the model, we performed a statistical analysis of the true values of P[sanssmem] during inference on the 10-objective task. We compared the results of models trained with vanilla GRPO and our method. Specifically, let the dataset size be M, and denote the memory component at step t of the m-th trajectory as τm(stmem), with the corresponding ground truth answer string being τm(sans). Figure 5 shows the grouped results of P[τm(sans)τm(stmem)], where the x-axis represents the group values of the conditional probability, and the y-axis shows the proportion of memory samples in that range relative to the total memory samples. This distribution illustrates the conditional probability distribution of the memory produced by the model. The line graph's y-axis represents the average accuracy of trajectories whose memory falls within each group, providing insight into the relationship between accuracy and the conditional probability.

条件概率分析。 为研究我们的奖励设计对模型的影响,我们在 10-objective 任务推理期间,对 P[sanssmem] 的真实值进行了统计分析。 我们比较了使用普通 GRPO 训练的模型和我们方法训练的模型的结果。 具体而言,令数据集大小为 M,并将第 m 条轨迹第 t 步的记忆组件记为 τm(stmem),对应的 ground truth 答案字符串记为 τm(sans) 图5展示了 P[τm(sans)τm(stmem)] 的分组结果,其中 x 轴表示条件概率的分组值,y 轴表示该范围内记忆样本占总记忆样本的比例。 这一分布展示了模型生成记忆的条件概率分布。 折线图的 y 轴表示记忆落入每个分组的轨迹的平均准确率,从而提供对准确率和条件概率之间关系的洞见。

Figure 4 displays the aggregated results of P[τm(sans)τm(stmem)] by step. The x-axis represents the step t of the memory, and the y-axis of the line graph shows the average conditional probability of the memory at step t across all M trajectories. This provides insight into how conditional probability evolves with the step. The histogram shows the proportion of memories at each step relative to the total number of trajectories.

图4按步骤展示了 P[τm(sans)τm(stmem)] 的聚合结果。 x 轴表示记忆所在的步骤 t,折线图的 y 轴表示所有 M 条轨迹中第 t 步记忆的平均条件概率。 这为理解条件概率如何随步骤演化提供了洞见。 直方图展示了每个步骤的记忆占总轨迹数的比例。

Ablation study. To validate the effectiveness of our design, we compared the performance of the memory-enabled model trained with vanilla GRPO and our model, keeping all other conditions identical. The only difference between the two models is the inclusion of a reward specifically for memory. The results are shown in the left panel of Figure 3. Additionally, we evaluated our method under different context retention settings: full context, retaining 1 or 3 interaction rounds. The results of these experiments are shown in the right panel of Figure 3.

消融研究。 为验证我们设计的有效性,我们在保持其他条件完全相同的情况下,比较了使用普通 GRPO 训练的带记忆模型和我们的模型的性能。 两个模型之间的唯一区别是是否包含专门针对记忆的奖励。 结果如图3左侧面板所示。 此外,我们在不同上下文保留设置下评估我们的方法:完整上下文、保留 1 轮或 3 轮交互。 这些实验结果如图3右侧面板所示。

5.5 Experimental Analyses

Our method demonstrates remarkable performance and generalization. As shown in Table 1, we achieve SOTA performance on tasks that are much more difficult than those in the training set, and our model continues to maintain leading performance when switched to a real-world web search environment that differs from the training setup. Additionally, as presented in Table 1, our model achieves SOTA performance while minimizing token consumption, thereby achieving the highest performance with the least resource usage. Furthermore, as seen in Figure 3, whether using full context, retaining 1 or 3 interaction rounds for inference, our method's performance remains consistently stable, showcasing strong generalization capabilities.

我们的方法展现出显著性能和泛化能力。表1所示,我们在比训练集难得多的任务上取得 SOTA 性能,并且当切换到与训练设置不同的真实 Web 搜索环境时,我们的模型仍然保持领先性能。 此外,如表1所示,我们的模型在最小化 token 消耗的同时取得 SOTA 性能,从而以最少资源使用实现最高性能。 进一步地,如图3所示,无论推理时使用完整上下文、保留 1 轮还是 3 轮交互,我们方法的性能都保持稳定,展现出强泛化能力。

The memory mechanism significantly reduces token consumption. As shown in Appendix Table, token consumption for tasks solved by MEM1, A-MEM, and our method, which all incorporate the memory mechanism, is noticeably lower than that of other baselines. Taking ReSearch as an example and comparing it with our method, when the task is relatively simple, such as a 2-objective task, the token consumption is only slightly higher than ours. However, as the complexity of the task increases, the gap between the two methods becomes more pronounced. By the time the task reaches 10 objectives, the number of tokens required by our method to solve a problem is approximately 1/3 of ReSearch's, with the token peak being 1/5. This is comparable to ReSearch's token consumption on a 4-objective task. Moreover, our method not only uses truncated contexts but also provides effective guidance on the memory content, resulting in even more compact contexts compared to other memory-related methods. As a result, token consumption in our method is the lowest among all baselines.

记忆机制显著降低 token 消耗。 如附录表所示,MEM1、A-MEM 和我们的方法都引入了记忆机制,它们解决任务的 token 消耗明显低于其他基线。 以 ReSearch 为例并与我们的方法比较,当任务相对简单时,例如 2-objective 任务,其 token 消耗只比我们的方法略高。 然而,随着任务复杂度增加,两种方法之间的差距变得更加明显。 当任务达到 10 个目标时,我们的方法解决一个问题所需的 token 数约为 ReSearch 的 1/3,token 峰值约为其 1/5。 这与 ReSearch 在 4-objective 任务上的 token 消耗相当。 此外,我们的方法不仅使用截断上下文,还对记忆内容提供有效指导,因此相比其他记忆相关方法产生更紧凑的上下文。 因此,我们方法的 token 消耗在所有基线中最低。

The result of grouped conditional probability analysis
图5:The result of grouped conditional probability analysis.

The ability of memory to provide strongly relevant information is crucial for task success. In the baseline methods, A-MEM generates memory based on a RAG approach. As mentioned in the introduction, the memory obtained by this method is not necessarily the most relevant for solving the task and contains a significant amount of redundancy. As a result, while A-MEM reduces token consumption compared to ReAct, its performance does not show a significant improvement. On the other hand, MEM1 generates memory by combining the model's summary with reasoning, creating a stronger link between memory generation and the task-solving process. This allows MEM1 to show a considerable improvement on long-horizon tasks. Furthermore, our method explicitly guides the model to retain the context that most strongly contributes to solving the problem, outperforming other memory-related baselines in all datasets.

记忆提供强相关信息的能力对任务成功至关重要。 在基线方法中,A-MEM 基于 RAG 方法生成记忆。 如引言所述,该方法获得的记忆不一定是解决任务最相关的信息,并且包含大量冗余。 因此,尽管 A-MEM 相比 ReAct 降低了 token 消耗,但其性能没有显著提升。 另一方面,MEM1 通过结合模型摘要和推理来生成记忆,在记忆生成与任务求解过程之间建立了更强联系。 这使 MEM1 在长程任务上表现出可观提升。 此外,我们的方法显式引导模型保留对解决问题贡献最大的上下文,在所有数据集上都优于其他记忆相关基线。

The number of context steps impacts performance on long-horizon tasks. As shown in the right panel of Figure 3, we present the results of inference using complete context, truncated in 1 step and in 3 steps. Overall, the performance differences between these methods fluctuate within an acceptable range, with the trend showing that the more context steps used, the better the performance on short-horizon tasks, but the weaker the performance on long-horizon tasks. This effect is particularly noticeable on long-horizon tasks. We believe that this aligns with the phenomenon of attention dilution caused by long contexts, which leads to performance degradation.

上下文步数会影响长程任务性能。图3右侧面板所示,我们展示了使用完整上下文、截断到 1 步和截断到 3 步进行推理的结果。 总体而言,这些方法之间的性能差异在可接受范围内波动,其趋势表明:使用的上下文步数越多,短程任务性能越好,但长程任务性能越弱。 这一影响在长程任务上尤其明显。 我们认为,这与长上下文导致的注意力稀释现象一致,而该现象会带来性能退化。

Our reward design positively contributes to improving the effective information content. Figure 3 presents a performance comparison between vanilla GRPO and our method. The results show that our reward design leads to an improvement in the model's performance. We also quantitatively analyzed the information content of the memory in the trajectories of both vanilla GRPO and our method. The bar graph in Figure 5 indicates that, compared to the baseline, our method’s probability distribution is more skewed toward higher values, which contributes to greater precision of responses, as confirmed by the line graph.

我们的奖励设计对提升有效信息含量有积极贡献。 图3展示了普通 GRPO 与我们方法之间的性能比较。 结果表明,我们的奖励设计提升了模型性能。 我们还定量分析了普通 GRPO 和我们方法的轨迹中记忆的信息含量。 图5中的柱状图表明,相比基线,我们方法的概率分布更偏向较高值,这有助于提高响应精度,并由折线图得到确认。

Additionally, the line graph in Figure 4 shows that, for the first 10 steps of 10-objective task, the mean probability of our method increases as the steps progress, whereas the baseline shows a decreasing trend. We believe this reflects the more effective organization of memory by our method compared to the baseline. After 10 steps, our method’s probability starts to decrease, which is reasonable given that the typical number of search steps in a 10-objective task is around 10. If the task is not completed by this point, it suggests that some information is difficult to find and is still being searched for. In contrast, the baseline experiences more difficulty in the first 10 steps, and as seen in the bar chart, only 20% of the search examples continue after the 10th step. We hypothesize that the few remaining examples that did not abandon exploration likely achieved relatively higher accuracy, which explains the continued increase in probability after the 10th step. Overall, both the final performance and the probability analysis validate that our reward design is effective and aligns with expectations.

此外,图4中的折线图显示,对于 10-objective 任务的前 10 步,我们方法的平均概率会随着步骤推进而上升,而基线呈下降趋势。 我们认为,这反映出相比基线,我们的方法对记忆进行了更有效组织。 10 步之后,我们方法的概率开始下降,这是合理的,因为 10-objective 任务的典型搜索步数约为 10。 如果任务在这一点仍未完成,说明某些信息难以找到,并且仍在继续搜索。 相比之下,基线在前 10 步遇到更多困难,而且如柱状图所示,只有 20% 的搜索样例在第 10 步后继续。 我们推测,少数没有放弃探索的剩余样例可能取得了相对更高的准确率,这解释了第 10 步后概率继续上升的现象。 总体而言,最终性能和概率分析都验证了我们的奖励设计是有效的,并且符合预期。

6. Conclusion

Our method optimizes memory management for agents by introducing a novel reward design that retains only relevant information, improving task performance and reducing computational costs. By integrating memory, reasoning, and tool invocation via reinforcement learning, we achieve superior performance, especially on long-horizon tasks, with efficient token consumption. Future work will focus on optimizing memory reward design and enhancing the scalability of self-memory method in more applications.

我们的方法通过引入一种只保留相关信息的新型奖励设计来优化智能体的记忆管理,从而提升任务性能并降低计算成本。 通过借助强化学习整合记忆、推理和工具调用,我们实现了更优性能,尤其是在长程任务上,同时保持高效的 token 消耗。 未来工作将聚焦于优化记忆奖励设计,并增强自记忆方法在更多应用中的可扩展性。