PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2609.27334v1cs.AI

TL;DR

Write-time memory fixes an artifact before the future task is known, risking irreversible information loss and delayed learning signals. JITMEM retains raw trajectories and curates task-conditioned payloads at read time, training the curator from immediate task rewards. Across three benchmarks, it outperforms all baselines by reported absolute success-rate margins.

  • Problem

    Write-time memory fixes a query-independent artifact before the future task is known, potentially discarding information that a later task could need.

    Source: §1 · p. 2

  • Method

    JITMEM retains raw trajectories and synthesizes a compact payload conditioned on retrieved traces and the current task, using same-task success to train the curator.

    Source: §1 · p. 2

  • Results

    JITMEM outperforms all baselines by 16.2, 16.3, and 3.9 absolute success-rate (SR) points on ALFWorld, WebShop, and τ 2-bench, respectively.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Ablations show that RL training learns to distill retrieved experience rather than acquiring standalone task-solving knowledge.

    Source: §5 · p. 10

    The BM25 retriever may become a bottleneck as the memory bank grows large and diverse, and the curator adds an extra LLM call per task.

    Source: §5 · p. 10

Abstract

from arXiv · show

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$-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.

1 Introduction

JITMEM defers memory curation until the downstream task is known, preserving raw experience so the curator can tailor what it extracts. This also makes curator learning depend on the current task’s immediate reward, and JITMEM outperforms the reported baselines across three benchmarks.

Motivation

  • The paper frames memory’s purpose as improving future task performance and asks when in the agent lifecycle memory should be shaped.Existing systems commonly distill completed trajectories into fixed artifacts before a future query is known.Source: §1 · p. 1
  • Write-time curation irreversibly discards details before the future task is known, while a fixed artifact must serve different downstream queries.The same trajectory may support distinct lessons, and which lesson matters depends on the task.Source: §1 · pp. 1–2

JITMEM: a read-time memory curator

  • JITMEM stores raw trajectories and, at read time, conditions a compact payload on both retrieved traces and the current task.The same stored trajectory can yield different payloads for different queries.Source: §1 · p. 2

Read-time curation simplifies credit assignment

  • Because the payload is consumed by the current task, its success reward immediately trains the curator without grouping related tasks to create a learning signal.This contrasts with learned write-time curation, where a storage decision may only be evaluated by a later matching query.Source: §1 · p. 2

Empirical validation and analysis

  • JITMEM outperforms all baselines by 16.2, 16.3, and 3.9 absolute success-rate (SR) points on ALFWorld, WebShop, and τ 2-bench, respectively.Even without training, read-time curation is competitive with or better than write-time methods; RL training further compounds the gains.Source: §1 · p. 2

2 Related Work

Prior work includes heuristic write-time artifacts, learned write-time memory policies, in-session context management, and test-time context construction. JITMEM differs by curating persistent episodic memory at read time, with the current task available to shape the payload.

Heuristic write-time memory

  • Heuristic write-time systems store query-independent artifacts before future tasks arrive, while ReasoningBank distills reasoning strategies and retrieves them by cosine similarity.Systems distill reflections, insights, workflows, skills, or other forms of memory; prediction-error signals can adapt what is stored but do not change the write-time trigger.Source: §2 · p. 3

Learned write-time memory

  • Learned write-time methods train memory operations, but SkillOS needs task grouping because rewards for storage decisions arrive only when a future query matches.JITMEM instead moves curation to read time, making the reward immediate and removing the need for task grouping.Source: §2 · p. 3

Learned in-session working memory

  • In-session memory methods optimize context during a single task, whereas JITMEM curates persistent episodic memory across tasks.The related approaches include context compression, rewriting working memory, injecting reminders, and step-level selection.Source: §2 · p. 3

Read-time and test-time context processing

  • Other test-time context methods either leave retrieved traces or lessons query-independent, use a fixed pool, or entangle curation and execution.Unlike MemHarness, JITMEM separates curator and executor, supporting transfer across executors while using a persistent streaming bank.Source: §2 · p. 3

3 Method

JITMEM retrieves raw trajectories, conditions a compact payload on the current task, and passes it to a frozen executor; only the resulting trajectory may be stored. The curator is trained from immediate task rewards, while training and evaluation use different memory-bank conditions.

Inference pipeline

  • At each task, JITMEM retrieves relevant traces, curates a task-adaptive payload, executes with a frozen agent, then considers the resulting trajectory for storage.The payload is ephemeral and is not stored; only the resulting trajectory is considered for insertion.Source: §3 · p. 4

