PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zeyuan Liu, Jeonghye Kim, Xufang Luo, Dongsheng Li, Yuqing Yang

arXiv:2602.23008v2cs.LGcs.AI

TL;DR

LLM agents often rely on prior knowledge and struggle to explore novel states in unfamiliar environments. EMPO^2 combines memory with on- and off-policy reinforcement learning, outperforming GRPO on ScienceWorld and WebShop and adapting rapidly to new tasks with memory.

  • Problem

    Current LLM agents rely heavily on pretrained knowledge and struggle when solving tasks requires discovering novel states through exploration.

  • Method

    EMPO^2 jointly updates the parametric policy with reinforcement learning and a self-generated non-parametric memory using complementary on- and off-policy updates.

  • Results

    EMPO^2 outperforms GRPO on ScienceWorld and WebShop, achieving a 128.6% improvement on ScienceWorld, and adapts rapidly to unfamiliar tasks with memory.

  • Takeaways & Limitations

    EMPO^2 is a promising framework for developing LLM agents with stronger exploration, training efficiency, generalization, and adaptability.

  • Takeaways & Limitations

    The study primarily uses Qwen2.5-7B-Instruct and simple similarity-based memory retrieval, leaving generality across models and retrieval mechanisms unresolved.

Abstract

from arXiv · show

Exploration remains the key bottleneck for large language model agents trained with reinforcement learning. While prior methods exploit pretrained knowledge, they fail in environments requiring the discovery of novel states. We propose Exploratory Memory-Augmented On- and Off-Policy Optimization (EMPO$^2$), a hybrid RL framework that leverages memory for exploration and combines on- and off-policy updates to make LLMs perform well with memory while also ensuring robustness without it. On ScienceWorld and WebShop, EMPO$^2$ achieves 128.6% and 11.3% improvements over GRPO, respectively. Moreover, in out-of-distribution tests, EMPO$^2$ demonstrates superior adaptability to new tasks, requiring only a few trials with memory and no parameter updates. These results highlight EMPO$^2$ as a promising framework for building more exploratory and generalizable LLM-based agents.

1 INTRODUCTION

LLM agents often rely on prior knowledge and struggle to explore environments where progress requires discovering novel states. EMPO2 addresses this limitation by jointly updating the policy and memory, improving exploration and performance on ScienceWorld and WebShop.

  • Motivation: LLM agents often exploit pretrained knowledge rather than systematically exploring, limiting performance when tasks require discovering novel states or acquiring new information.Many systems conduct only limited search within familiar distributions.
  • Motivation: External memory can help agents correct failed attempts without parameter updates, but its performance tends to saturate quickly with static parameters.Collecting experiences alone cannot fully overcome the limitations of fixed policies.
  • EMPO2: EMPO2 jointly updates parametric policy parameters through RL and non-parametric memory through interaction to enable broader exploration and adaptation.The memory updates also enhance the efficiency of parametric learning.
  • Experiments: EMPO2 is evaluated on ScienceWorld and WebShop against non-parametric and parametric offline and online RL approaches.Both environments require exploration to solve complex multi-step embodied reasoning tasks.
  • Results: 128.6% improvement on ScienceWorld and 11.3% improvement on WebShop over GRPO are achieved by EMPO2 in the reported experiments.The paper summarizes these gains as substantial improvements over prior algorithms.

2 PRELIMINARIES

The preliminaries define online RL as alternating policy rollouts and updates, formalize multi-step LLM–environment interaction, and introduce GRPO’s rollout-comparison mechanism without a value function.

  • Online RL: Online RL alternates between generating trajectories with the current policy πθ and optimizing that policy from the resulting rollouts.The process consists of a rollout phase followed by an update phase.
  • Policy Rollout: Policy rollout samples a task u, generates actions conditioned on the task and current state, and receives rewards and next states from the environment.The interaction continues over multiple timesteps, producing trajectories such as u, a1, r1, s1, a2, r2, …, sT.
  • Group Relative Policy Optimization: GRPO compares multiple rollouts of the same task and removes the value-function requirement used in PPO.For task u, πθ generates N trajectories, each assigned a return equal to the sum of rewards along its trajectory.
  • Group Relative Policy Optimization: GRPO assigns positive relative advantages to above-average-return trajectories and negative advantages to below-average-return trajectories.The relative advantage is normalized by the standard deviation of rollout returns.

