Technical Summary: Learning to Explore: Scaling Agentic Reasoning via Exploration-Aware Policy Optimization

Problem Statement

Recent advancements in agentic test-time scaling allow Large Language Models (LLMs) to gather environmental feedback before committing to final actions, enabling multi-step reasoning during deployment. However, existing methods suffer from a critical limitation: they employ undifferentiated exploration strategies. Current approaches typically entangle exploration and action selection within a single policy, lacking the ability to adaptively distinguish when exploration is truly required. Consequently, agents often engage in indiscriminate exploration even in well-understood states, accumulating low-value information and obscuring critical signals. Conversely, they may fail to explore sufficiently in unfamiliar states, leading to suboptimal decision-making. The core challenge is to teach agents to adaptively allocate interaction steps based on uncertainty, distinguishing between when information-seeking exploration is necessary and when direct goal-pursuit is sufficient.

Methodology

The paper proposes Exploration-Aware Policy Optimization (EAPO), a framework designed to enable agents to explore selectively at appropriate states and leverage acquired information for decision-making. The methodology consists of three primary components:

1. Exploration-and-Memory Reasoning Mode

To structure the reasoning process, the authors introduce a specific output format where the agent explicitly generates:

• Exploration Guidance ($e_t$): A structured strategy specifying what information to acquire and candidate actions to take.
• Accumulative Memory ($m_t$): A summary of task-relevant information extracted from past interactions and newly observed states.
• Executable Action ($a_t$): The final action taken.

This separation allows the agent to deliberate on environmental dynamics before committing to execution and maintains a persistent log of visited states, facilitating the reuse of acquired information.

2. Fine-Grained Reward Modeling via Variational Inference

To incentivize valuable exploration without relying on costly online rollouts, the authors design a reward function that evaluates the utility of exploratory actions.

• Bayesian Perspective: The method treats exploration and memory as latent variables. It defines a Bayesian exploratory reward based on the posterior probability of task success given the current state, exploration strategy, and memory.
• Variational Proxy: Since the true posterior is intractable, a learnable variational proxy $q_\phi$ is trained to approximate the distribution of optimal exploration-memory configurations.
• Reward Definition: The reward $R_{explore}$ is defined as the maximum of:
  1. The utility of immediate exploitation (using existing memory).
  2. The utility of proactive exploration (incorporating a newly observed state into memory), discounted by $\gamma^2$ to penalize latency and prevent over-conservatism.
• Total Reward: The final reward is a weighted combination of the task success signal, a format reward (ensuring structured output), and the exploratory reward.

3. Exploration-Aware Two-Stage Training Strategy

To address the challenges of credit assignment and policy optimization in the presence of exploration, EAPO employs a two-stage approach:

• Stage 1: Learning to Rollback (SFT): Agents are supervised fine-tuned on an expert dataset to acquire the capability of "rolling back" to previous states (e.g., clicking a "back" button) after an exploratory action leads to an undesirable state. This treats exploration as a reversible behavior rather than a terminal risk.
• Stage 2: Exploration-Aware Policy Optimization (GRPO): The authors modify Group Relative Policy Optimization (GRPO) by introducing visitation depth ($\kappa$) as a clustering criterion. Transitions are grouped not just by state, but by the number of times the agent has revisited that state. This ensures that exploratory actions are compared against other actions at similar informational stages, preventing the underestimation of exploration value that occurs when comparing early exploratory steps with final execution steps in the same group.

Key Contributions

1. Exploration-Aware Framework: A novel policy optimization method that explicitly teaches agents to distinguish between when to explore and when to execute, moving beyond uniform exploration strategies.
2. Variational Reward Modeling: A principled approach to quantifying the utility of exploratory actions using variational inference, avoiding the high computational costs of online trajectory sampling while providing stable reward signals.
3. Structured Reasoning & Memory: The introduction of explicit <explore> and <memory> tags that integrate information gathering and state summarization directly into the agent's reasoning loop.
4. Visitation-Depth Grouping: A mechanism to stabilize advantage estimation by clustering transitions based on how often a state has been visited, ensuring fair comparison between exploration and exploitation phases.
5. Rollback Capability: The integration of a supervised rollback mechanism that allows agents to recover from exploratory missteps, making exploration a safe and reversible process.

Experimental Results

The authors evaluated EAPO across four challenging environments: ALFWorld and WebShop (text-based), and AndroidWorld and OSWorld (GUI-based).

• Performance Gains: EAPO consistently outperformed strong baselines (including GRPO, DAPO, GiGPO, LAMER, and various closed-source models) across all environments. Improvements ranged from 20% to 60% over existing methods.
• Model Efficiency: A 2B-scale model trained with EAPO achieved performance comparable to or exceeding substantially larger general and agentic models (e.g., 72B parameter models), demonstrating that effective exploration strategies can compensate for parameter scale.
• Adaptive Behavior: Analysis showed that agents trained with EAPO exhibit adaptive exploration: they increase exploration degrees in uncertain states and converge to selective exploration as they learn the environment, avoiding indiscriminate trial-and-error.
• Generalization: Models trained on AndroidWorld generalized effectively to the unseen OSWorld environment with only slight performance degradation, suggesting the learned exploration policies are domain-invariant.
• Efficiency: The method incurred only about 30% additional training overhead (primarily from training the variational reward model) compared to standard GRPO, while avoiding the expensive online sampling required by alternative approaches. Inference time saw a modest increase (at most 10 additional steps on average) due to exploration, controlled by the discount factor.

Significance and Claims

The paper claims that EAPO represents a significant step forward in agentic test-time scaling. By explicitly modeling the utility of exploration and separating information-seeking from decision-making, the framework enables agents to:

• Reason about uncertainty: Agents learn to identify when the current state or memory is insufficient for a decision and proactively seek information.
• Optimize resource allocation: Agents avoid wasting interaction steps on redundant exploration in known states while effectively gathering critical data in novel scenarios.
• Achieve robust generalization: The ability to adaptively explore allows agents to handle out-of-distribution environments (e.g., updated UI layouts) without requiring environment-specific retraining.

The authors conclude that learning when and how to explore is a fundamental capability that can substantially improve decision quality and environment understanding, often more effectively than simply scaling model parameters. The work provides a principled framework for integrating exploration into the reasoning process of agentic LLMs, moving beyond heuristic or uniform exploration strategies.