Memory Bank

  • The bank stores complete, unabstracted trajectories and uses an executor-as-judge quality gate to retain trajectories judged successful.Preserving raw traces allows the curator to extract different information from the same experience for different tasks.Source: §3 · p. 4

Retrieval and Memory Curator

  • BM25 retrieves the top-k trajectories by task-description relevance, and the curator turns the current task plus retrieved traces into a compact guidance payload.The same retrieved trajectory can produce a different distillation for each task that retrieves it.Source: §3 · p. 4

Agent Executor

  • The frozen executor acts on the current task and curated payload, allowing the trained curator to serve multiple executors without retraining.The executor also serves as the LLM-as-judge for memory updates.Source: §3 · p. 4

Curator Training

  • The curator’s reward is the native task-success metric for the same task, making its reward immediate and avoiding delayed-return or task-grouping machinery.GRPO trains the curator using rewards from candidate payloads attempted by the frozen executor.Source: §3 · pp. 4–5

Training Bank and Evaluation Procedure

  • Training uses a fixed bank of successful base-executor trajectories, while deployment grows the bank online, creating a train/test distribution shift.Evaluation starts with an empty test bank by default; tasks share bank state within each batch, and results are averaged across different task orderings.Source: §3 · p. 5

4 Experiments

Experiments compare JITMEM with no-memory and write-time memory baselines across ALFWorld, WebShop, and τ 2-bench using three frozen executors. Results cover benchmark performance, executor transfer, efficiency, design ablations, and qualitative payloads.

4.1 Main Results

  • Across three frozen executors, reported results include 61.2 (+16.2) on ALFWorld and 32.8 vs. 16.5 (+16.3) SR on WebShop.Source: Table 1 · p. 6, §4.1 · p. 7
  • On ALFWorld, JITMEM-base as both executor and curator reaches 60.5 SR, versus 55.7 for ReasoningBank and 53.1 for SkillOS-base, supporting read-time curation gains.The pattern holds with Gemini-2.5-Pro on WebShop (61.0 vs. 40.2/41.0) and GPT-5.4 on τ 2-bench (75.6 vs. 71.7/66.4).Source: §4.1 · p. 6
  • A Qwen3-8B-trained curator transfers to stronger executors without retraining, coming within 1.4 SR points of a curator trained directly with GPT-5.4.Across Gemini-2.5-Pro and GPT-5.4 executors, the transferred curator improves over JITMEM-base and outperforms RL-trained SkillOS on Gemini-2.5-Pro.Source: Table 3 · p. 7, §4.1 · p. 7
  • JITMEM-base adds 1.9K input tokens over no memory, versus 10.7K for ReasoningBank and 13.4K for SkillOS-base, while reducing executor steps by 18.5%–21.9%.RL training further reduces JITMEM input tokens by 10.1%, output tokens by 13.0%, and steps by 12.1% versus JITMEM-base.Source: Table 4 · p. 7, §4.1 · p. 7

4.2 Ablation Study

  • Removing task conditioning lowers JITMEM-base by up to 3.1 on ALFWorld and 4.6 on WebShop; after RL training, the drops reach 11.4 and 10.4, respectively.The results indicate that RL learns to exploit the task signal rather than merely compress trajectories more effectively.Source: §4.2 · p. 8
  • Storing all trajectories with correctness labels lowers JITMEM-base by 1.5–2.9 on ALFWorld and 2.3–3.4 on WebShop; replacing raw traces with distilled items lowers it by 1.7–2.9 and 6.8–8.2, respectively.Correctness annotations do not fully suppress failed-trajectory noise, while query-independent write-time distillation can discard information needed at read time.Source: §4.2 · p. 8
  • On WebShop, staged bank refresh yields a 0.9 SR gain for GPT-5.4 and no change for Gemini-2.5-Pro; the cited passage does not state Qwen3-8B’s gain.Test-bank warm-starting changed SR by at most 1.3 across all three executors, within standard deviation; the curator handles early sparse retrieval without it.Source: Table 5 · p. 8, §4.2 · p. 9

4.3 Qualitative Analysis

  • For the same retrieved trace, the curator gives a hot-potato task heat/cool guidance but gives a newspaper-placement task target-location guidance; RL payloads also add environment-specific workflow steps.The trained curator’s desklamp-and-bowl workflow is absent from the untrained generic action sequence.Source: Figure 3 · p. 9, §4.3 · pp. 9–10, Figure 4 · p. 10

5 Conclusion

JITMEM stores raw trajectories and defers task-adaptive curation until read time, making curator learning an immediate single-task objective. Across three benchmarks, untrained read-time curation is competitive with or outperforms strong write-time baselines, and RL training further improves effectiveness, efficiency, and transfer; the design has practical limitations.