3 THE EXPLORATION PROBLEM OF LLM AGENTS

LLM agents must explore through direct interaction because pretrained knowledge may not match an environment’s rules or dynamics. Existing agents often fail to analyze unsuccessful actions and discover alternative behaviors, limiting online RL optimization.

  • Why exploration is necessary: Pretrained priors can conflict with environment rules, causing erroneous behavior and requiring adaptation through direct interaction and trial-and-error.Exploration seeks information beyond pretraining and may require atypical or counterintuitive actions.
  • Limitations of existing approaches: Prior alignment strategies use numerous golden trajectories, large-scale models, human engineering, or established simulation information, but achieve strong results mainly in constrained settings.These approaches include warm-start supervised fine-tuning and other forms of environment-specific support.
  • Core training challenge: Online RL without intrinsic exploration struggles to optimize effectively when agents must learn environment rules through trial and error.The paper therefore focuses on efficient online-RL training without embedding environment rules beforehand.
  • ScienceWorld example: In ScienceWorld, the agent must search for an unobserved red light bulb before activating it, but instead follows instructions literally and fails to locate the bulb.The example illustrates why agents need to analyze failure causes and explore alternative actions rather than repeat unchanged behavior.

4 METHOD

EMPO2 combines memory-augmented exploration with on- and off-policy optimization, using reflective tips from past rollouts and consolidating exploratory behavior through parameter updates. It also introduces intrinsic rewards for novel states to encourage exploration when extrinsic rewards are absent.

  • Memory and Tips: The agent reviews completed trajectories, generates reflective tips with its policy, and stores them in a memory buffer for subsequent retrieval.For episode i of task u ending at timestep t, tipi ∼πθ(st, u, tip-generation prompt).
  • Hybrid Learning Modes: EMPO2 uses memory-augmented prompting during rollouts and combines rollout and update modes into three hybrid learning configurations.The configurations are on-policy learning without memory, on-policy learning with memory, and off-policy learning.
  • Rollout Modes: During rollouts, memory-augmented prompting is selected with probability p, while prompting without memory is selected with probability 1 − p.Retrieved tips are selected by similarity search, and the number of retrieved tips is limited to 10.
  • Update Modes: Memory-augmented trajectories receive either on-policy updates with probability 1 − q or off-policy updates with probability q.Off-policy updates evaluate actions under the policy conditioned only on the current state and task, functioning as reward-guided knowledge distillation from tips-conditioned trajectories.
  • Intrinsic Rewards for Exploration: The intrinsic-reward mechanism stores distinct states, compares new states with prior entries using cosine similarity, and rewards sufficiently novel states.It encourages exploration without extrinsic reward and maintains policy entropy.

5 RELATED WORK

Related work spans data-driven, model-based, memory-augmented, reinforcement-learning, and exploration-focused approaches for improving LLM agents. These studies motivate EMPO^2’s combination of memory, online learning, and exploration mechanisms.

  • LLM Agents in Multi-Step Embodied Tasks: Data-driven approaches improve decision-making through data collection and imitation learning, while model-based agents construct world models, often using code generated by GPT-4.
  • Memory for LLM Agents: External-memory methods such as Reflexion and REMEMBERER enable LLMs to improve from past experiences without parameter updates.Reflexion stores verbal reflections, whereas REMEMBERER records observations, actions, rewards, and Q-values for retrieving similar cases as few-shot exemplars.
  • Learning by Knowledge Distillation: Context distillation trains a Student prompt from a Teacher prompt through offline supervised distillation, whereas the proposed method integrates reward-guided distillation into online RL.
  • RL for LLM Agents: RL optimizes LLM parameters using environment observations and rewards, with prior work spanning offline RL and recent online RL methods.Retrospex demonstrated offline RL for LLM agents, while later studies investigated online learning in real time.
  • Enhancing Exploration for Online RL: Online-RL exploration methods use intrinsic rewards, stored key states, or novelty mechanisms to encourage discovery of difficult states.Examples include count-based exploration, Random Network Distillation, Go-Explore, and its LLM extension Intelligent Go-Explore.

