cs.AISep 23, 2026

Just-in-Time Memory: Learning to Curate Task-Adaptive Memory for LLM Agents

Authors: Yefan Zhou, Yang Li, Zeyu Leo Liu, Semih Yavuz, Shafiq Joty

Organizations: Salesforce AI Research

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Agentic memory systems reuse past experience to improve future performance, yet most existing designs curate memory at write time: once a task is completed, its trajectory is distilled into a fixed artifact, such as a reflection, workflow, skill, or reasoning strategy, that is later retrieved by similarity. This forces the system to decide what is worth remembering before the future query is known, irreversibly discarding information and producing a query-independent summary that must serve many possible downstream tasks. Learning such a write-time curator is also difficult because the value of a storage decision may only become apparent when a relevant query arrives, potentially many tasks later, creating a long-horizon credit-assignment problem. We instead retain raw trajectories and defer curation until read time, when the current task is known. Given the retrieved traces and the new task, a memory curator synthesizes a compact, task-adaptive payload tailored to the immediate need. Because this payload is consumed on the same task, the curator can be trained directly from immediate task success, avoiding delayed utility signals and the need to artificially group related tasks. Across ALFWorld, WebShop, and τ2τ^2-bench, our Just-in-Time Memory (JitMem) consistently outperforms no-memory agents as well as heuristic and learned write-time memory methods, improving over the strongest baseline by 16.2, 16.3, and 3.9 absolute success-rate points, respectively. Notably, even an untrained curator is already competitive with or surpasses these baselines, showing that task-adaptive read-time curation itself is a major source of the gain; training the curator further compounds the improvement.

Figures & tables

Figure 1: Overview of JitMem . (a) Inference pipeline. Given the current task xtx_{t} , the retriever fetches raw trajectories from the memory bank. The curator distills them conditioned on xtx_{t} into a task-adaptive payload, which is injected into the executor’s context. After execution, an executor-as-judge assesses correctness, and successful trajectories are stored back to the bank. (b) Training pipeline. At each training step, a task is sampled and relevant trajectories are retrieved from a fixed training bank, then the curator generates multiple candidate payloads. The frozen executor attempts the task with each payload and returns the immediate task reward, which is used to update the curator via GRPO.
MethodsCuratorALFWorldWebShop
 SRScoreSR
Executor: Qwen3-8B
No Memory—47.9_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.2}}33.3_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.7}}\phantom{0}9.8_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.5}}
ReasoningBankQwen3-8B55.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 3.1}}35.4_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.1}}11.4_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}
MemPQwen3-8B49.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.7}}35.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}12.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.5}}
SkillOS-baseQwen3-8B53.1_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.5}}38.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}13.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.8}}
Table 1: Results on ALFWorld and WebShop across three frozen executors. Gains in red are relative to the strongest baseline in each block. denotes a prompted (untrained) curator; denotes an RL-trained curator. Mean and standard deviation over 3 runs with different task orderings.
MethodsCuratorτ2\tau^{2} -bench
 AirlineRetailTelecomMacro Avg.Micro Avg.
Executor: GPT-5.4
No Memory—65.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 5.7}}79.8_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 4.2}}50.2_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.0}}65.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 3.2}}65.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.8}}
ReasoningBankQwen3-8B67.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 5.7}}79.8_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 3.1}}46.5_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 5.4}}64.4_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.3}}63.8_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 3.0}}
ReasoningBankGPT-5.465.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 5.0}}84.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.2}}61.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 8.4}}70.4_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 4.1}}71.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 3.9}}
SkillOS-baseQwen3-8B68.5_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 4.6}}79.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.7}}52.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.1}}66.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.8}}66.3_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.3}}
Table 2: Results on τ2\tau^{2} -bench with GPT-5.4 as executor, reporting SR per domain and macro/micro averages. Mean and standard deviation over 4 runs with different task orderings.
Curator trainingQwen3-8BGPT-5.4
No training60.5_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.6}}79.3_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 3.6}}
Trained w/ Qwen3-8B exec.77.4_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.9}}86.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.7}}
Trained w/ GPT-5.4 exec.—88.1_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}
Transfer gap—1.4
Table 3: Executor transfer on ALFWorld (SR). Columns are test-time executors.
MethodIn. Tok. ( ↓\downarrow )Out. Tok. ( ↓\downarrow )Steps ( ↓\downarrow )
No Memory0 9.01.4017.8
ReasoningBank19.71.2616.2
SkillOS-base22.41.3616.9
JitMem -base10.91.0013.2
JitMem0 9.80.8711.6
Table 4: Average tokens (K) and steps per task on ALFWorld with GPT-5.4 executor.
Figure 2: Ablation of (a) training-free JitMem -base and (b) RL-trained JitMem on ALFWorld (top) and WebShop (bottom) across three executors. Each ablation bar removes one design choice, and the rightmost bar is the full method.
ExecutorScoreSR
Qwen3-8B  
JitMem61.1_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}32.8_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.7}}
w/ staged bank refresh61.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}35.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.3}}
w/ test bank warm-starting60.9_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.0}}32.5_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.5}}
Gemini-2.5-Pro  
JitMem61.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.8}}50.5_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.8}}
Table 5: Staged bank refresh and test bank warm-starting on WebShop.
Figure 3: Two tasks retrieve the same past experience but receive different curated payloads. The curator foregrounds state-change guidance for one task and placement-specific guidance for the other.
Figure 4: Comparison of untrained vs. RL-trained curator payloads on the same input. The untrained curator produces a generic action sequence; the trained curator recovers the environment-specific workflow (move to the desklamp, then examine the bowl with it).
Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