5 Conclusion

  • JITMEM preserves raw trajectories at write time and synthesizes task-adaptive payloads at read time, when the current task is known.This avoids committing prematurely to a single abstraction and makes curator learning an immediate single-step objective.Source: §5 · p. 10
  • Across ALFWorld, WebShop, and τ 2-bench, untrained read-time curation is competitive with or outperforms strong write-time baselines, while RL training further improves effectiveness, efficiency, and transfer.Ablations indicate that RL learns to distill retrieved experience rather than acquiring standalone task-solving knowledge.Source: §5 · p. 10
  • BM25 may bottleneck as memory banks grow large and diverse, each task requires an extra LLM call, and payload formats are fixed and hand-designed per benchmark.The authors propose stronger retrievers, joint payload-format optimization, and turn- or step-level adaptation as future work.Source: §5 · p. 10

A Experimental Details and Hyperparameters

The appendix specifies curator, executor, and judge prompts, baseline reproduction choices, experimental settings, and computing infrastructure. It also documents prompt-based memory distillation used in an ablation study.

A Experimental Details and Hyperparameters

  • The prompts guide curators to turn retrieved trajectories into concise, task-specific actionable advice, while executors use that advice to choose admissible actions.WebShop guidance covers product search and option selection; customer-service guidance prioritizes policy-compliant resolution steps, and the judge credits only completed outcomes confirmed by tool results.Source: A Experimental Details and… · pp. 14–16
  • The judge defines success as fully completing the requested action, with tool results confirming it and any required process evidenced in the transcript.Partial completion, escalation, or ambiguity is failure.Source: A Experimental Details and… · p. 17
  • The ablation distillation prompts request at most three non-overlapping, generalizable memory items rather than specific object names, locations, product IDs, prices, or properties.The prompts are included alongside the curator, executor, and judge prompts in Appendix A.Source: A Experimental Details and… · p. 17
  • WebShop no-memory reproduction reached 8.6 SR versus 9.8 reported; search guidance was calibrated so reproduced baselines stayed at or below reported results.The reproduction study also selected thinking mode for Qwen3-8B on ALFWorld and non-thinking mode on WebShop.Source: A Experimental Details and… · p. 18
  • Main results use three retrieved trajectories, evaluation batches of 10 on ALFWorld and WebShop and 5 on τ 2-bench, and executor limits of 30 or 200 turns, respectively.The history window is three; Qwen3-8B uses thinking mode on ALFWorld and non-thinking mode on WebShop.Source: A Experimental Details and… · p. 18
  • Experiments ran on one server with eight NVIDIA H200 GPUs; a Qwen3-8B curator training run took about 21 hours on ALFWorld and 27 hours on WebShop.The server had two Intel Xeon Platinum 8488C CPUs and 2 TB of system RAM.Source: A Experimental Details and… · p. 18

B Additional Results

Additional results examine retrieval-count sensitivity, design ablations, example payloads, and curator training curves. Performance is stable across the tested retrieval counts, while ablations show contributions from task adaptivity, successful-trajectory filtering, raw trajectory retention, and retrieved experience.

B Additional Results

  • Across WebShop executors, changing retrieved trajectories from k=3 to k=5 changes SR by less than 2 points in all cases, with differences within standard deviation for Qwen3-8B and GPT-5.4.The paper uses k=3 for main results for efficiency.Source: B Additional Results · p. 19, Table 8 · p. 20
  • Every JITMEM-base ablation degrades performance; removing raw trajectories causes up to −8.2 SR on WebShop, while removing retrieved trajectories from trained JITMEM causes up to −14.8 SR on ALFWorld and −15.2 SR on WebShop.Removing task adaptivity also degrades performance; staged bank refresh adds at most +2.8 SR, while test-bank warm-starting has negligible effect.Source: B Additional Results · p. 20
  • Example payloads turn retrieved experiences into task-specific guidance: an ALFWorld procedure, a WebShop search-and-attribute strategy, and an ordered MMS diagnostic plan.The τ 2-bench example retains the policy requirement for explicit user approval before a mutating tool call.Source: B Additional Results · p. 22
  • On ALFWorld and WebShop, validation success rises and executor turns fall over 100 GRPO steps, with stable training under a single task reward and no auxiliary content-quality reward, task grouping, or return shaping.The curves show training and validation metrics for the two benchmarks.Source: B Additional Results · p. 24, Figure 5 · p. 24, Figure 6 · p. 25
Loading 2609.27334v1…