Decoupling Exploration and Policy Optimization: Uncertainty Guided Tree Search for Hard Exploration

Paper: 2603.22273 Authors: Zakaria Mhammedi, James Cohan Categories: cs.LG

Abstract

Efficient autonomous exploration remains a major unsolved problem in reinforcement learning. The dominant paradigm addresses this by training agents with intrinsic motivation, maximizing a composite objective of extrinsic and intrinsic rewards. This paper argues that using RL machinery solely to expand state coverage is inefficient and proposes a new paradigm that explicitly separates exploration from exploitation. The method, called Go-With-Uncertainty (GowU), uses a tree-search strategy inspired by the Go-With-The-Winner (GWTW) algorithm, paired with epistemic uncertainty to systematically drive exploration. By removing the overhead of policy optimization during exploration, GowU explores an order of magnitude more efficiently than standard intrinsic motivation baselines on hard Atari benchmarks. The discovered trajectories are then distilled into deployable policies using supervised backward learning, achieving state-of-the-art scores by a wide margin on Montezuma’s Revenge, Pitfall!, and Venture. The authors further demonstrate the generality of their framework by solving MuJoCo Adroit dexterous manipulation and AntMaze tasks from image observations with sparse rewards, without expert demonstrations or offline datasets.

Key Contributions

• New exploration paradigm: Proposes decoupling exploration from policy optimization entirely, using uncertainty-guided tree search (GowU) instead of intrinsic-reward RL for the exploration phase
• Order-of-magnitude efficiency gains: GowU discovers high-reward trajectories using far fewer environment interactions than standard intrinsic motivation baselines
• State-of-the-art Atari results: Achieves 182,672 on Montezuma’s Revenge, 97,980 on Pitfall!, and 5,190 on Venture after policy distillation, surpassing all baselines by wide margins
• First sparse-reward visual solutions for MuJoCo tasks: Solves Adroit dexterous manipulation (door, hammer, relocate) and AntMaze directly from pixels without expert demonstrations, which has not been previously achieved

Methodology

Core Idea: Separate Exploration from Exploitation

The key insight is that while policy optimization is necessary for precise task execution, using RL to maximize intrinsic rewards for state coverage introduces unnecessary overhead. The intrinsic reward signal is non-stationary, making iterative policy updates inherently sample-inefficient. GowU instead treats sparse-reward environments as pure exploration problems first, then learns deployable policies from discovered trajectories in a separate phase.

Go-With-Uncertainty (GowU) Algorithm

GowU adapts the Go-With-The-Winner (GWTW) principle from randomized search theory. It maintains a population of particles (e.g., 32 per group) that explore the state space in parallel through a state-lineage tree:

1. Evolutionary rollout: Each particle interacts with its own environment instance for a random number of steps. If a particle enters a “dead” state (e.g., loses a life), it is pruned
2. Winner selection: Among surviving particles, the one with the highest cumulative reward is selected as the winner, with epistemic uncertainty (measured via RND) used as a tie-breaker. In sparse-reward settings, uncertainty is often the primary selection criterion
3. Particle redistribution: Dead particles are reset to the winner’s state, redirecting computational effort toward the most promising frontier
4. Group consolidation: Periodically, the entire population is collapsed to the single best particle’s state, ensuring subsequent search begins from the frontier of discovery
5. Failure recovery: If all particles die, they collectively roll back to ancestor states in the lineage tree, enabling exploration of alternative trajectories

Multiple groups (e.g., 4 for Atari) run in parallel with a shared uncertainty estimator, enabling implicit coordination where updates from one group discourage others from revisiting the same states.

Policy Distillation via Backward Learning

After exploration discovers high-reward trajectories, a backward learning curriculum distills them into deployable policies. The agent starts training near the end of the demonstration and progressively moves backward as it masters each segment. This creates a natural curriculum that significantly simplifies the RL problem.

Particle Policies

Individual particles use simple reward-free diversity policies based on bootstrapped zero-target networks. Importantly, ablation studies show that even fixed random policies per particle perform comparably to learned ensembles, confirming that the exploration power lies in the population management mechanisms, not individual particle behavior.

Results

Atari Hard-Exploration Benchmarks

GowU dramatically outperforms all baselines after policy distillation:

Game	GowU	Go-Explore	RND	MEME	BYOL-Hindsight
Montezuma’s Revenge	182,672	43,791	8,152	9,429	~14,517
Pitfall!	97,980	6,945	-3	7,821	~16,211
Venture	5,190	2,281	1,859	2,583	~2,328

Even the worst individual run across all games remains substantially above the best baseline. On Montezuma’s Revenge, GowU achieves over 4x the score of Go-Explore. On Pitfall!, the improvement is over 14x compared to Go-Explore.

MuJoCo Continuous-Control Tasks

GowU achieves near-perfect success rates on Adroit tasks from pixel observations with sparse rewards: 99.9% on hammer, 96.4% on door, and 93.9% on relocate. On AntMaze, the mean success rate is 86.3% with the best run reaching 95.1%. No prior method has achieved such results from pixels with sparse rewards and without expert demonstrations.

Ablation Studies

• Uncertainty is critical: Removing the uncertainty estimator causes exploration to fail entirely on Montezuma’s Revenge
• Group consolidation accelerates progress: Disabling it slows down exploration and increases memory usage
• Population management drives exploration: Fixed random policies per particle match learned ensemble policies, confirming that winner selection, pruning, and rollback are what matter
• Hyperparameter randomization provides robustness: Sampling parameters from reasonable ranges at each iteration is sufficient without precise tuning

Takeaways

This paper makes a compelling case that the dominant paradigm of using RL with intrinsic motivation for exploration is fundamentally inefficient. By cleanly separating exploration (via uncertainty-guided tree search) from exploitation (via backward learning), GowU achieves dramatic improvements across both discrete and continuous domains. The approach requires simulator access with reset capability during training, but produces policies that operate without any simulator at inference time. The framework is modular and agnostic to the specific uncertainty metric and particle policies, suggesting broad applicability. The authors note that this paradigm could extend to complex reasoning tasks where particles are LLM instances exploring reasoning chains, and to scaling up open-ended robotic learning in simulation.