Source-linked AI summary
Mem-$π$: Adaptive Memory through Learning When and What to Generate
Xiaoqiang Wang, Chao Wang, Hadi Nekoei, Christopher Pal, Alexandre Lacoste, Spandana Gella, Bang Liu, Perouz Taslakian
TL;DR
LLM agents remain limited by statelessness and cannot accumulate reusable experience across interactions. Mem-π generates context-specific guidance while learning when to provide it, consistently improving over retrieval-based and prior RL-optimized memory baselines across agentic benchmarks, including a 20% average relative gain over the base agent.
Problem
LLM agents remain stateless and cannot accumulate reusable experience across interactions.
Method
Mem-π uses a parametric memory policy to decide when to generate concise, context-specific guidance and what guidance to produce.
Results
20% relative average gain over the base agent was achieved across four benchmarks, with the WebArena gain approaching 50%.
Takeaways & Limitations
Across web navigation, terminal tool use, and embodied environments, adaptive generative memory consistently improves over retrieval-based and prior RL-optimized memory baselines.
Abstract
from arXiv · showhide
We present Mem-$π$, a framework for adaptive memory in large language model (LLM) agents, where useful guidance is generated on demand rather than retrieved from external memory stores. Existing memory-augmented agents typically rely on similarity-based retrieval from episodic memory banks or skill libraries, returning static entries that often misalign with the current context. In contrast, Mem-$π$ uses a dedicated language or vision-language model with its own parameters, separate from the downstream agent, to generate context-specific guidance for complex tasks. Conditioned on the current agent context, the model jointly decides when to produce guidance and what guidance to produce. We train it with a decision-content decoupled reinforcement learning (RL) objective, enabling it to abstain when generation would not help and otherwise produce concise, useful guidance. Across diverse agentic benchmarks spanning web navigation, terminal-based tool use, and text-based embodied interaction, Mem-$π$ consistently outperforms retrieval-based and prior RL-optimized memory baselines, achieving over 30% relative improvement on web navigation tasks.
1 Introduction
Mem-π addresses the limitations of retrieval-based memory by generating concise, context-specific guidance while learning whether generation is useful. It trains a separate generative memory policy with supervised and reinforcement learning, including abstention, and improves agent performance across diverse benchmarks.
- Retrieval-based memory reuses static stored experiences that can be irrelevant, partially aligned, overly specific, or unable to adapt to the agent’s current context.
- Mem-π models memory as a separate parametric policy πmem that jointly learns when to generate and what task-adaptive guidance to produce.The policy conditions on task instructions and environment observations and injects guidance into the downstream agent’s context.
- Mem-π distills offline experiences through supervised learning, then uses reinforcement learning with downstream task outcomes to refine memory generation and support abstention.The policy can emit either [GENERATE] with guidance or [ABSTAIN], while length-aware reward shaping discourages verbose or overly specific guidance.
- 20% relative gain over the base agent on average was achieved across four benchmarks, with the relative gain on WebArena approaching 50%.The evaluation covered WebArena, WorkArena, LifelongAgentBench, and ALFWorld, and adaptive generation consistently outperformed retrieval-based memory baselines.
- The decision-content decoupled objective separates learning whether to generate from learning what guidance to produce, addressing the imbalance caused by longer guidance sequences.
3 Experiments
Experiments across web navigation and embodied-task benchmarks show that Mem-π achieves strong performance through staged training, adaptive abstention, and context-specific guidance. It also improves efficiency and transfers to stronger agents while benefiting from visual observations.
- Benchmark results: Mem-π leads every WebArena sub-domain, with largest gains of +23.8 pp on Reddit and +28.2 pp on CMS.On WorkArena, it improves the base agent from 42.0% to 50.3% on average; on ALFWorld, it reaches 91.6%, a +6.3 pp improvement.
- Benchmark results: 35.0% on WebArena: Mem-π Stage 1 matches or surpasses Memory-R1 at 33.2% and MemRL at 34.0% without RL training.The result supports offline parametric knowledge as a strong initialization strategy.
- Ablations: +8.1 pp on WebArena overall: the full model improves over Stage 1, with largest jumps on CMS (+25.4 pp), Reddit (+4.8 pp), and Maps (+3.0 pp).ALFWorld gains a more modest +1.6 pp.
- Ablations: −6.8 pp on WebArena: unified single-stage training underperforms staged optimization, while removing Stage 1 initialization causes a 5.2 pp drop.The two stages optimize imitation and task-success rewards that can produce conflicting gradients when trained jointly.
- Multimodal analysis: 2.7 pp: the multimodal variant improves overall WebArena SR over text-only, with largest gains in CMS (+3.8 pp) and Shopping (+3.3 pp).GitLab shows the smallest gain at +0.9 pp, where code-centric content provides less visual signal.
- Adaptive memory behavior: 71% versus around 13%: Mem-π abstains frequently on easiest tasks but rarely on hardest tasks, where SR improvement peaks at +9.7 pp versus +1.3 pp on easiest tasks.This pattern indicates adaptive generation based on task difficulty rather than conservative abstention.
- Transfer and efficiency: 138 memory tokens per task: Mem-π uses 31% fewer tokens than Stage 1 and 38% fewer than Memory-R1 while improving Stage 1 from 35.0% to 43.1%.It also retains an advantage with GPT-5.4-mini on ALFWorld, yielding +6.3 pp versus RAG’s +1.8 pp.
4 Related Work
Prior agent-memory research has shifted from static pipelines toward learned memory operations and structured knowledge retrieval. Mem-π differs by treating memory as a generative policy that constructs task-adaptive guidance from parametric knowledge rather than retrieving fixed entries.
- Learning-based agent memory: Agent memory has progressed from static pipelines toward learned memory operations optimized jointly with downstream task outcomes.Some approaches distill interaction trajectories into structured rules, guidelines, or strategies retrieved at inference time.
- Learning-based agent memory: Existing learning-based systems remain retrieval-centric: they improve access to stored entries, while memory content is fixed when written.This limitation motivates modeling memory as a generative policy rather than relying solely on stored entries.
- Generative memory: Generative memory encodes experience into model parameters and produces useful information on demand, unlike retrieval-based memory that stores and retrieves external entries.Related approaches include parametric modules that internalize retrieval behavior into learnable parameters and methods that compress long contexts.
5 Conclusion
Mem-π formulates agent memory as a generative policy rather than retrieval over explicit memory entries. It internalizes reusable behavioral knowledge through distillation and jointly optimizes when to generate memory and what guidance to produce with a decision-content decoupled RL objective.
- Conclusion: Mem-π replaces retrieval over explicit memory entries with a generative memory policy πmem.The framework formulates agent memory as a generative policy rather than retrieval.
- Conclusion: Experience and adaptation distillation internalize reusable behavioral knowledge into model parameters and refine it using downstream task rewards.Mem-π combines experience distillation with adaptation distillation before further reward-based refinement.
- Conclusion: A decision-content decoupled RL objective jointly optimizes when to generate memory and what guidance to generate.The objective separates routing from content optimization through structured counterfactual advantages and per-token credit assignment.
A Experimental Details · A.1 Benchmark Details
The experimental evaluation spans web navigation, enterprise workflows, terminal-based tool use, and text-based embodied household interaction. Benchmark settings test task performance, cross-goal generalization, experience transfer, and memory distillation on unseen tasks.
- A.1 Benchmark Details: WebArena comprises 812 realistic web-navigation tasks across five fully functional web domains.Its domains include Shopping, CMS, GitLab, and Reddit, with tasks covering e-commerce, content management, software development, and social forum workflows.
- A.1 Benchmark Details: WorkArena evaluates enterprise workflows on ServiceNow across dashboard navigation, forms, list filtering/sorting, and knowledge or service-base tasks.It contains 33 atomic task templates, each instantiated with random seeds under a cross-goal generalization setting.
- A.1 Benchmark Details: LifelongAgentBench contains 1,396 terminal-based task instances across Database, Operating System, and Knowledge Graph environments, with experiments focused on Database and Operating System.The benchmark evaluates lifelong learning and experience reuse in interactive terminal environments.
- A.1 Benchmark Details: The Database subset contains 500 tasks evaluating 22 SQL-related skills, with correctness verified automatically through SQL engine outputs.Skills include querying, filtering, aggregation, joins, set operations, and data manipulation.
- A.1 Benchmark Details: The Operating System subset contains 500 tasks evaluating 29 Bash-command skills, with correctness determined from the final OS state.Skills span file operations, permissions, user management, text processing, compression, process inspection, and system monitoring.
- A.1 Benchmark Details: Both terminal subsets use a 7:3 train–test split with seed 42, comprising 350 training and 150 test tasks per subset.This evaluates whether Mem-π can distill reusable terminal-interaction experience and transfer it to unseen tasks within the same tool-use domain.
- A.1 Benchmark Details: ALFWorld converts household manipulation tasks into textual observations and actions, covering six task types that require navigation and object-state manipulation under partial observability.The setting uses 3,553 training tasks as the experience pool and evaluates transfer on 134 validation-unseen tasks rather than memorized trajectories.
A.2 Implementation Details
Mem-π is implemented with a two-stage training pipeline, task-derived experience hints, structured GRPO rollouts, and benchmark-specific verification. Its generation-length penalty applies only to generated memories and remains small relative to the binary task reward.
- Implementation stack: Mem-π uses PyTorch, Hugging Face Transformers, TRL, and vLLM, with special decision tokens for abstention and generation.The [ABSTAIN] and [GENERATE] tokens are initialized from semantically related vocabulary embeddings; symmetric averaging yields approximately 50% initial abstention probability.
- Experience bank and Stage 1: Stage 1 distills five JEF-Hinter-generated procedural hints per task from πagent training trajectories while excluding test-task hints.Training uses an autoregressive supervised objective with AdamW, learning rate 2×10−5, batch size 32, maximum length 2,048 tokens, and three epochs across 8× NVIDIA H100-80GB GPUs.
- Stage 2: Stage 2 applies decision-content decoupled GRPO using four rollout branches: one forced [ABSTAIN] branch and three [GENERATE] branches.Generated memories are capped at Lmax=256 tokens, with sampling temperature 1.0 and top_p 0.95; optimization uses AdamW at learning rate 1×10−6 and batch size 8.
- Evaluation protocol: Evaluation uses official or built-in verifiers for WebArena, WorkArena, LAB, and ALFWorld, with reported results averaged over three independent seeds.LAB correctness uses SQL execution and OS-state checks, while ALFWorld uses the environment’s terminal-condition checker.
- Length regularizer: The length penalty applies only to [GENERATE] branches, saturates at −0.1 when λlen=0.1, and is an order of magnitude smaller than the binary task reward.Generation uses Lmax=256, and prompt-leakage tokens are removed before measuring substantive memory length.
A.3 Abstention Training Dynamics
Mem-π’s abstention behavior converges quickly to benchmark-dependent rates, reflecting differences in task novelty and distribution breadth. Symmetric initialization is essential for rapid cold-start exploration and reaching stable abstention levels.
- Benchmark-dependent abstention: Final abstention rates vary by benchmark: 34% on WebArena, 28% on WorkArena, 21% on ALFWorld, and 39% on LAB.The higher LAB rate reflects more novel configurations, while ALFWorld’s lower rate reflects its narrow household-task distribution.
- Training dynamics: With symmetric initialization, abstention rates converge within approximately 100 training steps to stable, task-appropriate levels.Stage 2 training begins from a symmetrically initialized Stage 1 checkpoint with approximately 50% initial abstention probability.
- Training dynamics: Without symmetric initialization, the [ABSTAIN] token starts near zero probability, recovers slowly to approximately 21%, and never reaches the symmetric variant’s converged rate.This shows that balanced initialization of the decision tokens is critical for cold-start exploration.
A.4 Additional Case Studies
The case studies partition tasks into eight outcome regions, showing that Mem-π wins through context-specific generation or abstention while sometimes matching or underperforming retrieval. Across the regions, adaptive memory helps on 25 tasks, shares success with retrieval on 40, and fails alongside or behind baselines on 13.
- Region 001: In a Magento search-term task, Mem-π succeeds by following the query’s explicit top-three requirement instead of copying retrieval’s top-two hint.The generated guidance outputs the first three search terms, whereas RAG specifies reading only the first two rows.
- Region 101: In a storage-search task, Mem-π abstains, restoring the base agent’s broader search after retrieval narrows the search to unsuitable product families.This representative Region 101 case illustrates abstention recovering from RAG noise.
- Mem-π wins: 25 tasks fall in Mem-π-win regions: generation corrects retrieval’s mismatched specification, while abstention avoids retrieval noise.Regions 001 and 101 contain 15 and 10 tasks, respectively.
- Shared successes: 40 tasks are shared successes where memory is needed or all approaches solve, indicating Mem-π preserves retrieval-aligned behavior.Regions 011 and 111 contain 12 and 28 tasks, respectively.
- Failure modes: 13 tasks expose failure modes: Mem-π regresses on four tasks, loses to RAG on six, and both memory methods hurt on three.A further 87 tasks are unsolved by any method because of environment or tool limitations, not memory quality.