AUTOMEM: Automated Learning of Memory as a Cognitive Skill

Shengguang Wu, Hao Zhu, Yuhui Zhang, Xiaohan Wang, Serena Yeung-Levy

AUTOMEM treats memory management as a trainable skill, using meta-LLM loops to optimize agent scaffolds and memory proficiency.

How can we automate the optimization of an LLM's memory management strategy to improve performance on long-horizon tasks?

LLM agents in long-horizon tasks often fail because they lack a systematic way to decide what to store, when to retrieve, and how to organize their knowledge. AUTOMEM promotes file-system operations to first-class actions, then uses a meta-LLM to iteratively revise the agent's memory structure and fine-tune a dedicated memory specialist on the agent's own successful decisions. This approach doubles or triples performance on long-horizon games, bringing 32B open-weight models to the level of frontier proprietary systems without modifying the base model's task-action weights.

Paper Primer

The framework optimizes memory along two axes: structure and proficiency. The structure loop acts as a code reviewer, analyzing full episode traces to revise prompts and file schemas, while the proficiency loop uses the agent's own history to train a memory specialist: a LoRA-adapted model that handles memory operations while the base model remains frozen for task actions.

Memory management is a higher-leverage optimization axis than model scale.

The optimized 32B agent significantly outperforms the 72B base model on all three tested environments (Crafter, MiniHack, NetHack). Progression rates improved by ~2×–4× across environments compared to the base agent.

Why treat memory as a learnable skill rather than a fixed architectural module?

Fixed memory mechanisms often fail to adapt to the specific information-management needs of long-horizon tasks; by making memory operations first-class actions, the model gains the flexibility to decide what to record and how to organize its own external records.

Why is trajectory-level review necessary for this optimization?

In long-horizon tasks, the consequences of a memory mistake—such as failing to record a map coordinate—may remain hidden for thousands of steps, making per-step reward signals insufficient for identifying the root cause of failure.

Memory management should be treated as an independently learnable skill; researchers can achieve frontier-level performance on long-horizon tasks by automating the iterative refinement of memory scaffolds and training specialized memory models.

Metamemory as a Trainable Skill

LLMs lack innate memory management; AUTOMEM trains memory as a cognitive skill.

LLMs treat external memory as a static module, yet long‑horizon tasks routinely exceed the context window. A single memory mistake can remain hidden for thousands of steps, making human review of full trajectories infeasible. This gap motivates a learned, trainable memory skill.

Metamemory is the ability to decide what information to store, when to retrieve it, and how to organize the stored knowledge—essentially a learned strategy for managing one’s own memory.

We introduce AUTOMEM, which treats memory management as a trainable skill. The model’s action space is extended with file‑system primitives (read, write, search, append, create), allowing the same forward pass that selects a task action to also select a memory action. AUTOMEM improves the skill along two axes—scaffold structure (prompts, file schemas, action vocabularies) and proficiency (the model’s parametric ability to choose good memory actions)—via two sequential outer loops: a structure‑revision loop that iterates on the agent’s memory scaffold, and a proficiency‑training loop that fine‑tunes a dedicated memory model on its own successful memory decisions.

**Figure 1.** Memory skill optimization with Qwen2.5-32B-Instruct. Starting from a base agent equipped with file-system memory (v0), AUTOMEM progressively improves performance through memory scaffold optimization (v0–v5/v4/v2), followed by memory proficiency training (+train) that yields further gains on top of the optimized scaffold.

Memory management is a cognitive skill that can be automated.

The AUTOMEM Framework

AUTOMEM automates memory‑skill learning through two outer loops that refine structure and proficiency.

LLMs inherit generic context windows but lack a systematic, trainable strategy for managing external memory, leading to redundant logs and missed information in long‑horizon tasks.

Memory scaffolding is the mutable file‑system layout (which files exist, their schemas, and how the agent reads/writes them) that defines the space of possible memory decisions.

Step 1 appends “(0,0)” to

Step 2 appends “(0,1)”.

Step 3 appends another “(0,0)”, creating a duplicate.

Scaffold revision replaces

After revision the file holds only two unique entries, reducing read‑time and storage.

Deduplication shrinks the searchable memory without losing information, illustrating how a simple schema change can dramatically improve downstream decision speed.

How does “memory scaffolding” differ from static prompt templates that list file names?

Scaffolding is mutable code‑level structure that the agent can read and write during execution; a static prompt merely describes a fixed set of files and cannot be altered by the agent itself.

Memory proficiency is the parametric ability of the model to issue high‑quality LOG and PLAN actions given a fixed scaffold.

Why not simply finetune the entire base model on the same memory data?

Finetuning the whole model would entangle memory behavior with action generation, risking degradation of the already‑good gameplay policy. The specialist isolates the memory skill, guaranteeing additive gains without trade‑offs.

