PaperlayerView exact revision on arXiv

Source-linked AI summary

On the Importance of Exploration for Generalization in Reinforcement Learning

Yiding Jiang, J. Zico Kolter, Roberta Raileanu

arXiv:2306.05483v1cs.LG

TL;DR

Deep RL generalization methods have largely focused on representation learning, while exploration’s role on common benchmarks has lacked convincing empirical demonstration. The paper proposes EDE, which uses an ensemble of Q-value distributions to explore states with high epistemic uncertainty. It reports state-of-the-art results for a value-based approach on Procgen and Crafter.

  • Problem

    Deep RL generalization approaches have mostly focused on representation learning, and a convincing empirical demonstration of exploration’s role on common generalization benchmarks has been lacking.

    Source: §1 · p. 1

  • Method

    EDE uses an ensemble of Q-value distributions to encourage exploration of states with high epistemic uncertainty.

    Source: §1 · p. 2

  • Results

    The paper reports that EDE is the first model-free value-based method to achieve state-of-the-art performance on Procgen and Crafter in sample efficiency and generalization.

    Source: §1 · p. 2, §7 · p. 9

  • Takeaways & Limitations

    Exploring training environments can help agents acquire knowledge useful for decisions in other, potentially unseen, environments.

    Source: §1 · p. 2

    The paper focuses on exploration as one approach to CMDP generalization and does not claim it is the only approach.

    Source: §3 · p. 5

Abstract

from arXiv · show

Existing approaches for improving generalization in deep reinforcement learning (RL) have mostly focused on representation learning, neglecting RL-specific aspects such as exploration. We hypothesize that the agent's exploration strategy plays a key role in its ability to generalize to new environments. Through a series of experiments in a tabular contextual MDP, we show that exploration is helpful not only for efficiently finding the optimal policy for the training environments but also for acquiring knowledge that helps decision making in unseen environments. Based on these observations, we propose EDE: Exploration via Distributional Ensemble, a method that encourages exploration of states with high epistemic uncertainty through an ensemble of Q-value distributions. Our algorithm is the first value-based approach to achieve state-of-the-art on both Procgen and Crafter, two benchmarks for generalization in RL with high-dimensional observations. The open-sourced implementation can be found at https://github.com/facebookresearch/ede .

1 Introduction

Generalization in deep RL remains challenging, particularly for value-based methods, while most existing approaches focus on representation learning. The paper argues that exploration can also support generalization by helping agents learn information useful in other environments.

1 Introduction

  • Existing deep RL algorithms struggle to generalize across unseen contextual MDPs, especially value-based methods, despite competitive performance in singleton MDPs.Most approaches treat generalization as a representation-learning problem and apply supervised-learning regularization, neglecting agents’ ability to collect data through exploration.Source: §1 · p. 1
  • A convincing empirical demonstration that exploration supports generalization on common RL benchmarks had been lacking outside niche areas such as task-agnostic RL and meta-RL [122] [64].The paper identifies this as an unresolved connection and studies exploration as a factor influencing generalization.Source: §1 · p. 1
  • Exploration can gather information about parts of an environment that are not needed for its training-optimal policy but may be useful at test time.Trajectories suboptimal in one MDP may be optimal in another MDP from the same family.Source: Figure 1 · p. 2, §1 · p. 2
  • EDE uses an ensemble of Q-value distributions to encourage exploration of states with high epistemic uncertainty, distinguishing it from irreducible aleatoric uncertainty.The paper reports state-of-the-art performance for a model-free value-based method on Procgen and Crafter in sample efficiency and generalization.Source: §1 · p. 2
  • The paper identifies exploration as a factor in generalization, proposes an epistemic-uncertainty-based method, and reports state-of-the-art performance on Procgen and Crafter.These are the paper’s stated contributions.Source: §1 · p. 2

2 Background

The background defines the RL and contextual-MDP setting in which a policy is evaluated across a distribution of environments. It frames generalization in terms of test return under finite interaction and explains the value-based and distributional quantities used in the paper.