6 EXPERIMENTS

Experiments evaluate EMPO2 on ScienceWorld and WebShop with Qwen2.5-7B-Instruct, measuring the trained model without memory at test time. EMPO2 outperforms baselines and benefits from combining memory-based exploration with on- and off-policy learning.

  • Experimental Setup: Experiments use Qwen2.5-7B-Instruct on ScienceWorld and WebShop, evaluating trained models without memory at test time.ScienceWorld requires long-term multi-step planning, hypothesis testing, outcome interpretation, and exploration; WebShop requires searching, navigating, and purchasing products.
  • ScienceWorld Results: On ScienceWorld, EMPO2 delivers the strongest reported performance, improving over Qwen2.5-7B-Instruct, Reflexion, Retrospex, and GRPO.ScienceWorld uses five variants for training and 20 unseen test variants; failed tasks yield returns from -100 to 100.
  • Adaptation in New Tasks with Memory Updates: 136% average improvement across three scenarios within 10 steps demonstrates EMPO2’s rapid adaptation when memory is enabled.Without memory, EMPO2 already performs better than GRPO on novel tasks; GRPO is more variable and sometimes fails to adapt.
  • WebShop Results: On WebShop, EMPO2 surpasses all baselines, achieving higher scores and success rates than GiGPO.The comparison averages results over three random seeds and reports mean score and mean success rate (%).
  • Ablation Analysis: Removing either off-policy learning or on-policy learning with memory produces suboptimal learning, showing their complementary roles.On-policy updates support stable learning, while off-policy updates enable reasoning as if guided by additional tips; combining them improves convergence and final performance.

7 CONCLUSION

The conclusion presents EMPO2 as a memory-augmented RL method that combines on- and off-policy learning to improve exploration, training efficiency, stability, and few-shot adaptation. It also identifies simple memory retrieval and limited model-family coverage as directions for improvement.

  • Contributions: EMPO2 enhances exploration in parametric RL through non-parametric memory updates.The method is presented as a novel RL approach for improving exploration.
  • Contributions: EMPO2 integrates on-policy and off-policy learning to improve training efficiency and stability.Experiments report gains in training efficiency on ScienceWorld and WebShop.
  • Contributions: EMPO2 adapts rapidly to new domains in a few-shot manner by incorporating additional memory.This capability is demonstrated through adaptation to new domains with added memory.
  • Limitations and future work: The current memory implementation uses simple similarity-based retrieval, while more advanced mechanisms may further improve performance.The conclusion identifies retrieval sophistication as a limitation and future improvement direction.
  • Limitations and future work: Experiments primarily use Qwen2.5-7B-Instruct, motivating evaluation across broader model families and sizes.Scaling to larger models is specifically identified as a direction for studying generality and robustness.

ETHICS STATEMENT … C DETAILED EXPLANATION OF IMPORTANCE SAMPLING RATIOS IN POLICY UPDATES

The paper documents responsible benchmark-only evaluation, describes EMPO2’s memory-augmented rollout and update procedure, specifies its prompting setup, and explains three importance-sampling scenarios with practical stability considerations.

  • ETHICS STATEMENT: The study uses publicly available ScienceWorld and WebShop benchmarks without private or sensitive data, follows dataset licenses and community standards, and collects or modifies no additional data.
  • ETHICS STATEMENT: Because online RL may be misapplied in safety-critical contexts, the study is confined to benchmarks and recommends careful scrutiny for real-world LLM applications.
  • A PSEUDO CODE: Compared with GRPO, EMPO2 adds a memory buffer, tip retrieval and addition, two rollout modes, and two policy-update modes.
  • A PSEUDO CODE: Each training iteration samples tasks, initializes identical environments, mixes prompting without memory and memory-augmented prompting, retrieves tips when applicable, and appends generated tips to memory.
  • B PROMPTS: The prompts use dynamically filled progress information, ScienceWorld demonstrations and exploration guidance, while WebShop prompts provide task, history, observations, and admissible-action instructions.
  • C DETAILED EXPLANATION OF IMPORTANCE SAMPLING RATIOS IN POLICY UPDATES: The importance-sampling ratio ρθ compares the current policy’s action probability with the old policy’s probability, with calculation determined by tip use during rollout and update.
  • C DETAILED EXPLANATION OF IMPORTANCE SAMPLING RATIOS IN POLICY UPDATES: Regular on-policy updates use identical (st, u) conditioning, whereas on-policy updates with tips consistently condition both policies on (st, u, tips t).
  • C DETAILED EXPLANATION OF IMPORTANCE SAMPLING RATIOS IN POLICY UPDATES: Off-policy updates sample actions with tips but recompute current-policy log-probabilities without tips, enabling guidance internalization while introducing distribution-shift instability and bias–variance trade-offs.

