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 373 papers

Explore similar work

CardsList

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

    Jul 15, 2026Eric Hanchen Jiang, Zhi Zhang, Yuchen Wu +11Persistent Memory for Language ModelsMemory Management for LLM Agents

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

  2. Rosetta Memory: Adaptive Memory for Cross-LLM Agents

    Jun 5, 2026Hao Yang, Shiqi Shen, Haoxuan Li +3LLM Agent MemoryPersistent Memory for Language Models

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

  3. TrawMem: Reframing Agent Memory as Task-Oriented Role-Adaptive Workspaces

    Oct 27, 2026Mingyu Zhang, Zixu Li, Zhiheng Fu +3LLM Agent MemoryMemory Management for LLM Agents

    Shandong University · National University of Singapore