2 Background

  • A value-based method derives its policy from the state-action value function, while distributional RL models a distribution over returns rather than a single value estimate.The Q-function is defined as expected return conditional on the starting state and action, and the greedy policy selects the action with the largest Q-value.Source: §2 · p. 3
  • A CMDP consists of MDPs sharing state and action spaces but differing in rewards, transitions, and initial-state distributions, with a structured distribution over environments.Training uses a finite or infinite collection of MDPs drawn from that distribution.Source: §2 · p. 3
  • The objective in a CMDP is to find one policy maximizing expected discounted return across trajectories and MDPs drawn from the environment distribution.The paper calls this quantity the test return.Source: §2 · p. 3
  • With finite interactions in training environments, generalization concerns the performance a learning algorithm achieves on the test return; the setting assumes a policy can be optimal across the environment distribution.If the training set is infinite, the challenge is sample-efficient optimization; with a finite set, it also includes generalization to unseen environments.Source: §2 · p. 3

3 Generalization in a Tabular CMDP

In a tabular CMDP, exploration strategy affected test performance even when methods quickly solved the training environment. The experiments motivate the paper’s view that exploration can help agents recover in test environments by exposing them to states rarely visited during training.

3 Generalization in a Tabular CMDP

  • The toy CMDP illustrates that exploring training environments can help both find their rewards and support decisions in new environments at test time.In Procgen-like settings, exploring states rarely visited during training may help an agent recover from suboptimal states at test time.Source: §3 · pp. 3–4
  • The experiment used a 5 × 5 grid with deterministic transitions, a fixed training start, and a different test start; only the initial-state distribution differed.The terminal state rewarded +2, other steps rewarded −0.04, and episodes were capped at 250 steps.Source: §3 · p. 3
  • The comparison used tabular Q-learning with ε-greedy exploration and uncertainty-driven UCB exploration, avoiding function-approximation confounds.UCB selects actions using Q-values plus a bonus based on log(t)/N(s,a), with c as the exploration coefficient.Source: §3 · p. 4
  • Greedy exploration generalized worst and had the highest variance, while UCB reliably found the test-optimal policy with negligible final-return variance.All three methods quickly achieved optimal training return; similar observations came from simulations with changing dynamics.Source: §3 · p. 4, Figure 2 · p. 4
  • The simple UCB extension used for the didactic example does not retain the bandit algorithm’s regret guarantee.The authors describe it as effective for illustration and distinguish it from a more complicated MDP extension.Source: §3 · p. 4
  • The paper does not claim exploration is the only way to improve RL generalization; it describes exploration as complementary to other approaches.The stated focus is exploration for generalization in contextual MDPs.Source: §3 · p. 5

4 Exploration via Distributional Ensemble

EDE combines distributional Q-value estimates and deep ensembles to target epistemic uncertainty for exploration. Its exploration choices include UCB and Thompson sampling, with actor-specific coefficients used to diversify UCB experience.

4 Exploration via Distributional Ensemble

  • EDE combines deep ensembles with distributional RL to estimate epistemic uncertainty and encourage exploration of states with high reducible uncertainty.The method builds on Clements et al. [19] but replaces MAP sampling with a deep ensemble; ensemble heads use different minibatches and random initializations.Source: §4 · pp. 5–7
  • Distributional Q-value estimates distinguish epistemic uncertainty from aleatoric uncertainty, which reflects irreducible stochasticity in the environment.The method treats uncertainty in expected Q-values as the quantity relevant to exploration.Source: §4 · pp. 5–6
  • EDE uses UCB or Thompson sampling to turn its uncertainty estimates into exploration policies, with UCB performing better than Thompson sampling on Procgen with parallel workers.The preferred choice depends on the task and setting.Source: §4 · pp. 6–7
  • On Procgen, EDE greatly outperforms baselines on median and IQM and is competitive with state-of-the-art policy optimization methods on mean and optimality gap.For the chosen evaluation configuration, optimality gap equals 1-mean.Source: Figure 4 · p. 7
  • For UCB, EDE assigns different exploration coefficients to actors to diversify collected data; this technique is called temporally equalized exploration (TEE).TEE’s coefficient-distribution shape is controlled by λ and α.Source: §4 · pp. 6–7

5 Experiments

On Procgen, EDE improves on value-based baselines and is competitive with a leading policy-optimization method; on Crafter, it scores significantly higher than the presented baselines. The experiments also indicate that epistemic uncertainty and diverse training data both matter for Procgen performance.