SourceALFWorldWebShop
SRScoreSR
Executor: Qwen3-8B
SkillOS Reported47.9_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.2}}33.3_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.7}}9.8_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.5}}
Reproduced (non-thinking)34.5_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.3}}36.4_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.2}}8.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.1}}
Reproduced (thinking)42.1_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.0}}29.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.3}}4.4_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.7}}
Executor: Gemini-2.5-Pro
Appendix
Table 6: No-memory baseline reproduction. “Reported” from Ouyang et al. (2026a) ; “Reproduced” from our runs. In all cases, our reproduced performance is at or below the reported results, ensuring gains are measured conservatively.
HyperparameterValue
Optimization
Base policy (curator)Qwen3-8B (non-thinking)
Advantage estimatorGRPO (no std. normalization)
Learning rate1×10−61\times 10^{-6}
LR scheduleconstant with warmup
Warmup steps55 (ratio 0.050.05 of 100100 steps)
Appendix
Table 7: RL (GRPO) training hyperparameters.
ExecutorScoreSR
Qwen3-8B  
k=3k=361.1_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}32.8_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.7}}
k=5k=560.5_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.4}}32.8_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.3}}
Gemini-2.5-Pro  
k=3k=361.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.8}}50.5_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.8}}
k=5k=559.2_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.4}}48.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.7}}
Appendix
Table 8: Ablation of the number of retrieved trajectories kk on WebShop.
MethodsCuratorALFWorldWebShop
 SRScoreSR
Executor: Qwen3-8B
No Memory—47.9_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.2}}33.3_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.7}}\phantom{0}9.8_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.5}}
ReasoningBankQwen3-8B55.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 3.1}}35.4_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 1.1}}11.4_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}
MemPQwen3-8B49.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.7}}35.7_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}12.0_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.5}}
SkillOS-baseQwen3-8B53.1_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 2.5}}38.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.9}}13.6_{\,{\color[rgb]{0.5,0.5,0.5}\scriptscriptstyle 0.8}}
Appendix
Table 9: Consolidated ablation and baseline results on ALFWorld and WebShop across three executors. Each block reports baselines, JitMem -base ablations ( blue ), and JitMem ablations ( blue ). denotes prompted curator; denotes RL-trained. Mean ±\pm std over 3 runs.
Figure 5: GRPO training curves on ALFWorld (Qwen3-8B executor). Top row: training SR and executor turns. Bottom row: the same metrics on validation. SR rises and turns fall steadily over 100 steps.
Figure 6: GRPO training curves on WebShop (Qwen3-8B executor). Top row: training SR, training score, training executor turns. Bottom row: the same three metrics on validation.
More in this topic

LLM Agent Memory

Large language model agent memory lets agents retain, organize, retrieve, and update information across interactions and long-running tasks.

See all 368 papers

Explore similar work

CardsList
Jul 15, 2026

Memory as a Controlled Process: Learned Adaptive Memory Management for LLM Agents

