[ICLR 2025] EMPO2: Escaping the Exploration Trap in LLM Agents via Memory-Augmented Distillation

TL;DR

Reinforcement Learning (RL) for LLM agents typically fails when the agent needs to "explore" rather than just "reason" from its pre-training. EMPO2 (Exploratory Memory-Augmented On- and Off-Policy Optimization) solves this by using a self-generated memory buffer as a scaffold. It doesn't just use memory as a crutch; it uses off-policy distillation to bake that "memory-guided wisdom" directly into the model's weights, resulting in a 128.6% performance boost on complex tasks like ScienceWorld.

The Problem: Why LLM Agents Stop Exploring

Most current LLM agents suffer from Exploitation Bias. When faced with a task like "Turn on the red light bulb" in a room where no bulb is visible, standard RL agents (like GRPO) repeatedly try to "focus on the bulb" because their pre-training tells them that's the next logical step. They fail to understand that they first need to seek the bulb in another room.

Existing solutions are bifurcated:

1. Non-parametric Memory (e.g., Reflexion): Excellent for short-term fixes, but the model "forgets" the lessons if the memory is removed.
2. Parametric RL (e.g., PPO/GRPO): Permanently changes weights but struggles to discover successful trajectories in hard-exploration environments.

Methodology: The "Memory Scaffold"

EMPO2 introduces a hybrid architecture that treats exploration and optimization as two interleaved cycles.

1. Self-Generated Memory Tips

The agent uses its own reasoning capability to review past failures and successes, generating verbal "tips" (e.g., "The circuit was missing a battery connection"). These are stored in a vector database.

2. The Hybrid Rollout

During training, the agent operates in two modes:

• Mode A (Naive): Standard prompting.
• Mode B (Memory-Augmented): Prompting with retrieved tips.

3. Off-Policy Internalization (The Secret Sauce)

This is where the magic happens. When an agent succeeds using a "tip," EMPO2 uses Off-Policy Updates. It updates the "Naive" policy to produce the same high-reward actions even when the tip is absent. This effectively distills the external guidance into the model's permanent parameters.

Figure 1: The dual-update paradigm where non-parametric exploration bootstraps parametric learning.

Experiments: Breaking the Ceiling

The researchers tested EMPO2 on ScienceWorld and WebShop.

• Performance: In ScienceWorld, EMPO2 achieved an average score of 75.9, obliterating the GRPO baseline of 33.2.
• OOD Adaptability: Perhaps more impressively, once trained, EMPO2 showed a remarkable ability to adapt to entirely new tasks. By using a few trials to populate its memory, it could solve unseen biology or electricity tasks without further weight updates.

Table 1: Comparison across 19 ScienceWorld tasks showing EMPO2's dominance.

Critical Insight: Stabilization is Key

Off-policy training with LLMs is notoriously unstable. If a "tip" makes a previously low-probability action highly likely, the Importance Sampling ratio $({\rho }_{h}eta)$ can explode, leading to NaN gradients. The authors introduced a Token Masking mechanism that ignores advantage terms for tokens whose probability falls below a threshold $\delta$, ensuring the distillation process remains stable.

Conclusion & Future Outlook

EMPO2 proves that we don't need a "perfect" teacher to train better agents—the agent can be its own teacher. By structuring the interaction between memory and weights, we can create agents that explore like humans: they use external notes to learn, but eventually, the knowledge becomes second nature.

Limitations: The current system relies on a simple cosine-similarity search for memory. Future iterations could benefit from more semantic, hierarchical memory structures to handle even longer-horizon tasks in worlds like Minecraft or complex codebase navigation.