5.1 Procgen

  • EDE achieves almost 3 times the median score of QR-DQN and more than 2 times its IQM on Procgen, while remaining competitive with IDAAC on mean and optimality gap.The evaluation uses the standard 25M-step, easy-mode benchmark and aggregates min-max normalized scores across games and five seeds.Source: §5.1 · p. 7
  • On Procgen, combining epistemic uncertainty with diverse data performs best: QR-DQN+TEE trails QR-DQN+UCB, which in turn trails EDE.The comparisons show that either ingredient alone is insufficient to match their combination.Source: §5.1 · pp. 7–8
  • Methods targeting total uncertainty, including Bootstrapped DQN and DQN+UCB, perform worse than DQN, while QR-DQN+UCB outperforms QR-DQN and DQN.The authors attribute the contrast to the importance of using distributional RL to measure epistemic uncertainty in CMDPs.Source: §5.1 · p. 8
  • On Procgen, ϵz-greedy and NoisyNet remain below EDE; UA-DQN also performs worse than QR-DQN+UCB, and none of the compared methods significantly exceeds another on IQM.The exploration-baseline comparison is shown in Figure 5.Source: §5.1 · p. 8, Figure 5 · p. 8

5.2 Crafter

  • On Crafter, EDE scores significantly higher than the presented baselines, including DreamerV2, and is competitive with LSTM-SPCNN.The comparison follows 1M training steps over 10 runs; Crafter requires generalization across procedurally generated training environments, not to new environments.Source: §5.2 · pp. 8–9, Table 1 · p. 8

6 Related Works

Prior generalization work spans regularization, representation learning, and other approaches, while exploration methods have often targeted sample efficiency in singleton MDPs or sparse-reward settings. This paper positions exploration as a route to faster training and better generalization for value-based methods in contextual MDPs.

6 Related Works

  • Prior RL generalization work includes regularization, representation learning, information-theoretic approaches, curriculum learning, planning, and diverse policies.Raileanu and Fergus [89] identify value overfitting in CMDPs, but their policy-value decoupling approach cannot be applied to value-based methods whose policy is defined by the Q-function.Source: §6 · p. 9
  • Common exploration approaches include noise, bootstrapping, ensembles, uncertainty estimation, and distributional RL, often targeting sample efficiency or temporally extended exploration in singleton MDPs.Other methods focus on sparse-reward CMDPs, while task-agnostic RL explores without reward before learning a downstream task.Source: §6 · p. 9
  • The paper focuses on exploration as a means to improve value-based agents’ training speed and generalization to unseen MDPs.It also draws on distributional RL as a way to avoid value overfitting in CMDPs [89].Source: §6 · p. 9

7 Conclusion

The paper argues that exploration can improve both policy learning in training environments and knowledge for decisions in unseen environments. It proposes EDE, which targets epistemic uncertainty, and reports state-of-the-art performance on Procgen and Crafter; its ensemble increases computational cost.

7 Conclusion

  • EDE is the first value-based method reported to achieve state-of-the-art performance on both Procgen and Crafter.These are benchmarks for generalization in RL with high-dimensional observations.Source: §7 · p. 9
  • Tabular experiments indicate that exploration can help find training-environment policies and acquire knowledge useful in unseen environments.The conclusion suggests exploration is important for RL algorithms trained and tested on CMDPs.Source: §7 · p. 9
  • EDE encourages exploration of states with high epistemic uncertainty using deep ensembles and distributional RL to distinguish epistemic from aleatoric uncertainty.Source: §7 · p. 9
  • EDE is more computationally expensive because of its ensemble, and the authors expect it could benefit from more efficient uncertainty estimation.Source: §7 · pp. 9–10

A Sample Complexity of Q-Learning

The analysis relates Q-learning sample complexity to how well a behavioral policy visits state-action pairs. It also uses a pruned sub-MDP argument to explain faster learning near training initial states, while noting the limits of the analysis’s continuous-MDP proxy.

Sample-complexity analysis

  • Q-learning sample complexity is inversely proportional to the visitation probability of the hardest-to-reach state, so more uniform exploration can reduce Q-value error across state-action pairs.The analysis connects lower Q-value error to generalization across initial state distributions.Source: A Sample Complexity of Q-Learning · p. 17

Effect of resetting

  • Because resets concentrate occupancy near the initial-state distribution, Q-values farther away may converge more slowly, explaining weaker performance from those states.The stationary occupancy is independent of the initial distribution, but may remain close to the reset distribution when exploration is inefficient.Source: A Sample Complexity of Q-Learning · p. 17

Analysis scope

  • The continuous-MDP proxy can be ergodic when the original episodic MDP is, but generally does not preserve its optimal policy, though it may be close for common MDPs.The sample-complexity argument is framed around Q-learning with a fixed behavioral policy and an ergodic environment.Source: A Sample Complexity of Q-Learning · p. 17

