AutoMem: Automated Learning of Memory as a Cognitive Skill

Technical Breakdown§

AutoMem addresses the challenge of memory management in LLM agents operating over long horizons. The core idea is to treat memory operations (read, write, delete, etc.) as first-class actions, and then automate the optimization of both (1) the memory structure (how memory files are organized, schemas, prompts) and (2) the memory proficiency of the model.

Two-Loop Architecture§

• Outer Loop (Structure Optimization): A strong LLM (e.g., GPT-4) reviews complete trajectories of a base agent (e.g., Llama 32B) and proposes revisions to the memory prompts, file schemas, and action definitions. This loop runs iteratively, improving the scaffolding that governs memory interactions.

• Inner Loop (Proficiency Tuning): From many episodes, episodes where the agent made particularly good memory decisions (e.g., correctly retrieving relevant info, organizing files logically) are identified. These are used as supervised fine-tuning data to update the base model's weights directly, sharpening its memory usage.

Memory Actions§

The agent's action space is extended with file-system operations:

memory_actions = ['read', 'write', 'append', 'delete', 'create_file', 'list_dir', 'search']

These actions are executed on a structured memory directory that the agent can reorganize. For example, an agent can create a file goal.txt and later read it to recall its current goal.

Formalization§

The overall objective is to maximize task reward $R$ by optimizing two components:

• Memory structure parameters $\theta_s$ (prompts, schemas)
• Model weights $\theta_m$

$$ \max_{\theta_s, \theta_m} \mathbb{E}_{\tau \sim \pi(\cdot|\theta_s,\theta_m)} [R(\tau)] $$

Where $\tau$ is a trajectory of (observation, memory_action, task_action) steps.

Training Signal for Inner Loop§

Given a set of trajectories, AutoMem identifies positive memory episodes using a heuristic: if the agent later performs a task action that depends on a previously stored memory piece, that storage decision is considered good. These episodes are extracted and used for SFT (supervised fine-tuning) on the base model.

Implementation Sketch§

# Pseudocode for memory action execution
class MemoryManager:
    def __init__(self, base_dir):
        self.base_dir = base_dir
    
    def execute(self, action_type, **kwargs):
        if action_type == 'write':
            path = os.path.join(self.base_dir, kwargs['filename'])
            with open(path, 'w') as f:
                f.write(kwargs['content'])
        elif action_type == 'read':
            path = os.path.join(self.base_dir, kwargs['filename'])
            return open(path).read()
        # ... other actions

# Outer loop: revise memory prompt
revised_prompt = strong_llm_review(trajectory, current_prompt)

# Inner loop: fine-tune model on good memory sequences
for memory_seq in positive_memory_episodes:
    train_step(model, memory_seq)

Experimental Results§

AutoMem was evaluated on three procedurally generated games: Crafter, MiniHack, and NetHack. By optimizing only memory (task actions unchanged), the base 32B open model improved ~2x-4x, matching or surpassing frontier closed models like Claude Opus 4.5 and Gemini 3.1 Pro Thinking. This demonstrates that memory management is a separable, high-leverage skill.