PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Xingyuan Hua, Sheng Yue, Ju Ren

arXiv:2605.08978v2cs.AI

TL;DR

Existing agentic models often struggle to explore effectively in complex, long-horizon settings. EAPO learns when exploration is useful through exploration-aware rewards and optimization, improving performance across four environments and outperforming existing methods by 20%–60%.

  • Problem

    Agentic models struggle in complex, long-horizon settings because goal-oriented training limits generalization and environmental information gathering, while optimization can mix exploratory and execution actions.

  • Method

    EAPO uses exploration-and-memory reasoning, Bayesian exploratory rewards, and exploration-aware grouping to distinguish information gathering from task execution.

  • Results

    EAPO consistently outperforms existing methods by 20%–60% across four environments, with particularly strong gains in complex long-horizon GUI control tasks.

  • Takeaways & Limitations

    EAPO enables agents to learn when and how to explore, improving decision-making across embodied, shopping, and web/mobile GUI tasks.

  • Takeaways & Limitations

    EAPO relies on manually specified structured exploration and memory representations, which may restrict strategy expressiveness and adaptability to flexible information-acquisition tasks.

Abstract

from arXiv · show

Recent advancements in agentic test-time scaling allow models to gather environmental feedback before committing to final actions. A key limitation of existing methods is that they typically employ undifferentiated exploration strategies, lacking the ability to adaptively distinguish when exploration is truly required. In this paper, we propose an exploration-aware reinforcement learning framework that enables LLM agents to adaptively explore only when uncertainty is high. Our method introduces a fine-grained reward function via variational inference that explicitly evaluates exploratory actions by estimating their potential to improve future decision-making, together with an exploration-aware grouping mechanism that separates exploratory actions from task-completion actions during optimization. By targeting informational gaps, this design allows agents to explore selectively and transition to execution as soon as the task context is clear. Empirically, we demonstrate that our approach achieves consistent improvements across a range of challenging text-based and GUI-based agent benchmarks. Code is available at https://github.com/HansenHua/EAPO-ICML26 and models are available at https://huggingface.co/hansenhua/EAPO-ICML26.

1. Introduction

Current agentic models struggle in complex, long-horizon tasks, while existing test-time scaling often explores indiscriminately because exploration and action selection are entangled. EAPO explicitly supports state-appropriate exploration and improves decision-making across four environments with limited training overhead.

  • Motivation: Current agentic models often struggle in complex, long-horizon settings, including web navigation, scientific research, and embodied agentic tasks.
  • Problem: Existing test-time scaling methods entangle exploration with action selection, causing indiscriminate exploration in well-understood states and accumulation of low-value information.
  • Method: EAPO teaches agents to explore at appropriate states through an exploration-and-memory reasoning mode that generates exploration guidance and summarizes newly observed states.
  • Evaluation: EAPO is evaluated across 4 challenging environments spanning embodied tasks, online shopping, and web/mobile GUI control, enhancing decision-making across all environments.
  • Results: 20%–60%: EAPO consistently outperforms existing methods, particularly on complex long-horizon GUI control tasks.EAPO incurs only about 30% additional training overhead, enables a 2B-scale model to outperform most substantially larger models, and generalizes to unseen scenarios without additional fine-tuning.

2. Related Work

Related work studies test-time scaling for richer reasoning and decision-making, reinforcement learning for LLM agents, and exploration–exploitation trade-offs during policy optimization. However, existing exploration methods primarily improve training-time diversity rather than enabling explicit, adaptive exploration at test time in unfamiliar states.

  • LLM Test-Time Scaling: Test-time scaling branches multiple reasoning trajectories and selects or aggregates final answers to diversify intermediate reasoning during inference.This line extends LLM reasoning and decision-making beyond single-step generation through prompting strategies.
  • RL for LLM Agents: Reinforcement learning encourages LLM agents to explore diverse reasoning chains under the guidance of verifiable rewards.
  • RL for LLM Agents: Policy-optimization methods balance exploration and exploitation by encouraging diverse action selection and preventing premature convergence.
  • RL for LLM Agents: Existing approaches mainly enhance exploration during training rather than enabling explicit, adaptive exploration at test time in unfamiliar states.

3. Preliminaries

