Source-linked AI summary
Model-Based Active Exploration
Pranav Shyam, Wojciech Jaśkowski, Faustino Gomez
TL;DR
Efficient exploration in reinforcement learning remains difficult because reactive methods reward agents only after fortuitous encounters with novelty. MAX actively plans toward novel transitions using ensemble-model disagreement as a Bayesian novelty estimate. It outperforms reactive exploration in discrete and continuous settings, while its derivation omits a second-order utility effect and it trades computational efficiency for data efficiency.
Problem
Efficient exploration in large, high-dimensional reinforcement-learning environments remains unsolved, while reactive novelty rewards can cause inefficient over-commitment.
Method
MAX uses a bootstrap ensemble of learned forward models and plans in a surrogate exploration MDP toward transitions whose predicted futures disagree.
Results
MAX is significantly more efficient than reactive baselines in discrete environments, reaches the far end of Ant Maze in 40 episodes and 12k steps, and yields superior exploited performance in Half Cheetah.
Takeaways & Limitations
MAX provides a generic active-exploration method that distinguishes learnable from unlearnable unknowns and supports task-agnostic model reuse.
Takeaways & Limitations
The derivation omits second-order effects whereby training on encountered transitions changes the utility of remaining transitions, and MAX is less computationally efficient than the baselines.
Abstract
from arXiv · showhide
Efficient exploration is an unsolved problem in Reinforcement Learning which is usually addressed by reactively rewarding the agent for fortuitously encountering novel situations. This paper introduces an efficient active exploration algorithm, Model-Based Active eXploration (MAX), which uses an ensemble of forward models to plan to observe novel events. This is carried out by optimizing agent behaviour with respect to a measure of novelty derived from the Bayesian perspective of exploration, which is estimated using the disagreement between the futures predicted by the ensemble members. We show empirically that in semi-random discrete environments where directed exploration is critical to make progress, MAX is at least an order of magnitude more efficient than strong baselines. MAX scales to high-dimensional continuous environments where it builds task-agnostic models that can be used for any downstream task.
1. Introduction
MAX reframes exploration as an active planning problem: an ensemble of forward models estimates novelty through predictive disagreement and directs the agent toward informative transitions. It improves exploration efficiency in discrete settings and scales to task-agnostic modeling in continuous environments.
- Motivation: Reactive exploration rewards agents after they encounter novelty, causing over-commitment as novelty fades and the reward must be unlearned.MAX instead treats exploration as seeking novelty based on predicted action consequences.
- Motivation: Bayesian exploration measures transition novelty through disagreement among probable models’ next-state predictions.This provides an internal estimate for selecting action sequences likely to produce interesting transitions.
- Method: MAX uses an ensemble of learned forward dynamics models to identify learnable uncertainty and plans visits that resolve predictive conflicts.Random environmental noise appears as confusion shared by all models rather than disagreement, so it does not drive exploration.
- Method: In discrete environments, MAX evaluates novelty with Jensen-Shannon Divergence, while continuous environments use Jensen-Rényi Divergence because JSD is intractable there.The divergence compares predicted next-state distributions across ensemble models.
- Scope: For pure exploration, MAX collects task-agnostic data that can later support multiple downstream tasks, including tasks unknown during exploration.The paper focuses on exploration without external reward followed by exploitation.
- Results: In Ant Maze, MAX reaches the far end of the U-shaped maze in 40 episodes and 12k steps, while reactive baselines reach only about halfway.The continuous-environment evaluation demonstrates scaling to high-dimensional exploration.
2. Model-Based Active Exploration
MAX turns pure exploration into an internal MDP whose utility rewards predicted novelty, using Bayesian model disagreement to plan informative transitions. A bootstrap ensemble approximates transition uncertainty, with Jensen-Shannon or Jensen-Rényi divergence supporting discrete and continuous environments.
- Pure exploration seeks to learn an accurate transition model while minimizing the number of observed state transitions, independently of external reward.
- Bayesian information gain reduces to Jensen-Shannon divergence among predicted next-state distributions, allowing novelty to be estimated before transitions occur.The resulting utility measures disagreement among possible transition models weighted by their probability.
- MAX defines an internal exploration MDP whose utility is the expected information gained from transitions rather than external reward.The exploration MDP starts from the agent’s current state and uses utility as its learning signal.
- Bootstrap Ensemble Approximation: A bootstrap ensemble of independently trained transition models approximates the posterior over environments using different subsets of the agent’s transition history.Agreement reflects data-supported generalization, whereas disagreement indicates unknowns not warranted by observed data.
- Large Continuous State Spaces: For continuous state spaces, MAX replaces intractable mixture-of-Gaussians Shannon-entropy calculations with Jensen-Rényi divergence using quadratic Rényi entropy.At α = 2, quadratic Rényi entropy has a closed-form solution for a mixture of Gaussians.
3. Experiments
Experiments evaluate MAX against reactive and active baselines in difficult discrete and high-dimensional continuous environments. MAX explores efficiently, distinguishes uncertainty from risk, and supports directed exploration and downstream task performance.
- 3.1. Discrete Environment: MAX explores 100% of Chain transitions in around 15 episodes, while baseline methods reach 40% in 60 episodes.The randomized chain requires exploiting environment structure because random exploration rarely reaches the far-right states.
- 3.1. Discrete Environment: In a stochastic-trap Chain variant, MAX slows down but still explores the transitions, suggesting it can distinguish environmental risk from uncertainty.The trap modifies the left-most state with stochastic dynamics.
- 3. Experiments: The experiments compare MAX with TVAX, JDRX, PERX, and random exploration across discrete and continuous settings.The baselines omit at least one feature of MAX, including active planning or uncertainty-aware utility estimation.
- 3.2. Continuous Environments: On Half Cheetah, active methods outperform reactive methods, while MAX is better than TVAX because it uses principled trajectory sampling and utility evaluation.Running is easy for random exploration, whereas flipping requires directed active exploration.
- 3.2. Continuous Environments: In Continuous Mountain Car, MAX follows spiral trajectories that progressively fill gaps between earlier paths while mapping state-space uncertainty.Colors encode average uncertainty over actions, and dotted lines represent agent trajectories.
4. Discussion
The discussion frames exploration and model bias as challenges for model-based RL, while noting assumptions and computational trade-offs in MAX's derivation and implementation.
- Model-based RL can be more efficient and general than model-free RL, but model bias may produce overoptimistic policies that fail in the environment.Bias can arise from improper generalization or poor exploration, and strong policy search may exploit model degeneracies.
- The derivation assumes that a policy's utility is the average utility of its probable transitions.This assumption omits how training on some transitions can change the utility of remaining transitions.
- In the Chain environment, the omitted second-order effect caused MAX to plan loops between pairs of uncertain states rather than visit many different uncertain states.
- MAX is less computationally efficient than the paper's baselines because it trades computational efficiency for data efficiency.
5. Related Work
Related work spans intrinsic rewards, Bayesian experimental design, directed exploration, and model-based RL, with prior approaches often limited by setting or approximation.
- MAX is related to adversarial curiosity frameworks in which modules propose experiments that generate disagreement about predicted outcomes.
- Earlier exploration methods used prediction errors, learning progress, intrinsic rewards, and directed or undirected exploration signals.
- Bayesian experimental design provides a framework for sequential experiments, but is usually restricted to linear models with Gaussian assumptions.
- Sun et al. developed Bayesian curiosity-driven exploration using learning progress and dynamic programming, but evaluated it only in small tabular MDPs with Dirichlet transition priors.
- Other Bayesian-inspired methods used Bayesian neural networks, information gain, variational inference, or pseudo-count-based exploration bonuses.
- Deep-learning model-based RL has shown promise in high-dimensional environments, while learning accurate models and exploiting them for downstream tasks remains active research.
6. Conclusion
The conclusion presents MAX as a model-based pure-exploration algorithm that actively seeks learnable unknowns while distinguishing them from unlearnable unknowns.
- MAX distinguishes learnable unknowns from unlearnable unknowns and searches for policies that actively seek the learnable ones.
- An ensemble of models lets MAX simulate and evaluate exploration policies, enabling direct optimization without actual environmental interaction.
- Experiments in difficult discrete and high-dimensional continuous environments indicate that MAX is a powerful generic exploration method.
A. Policy Evaluation
The policy-evaluation formulation defines information gain over policy-induced transitions, expands it through action utility, and reduces it to an expectation.
- Information gain IG(π) is defined for a policy π over the transitions occurring when π acts as the behavioural policy in the external MDP.
- The formulation expands policy information gain using an action utility function u(s, a).
- The action utility u(s, a) quantifies the net utility of taking action a from state s.
- The expanded formulation is reduced to an expectation over the relevant transitions.
B. Action Evaluation
MAX evaluates actions through an exploration MDP whose utility is derived from model uncertainty, using ensemble predictions to quantify novelty and guide exploration.
- The exploration utility is expanded from the Bayesian formulation using KL-divergence identities and entropy manipulations.
- The entropy difference between an average next-state distribution and the average component entropy is identified as Jensen-Shannon Divergence.JSD is also termed the Information Radius.
- The stochastic trap experiment shows that random environmental outcomes increase model disagreement near the initial state and can make MAX struggle despite its uncertainty-risk distinction.
- In the chain setting, three neural forward models predict categorical next-state distributions, and open-loop exploration policies are searched with MCTS after a random warm-up.The implementation uses a three-episode warm-up, 150 training iterations, and 25 MCTS rounds with five random trajectories per expansion step.
C.3. Baselines
The baselines compare reactive exploration bonuses with stochastic value over-estimation, using tuned neural-network configurations and multiple random seeds.
- Exploration Bonus DQN temporarily overestimates first-visit transitions by assigning them an extra bonus reward.For the simple environment, the bonuses were provided by an oracle.
- Bootstrapped DQN maintains multiple Q-value heads and randomly selects one before each episode, producing stochastic over-estimation in novel states.
- Both baselines disable ε-greedy exploration so their exploration depends solely on the respective method.Their Q-value functions use multilayer fully connected neural networks.
- All neural networks use Glorot-initialized, tanh-activated fully connected layers whose depth and width are tuned.
- Each hyper-parameter configuration is evaluated with 5 random seeds and ranked by the median area under the exploration curve.
C.4. Supplementary Figures
Supplementary analyses show that MAX’s utility tracks novel transitions, its ensemble learns the chain quickly, and its performance generally improves with more planning resources.
- High exploration-policy utility correlates with encounters with novel transitions during exemplary episodes.Figure 5 marks novel-transition encounters with red points.
- Less than 20 episodes are required for the ensemble models to correctly learn all transitions in the analyzed run.
- The ensemble models nearly agree on training-set transitions and disagree mainly on unobserved transitions, driving exploration toward those transitions.
- Figure 6 shows MAX planning for uncertain, unvisited transitions in the final two episodes and minimizing their corresponding uncertainties.
- More trajectories, ensemble models, and planning iterations generally improve performance, although MAX remains efficient with minimal settings.Figure 7 varies ensemble size, trajectories per MCTS iteration, and planning iterations; curves summarize medians over 5 seeds.
D.1. Numerically Stable Jensen-Rényi Divergence Implementation
The implementation uses a LOG-SUM-EXP trick to calculate the entropy of the mean in the Jensen-Rényi divergence computation.
- MAX uses LOG-SUM-EXP to calculate the entropy of the mean.
D.2. Experimental Setup
The continuous-environment setup normalizes model inputs and targets, repeatedly retrains models and policies during exploration, and uses ensemble-generated data for exploitation.
- 256 warm-up steps preceded exploration, after which models and SAC policies were retrained from scratch every 25 steps.SAC also received 100 steps of off-policy training using current exploration experience.
- During exploitation, the model ensemble was trained for 200 epochs every 2000 exploration steps, generating recursive 100-step rollouts for SAC training.Rollouts used 250 episodes, length 100, with 128 actors.
- Models predicted normalized state deltas rather than raw next states, and utility computations used normalized scales for stability.
- Table 4 lists the hyperparameters used for models in continuous environments.
- The running reward was defined as r_t = v_x, where v_x denotes velocity along the x axis at time t.