D EXPERIMENTS DETAILS … E QUALITATIVE ANALYSIS ON TIPS

The experiments standardize Retrospex evaluations around Qwen2.5-7B-Instruct, remove selected human heuristics in ScienceWorld, and implement EMPO2 as a multi-step extension of GRPO with memory integration. ScienceWorld and WebShop use separate implementations and shared algorithm-specific training configurations.

  • D.1 RETROSPEX: Retrospex experiments standardize both environments on Qwen2.5-7B-Instruct using offline trajectories and supervised fine-tuning.The original Retrospex models were Flan-T5-Large for ScienceWorld and Llama-3-8B-Instruct for WebShop.
  • D.1 RETROSPEX: ScienceWorld evaluation removes Retrospex’s human-designed “focus on” and “teleport” heuristics to better reflect inherent agent capabilities.The original procedure skips some focus actions and replaces step-by-step go-to actions with direct teleport moves.
  • D.2 ONLINE RL: SCIENCEWORLD: EMPO2 builds on verl’s GRPO implementation and extends single-response rollouts to multi-step environment interaction.Episodes continue until reaching the maximum episode length or receiving an environment termination signal.
  • D.2 ONLINE RL: SCIENCEWORLD: The ScienceWorld memory buffer resets per buffer, stores nonduplicate content, and evicts the oldest entries after exceeding 1,000 memories.Stored entries include a counter, key, content, and score.
  • D.2 ONLINE RL: SCIENCEWORLD: ScienceWorld memory retrieval filters candidates by cosine similarity above 0.5 and returns the top 10 by score.Memory retrieval is used during on-policy-with-memory and off-policy phases.
  • D.2 ONLINE RL: SCIENCEWORLD: ScienceWorld online RL uses 32-token responses per step, 4,500 total tokens, 30-step episodes, and a 1 × 10−6 actor learning rate.GRPO additionally uses group size 8, mini-batch size 16, and a KL-divergence coefficient of 0.0.
  • D.3 ONLINE RL: WEBSHOP: WebShop uses the verl-agent GRPO framework with the same EMPO2 modifications, while all online algorithms share a 512-token response limit, 15-step episodes, and a 1 × 10−6 actor learning rate.The configuration also specifies group size 8, rollout temperature 1.0, validation temperature 0.4, mini-batch size 64, and KL coefficient 0.01.

E.1 MORE EXAMPLES OF GENERATED TIPS · E.2 EFFECTS OF TIPS ON EXPLORATION BEHAVIOR · F MORE ABLATION STUDY