The paper formulates agentic tasks as finite-horizon Markov decision processes optimized for expected discounted cumulative reward. It also specializes this formulation to GUI agents and introduces GRPO as a practical training method based on grouped trajectories and token-level advantages.

  • Agentic Tasks: Agentic tasks are modeled as an MDP ⟨S, A, P, T, R, µ, γ⟩ with state, action, transition, horizon, reward, initial-state, and discount components.The policy πθ(a|s) defines the action distribution, while P(s′|s, a) specifies transition probabilities.
  • Agentic Tasks: The learning objective is to maximize expected discounted cumulative reward.
  • GUI-based Agentic Tasks: For GUI tasks, states pair visual UI contexts with task descriptions, actions execute taps, swipes, or text entry, and transitions capture application navigation logic.The horizon limits interaction steps, and an LLM policy governs behavior at each step.
  • GUI-based Agentic Tasks: GUI-task rewards are often sparse binary signals indicating whether the agent completes the task.
  • Group Relative Policy Optimization (GRPO): GRPO generates grouped trajectories for the same task, concatenates each trajectory’s tokens into a complete action, and optimizes using token-level importance weights and advantages.The group contains G trajectories, λ is a hyperparameter, and R(g, ai,t) denotes token reward-to-go.

4. Exploration and Memory Mode

The method structures agent outputs around explicit exploration strategies and accumulated memories, enabling agents to gather and reuse task-relevant information during execution. It then assigns exploration-aware rewards through a variational formulation that evaluates exploratory utility while discouraging unnecessary exploration and avoiding costly online rollouts.

  • Structured exploration and memory: The agent maintains an exploration strategy and accumulated task-relevant memory, initialized as empty strings and updated throughout execution.The strategy specifies needed information and candidate actions, while memory summarizes information extracted from past interactions.
  • Structured exploration and memory: At each step, the model receives the task, current state, preceding exploration, and prior memory, then outputs an exploration strategy, updated memory, and executable action.The structured inputs and outputs are represented as ˜st = [g; st; et−1; mt−1] and ˜at = [et; mt; at].
  • Exploration-aware reward: The variational density-based reward provides stable, efficient utility estimates that are robust to policy stochasticity, decouple estimation from environment interaction, and eliminate expensive online rollouts.A learnable proxy qϕ(e, m|s) approximates the intractable success-conditioned posterior, and KL-regularized policy optimization trains it.
  • Exploration-aware reward: The reward model treats learning to explore as proactively enriching memory with information useful for successful future trajectories.The posterior exploration-memory distribution conditioned on success characterizes the utility of exploration strategies and memory states from the current state.
  • Exploration-aware reward: A discounted exploratory-gain term encourages exploration only when its anticipated utility outweighs the latency cost of acquiring and synthesizing new information.The discount factor γ2 also prevents redundant continued exploration after sufficient information has been acquired for a correct decision.
  • Exploration-aware reward: The total transition reward combines exploratory, format, and task-success components, with the latter two enforcing structured outputs and episodic task completion.The format reward is binary for template compliance, while the task reward is a binary signal for reaching the task goal.

5. Exploration-Aware Training

Section 5 presents EAPO, a two-stage training algorithm that makes exploration-aware policy optimization practical by learning rollback behavior and separating transitions by exploration stage. This addresses reward-attribution failures and misleading advantage estimates caused by mixing exploratory and task-execution actions.

  • Motivation: Standard training underestimates exploratory-action utility because irreversible downstream consequences prevent reliable attribution of future success to information gained during exploration.The agent lacks rollback capability, so exploration rewards become entangled with later consequences.
  • Rollback Training: EAPO first uses supervised fine-tuning on expert rollback transitions to teach the agent to recover previous states.Transitions are collected by prompting a teacher LLM and retained when the generated action successfully recovers the prior state.
  • Rollback Training: Rollback training enables the agent to treat exploration as reversible behavior rather than a terminal risk.The learned capability allows reliable recovery of the previous state.
  • Exploration-Aware Grouping: EAPO clusters transitions using environment states and visitation depth because complete-state grouping over exploration histories and memories is infeasible.Visitation depth counts how often the agent revisits the same state during exploration.
  • Exploration-Aware Grouping: Separating exploration stages prevents discounted exploratory returns from being compared directly with final-action returns, which would underestimate exploration value and reduce exploration.The resulting algorithm is named EAPO and is specified in Algorithm 1.

6. Experiment

Experiments evaluate EAPO across text-based and GUI-based environments, finding consistent baseline improvements and adaptive exploration behavior. Analyses further examine parameters, ablations, generalization, and runtime efficiency.

  • Environments: Experiments cover four environments across text-based ALFWorld and WebShop and GUI-based AndroidWorld and OSWorld.The study uses two domains and four environments.
  • Comparative Results: EAPO consistently outperforms baselines in all four environments, often by a significant margin.The method explicitly learns exploration-aware policies through reward modeling and grouping mechanisms.
  • Key Parameters: Increasing γ generally promotes exploration by preserving rewards from information-gathering actions, although overly large values may harm task performance.The discount factor is varied from 0.5 to 1.0, alongside group size from 4 to 32 and KL coefficient from 0.005 to 1.0.
  • Exploration Learning: During training, larger step-level groups and exploration degree initially increase, then stabilize or converge as agents learn when exploration is useful.These trends indicate increasing information acquisition early in training followed by learned discrimination of exploration needs.
  • Case Studies: At test time, EAPO explores uncertain environments by proposing candidate actions, observing resulting states, summarizing them, and storing memories, sometimes across multiple steps.The behavior is demonstrated through an OSWorld case study.
  • Generalization: EAPO generalizes from AndroidWorld to unseen OSWorld environments by explicitly modeling exploration and memory without environment-specific retraining.The reported mechanism disentangles exploratory reasoning from action execution and maintains structured memory of visited states.
  • Run Time: EAPO adds less than a 15% increase in training time, while its inference overhead is at most 10 additional average steps.Training overhead mainly comes from variational-distribution training; inference cost reflects exploration moderated by γ.

