Skip to content

Exploratory Memory-Augmented LLM Agent via Hybrid On- and Off-Policy Optimization

🕒 Published (v1): 2026-01-01 · Source: ICLR · Venue: ICLR 2026 · link

Why this paper was selected

Hybrid on/off-policy RL for exploratory memory-augmented LLM agents

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

EMPO² is a hybrid RL framework for LLM agents that combines parametric (gradient) updates with non-parametric (external memory) updates to overcome the exploration bottleneck in online RL. The agent generates self-reflective tips from past rollouts, uses them to guide future exploration, and then distills the learned behaviors back into model parameters via both on- and off-policy losses. On ScienceWorld and WebShop, it achieves 128.6% and 11.3% improvements over GRPO respectively.

Problem

Current LLM agents trained with online RL (e.g., GRPO) exploit pretrained knowledge but fail to explore novel states — they stagnate when tasks require discovering objects or strategies not in the model's prior. Pure non-parametric methods (Reflexion) saturate quickly because fixed parameters cannot diversify collected experiences. The gap is a principled mechanism that lets memory bootstrap parametric learning rather than substituting for it.

Method

EMPO² (Exploratory Memory-Augmented On- and Off-Policy Optimization) operates in two rollout modes and two update modes, yielding three hybrid learning configurations:

Rollout modes: - Without memory (prob. \(1-p\)): standard policy rollout \(a_{t+1} \sim \pi_\theta(\cdot \mid s_t, u)\) - With memory (prob. \(p\)): memory-augmented rollout \(a_{t+1} \sim \pi_\theta(\cdot \mid s_t, u, \text{tips}_t)\), where \(\text{tips}_t = \text{Retr}(s_t; \mathcal{M})\) retrieves up to 10 relevant tips via embedding similarity

Memory update: After each episode, the policy \(\pi_\theta\) itself generates a reflective tip \(\text{tip}_i \sim \pi_\theta(s_T, u, \text{tip-generation prompt})\) stored in buffer \(\mathcal{M}\).

Update modes for memory-augmented rollouts: - On-policy (prob. \(1-q\)): importance ratio uses the tips-conditioned prompt - Off-policy (prob. \(q\)): stored log-probs under tips are replaced with log-probs under the no-tips prompt — effectively reward-guided knowledge distillation where the tips-conditioned policy acts as teacher and the vanilla policy is the student

The off-policy update is stabilized via a masking mechanism that zeros the advantage for tokens where \(\pi_\theta(a_t^{(i)} \mid s_t^{(i)}, u) < \delta\), preventing gradient blow-up from low-probability tokens. The modified GRPO loss is:

\[\mathcal{L} = \mathbb{E}\left[\frac{1}{NT}\sum_{i,t} \min\!\left(\rho_\theta^{(i,t)} \hat{A}_t^{(i)},\, \text{clip}(\rho_\theta^{(i,t)}, 1\pm\epsilon)\hat{A}_t^{(i)}\right) \cdot \mathbf{1}_{\pi_\theta \geq \delta}\right] - \beta D_{\mathrm{KL}}(\pi_\theta \| \pi_{\mathrm{ref}})\]

An intrinsic reward \(r_{\text{intrinsic}} = \frac{1}{n}\) (where \(n\) = number of similar past states by cosine similarity) encourages visiting novel states when extrinsic reward is absent.

Key Contributions

  • Unified dual-update paradigm: simultaneous parametric (RL gradient) and non-parametric (memory) updates during training
  • Self-generated tips: the evolving policy \(\pi_\theta\) produces its own reflective guidance, avoiding reliance on fixed or external annotators
  • Hybrid on/off-policy GRPO extension: off-policy update acts as reward-gated knowledge distillation, enabling the model to internalize memory benefits so inference requires no memory
  • Token-level masking to stabilize off-policy training and prevent gradient divergence
  • Novelty-based intrinsic reward maintaining policy entropy throughout training
  • OOD generalization: trained model uses memory to quickly adapt to unseen tasks with no weight updates

Results

  • ScienceWorld (Qwen2.5-7B-Instruct, 19 tasks): EMPO² avg return 75.9 vs. GRPO 33.2 vs. Reflexion 17.1 vs. Retrospex 33.8 vs. naive −61.3 — a 128.6% improvement over GRPO
  • 7 out of tasks that began with negative reward reached score of 100
  • WebShop: EMPO² 88.3 vs. GRPO 79.3 — an 11.3% improvement over GRPO
  • OOD (no weight updates, few trials): EMPO² 88.1 vs. Retrospex 58.1 on in-domain; on OOD biology/electricity tasks, EMPO² shows strong positive transfer where GRPO degrades (e.g., GRPO −21.4 / −6.2 vs. EMPO² positive scores)
  • GRPO converges prematurely on the power-component task; EMPO² continues improving and solves it

Limitations

  • Memory retrieval cost and storage grow with interaction; the paper does not report inference-time latency or memory size experiments
  • Tip quality depends on the evolving policy — early in training, generated tips may be noisy or misleading
  • Evaluated only on text-based embodied environments (ScienceWorld, WebShop); generalization to visual or tool-use domains is untested
  • Two hyperparameters \(p\) (memory rollout probability) and \(q\) (off-policy update probability) require tuning; their interaction is addressed only via ablations
  • Results use a single base model (Qwen2.5-7B-Instruct); scalability to larger or different model families is not shown

Relevance to Agentic AI / LLM Agents

EMPO² directly addresses the core bottleneck in training LLM agents with RL: the exploration–exploitation imbalance that causes standard GRPO to stagnate. The framework's insight — that non-parametric memory can bootstrap parametric learning rather than replace it — offers a principled alternative to warm-start SFT and closes the gap between verbal RL methods (Reflexion) and online RL. The off-policy distillation mechanism is particularly relevant for agentic settings where inference-time memory may be unavailable or expensive, since it internalizes exploratory behaviors into model weights. This work sits at the intersection of memory-augmented agents, exploration in RL, and online self-improvement, making it a strong reference for research on sample-efficient agent training without human-curated trajectories.