The framework runs two sequential outer loops: the first rewrites the scaffold; the second trains a dedicated memory specialist, each using a meta‑LLM as a high‑level optimizer.

LOG: decide whether the just‑observed environment response should be recorded; choose file, mode (append, create, overwrite), and content.

PLAN: query the file system for entries relevant to the current decision; read whole files or tail sections, then produce the next world action.

Outer‑loop 1 scaffold revision pseudocode

Outer‑loop 2 memory‑proficiency training engine

Together the two outer loops turn memory management from a hard‑coded utility into a trainable skill that can be systematically improved on any long‑horizon task.

Empirical Evaluation

AUTOMEM boosts progression and cuts wasted actions across three long‑horizon game benchmarks.

BALROG Benchmarks are a suite of procedurally generated long‑horizon games used to evaluate how well agents manage external memory while pursuing complex tasks.

Scaffold optimization more than doubles progression on all three BALROG benchmarks.

Crafter improves from 25.0% to 47.27% (×1.89), MiniHack from 7.5% to 27.5% (×3.67), NetHack from 0.42% to 1.57% (×3.74).

**Figure 2.** Long-horizon game environments for evaluating memory skills. All three environments are stochastic worlds, making each episode unique and minimizing the influence of pretraining knowledge. Crafter is an open-world survival game with crafting, combat, and resource management. MiniHack presents focused puzzle, navigation and combat tasks within the NetHack engine. NetHack is among the most complex games: episodes span $10^4-10^5$ turns with a vast exploration space, taking human players typically years to master.

**Figure 4.** Effect of scaffold optimization on gameplay and memory behavior (v0 → v5 for Crafter, v0 → v4 for MiniHack, v0 → v2 for NetHack). Left: the unproductive game action rate (fraction of steps that are either stuck or oscillating) drops 32–65% across all three environments. Right three panels: memory operations become more efficient and better targeted—redundant writes drop sharply (-68 to -83%), the empty-search rate (memory SEARCHes returning nothing) falls (-13 to -50%), and per-step input context shrinks (-3 to -30%) as leaner memory compresses what the model must attend to. All values are v0 vs. final scaffold version; lower is better in every panel.

**Figure 6.** Qualitative behavior across the two optimization stages. Each row is an environment; its three cells show the base scaffold v0, the evolved scaffold, and the + trained specialist on a representative evaluation episode. The note under each cell reports the agent's behavior and its progression rate (prog.); for NetHack it also lists the dungeon level (Dlvl, the agent's depth) and the experience level and points (Xp). Crafter: the base agent only gathers wood; the evolved scaffold crafts stone tools, builds a furnace, and mines iron; and the trained specialist reaches that same crafting tier while also feeding itself (9 -> 55 -> 59% progression). MiniHack Corridor-R3, a task that requires navigating branching corridors to reach a goal staircase: neither the base nor the evolved scaffold reaches the staircase, whereas the trained specialist solves the task (0 -> 0 -> 100%) NetHack: the base agent dies at experience level 1 within a few hundred steps; the evolved scaffold survives thousands of steps and reaches experience level 2; and the trained specialist survives far longer and reaches experience level 4 (0 -> 1.85 -> 2.42% progression).

**Table 1.** Performance comparison of various agents on Crafter, MiniHack, and NetHack benchmarks. The table categorizes agents into "Frontier proprietary", "Open-weight", "Qwen2.5-32B-Instruct with basic context-management baselines", and "Qwen2.5-32B-Instruct with AUTOMEM (ours)".

**Table 2.** LOG-phase memory writes per SEARCH on the evolved scaffold with base model vs. + trained memory specialist (lower = more retrieval before writing). The trained specialist checks memory before writing in every environment.

AUTOMEM significantly reduces unproductive agent actions across all tested environments.

Related Work and Context

We situate AUTOMEM among memory‑augmented LLM work and recap its impact.

External memory mechanisms extend the fixed context window of LLMs, enabling agents to retrieve or store information beyond the immediate prompt.

RAG pairs a language model with a document retriever so the model can look up relevant passages at inference time instead of relying solely on its internal parameters.

Beyond RAG, systems such as MemGPT, Generative Agents, A‑MEM, MemoryBank, MemLLM, and Self‑Notes explore paging, timestamped streams, active retain/forget decisions, persistent stores, and interleaved reasoning‑memory tokens—each touching a different point on the “memory proficiency” axis.

Cognitive‑science concepts of metamemory and the Extended‑Mind thesis inspire treating memory management as a monitor‑control loop, a view instantiated in CoALA and MetaMem and operationalized in AUTOMEM.