7. Limitation and Discussion

EAPO teaches agents to explore at appropriate states by modeling exploration utility and using exploration-aware optimization to distinguish beneficial exploration from avoidable exploration. However, its manually specified structured exploration and memory representations may restrict expressiveness and adaptability.

  • Contributions: EAPO models exploration utility and uses exploration-aware optimization to guide agents toward exploring at appropriate states.The method is designed to help agents leverage acquired information for decision-making.
  • Effectiveness: EAPO enables agents to distinguish when exploration is beneficial and when it should be avoided.Extensive experiments on agentic tasks corroborate the effectiveness of exploration at test-time and EAPO’s superiority.
  • Limitations: EAPO relies on manually specified structured exploration and memory representations throughout training, potentially limiting strategy expressiveness and task adaptability.This limitation is especially relevant for tasks requiring more flexible forms of information acquisition.

Impact Statement

The work advances principled exploration mechanisms for agentic large language models, while noting that increased exploration capability may heighten autonomy and real-world risks without sufficient constraints.

  • Contribution: The framework enables agentic large language models to learn when and how to explore.It advances the understanding and design of exploration mechanisms through a principled approach.
  • Broader implications: More capable exploration may increase agent autonomy and effectiveness but could amplify unintended behaviors and risks if misused or insufficiently constrained.These concerns are especially relevant when deploying such models in complex real-world environments.

A. Derivation

The derivation formulates memory and exploration as latent variables, obtaining KL-based objectives and a variational lower bound on success probability. Under a regularized soft policy and sparse binary rewards, it characterizes the optimal distribution and simplifies the Q-function to the terminal reward.

  • KL-based objective: The KL objective decomposes into expected negative success log-likelihood, KL divergence between q(e, m|s) and p(e, m|s), and log p(success|s).This follows the stated expansion in Eq. (16).
  • Variational lower bound: Treating memory and exploration as latent variables yields a variational lower bound on πθ(success|s), with expected log policy success minus KL(q(m, e|s)∥πθ(m, e|s)).The bound is derived using the ELBO framework.
  • Optimal policy: The regularized soft policy gives the optimal exploration-memory distribution as q(e, m|s) ∝ πref(e, m|s) exp(Q(s, e, m)).Q(s, e, m) is defined through cumulative trajectory reward under q and πθ.
  • Sparse-reward simplification: For sparse binary rewards, the Q-function equals the final-round reward, Q(st, et, mt) = rT, and Jensen’s inequality introduces a hyperparameter β in the subsequent objective.The derivation then substitutes this result into the objective to obtain its final form.

B. Alternative Reward Function … C.2. Baselines

The alternative online exploratory reward combines direct task progress with the utility of refined exploratory rollouts, while experiments show that more trajectory samples improve estimation. The evaluation spans four text- and GUI-based environments and compares EAPO with six established baselines.

  • B. Alternative Reward Function: R1 rewards actions according to their direct rollout after commitment, favoring correct actions that move the agent toward the target state.This supports goal-directed behavior and efficient completion when sufficient information is available.
  • B. Alternative Reward Function: R2 evaluates a refined rollout after the agent explores future states, capturing informative actions that are not immediately correct.The refined action influences the subsequent decision process through a trajectory that returns to the original state.
  • B. Alternative Reward Function: The online reward’s exploration degree is slightly lower than the trained reward because the current policy may underestimate exploratory actions’ long-term information gain.Figure 3 compares the exploration degree of EAPO with the alternative online reward.
  • B. Alternative Reward Function: Performance consistently improves as sampled trajectories increase from 1 to 10, while EAPO performs comparably to multi-trajectory sampling.More trajectories provide more accurate action-utility estimates and stronger supervision for policy optimization.
  • C.1. Environments: The method is evaluated across four environments in two areas: ALFWorld and WebShop for text-based interaction, and AndroidWorld and OSWorld for GUI-based tasks.AndroidWorld contains 116 dynamic tasks across 20 Android apps, while OSWorld contains 369 web and desktop tasks.
  • C.2. Baselines: EAPO is compared with six baselines implemented from publicly available implementations.The baselines include Min-p, OverRIDE, GRPO, DAPO, GiGPO, and LAMER.
  • C.2. Baselines: The baselines cover dynamic decoding, critic-free group-based reinforcement learning, multidimensional credit assignment, and meta-reinforcement learning with sampled trajectories.Specifically, Min-p and OverRIDE modify decoding; GRPO, DAPO, and GiGPO use group-based optimization; LAMER adapts through in-context interaction with sampled experiences.