Generated tips evolve with task progress by identifying incomplete trajectories, incorrect actions, and missing connections across ScienceWorld tasks. Memory-augmented prompting helps agents use failed attempts as guidance, avoid repeating mistakes, and explore more effective behaviors.

  • E.1 MORE EXAMPLES OF GENERATED TIPS: Tips evolved as trajectories progressed, highlighting incomplete tasks, incorrect configurations, and missed circuit-construction steps.Examples included focusing on a red light bulb without completing the task and connecting a battery incorrectly.
  • E.1 MORE EXAMPLES OF GENERATED TIPS: 7.0 score accompanied a tip noting that the green wire and light bulb were connected but the battery was not found.The trajectory remained incomplete because powering the green light bulb required the missing battery connection.
  • E.1 MORE EXAMPLES OF GENERATED TIPS: 13.0 score accompanied a tip identifying a missing battery connection after connecting the green wire to the green light bulb.The task was not fully completed despite the partial circuit connection.
  • E.1 MORE EXAMPLES OF GENERATED TIPS: 83.0 score accompanied tips describing an unpowered green light bulb and repeatedly closed workshop doors.The examples show that tips also captured environmental and action-level obstacles after circuit progress.
  • E.1 MORE EXAMPLES OF GENERATED TIPS: 80.0 score accompanied a tip recognizing that green paint had been created by mixing yellow and blue paint, making a focus action unnecessary.Other examples documented ambiguous pouring actions and failures to locate or use the required paint materials.
  • E.2 EFFECTS OF TIPS ON EXPLORATION BEHAVIOR: Without memory, agents tend to repeat mistakes because they cannot incorporate feedback from previous failures into later attempts.Memory-augmented prompting lets agents consult unsuccessful attempts, avoid repeating errors, explore novel behaviors, expand search capabilities, and improve learning performance.

F.1 MODE SELECTION PROBABILITY … G.1 COST ANALYSIS OF MEMORY-AUGMENTED ROLLOUTS

EMPO² remains effective across broad hyperparameter settings, with memory rollout and off-policy update probabilities balancing exploration, internalization, stability, and robustness. Intrinsic rewards support continued learning, while memory improves exploration at a measurable rollout-time cost.

  • F.1 MODE SELECTION PROBABILITY: EMPO² ablates memory-rollout probability p and off-policy update probability q to assess their effects during rollout and update phases.The study systematically investigates both hyperparameters.
  • F.1 MODE SELECTION PROBABILITY: p = 0.4 and p = 0.7 accelerate initial learning, whereas p = 0.7 fluctuates later and p = 0.25 provides stable convergence across diverse tasks.At p = 0.1, performance degrades significantly because EMPO² effectively collapses to GRPO.
  • F.1 MODE SELECTION PROBABILITY: q = 0.3 and q = 0.95 underperform, while q = 0.85 enables faster early exploration than the default q = 2/3.Large q overemphasizes distillation, whereas small q slows knowledge internalization; the default prioritizes overall robustness.
  • F.1 MODE SELECTION PROBABILITY: The default hyperparameters provide a balanced configuration that generalizes across multiple tasks without task-specific tuning and remains adaptable to further optimization.These results indicate effective performance across a broad hyperparameter range.
  • F.2 ROLE OF INTRINSIC REWARD: Intrinsic-reward variants with 0.5×, 2×, removal, or replacement by an RND bonus are compared against the full EMPO² method.The ablation examines how intrinsic-reward design affects the proposed algorithm.
  • F.2 ROLE OF INTRINSIC REWARD: Removing the intrinsic reward causes learning to plateau lower, while 0.5× slows convergence, 2× introduces minor instabilities, and reward-using variants reach similar final performance.The RND-based variant also converges to a similar final-performance level.
  • G.1 COST ANALYSIS OF MEMORY-AUGMENTED ROLLOUTS: 50.4 seconds per iteration, or about 19% of total rollout time, is added by memory operations including tip generation, retrieval, and storage.Experiments use Qwen2.5-7B-Instruct on 8 A100 40GB GPUs; tip generation and storage account for a substantial portion of the cost.

G.2 COST ANALYSIS OF TOTAL TRAINING TIME

EMPO2’s training-time cost is driven mainly by memory-augmented rollouts and response length. Despite these costs, EMPO2 remains substantially more efficient than GRPO when performance is evaluated against training time.

  • Cost factors: 19% of total rollout time is attributed to the memory component, incurred with a 25% probability under memory-augmented prompting.The memory component’s expected cost is determined by its rollout-time share and selection probability.
  • Cost factors: Increasing response length slows rollouts, which constitute a major portion of total LLM-based RL training cost.Response length is identified as the second primary factor influencing training time relative to GRPO.
  • Time–performance comparison: EMPO2 exhibits substantially higher efficiency than GRPO when performance is plotted against training time on ScienceWorld’s power-component task.Figure 13 uses training time as the x-axis for a fair comparison with GRPO.
Loading 2602.23008v2…