Why so pessimistic? Estimating uncertainties for offline RL through ensembles, and why their independence matters.

1 INTRODUCTION . Offline reinforcement learning ( RL ) , also referred to as batch RL , is the setting where we are provided with a dataset of interactions with a Markov Decision Process ( MDP ) , and the goal is to learn an effective policy without further interactions with the MDP . Offline RL holds the promise of dataefficiency through data reuse and improved safety due to minimizing the need for policy rollouts . As a result , offline RL has been a subject of significant renewed interest in the machine learning literature ( Levine et al. , 2020 ) . One common approach to offline RL , known as model-free , uses value estimation through approximate dynamic programming ( ADP ) . The predominant algorithmic philosophy with most success in ADP-based offline RL is to limit obtained policies to the support set of the available offline data , with the intuition being that such constraints would reduce inaccurate value estimates since the actions chosen by the policy are close to the observed data . A large variety of methods have been developed for enforcing such constraints , examples of which include regularizing policies with behavior cloning objectives ( Kumar et al. , 2019 ; Fujimoto & Gu , 2021 ) , only performing updates on actions observed in ( Peng et al. , 2019 ; Nair et al. , 2020 ; Wang et al. , 2020 ; Ghasemipour et al. , 2021 ) or close to ( Fujimoto et al. , 2019 ) the offline dataset , and regularizing to lower the estimated value of actions not seen in the dataset ( Wu et al. , 2019 ; Kumar et al. , 2020 ; Kostrikov et al. , 2021 ) . Support constraint forms of regularization correspond to an overly pessimistic assumption that any action outside the provided data leads to highly negative outcomes ( Buckman et al. , 2020 ) . Instead , it would be preferable if we could place more trust into the predictions of value networks beyond the training dataset . Indeed in adjacent fields of A.I. , for different tasks of interest , researchers have spent significant effort developing architectures with built-in inductive biases to help neural networks not only obtain strong accuracy on the training set , but perform well beyond the data they were trained on : convolutional and residual networks for computer vision ( He et al. , 2016 ) , recurrent networks and transformers for NLP ( Vaswani et al. , 2017 ) , graph neural networks for graph-based data ( Li et al. , 2015 ) , and Nerf and Siren for effective signal modelling ( Sitzmann et al. , 2020 ; Mildenhall et al. , 2020 ) . However , trusting the predictions of neural networks beyond the training data is a challenging ordeal . Many works have demonstrated that neural networks are prone to making high-confidence incorrect predictions outside the training set , even on i.i.d . test data ( Guo et al. , 2017 ) . This is a particularly major issue for value estimation , since errors rapidly accumulate through dynamic programming and lead to catastrophic results . Thus , in order to trust our models beyond the training set , a potentially promising approach would be to obtain high quality uncertainty estimates on predicted action values . In current supervised learning literature , “ Deep Ensembles ” and their more efficient variants have been shown to be the most effective method for uncertainty estimation ( Ovadia et al. , 2019 ) . In this work , we aim to transfer the success of ensembles for uncertainty estimation to the setting of offline RL . We begin by presenting MSG , an actor-critic offline RL algorithm leveraging an ensemble of Q-value functions . MSG trains N Q-value functions completely independent of one-another , and updates an actor network with respect to a lower confidence bound ( LCB ) on action values obtained from the ensemble . By referring to the literature on infinite-width neural networks , we theoretically demonstrate the critical importance of independence in Q-functions , a deviation from standard practice in RL which more often uses the full ensemble to compute identical target values for training each ensemble member . Our theoretical predictions are corroborated by experiments in a toy MDP , and their relevance to practical settings is also verified through benchmark experiments . In established benchmarks for offline RL , we demonstrate that MSG matches , and in the more challenging domains , significantly exceeds the prior state-of-the-art . Inspired by this success , we investigate whether the performance of MSG can be recovered through modern efficient approaches to ensembling . While we demonstrate that efficient ensembles continue to outperform current stateof-the-art batch RL algorithms , they can not recover the performance of MSG using full deep ensembles . We hope that our work highlights some of the unique challenges of uncertainty estimation in reinforcement learning , and the significant impact it can have on the performance of algorithms . We also hope that our work encourages stronger engagement from researchers specializing in uncertainty estimation techniques – who typically use computer vision and NLP tasks as benchmarks – to use offline RL as an additional testbed . Our results reveal that offline RL presents unique challenges not seen in standard supervised learning , due to the accumulation of errors over multiple back ups , and can be a valuable domain for testing novel uncertainty estimation techniques . 2 RELATED WORK . Uncertainty estimation is a core component of RL , since an agent only has a limited view into the mechanics of the environment through its available experience data . Traditionally , uncertainty estimation has been key to developing proper exploration strategies such as upper confidence bound ( UCB ) and Thompson sampling ( Lattimore & Szepesvári , 2020 ) , in which an agent is encouraged to seek out paths where its uncertainty is high . Offline RL presents an alternative paradigm , where the agent must act conservatively and is thus encouraged to seek out paths where its uncertainty is low ( Buckman et al. , 2020 ) . In either case , proper and accurate estimation of uncertainties is paramount . To this end , much research has been produced towards the end of devising provably correct uncertainty estimates ( Thomas et al. , 2015 ; Feng et al. , 2020 ; Dai et al. , 2020 ) , or , at least , bounds on uncertainty that are good enough for acting either exploratory ( Strehl et al. , 2009 ) or conservatively ( Kuzborskij et al. , 2021 ) . However , these approaches require exceedingly simple environment structure , typically either a finite discrete state and action space or linear spaces with linear dynamics and rewards . While theoretical guarantees for uncertainty estimation are more limited in practical situations with deep neural network function approximators , a number of works have been able to achieve practical success , for example using deep network analogues for count-based uncertainty ( Ostrovski et al. , 2017 ) , Bayesian uncertainty ( Ghavamzadeh et al. , 2016 ; Yang et al. , 2020 ) , and bootstrapping ( Osband et al. , 2019 ; Kostrikov & Nachum , 2020 ) . Many of these methods employ ensembles . In fact , in continuous control RL , it is common to use an ensemble of two value functions and use their minimum for computing a target value during Bellman error minimization ( Fujimoto et al. , 2018 ) . A number of works in offline RL have extended this to propose backing up minimums or lower confidence bound estimates over larger ensembles ( Kumar et al. , 2019 ; Wu et al. , 2019 ; Agarwal et al. , 2020 ) . In our work , we continue to find that ensembles are extremely useful for acting conservatively , but the manner in which ensembles are used is critical . Specifically our proposed MSG algorithm advocates for using independently learned ensembles , without sharing of target values , and this import design decision is supported by empirical evidence . The widespread success of ensembles for uncertainty estimation in RL echoes similar findings in supervised deep learning . While there exist proposals for more technical approaches to uncertainty estimation ( Li et al. , 2007 ; Neal , 2012 ; Pawlowski et al. , 2017 ) , ensembles have repeatedly been found to perform best empirically ( Lee et al. , 2015 ; Lakshminarayanan et al. , 2016 ) . Much of the active literature on ensembles in supervised learning is concerned with computational efficiency , with various proposals for reducing the compute or memory footprint of training and inference on large ensembles ( Wen et al. , 2020 ; Zhu et al. , 2019 ; Havasi et al. , 2020 ) . While these approaches have been able to achieve impressive results in supervised learning , our empirical results suggest that their performance suffers significantly in an offline RL setting compared to deep ensembles , and even naive Multi-Head ensembles ( Lee et al. , 2015 ; Osband et al. , 2016 ; Tran et al. , 2020 ) which are not considered to be as effective in the supervised learning setting ( Havasi et al. , 2020 ) . 3 METHODOLOGY . Notation Throughout this work , we represent Markov Decision Process ( MDP ) as M = 〈S , A , r , P , γ〉 , with state space S , action space A , reward function r : S ×A → R , transition dynamics P , and discount γ . In offline RL , we assume access to a dataset of interactions with the MDP , which we will represent as collection of tuples D = { ( s , a , s′ , r , term ) } N , where t is an indicator variable that is set to True when s′ is a terminal state . Two neural networks with an identical architecture , trained in an identical manner , will make different predictions outside the training set when their weights are initialized with different random draws from the initial weight distribution . So a natural question is , “ Which network ’ s predictions can we trust ? ” An answer to this question can be motivated by referring to the literature on infinite-width networks . When performing mean-squared error regression in the infinite-width regime , informally speaking , an intriguing property is that the distribution of predictions on unseen data is given by a Gaussian Process whose kernel function is solely defined by , 1 ) the architecture , and 2 ) the choice of initial weight distribution ( Lee et al. , 2019 ) . Hence , to leverage the built-in architectural inductive biases of value networks for offline RL , we can train policies to act with respect to the lower-confidence bound ( LCB ) of the derived Gaussian Process . In what follows , we present our proposed algorithm , MSG , which leverages ensembles to approximate the LCB . Additionally , using the literature on infinite-width networks , we demonstrate the theoretical advantage of forming ensembles in the manner we propose , which deviates from the current use of ensembles in the reinforcement learning literature . 3.1 MODEL STANDARD-DEVIATION GRADIENTS ( MSG ) . MSG follows an actor critic setup where in each iteration , we first estimate the Q-values of the current policy , and subsequently optimize the policy through gradient ascent on the lower confidence bound of action value estimates . Policy Evaluation At the beginning of training , we create an ensemble ofN Q-functions by taking N samples from the initial weight distribution . Throughout training , the loss for the Q-functions is the standard least-square Q-evaluation loss , L ( θi ) = E ( s , a , r , s′ , t ) ∼D [ ( Qθi ( s , a ) − yi ( r , s′ , term , π ) ) 2 ] ( 1 ) yi : = r + ( 1− term ) · γ · Ea′∼π ( s′ ) [ Qθ̄i ( s ′ , a′ ) ] ( 2 ) where θi , θ̄i denote the parameters and target network parameters for the ith Q-function , and term = 1 [ s′ is terminal ] . In practice , the expectation in equation 1 is estimated by a minibatch , and the expectation in equation 2 is estimated with a single action sample from the policy . After every update to the Q-function parameters , their corresponding target parameters are updated to be an exponential moving average of the parameters in the standard fashion . A key factor to note is that , in contrast to the typical usage of ensembles in actor critic algorithms , each ensemble member ’ s update ( and yi ) is completely independent from other ensemble members . As we will present below , this an important algorithmic choice , both theoretically ( Theorem 4.1 ) and empirically ( section 5.3.2 ) . Policy Optimization As described above , the choice of architecture and weight initialization induces a distribution on predictions . Having approximated this distribution using an ensemble of value networks , we can optimize the policy with respect to a lower confidence bound on action values . Specifically , our proposed policy optimization objective in MSG is , max π Es∼D , a∼π ( s ) [ µ ( s , a ) + β · σ ( s , a ) ] where µ ( s , a ) = mean i [ Qθi ( s , a ) ] , σ ( s , a ) = std i [ Qθi ( s , a ) ] where β ∈ R is a hyperparameter that trades off conservatism and optimism . As our problem setting is that of offline RL requiring conservatism , we use β ≤ 0 .

Summary: The paper agues that methods, which they call "support constraint", consisting in assuming that any action outside the training yields negative outcomes, lead the agent's training to be overly pessimistic. In contrast, they propose MSG an algorithm where a lower bound of the value function is estimate from ensembles of networks. They provide some theoretical arguments on infinite-width neural networks and validate their methods on many applications, ranging from illustrative toy examples to classic continuous state-action spaces offline RL benchmarks.