Uni mem learns to manage long agent memory with new rl technique
Uni-Mem: Learning a Long-Horizon Agentic Memory Manager with Unified Online Policy Optimization
Computation and Language
Summary
Managing long-term memory for intelligent agents can be tricky because it involves deciding what to remember and what to look up later. The authors study a new way to teach an agent both how to build and use its memory using a single learning method. To make this work better, they create new techniques that separate learning steps and assign feedback more precisely to the parts responsible. Their approach improves memory handling in tests compared to previous methods that only partially learned these tasks.
What this means in practice
- •For autonomous robot developers: Improve robot memory by jointly optimizing how they store and recall information during long missions using a unified reinforcement learning policy.
- •For virtual assistant engineers: Enable virtual assistants to better remember and retrieve relevant context over long conversations by training a single memory manager.
Authors
Zhiqian Xia, Jielong Tang, Ling-An Zeng, Shixiong Kai, Rong Yao, Hao Kong, Xiao Dong, Wei Wang, Mingxuan Yuan, Siyu Xia, Haifeng Xia
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.