Sub-MDP argument

  • A sub-MDP prunes state-action pairs visited less often than a threshold, and can have better sample complexity than the full MDP.If the optimal policy’s visited pairs remain in the pruned set, the sub-MDP recovers the optimal policy from the specified initial distribution.Source: A Sample Complexity of Q-Learning · p. 18

Sample-complexity analysis

  • UCB-style exploration targets low-count state-action pairs, which can increase minimum visitation and improve sample complexity; ε-greedy does not explicitly target visitation and may poorly cover large spaces.The text notes that episodic RL state-of-the-art algorithms generally use UCB-based exploration.Source: A Sample Complexity of Q-Learning · p. 18

B Case Study: Simplified Bigfish

The simplified Bigfish example illustrates how exploring a rarer, suboptimal training trajectory can provide useful behavior for a different test environment. The accompanying experiments compare exploration strategies, policy-gradient and full-trajectory updates, and Sarsa(λ) across generalization settings.

B Case Study: Simplified Bigfish

  • In Bigfish, uncertainty-aware exploration can sample a rarer, suboptimal training trajectory that shares a state with the test trajectory, helping the agent choose an action that avoids death and earns reward.The example is simplified; the authors note that real state sequences need not match exactly and that not all CMDPs have this structure.Source: B Case Study: Simplified Bigfish · pp. 19–20

C.1 Generalization to Different Initial States

  • The C.1 comparison measures mean test return across exploration hyperparameters in a grid setting where train and test differ only in their initial-state distributions.Source: Figure 8 · p. 20, §C.1 · p. 20

C.2 Policy gradient

  • Across exploration hyperparameters, UCB and policy gradient are relatively robust, while ε-greedy varies greatly and remains worse even at very large ε.Policy gradient performs between Greedy and UCB, with better convergence and lower variance than Greedy but worse performance than UCB.Source: Figure 8 · p. 20, §C.2 · p. 21

C.3 Full trajectory update and Sarsa(λ)

  • Both policy gradient and Sarsa(λ) outperform Greedy, but the authors suggest whole-trajectory updates may help propagate learning to low-density states; their remaining gap from UCB suggests distinct benefits from exploration.The proposed mechanism contrasts whole-trajectory updates with off-policy updates that require repeated sampling of a state-action pair.Source: §C.3 · pp. 21–22

C.3.1 Details of Sarsa(λ)

  • With its best hyperparameter configuration, Sarsa(λ) outperforms Greedy and policy gradient, but remains below UCB and has much larger variance.Across exploration coefficients, larger values tend to yield higher average return and lower test-performance variance.Source: §C.3.1 · p. 22

C.4 Generalization to Different Dynamics

A windy-grid experiment examines whether exploration learned under familiar dynamics helps when test-time transitions change. The section also describes QR-DQN’s quantile representation and EDE’s exploration and ensemble-update procedures.

C.4 Generalization to Different Dynamics

  • Test-time wind makes the optimal expected return generally worse and more variable, while the authors report results similar to the preceding initial-state-shift experiment.The agent starts at (0, 3); at test time, a gust blows it two units in the y-direction with 40% probability.Source: §C.4 · p. 22, Figure 11 · p. 23, Figure 12 · p. 23
  • Exploration is presented as useful under dynamics shifts without assumptions about the perturbation type or magnitude and without simulator access.The wind perturbation is used to illustrate this setting.Source: §C.4 · pp. 22–23

D.1 Quantile Regression DQN

  • QR-DQN represents each state-action value distribution with predicted quantiles and learns them using quantile regression.Its target distribution is bootstrapped through the distributional Bellman operator, and quantile Huber loss stabilizes optimization.Source: §D.1 · pp. 23–24

E Algorithmic Details

  • EDE estimates epistemic uncertainty and expected Q-values from ensemble value distributions, then selects actions using Thompson sampling or UCB.Each ensemble head uses an independently sampled replay minibatch, while one randomly selected head alone updates the shared feature extractor on each update.Source: E Algorithmic Details · p. 25
  • EDE processes multiple minibatches per update, increasing computational cost; the authors report that deep ensembles are important for accurate uncertainty estimation in CMDPs.They report EDE with five members runs about 2.4 times slower than base QR-DQN.Source: E Algorithmic Details · p. 25

F Procgen Results

