Uncertainty Weighted Offline Reinforcement Learning

1 INTRODUCTION . Deep reinforcement learning ( RL ) has seen a surge of interest over the recent years . It has achieved remarkable success in simulated tasks ( Silver et al. , 2017 ; Schulman et al. , 2017 ; Haarnoja et al. , 2018 ) , where the cost of data collection is low . However , one of the drawbacks of RL is its difficulty of learning from prior experiences . Therefore , the application of RL to unstructured real-world tasks is still in its primal stages , due to the high cost of active data collection . It is thus crucial to make full use of previously collected datasets whenever large scale online RL is infeasible . Offline batch RL algorithms offer a promising direction to leveraging prior experience ( Lange et al. , 2012 ) . However , most prior off-policy RL algorithms ( Haarnoja et al. , 2018 ; Munos et al. , 2016 ; Kalashnikov et al. , 2018 ; Espeholt et al. , 2018 ; Peng et al. , 2019 ) fail on offline datasets , even on expert demonstrations ( Fu et al. , 2020 ) . The sensitivity to the training data distribution is a well known issue in practical offline RL algorithms ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; 2020 ; Peng et al. , 2019 ; Yu et al. , 2020 ) . A large portion of this problem is attributed to actions or states not being covered within the training set distribution . Since the value estimate on out-of-distribution ( OOD ) actions or states can be arbitrary , OOD value or reward estimates can incur destructive estimation errors that propagates through the Bellman loss and destabilizes training . Prior attempts try to avoid OOD actions or states by imposing strong constraints or penalties that force the actor distribution to stay within the training data ( Kumar et al. , 2019 ; 2020 ; Fujimoto et al. , 2019 ; Laroche et al. , 2019 ) . While such approaches achieve some degree of experimental success , they suffer from the loss of generalization ability of the Q function . For example , a state-action pair that does not appear in the training set can still lie within the training set distribution , but policies trained with strong penalties will avoid the unseen states regardless of whether the Q function can produce an accurate estimate of the state-action value . Therefore , strong penalty based solutions often promote a pessimistic and sub-optimal policy . In the extreme case , e.g. , in certain benchmarking environments with human demonstrations , the best performing offline algorithms only achieve the same performance as a random agent ( Fu et al. , 2020 ) , which demonstrates the need of robust offline RL algorithms . In this paper , we hypothesize that a key aspect of a robust offline RL algorithm is a proper estimation and usage of uncertainty . On the one hand , one should be able to reliably assign an uncertainty score to any state-action pair ; on the other hand , there should be a mechanism that utilizes the estimated uncertainty to prevent the model from learning from data points that induce high uncertainty scores . The first problem relates closely to OOD sample detection , which has been extensively studied in the Bayesian deep learning community . ( Gal & Ghahramani , 2016a ; Gal , 2016 ; Osawa et al. , 2019 ) , often measured by the epistemic uncertainty of the model . We adopt the dropout based approach Gal & Ghahramani ( 2016a ) , due to its simplicity and empirical success in practice . For the second problem , we provide an intuitive modification to the Bellman updates in actor-critic based algorithms . Our proposed Uncertainty Weighted Actor Critic ( UWAC ) is to simply down weigh the contribution of target state and action pairs with high uncertainty . By doing so , we prevent the Q function from learning from overly optimistic targets that lie far away from training data distribution ( high uncertainty ) . Empirically , we first verified the effectiveness of dropout uncertainty estimation at detecting OOD samples . We show that the uncertainty estimation makes intuitive sense in a simple environment . With the uncertainty based down weighting scheme , our method significantly improves the training stability over our chosen baseline ( Kumar et al. , 2019 ) , and achieves state-of-the-art performance in a variety of standard benchmarking tasks for offline RL . Overall , our contribution can be summarized as follows : 1 ) We propose a simple and efficient technique ( UWAC ) to counter the effect of OOD samples with no additional loss terms or models . 2 ) We experimentally demonstrate the effectiveness of dropout uncertainty estimation for RL . 3 ) UWAC offers a novel way for stabilizing offline RL . 4 ) UWAC achieves SOTA performance on common offline RL benchmarks , and obtains significant performance gain on narrow human demonstrations . 2 RELATED WORK . In this work , we consider offline batch reinforcement learning ( RL ) under static datasets . Offline RL algorithms are especially prone to errors from inadequate coverage of the training set distribution , distributional shifts during actor critic training , and the variance induced by deep neural networks . Such error have been extensively studied as ” error propagation ” in approximate dynamic programming ( ADP ) ( Bertsekas & Tsitsiklis , 1996 ; Farahmand et al. , 2010 ; Munos , 2003 ; Scherrer et al. , 2015 ) . Scherrer et al . ( 2015 ) obtains a bound on the point-wise Bellman error of approximate modified policy iteration ( AMPI ) with respect to the supremum of the error in function approximation for an arbitrary iteration . We adopt the theoretical tools from ( Kumar et al. , 2019 ) and study the accumulation and propagation of Bellman errors under the offline setting . One of the most significant problems associated with off-policy and offline RL is the bootstrapping error ( Kumar et al. , 2019 ) : When training encounters an action or state unseen within the training set , the critic value estimate on out-of-distribution ( OOD ) samples can be arbitrary and incur an error that destabilizes convergence on all other states ( Kumar et al. , 2019 ; Fujimoto et al. , 2019 ) through the Bellman backup . Yu et al . ( 2020 ) trains a model of the environment that captures the epistemic uncertainty . The uncertainty estimate is used to penalize reward estimation for uncertain states and actions , promoting a pessimistic policy against OOD actions and states . The main drawback of such a model based approach is the unnecessary introduction of a model of the environment – it is often very hard to train a good model . On the other hand , model-free approaches either train an agent pessimistic to OOD states and actions ( Wu et al. , 2019 ; Kumar et al. , 2020 ) or constrain the actor distribution to the training set action distribution ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Wu et al. , 2019 ; Jaques et al. , 2019 ; Fox et al. , 2015 ; Laroche et al. , 2019 ) . However , the pessimistic assumption that all unseen states or actions are bad may lead to a sub-optimal agent and greatly reduce generalization to online fine-tuning ( Nair et al. , 2020 ) . Distributional constraints , in addition , rely on approximations since the actor distribution is often implicit . Such approximations cause practical training instability that we will study in detail in section 5.4 . We propose a model-free actor-critic method that down-weighs the Bellman loss term by inverse uncertainty of the critic target . Uncertainty estimation has been implemented in model-free RL for safety and risk estimation ( Clements et al. , 2019 ; Hoel et al. , 2020 ) or exploration ( Gal & Ghahramani , 2016a ; Lines & Van Der Wilk ) , through ensembling ( Hoel et al. , 2020 ) or distributional RL ( Dabney et al. , 2018 ; Clements et al. , 2019 ) . However , distributional RL works best on discrete action spaces ( Dabney et al. , 2018 ) and require additional distributional assumptions when extended to continuous action spaces ( Clements et al. , 2019 ) . Our approach estimates uncertainty through Monte Carlo dropout ( MC-dropout ) ( Srivastava et al. , 2014 ) . MC-dropout uncertainty estimation is a simple method with minimal overhead and has been thoroughly studied in many traditional supervised learning tasks in deep learning ( Gal & Ghahramani , 2016a ; Hron et al. , 2018 ; Kingma et al. , 2015 ; Gal & Ghahramani , 2016b ) . Moreover , we observe experimentally that MC-dropout uncertainty estimation behaves similarly to explicit ensemble models where the prediction is the mean of the ensembles , while being much simpler ( Lakshminarayanan et al. , 2017 ; Srivastava et al. , 2014 ) . The most relevant to our work are MOPO ( Yu et al. , 2020 ) , BEAR ( Kumar et al. , 2019 ) , and CQL ( Kumar et al. , 2020 ) . MOPO approaches offline RL from a different model-based paradigm , and obtains competitive results on some tasks with wide data distribution . However , due to the model-based nature , MOPO achieves limited performance on most other benchmarks due to the performance of the model being limited by the data distribution . On the other hand , BEAR and CQL both use actorcritic and do not suffer from the above problem . We use BEAR ( discussed in section 3.2 ) as our baseline algorithm and achieve significant performance gain through dropout uncertainty weighted backups . CQL avoids OOD states/actions through direct Q value penalty on actions that leads to OOD unseen states within the training set . However the penalty proposed by CQL 1 ) risks hurting Q estimates for ( action , state ) pairs that are not OOD , since samples not seen within the dataset can still lie within the true dataset distribution ; 2 ) limits the policy to be pessimistic , which may be hard to fine-tune once on-policy data becomes available . Additionally our method is not limited to BEAR and can apply to other actor-critic methods like CQL . We leave such exploration to future works . 3 PRELIMINARIES . 3.1 NOTATIONS . Following Kumar et al . ( 2019 ) , we represent the environment as a Markov decision process ( MDP ) comprising of a 6-tuple ( S , A , P , R , ρ0 , γ ) , where S is the state space , A is the action space , P ( s′|s , a ) is the transition probability distribution , ρ0 is the initial state distribution , R : S×A → R is the reward function , and γ ∈ ( 0 , 1 ] is the discount factor . Our goal is to find a policy π ( s|a ) from the set of policy functions Π to maximize the expected cumulative discounted reward . Standard Q-learning learns an optimal state-action value function Q∗ ( s , a ) , representing the expected cumulative discounted reward starting from s with action a and then acting optimally thereafter . Q-learning is trained on the Bellman equation defined as follows with the Bellman optimal operator T defined by : T Q ( s , a ) : = R ( s , a ) + γEP ( s′|s , a ) [ max a′ Q ( s′ , a′ ) ] ( 1 ) In practice , the critic ( Q function ) is updated through dynamic programming , by projecting the target Q estimate ( T Q ) into Q ( i.e . minimizing Bellman Squared Error E [ ( Q− T Q ) 2 ] ) . Since maxa′ Q ( s ′ , a′ ) in generally intractable in continuous action spaces , an actor ( πθ ) function is learned to maximize the critic function ( πθ ( s ) : = arg maxaQ ( s , a ) ) ( Haarnoja et al. , 2018 ; Fujimoto et al. , 2018 ; Sutton & Barto , 2018 ) . In the context of offline reinforcement learning , naively performing maxa′ Q ( s′ , a′ ) in equation 1 may result in an a′ unseen within the training dataset ( OOD ) , and resulting in aQ estimate with very large error that can propagate through the Bellman bootstrapping and destabilize training on other states ( Kumar et al. , 2019 ) .

Summary: This paper proposes to use an uncertainty-weighted objective for offline RL with BEAR (Kumar et al.) that penalizes the MMD distance between the learned policy and the previous policy. The uncertainty weighted objective weights the policy improvement objective with the variance in the Q-function, where this variance primarily represents aleatoric or intrinsic uncertainty, not the epistemic or belief uncertainty. They show that their method performs reasonably better than prior methods on the D4RL datasets and show that the learned Q-values are better than BEAR.