C.3. Implementation Details … D.7.5. EXPLORATION DEGREE

The paper details a consistent implementation and agent interaction protocol, then evaluates EAPO across comparisons, transfer, hyperparameter studies, convergence, reward components, grouping, and exploration behavior. These results indicate adaptive exploration improves performance while learning when to stop exploring.

  • C.3. Implementation Details: The reward model shares the policy model, with identical hyperparameters across datasets, and experiments use PyTorch 2.8.0, verl, Ubuntu 22.04.4 LTS, and 8 NVIDIA H800 GPUs.The implementation is built on the open-source verl framework.
  • C.4. Pseudocode: EAPO alternates supervised fine-tuning and reinforcement learning by sampling trajectories, grouping state-action transitions, optimizing the reward model, and updating the policy.Algorithm 1 initializes rollback data, a reward model, and policy network before these optimization stages.
  • C.5. Instruction Template: The instruction template combines fixed system prompts, basic task instructions, and action guidance requiring one correctly formatted action based on the goal, UI state, and action history.The protocol explicitly separates reasoning, exploration, memory, and action outputs into tagged sections.
  • D. Additional Results; D.1. Comparison with More Models; D.2. Comparison with Training Methods: A 2B EAPO model outperforms larger general LLMs, while training comparisons report improvements of 20%−60%, faster convergence, and greater stability, especially on complex GUI tasks.The reported findings attribute these gains to obtaining dynamic information through adaptive exploration.
  • D.3. Generalization: Models trained on AndroidWorld with EAPO retain strong performance across OSWorld task domains, with only slight degradation compared with models trained directly on OSWorld.This cross-environment evaluation tests transfer to unseen scenarios.
  • D.4. Impact of Discount γ: Increasing γ from 0.5 to 1.0 encourages more exploration but can cause redundant information and longer trajectories; the method sets γ = 0.9 as a trade-off.Average steps increase with γ, especially as γ approaches 1.0.
  • D.5. Impact of Sampling Group Size G; D.6. Impact of KL loss coefficency λ: Larger sampling groups generally stabilize optimization but eventually saturate with added computational overhead, while increasing λ first improves and then degrades performance.Small λ permits aggressive deviation from the reference model and unstable updates; stronger KL regularization stabilizes training and preserves prior knowledge.

D.7.6. ABLATION · D.7.7. RUN-TIME · E. Case Study

The ablations show that exploratory reward, exploration-aware grouping, and format reward are essential for effective, stable exploration and downstream task completion. EAPO incurs manageable runtime overhead from its additional reward-model and exploration-related computation, while the case study demonstrates exploration, rollback, and successful execution.

  • D.7.6. ABLATION: Without exploratory reward, agents cannot evaluate exploratory-action usefulness, causing ineffective exploration, premature convergence, and degraded performance across all environments.The passage also notes that EAPO loses competitiveness without SFT because exploration accuracy and rollback capability are coupled.
  • D.7.6. ABLATION: Without exploration-aware grouping, exploratory and task-execution actions mix, causing exploration degree and task performance to rise initially before collapsing.Exploratory actions become increasingly underestimated during training, impairing long-term optimization.
  • D.7.6. ABLATION: Without format reward, agents fail to follow the predefined structure for exploration signals and memory updates, limiting information reuse for decision-making and task completion.This prevents effective organization, storage, and reuse of information obtained through exploration.
  • D.7.7. RUN-TIME: EAPO’s total runtime is approximately twice that of other group-based methods because it trains an additional reward model, but remains practically manageable and scalable.The comparison uses the same model size on 8 NVIDIA H800 GPUs.
  • D.7.7. RUN-TIME: GRPO optimization time remains unchanged from the baseline; added computation mainly comes from reward-model training, transition grouping, and reward inference for policy-advantage estimation.These components account for the method’s additional computational cost.
  • D.7.7. RUN-TIME: A discount factor γ limits exploration-related inference overhead to a modest increase of at most 10 in average step.The discount reflects that exploration benefits require observing a new state and subsequently using the acquired information.
  • E. Case Study: In the OSWorld case study, the agent explores at Step 5, detects an incorrect state and rolls back at Step 6, then selects the correct action at Step 7.The agent incorporates the acquired information to progress toward successful task completion.
Loading 2605.08978v2…