Automated design frameworks (ADAS, AFlow, DSPy, PromptBreeder, APE) and memory‑evolution systems (MemEvolve, EvolveMem) automate architecture or prompt search; AUTOMEM differs by jointly optimizing a scaffold and a memory‑specialist through meta‑LLM‑driven outer loops that analyze entire long‑horizon trajectories.

Game‑based agents (Voyager, DEPS, NetPlay) showcase long‑horizon challenges; methods such as ReAct, Reflexion, Inner Monologue, and ExpeL focus on reasoning or planning, whereas our work isolates memory management as the primary lever for improvement.

We show that treating memory management as a trainable skill—optimizing both scaffold structure and proficiency—delivers substantial performance improvements across three game environments without altering the base model’s weights.

Implementation Details and Limitations

Implementation configurations and scaffold evolution details for reproducibility.

This appendix records the concrete implementation settings used for all experiments.

We run the BALROG harness with two minor tweaks: autopickup=True for MiniHack and a custom fallback action that searches the task’s action space for NetHack and MiniHack (instead of the default esc) while keeping Noop for Crafter.

The fallback ensures the LLM never stalls on an unparseable command, preserving episode continuity.

All environments share a 64×64 world and a 9×9 agent view; Crafter uses dense rewards, `unique_items`=True, and a 2000‑step episode limit.

MiniHack runs eight tasks (Boxoban‑Hard, Boxoban‑Medium, MazeWalk‑9×9, MazeWalk‑15×15, Corridor‑R3, CorridorBattle‑Dark, Quest‑Easy, Quest‑Medium) with 100‑step episodes and a –0.01 step penalty.

NetHack episodes last up to 100 000 steps and abort after 150 steps without progress; in all three environments `skip_more`=True suppresses the “–More–” prompt.

Evaluation uses a fixed seed list [42…51]; we run 10 Crafter episodes, 5 per MiniHack task (40 total), and 5 NetHack episodes per seed, guaranteeing direct comparability across scaffold and memory‑trained agents.

Outer‑loop #1 optimizes the scaffold via Claude Opus 4.6 (‑effort max); a revision is accepted only if its average progression on the fixed eval seeds strictly exceeds the prior iteration.

If a revision fails, we allow a single retry with the failed log supplied; a second failure restarts the meta‑loop from a clean session.

Outer‑loop #2 proceeds in three stages: (a) collect training data, (b) run a Claude Opus 4.7‑driven data‑engine, and (c) LoRA finetune and deploy a two‑model inference architecture.

During stage (a) the base model plays 100 Crafter episodes, 50 per MiniHack task (400 total), and 50 NetHack episodes, using random seeds that are disjoint from the evaluation set.

Stage (b) iterates: (i) set dataset composition, (ii) select examples from the trace pool, (iii) pick a LoRA configuration; advisory per‑environment priors are provided but the engine may deviate.

After each engine pass we clean code‑block wrappers, discard examples lacking memory operations, and trim actions, yielding final training sets of 1 597 Crafter, 444 MiniHack, and 800 NetHack examples.

Stage (c) trains LoRA adapters with `cutoff_len`=16384, bf16, AdamW, cosine LR schedule (warmup 0.05) on two GPUs via DeepSpeed ZeRO‑3; hyperparameters differ per environment as listed:

Crafter: rank 256, alpha 512, batch 32, attention‑only modules, LR 5e‑5, dropout 0, epochs 4. MiniHack: rank 128, alpha 256, batch 16, LR 5e‑5, dropout 0, epochs 3. NetHack: rank 256, alpha 512, batch 32, LR 5e‑5, dropout 0, epochs 1.

Section B details how Outer‑loop #1 iteratively reshapes the memory scaffold for each environment.

Crafter scaffold changes: v1 pre‑loads `game_knowledge`.txt and introduces a goals.txt template with low‑health/food warnings; v2 logs inventory deltas and adds a 22‑item achievement checklist; v3 escalates craft‑failure warnings and lists buildable items when a table or furnace appears; v4 predicts the next resource to gather and flags nearby threats; v5 blocks impossible crafts, rejects repeated Noop actions, and warns when inventory stagnates.

MiniHack scaffold changes: v1 tracks recent actions and warns on back‑and‑forth movement, adding per‑task rule sheets; v2 records visited cells, reports unexplored directions, and prepends single‑word direction replies; v3 adds a staircase‑directed planning prompt and shortens direction queries; v4 fixes map parsing, correctly handles impassable symbols, and routes around hazards.

NetHack scaffold changes: v1 introduces <|`UPSERT_MAP`|> to overwrite duplicate tiles, trims the action log, and seeds strategy.txt with safety rules; v2 issues high‑priority descent directives on staircases, maintains `current_status`.txt and inventory.txt each step, and warns on back‑and‑forth movement.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers