Shangtong Gui

Author directory

2026

Long-context LLMs and Retrieval-Augmented Generation (RAG) systems expand accessible context, but still defer state tracking and evidence consolidation to query time. This passive paradigm is brittle in streaming settings where facts are updated and answers depend on latent states rather than explicit spans. We introduce Unified Memory Agent (UMA), an end-to-end reinforcement-learning framework that trains a single policy to maintain external memory and answer questions. UMA updates a structured memory bank via CRUD operations and combines it with raw-context retrieval for QA. To jointly optimize memory maintenance and question answering, we propose Task-Stratified GRPO, which normalizes heterogeneous trajectories separately for more effective credit assignment. We also introduce Ledger-QA, a controlled benchmark for long-horizon state tracking over accumulated updates. Experiments on Ledger-QA and 13 TTL/AR benchmarks show that UMA improves long-horizon tracking, achieves the best average TTL/AR score among compared methods, and remains competitive on retrieval-centric benchmarks. Our results highlight learned proactive memory management as a complement to longer context windows and query-time retrieval.