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

TL;DR

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.

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

Figures (5)

• Figure 1: The self-memory inference process of our method, which only uses the previous step interaction for next step input with <mem> action.
• Figure 2: Overview of MemPO. At step $t$ of any trajectory $\tau_i$, the context is represented as $\{s_t^{mem}, s_t^{think}, s_t^{call}, s_t^{resp}\}$. The memory reward $R^M$ is calculated using conditional probabilities and contributes to the advantage $A^M$. The final advantage is the sum of $A^M$ and the trajectory-level advantage $A^T$. During inference, only the previous step’s content is used as context, discarding earlier information.
• Figure 3: The results of ablation study.
• Figure 4: The result of step conditional probability analysis.
• Figure 5: The result of grouped conditional probability analysis.