Large Language Model (LLM) agents increasingly rely on external memory systems to accumulate experience across tasks. Yet nearly all existing approaches, from graph-structured memories to reflective insight stores, access memory through fixed, hand-designed heuristics. We argue that this static view of memory is a core bottleneck for agentic learning because optimal memory behavior is fundamentally context-dependent. The early stages of the tasks, benefit from minimal retrieval because memory is sparse; recurring goal types benefit from plan reuse rather than generic nearest-neighbor lookup; stuck agents benefit from re-retrieval with alternative queries; and across long task streams, the memory store itself must be consolidated and pruned to remain useful. We present Memory as a Controlled Process (MemCon), a framework that models memory operations as a Markov Decision Process and learns an online policy that adaptively decides when, what, and how much to retrieve, when to inject a distilled plan, and when to consolidate or forget. MemCon is backend-agnostic: it wraps any existing memory implementation, learns from task-by-task binary feedback with no pretraining and no additional LLM calls, and uses a lightweight tabular contextual bandit with UCB exploration that converges within tens of tasks. Across 6 benchmarks, 3 agent frameworks, and 3 LLM backbones, MemCon consistently outperforms multiple memory baselines by up to 15.2 points in task success while reducing token consumption by 5--20%.

Eric Hanchen Jiang, Zhi Zhang, Yuchen Wu +11

University of California Los Angeles · University of Washington · Northwestern University

Persistent Memory for Language ModelsMemory Management for LLM AgentsAgent Memory ManagementMemory-Augmented Language Model Agents
AskPDF
Jun 5, 2026

Rosetta Memory: Adaptive Memory for Cross-LLM Agents

Memory is the key component for transforming a stateless LLM into a persistent, evolving agent through experience accumulation, long-horizon planning, and continual self-improvement. Existing memory systems typically take the LLM as the center and design memory operations tailored to a specific backbone. In practice, however, users frequently switch between LLMs, for example using Claude for coding and GPT for writing across tasks, or routing different steps to different backbones within a single task for cost-effective trade-offs. As a result, memory written by one model often needs to be consumed by another. Making upstream memory effectively adapt to and activate downstream LLMs remains a critical yet underexplored problem. To bridge this gap, we shift the perspective from LLM-centric memory design to memory-centric LLM adaptation. Specifically, we approach the above upstream-downstream memory adaptation problem from both the write and read sides, and design two profile-conditioned operators that are jointly trained to optimize how memory is stored and presented for better task completion. To ensure the learned operators generalize across a broad set of LLMs, we propose a minimum-gain sampling curriculum that prioritizes the least-served LLMs during training. To better measure the operators' actual contribution rather than the LLM's own capability, we design a performance-gap reward that compares against a naive memory baseline. Experiments on HotpotQA, 2WikiMultihopQA, and MuSiQue demonstrate that our model consistently outperforms baselines and remains robust under unseen-model replacement.

Hao Yang, Shiqi Shen, Haoxuan Li +3

Gaoling School of Artificial Intelligence, Renmin University of China · Weixin, Tencent · Institute for Artificial Intelligence, Peking University

LLM Agent MemoryPersistent Memory for Language ModelsAgent Memory ManagementMemory-Augmented Language Model AgentsMemory Management for LLM Agents
AskPDF
Oct 23, 2026

FlexA-Mem: Budget-Adaptive Agentic Memory for LLM Agents

Large language model (LLM) agents work in scenarios with very different memory needs. A question about past conversations needs the exact facts it names, while a multi-step task needs a trace of the steps already taken together with details that the next step cannot name. Existing memories are each designed and tested for one kind of scenario. Retrieval memories such as A-Mem miss steps that the current query does not express, summaries drop exact strings, and pairing the two fixes the split of the budget in advance. We propose FlexA-Mem, a general memory that writes each unit of history once into an immutable store and renders a fresh memory for every call. The rendered memory is a compact timeline of records plus detail that a fixed rule fits into the remaining room without an LLM call. The detail is either relevant units as raw text or excerpts or the verbatim pieces that best match the query. With one store, FlexA-Mem serves both kinds of scenario. On AppWorld challenge tasks at 4k, it outperforms the second-best method by 11.5%, and on 147 unused tasks with predictions fixed in advance it outperforms A-Mem with a summary by 12.2%. On LongMemEval and LoCoMo, it matches a dedicated keyword retriever that writes no memory, while memories built on text written by an LLM fall far behind at tight budgets. Ablations show that the timeline carries AppWorld completion, whereas recovery mainly saves steps there and decides accuracy on exact facts. Code is available at https://github.com/WujiangXu/Flex-Amem

Wujiang Xu, Jiaojiao Han, Xi Zhu +5

Rutgers University · Independent Researcher · Stanford University +2

Memory Management for LLM AgentsLong-Term Conversational MemoryLong-Horizon LLM Agents
AskPDF