EDE’s Procgen results vary across games: it is more sample-efficient on a subset, while on others it is comparable to policy optimization or fails to train. The accompanying analysis examines how training performance relates to test performance.

F Procgen Results

  • EDE is more sample-efficient on a subset of Procgen games, but is on par with policy optimization or fails to train on others.Hyperparameters for IDAAC are tuned per game, whereas EDE’s are tuned only on bigfish; the individual-game figure notes that chaser, maze, heist, and miner remain challenging for value-based methods.Source: F Procgen Results · p. 27, Table 2 · p. 27, Figure 13 · p. 28

F.1 Additional Ablation

  • DQN+UCB+TEE performs much worse than the other ablations; the authors attribute this to the unfavorable performance of DQN+UCB.The supplied figure passage does not report numerical scores.Source: Figure 14 · p. 28

G Additional Analysis

  • The method improves train and test performance on two CMDP benchmarks but does not fully close the generalization gap.The paper examines other potential sources of poor generalization, particularly for value-based methods.Source: G Additional Analysis · p. 29

G.1 Alternative hypothesis

  • Across many Procgen environments, algorithm training performance is strongly correlated with test performance, suggesting better optimization may also improve generalization.The analysis measures Spearman rank correlation across seven representative algorithms.Source: §G.1 · p. 29
  • Positive-correlation games show an almost linear train-test relationship with slope near one, while negative-correlation games are often saturated or too challenging for current methods.The analysis uses PPO-based algorithms and notes that the pattern may differ for value-based methods; PLR is an outlier in both plots.Source: §G.1 · pp. 29–30

G.2 Other Sources of Poor Generalization

The analysis finds that single-environment sample efficiency and partial observability are not the main explanations for value-based methods’ poor CMDP generalization. It identifies value overfitting and uncertainty estimation as relevant, alongside the role of exploration.

G.2 Other Sources of Poor Generalization

  • On single Procgen environments, DQN and PPO show no significant average performance gap, and DQN is usually slightly more sample-efficient.The comparison covers three seeds across ten environments where both methods score above 0.5.Source: §G.2 · p. 30
  • Providing task IDs to the Q-network does not improve training performance, suggesting partial observability may not be the main problem in these environments.The authors note that optimization can remain challenging without proper exploration even when task identity is available.Source: §G.2 · p. 30
  • Distributional RL can alleviate value overfitting by modeling multiple possible Q-values for a state, while the largest reported gains come from more sophisticated exploration.The distribution captures aleatoric uncertainty associated with unobserved contexts in CMDPs.Source: §G.2 · p. 30

H Uncertainty

  • Using only epistemic uncertainty for exploration performs significantly better than using both epistemic and aleatoric uncertainty on training and test performance.The authors describe epistemic uncertainty as reducible through data and aleatoric uncertainty as intrinsic to the environment.Source: H Uncertainty · p. 31
  • Three ensemble heads achieve comparable performance, while the authors report that deep ensembles estimate epistemic uncertainty better than MAP sampling in their comparison.The comparison uses three-member ensembles and is reported as an empirical finding; the authors identify further understanding as future work.Source: H Uncertainty · p. 31

I Architecture and Hyperparameters

The paper specifies the Procgen and Crafter configurations and reports sensitivity tests for EDE’s hyperparameters and feature-extractor updates. It also reports the hardware used for the experiments.

I.1 Procgen

  • Procgen uses an IMPALA ResNet feature extractor and M MLP heads, with hyperparameters adapted from prior work except for new algorithmic components.The Procgen comparison aims to keep hyperparameters as similar as possible while making changes for fairness; EDE-specific values are selected on bigfish.Source: §I.1 · p. 32

I.2 Crafter

  • Crafter uses the Rainbow experiment configuration with QR-DQN and EDE, a larger minibatch, and no prioritized replay.EDE uses Thompson sampling rather than UCB in this single-thread setting and does not use TEE.Source: §I.2 · p. 33

J Sensitivity Study

  • Test performance is slightly sensitive to φ but not sensitive to ensemble size M; training the feature extractor with every ensemble member’s gradient lowers performance.The study reports these sensitivity results on selected games, using one seed.Source: J Sensitivity Study · p. 34, Figure 17 · p. 35
  • The authors observe no significant performance difference across the tested TEE values of α and λ.The experiments were conducted on 2080 and V100 GPUs and took approximately 250 GPU days.Source: J Sensitivity Study · p. 34, K Hardware · p. 34, Figure 18 · p. 35
Loading 2306.05483v1…