Uncertainty-Aware Reward-Free Exploration with General Function Approximation
Junkai Zhang, Weitong Zhang, Dongruo Zhou, Quanquan Gu
TL;DR
The paper tackles the challenge of reward-free exploration under general function approximation by introducing GFA-RFE, which employs an uncertainty-aware intrinsic reward and variance-weighted learning to address heterogeneity in collected data. The method leverages a two-phase framework (exploration without rewards, followed by planning with rewards) and uses the generalized eluder dimension to characterize sample complexity. Theoretical results show a near-optimal sample complexity of $\tilde{O}(H^2 \log N_{\mathcal{F}}(\epsilon) \dim(\mathcal{F}) / \epsilon^2)$ for achieving $\epsilon$-optimal policies, improving upon prior reward-free approaches. Empirically, GFA-RFE demonstrates strong performance on the DeepMind Control Suite within URL benchmarks, matching or surpassing state-of-the-art unsupervised exploration methods, and confirms the value of uncertainty-aware weighting in both exploration and planning.
Abstract
Mastering multiple tasks through exploration and learning in an environment poses a significant challenge in reinforcement learning (RL). Unsupervised RL has been introduced to address this challenge by training policies with intrinsic rewards rather than extrinsic rewards. However, current intrinsic reward designs and unsupervised RL algorithms often overlook the heterogeneous nature of collected samples, thereby diminishing their sample efficiency. To overcome this limitation, in this paper, we propose a reward-free RL algorithm called \alg. The key idea behind our algorithm is an uncertainty-aware intrinsic reward for exploring the environment and an uncertainty-weighted learning process to handle heterogeneous uncertainty in different samples. Theoretically, we show that in order to find an $ε$-optimal policy, GFA-RFE needs to collect $\tilde{O} (H^2 \log N_{\mathcal F} (ε) \mathrm{dim} (\mathcal F) / ε^2 )$ number of episodes, where $\mathcal F$ is the value function class with covering number $N_{\mathcal F} (ε)$ and generalized eluder dimension $\mathrm{dim} (\mathcal F)$. Such a result outperforms all existing reward-free RL algorithms. We further implement and evaluate GFA-RFE across various domains and tasks in the DeepMind Control Suite. Experiment results show that GFA-RFE outperforms or is comparable to the performance of state-of-the-art unsupervised RL algorithms.