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

AI summary

Overview

Research area: LLM agent memory systems — specifically, when in an agent's lifecycle past experience should be distilled and how to train that distillation process with reinforcement learning.

Technical level: Advanced. The paper assumes familiarity with LLM agents, retrieval-augmented context construction, reinforcement learning (specifically GRPO), and benchmark evaluation of multi-step agents.

Scope: The paper proposes JitMem, a memory system that stores raw past trajectories losslessly and uses an RL-trained curator to synthesize a task-specific memory payload at read time, evaluated on ALFWorld (140 test tasks), WebShop (500 test instances), and τ²-bench.

What This Paper Is About

Most agent memory systems decide what to remember at write time: after a task ends, they compress its trajectory into a fixed artifact — a reflection, workflow, skill, or reasoning strategy — that is later retrieved by similarity. This forces the system to guess what will matter before the future task is even known, and it makes training hard because the reward for a storage decision only arrives when some later task happens to retrieve that artifact. The authors instead keep raw trajectories in the bank and defer curation to read time, when the current task is known, so a memory curator can produce a payload tailored to the immediate problem and be trained directly on that same task's success.

Key Contributions

1. Read-time curation enables task-adaptive memory. By deferring curation until the query is known, the same stored trajectory can yield different payloads for different downstream tasks — a property the authors argue write-time curators structurally cannot provide, since their artifact is fixed and query-independent.

2. Read-time curation simplifies credit assignment. Because the curated payload is consumed on the very task it was produced for, the reward is immediate, collapsing credit assignment to a single step and removing the task-grouping scaffolding that learned write-time curators such as SkillOS require.

3. JitMem, a read-time memory curator. The system stores raw trajectories losslessly, retrieves the top-k by BM25 over task descriptions, and synthesizes task-conditioned payloads via a curator trained with GRPO over a persistent streaming memory bank, with the executor frozen throughout.

4. Empirical validation and analysis. Across ALFWorld, WebShop, and τ²-bench, JitMem outperforms no-memory agents, heuristic write-time methods, and an RL-trained write-time curator. Ablations isolate task-conditioned curation, quality-filtered storage, and raw-trajectory retention as independent contributors.

Main Findings

• Main results over the strongest baseline. JitMem improves over the strongest baseline by 16.2, 16.3, and 3.9 absolute success-rate points on ALFWorld, WebShop, and τ²-bench respectively. Against the RL-trained write-time curator SkillOS with the same Qwen3-8B base, JitMem reaches 77.4 vs. 61.2 on ALFWorld and 32.8 vs. 16.5 SR on WebShop with a Qwen3-8B executor. With Gemini-2.5-Pro as executor, JitMem reaches 86.2 vs. 80.2 on ALFWorld and 50.5 vs. 41.3 on WebShop.

• Untrained read-time curation is already competitive. On WebShop, untrained JitMem-gemini reaches 61.0 SR versus 41.0 for SkillOS when both use Gemini-2.5-Pro as curator and executor. With GPT-5.4 as executor on ALFWorld, JitMem-base with a Qwen3-8B curator (79.3) outperforms ReasoningBank (77.9) and SkillOS-gpt (70.0), both of which use GPT-5.4 — suggesting the gain comes from task-adaptive read-time curation rather than curator model strength. The paper reports that read-time variants generally outperform write-time counterparts with the same curator model, though on WebShop with a Qwen3-8B executor JitMem-base (32.5 score, 11.7 SR) does not beat the write-time baselines.

• The learned curator transfers across executors. Trained once with Qwen3-8B as executor, the curator transfers to GPT-5.4 within 1.4 SR points of a curator trained directly with GPT-5.4 (86.7 vs. 88.1 on ALFWorld). Across the Gemini-2.5-Pro and GPT-5.4 executors, the transferred curator improves over JitMem-base by +6.2/+7.4 on ALFWorld and +6.1 on WebShop, and outperforms RL-trained SkillOS on Gemini-2.5-Pro.

• Payloads are more compact and reduce execution steps. Relative to write-time methods, JitMem's payload reduces input tokens by 50.3%–56.3% and executor steps by 28.4%–31.4%. On ALFWorld with GPT-5.4, JitMem-base adds only 1.9K input tokens over no memory versus 10.7K for ReasoningBank and 13.4K for SkillOS-base, while cutting executor steps by 18.5%–21.9%. RL training reduces input tokens by 10.1%, output tokens by 13.0%, and steps by 12.1% over JitMem-base.

• Gains concentrate where procedural guidance matters. On τ²-bench, per-domain gains are largest on Telecom (+11.0). On Airline and Retail, none of the memory methods improves over the no-memory agent beyond variance, and JitMem variants remain on par with baselines. The authors suggest read-time curation is most valuable when tasks require synthesizing procedural guidance rather than simple fact retrieval.

• Each design choice contributes independently. Removing the current task from the curator's input lowers JitMem-base by up to 3.1 on ALFWorld and 4.6 on WebShop, and lowers the RL-trained JitMem by up to 11.4 and 10.4 respectively — indicating RL learns to exploit the task signal rather than just compress. Replacing quality filtering with storing all trajectories plus correctness labels drops JitMem-base by 1.5–2.9 on ALFWorld and 2.3–3.4 on WebShop. Replacing raw traces with ReasoningBank-style write-time distillations drops JitMem-base by 1.7–2.9 on ALFWorld and 6.8–8.2 on WebShop. Forcing the retriever to return nothing drops the trained JitMem by up to 14.8 on ALFWorld and 15.2 on WebShop, confirming gains are grounded in distilling retrieved experience.

• Bank management choices matter little. A staged bank refresh (discarding the original training bank after 100 GRPO steps, rebuilding it with the trained curator, then training 50 more steps) improves SR by 2.8 for Qwen3-8B and 0.9 for GPT-5.4, with no change for Gemini-2.5-Pro — modest relative to the added training cost. Warm-starting the test bank with 100 training trajectories changes SR by at most 1.3 and stays within standard deviation.

Methodology in Plain English

The setting is a stream of tasks arriving one at a time. The system has four parts: a memory bank holding complete past trajectories, a retriever, a memory curator, and a frozen executor. Only the curator is trainable.

For each new task, the retriever uses BM25 over task descriptions (not trajectory contents) to fetch the top-k raw trajectories. The curator reads those trajectories alongside the current task and writes a short natural-language briefing — the memory payload — that highlights relevant past experience and specific guidance. That payload is prepended to the executor's prompt; the executor never sees the raw trajectories. After execution, the executor model also acts as a judge: only trajectories it deems successful are appended to the bank. The payload itself is ephemeral and is not stored.

Training uses GRPO. The curator generates a group of candidate payloads for a sampled task, the frozen executor attempts the task with each one, and each payload receives the benchmark's native reward (binary success on ALFWorld and τ²-bench, a continuous score on WebShop). Advantages are computed as each reward minus the group mean — the authors omit standard-deviation normalization — and the curator is updated from those advantages without a value network. The curator is initialized from Qwen3-8B with thinking mode disabled, trained for 100 steps at learning rate 1e-6 with batch size 32 and group size 8, using Qwen3-8B as the training executor for efficiency.

To keep learning stable, the training bank is built once by running the base executor on the training set and keeping successful trajectories with ground-truth labels, then held fixed. At test time the bank starts empty and grows as tasks are solved, so the paper reports a natural cold-start effect. Evaluation uses a batched streaming protocol where tasks in a batch share a bank state that is updated after each batch, and results are averaged over multiple runs with different random task orderings (3 runs for the ALFWorld/WebShop table, 4 for τ²-bench).

Why This Matters

Impact on research. The paper reframes a design question that most agent-memory work treats as settled: when should memory be shaped? It argues that deferring curation to read time removes two structural costs (irreversible information loss and query-independent artifacts) and, more importantly, converts a long-horizon credit-assignment problem into a single-step one. That reframing matters for anyone training memory policies, because it removes the need for task grouping or delayed-return machinery that prior learned write-time curators rely on. The untrained-curator results also suggest evaluation practice should separate the contribution of when memory is curated from the contribution of how strong the curator model is.

Real-world applications:

• Household and embodied assistants. ALFWorld-style tasks involve state transitions such as heating, cooling, or placing objects, where a past trajectory can teach different lessons — a state-change pattern for one task and an object-placement strategy for another.
• Web shopping and transactional agents. WebShop tests product purchase across 500 test instances, where compact guidance synthesized from successful purchase trajectories can reduce the number of browsing and comparison steps.
• Customer-service tool-use agents. τ²-bench covers airline, retail, and telecom domains requiring multi-turn dialogue, tool calls, and policy enforcement, with the largest gains reported on the most policy-verification-heavy domain (Telecom).
• Multi-model deployments. Because curation is decoupled from execution, one trained curator can be reused across different executor models, so an organization upgrading its executor does not have to retrain memory from scratch.

**Industry relevance

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