Pessimistic Bootstrapping for Uncertainty-Driven Offline Reinforcement Learning

1 INTRODUCTION . Deep Reinforcement Learning ( DRL ) ( Sutton & Barto , 2018 ) achieves remarkable success in a variety of tasks . However , in most successful applications , DRL requires millions of interactions with the environment . In real-world applications such as navigation ( Mirowski et al. , 2018 ) and healthcare ( Yu et al. , 2019 ) , acquiring a large number of samples by following a possibly suboptimal policy can be costly and dangerous . Alternatively , practitioners seek to develop RL algorithms that learn a policy based solely on an offline dataset , where the dataset is typically available . However , directly adopting online DRL algorithms to the offline setting is problematic . On the one hand , policy evaluation becomes challenging since no interaction is allowed , which limits the usage of on-policy algorithms . On the other hand , although it is possible to slightly modify the off-policy value-based algorithms and sample solely from the offline dataset in training , such modification typically suffers from a significant performance drop compared with their online learning counterpart ( Levine et al. , 2020 ) . An important reason for such performance drop is the so-called distributional shift . Specifically , the offline dataset follows the visitation distribution of the behavior policies . Thus , estimating the Q-functions of the corresponding greedy policy with the offline dataset is biased due to the difference in visitation distribution . Such bias typically leads to a significant extrapolation error for DRL algorithms since the estimated Q-function tends to overestimate the out-of-distribution ( OOD ) actions ( Fujimoto et al. , 2019 ) . To tackle the distributional shift issue in offline RL , previous successful approaches typically fall into two categories , namely , policy constraint ( Kumar et al. , 2019 ; Fujimoto & Gu , 2021 ) and conservatism ( Kumar et al. , 2020 ; Yu et al. , 2020 ; 2021 ) . Policy constraint aims to restrict the learned policy to be close to the behavior policy , thus reducing the extrapolation error in policy evaluation . Conservatism methods seek to penalize the Q-functions for OOD actions in policy evaluation and hinge on a gap-expanding property to regularize the OOD behavior . Nevertheless , since policy constraints explicitly confine the policy to be close to the behavior policy , it tends to be easily affected by the non-optimal behavior policy . Meanwhile , although the conservatism algorithms such as CQL ( Kumar et al. , 2020 ) do not require policy constraints , CQL equally penalizes the OOD actions and lacks a precise characterization of the OOD data , which can lead to overly conservative value functions . To obtain a more refined characterization of the OOD data , uncertainty quantification is shown to be effective when associated with the model-based approach ( Yu et al. , 2020 ; Kidambi et al. , 2020 ) , where the dynamics model can be learned in static data thus providing more stable uncertainty in policy evaluation . Nevertheless , model-based methods need additional modules and may fail when the environment becomes high-dimensional and noisy . In addition , the uncertainty quantification in model-free RL is more challenging since the Q-function and uncertainty quantifier need to be learned simultaneously ( Yu et al. , 2021 ) . To this end , we propose Pessimistic Bootstrapping for offline RL ( PBRL ) , an uncertainty-driven model-free algorithm for offline RL . To acquire reliable Q-function estimates and their corresponding uncertainty quantification , two components of PBRL play a central role , namely , bootstrapping and OOD sampling . Specifically , we adopt bootstrapped Q-functions ( Osband et al. , 2016 ) for uncertainty quantification . We then perform pessimistic Q-updates by using such quantification as a penalization . Nevertheless , solely adopting such penalization based on uncertainty quantification is neither surprising nor effective . Training the Q-functions based solely on the offline dataset does not regularize the OOD behavior of the Q-functions and suffers from the extrapolation error . To this end , we propose a novel OOD sampling technique as a regularizer of the learned Q-functions . Specifically , we introduce additional OOD datapoints into the training buffer . The OOD datapoint consists of states sampled from the training buffer , the corresponding OOD actions sampled from the current policy , and the corresponding OOD target based on the estimated Q-function and uncertainty quantification . We highlight that having such OOD samples in the training buffer plays an important role in both the Q-function estimation and the uncertainty quantification . On the one hand , OOD sampling controls the OOD behavior in training , which guarantees the stability of the trained bootstrapped Q-functions . On the other hand , a stable estimation of the bootstrapped Qfunctions further ensures that the uncertainty quantification is reliable . We further show that PBRL is provably efficient under the linear MDP assumptions . We highlight that PBRL exploits the OOD state-action pairs by casting a more refined penalization over OOD data points , allowing PBRL to acquire better empirical performance than the policyconstraint and conservatism baselines . As an example , if an action lies close to the support of offline data but is not contained in the offline dataset , the conservatism ( Kumar et al. , 2020 ) and policy constraint ( Fujimoto et al. , 2019 ) methods tend to avoid selecting it . In contrast , PBRL tends to assign a small Q-penalty for such an action as the underlying epistemic uncertainty is small . Hence , the agent trained with PBRL may consider such action if the corresponding value estimate is high , yielding better performance than the policy constraint and conservatism baselines . Our experiments on the D4RL benchmark ( Fu et al. , 2020 ) show that PBRL provides reasonable uncertainty quantification and yields better performance compared to the state-of-the-art algorithms . 2 PRELIMINARIES . We consider an episodic MDP defined by the tuple ( S , A , T , r , P ) , where S is the state space , A is the action space , T ∈ N is the length of episodes , r is the reward function , and P is the transition distribution . The goal of RL is to find a policy π that maximizes the expected cumulative rewards E [ ∑T−1 i=0 γ tri ] , where γ ∈ [ 0 , 1 ) is the discount factor in episodic settings . The corresponding Q-function of the optimal policy satisfies the following Bellman operator , T Qθ ( s , a ) : = r ( s , a ) + γEs′∼T ( ·|s , a ) [ max a′ Qθ− ( s ′ , a′ ) ] . ( 1 ) where θ is the parameters of Q-network . In DRL , the Q-value is updated by minimizing the TDerror , namely E ( s , a , r , s′ ) [ ( Q−T Q ) 2 ] . Empirically , the target T Q is typically calculated by a separate target-network parameterized by θ− without gradient propagation ( Mnih et al. , 2015 ) . In online RL , one typically samples the transitions ( s , a , r , s′ ) through iteratively interacting with the environment . The Q-network is then trained by sampling from the collected transitions . In contrast , in offline RL , the agent is not allowed to interact with the environment . The experiences are sampled from an offline dataset Din = { ( sit , ait , rit , sit+1 ) } i∈ [ m ] . Naive off-policy methods such as Q-learning suffer from the distributional shift , which is caused by different visitation distribution of the behavior policy and the learned policy . Specifically , the greedy action a′ chosen by the target Q-network in s′ becomes OOD-action since ( s′ , a′ ) is not contained in Din . The value functions evaluated on OOD actions have significant extrapolation errors . Such errors can be further amplified through propagation and potentially diverges during training . In what follows , we introduce PBRL , which tackles such a challenge by uncertainty quantification and OOD sampling . 3 PESSIMISTIC BOOTSTRAPPING FOR OFFLINE RL . 3.1 UNCERTAINTY QUANTIFICATION WITH BOOTSTRAPPING . In PBRL , we maintain K bootstrapped Q-functions in critic to quantify the epistemic uncertainty . Formally , we denote Qk as the k-th Q-function in the ensemble . Qk is updated by fitting the following target T̂ Qkθ ( s , a ) : = r ( s , a ) + γÊs′∼P ( ·|s , a ) , a′∼π ( ·|s ) [ Qkθ− ( s ′ , a′ ) ] , ( 2 ) where we denote T̂ as the empirical Bellman operator , which estimates the expectation E [ Qkθ− ( s ′ , a′ ) | s , a ] empirically based on the offline dataset . We adopt such an ensemble technique from Bootstrapped DQN ( Osband et al. , 2016 ) , which is initially proposed for the online exploration task . Intuitively , the ensemble forms an estimation of the posterior distribution of the estimated Qfunctions , which yields similar value on areas with rich data and diversely on areas with scarce data . Indeed , the deviation among the bootstrapped Q-functions yields an epistemic uncertainty estimation , which we aim to utilize as a penalization in estimating the Q-functions . Specifically , we introduce the following uncertainty quantification U ( s , a ) based on the Q-functions { Qk } k∈ [ K ] , U ( s , a ) : = Std ( Qk ( s , a ) ) = √ 1 K ∑K k=1 ( Qk ( s , a ) − Q̄ ( s , a ) ) 2 . ( 3 ) Here we denote Q̄ as the mean among the ensemble of Q-functions . From the Bayesian perspective , such uncertainty quantification yields an estimation of the standard deviation of the posterior of Q-functions . To better understand the effectiveness of such uncertainty quantification , we illustrate with a simple prediction task . We use 10 neural networks with identical architecture and different initialization as the ensemble . We then train the ensemble with 60 datapoints in R2 plane , where the covariate x is generated from the standard Gaussian distribution , and the response y is obtained by feeding x into a randomly generated neural network . We plot the datapoints for training and the uncertainty quantification in Fig . 1 ( a ) . As shown in the figure , the uncertainty quantification rises smoothly from the in-distribution datapoints to the OOD datapoints . In offline RL , we perform regression ( s , a ) → T̂ Qk ( s , a ) in Din to train the bootstrapped Qfunctions , which is similar to regress x → y in the illustrative task . The uncertainty quantification allows us to quantify the deviation of a datapoint from the offline dataset , which provides more refined conservatism compared to the previous methods ( Kumar et al. , 2020 ; Wu et al. , 2019 ) . 3.2 PESSIMISTIC LEARNING . We now introduce the pessimistic value iteration based on the bootstrapped uncertainty quantification . The idea is to penalize the Q-functions based on the uncertainty quantification . To this end , we propose the following target for state-action pairs sampled from Din , T̂ inQkθ ( s , a ) : = r ( s , a ) + γÊs′∼P ( ·|s , a ) , a′∼π ( ·|s ) [ Qkθ− ( s ′ , a′ ) − βin Uθ− ( s′ , a′ ) ] , ( 4 ) where Uθ− ( s′ , a′ ) is the uncertainty estimation at ( s′ , a′ ) based on the target network , and βin is a tuning parameter . In addition , the empirical mean Ês′∼P ( ·|s , a ) , a′∼π ( ·|s ) is obtained by sampling the transition ( s , a , s′ ) fromDin and further sampling a′ ∼ π ( · | s′ ) from the current policy π . We denote T̂ inQkθ ( s , a ) as the in-distribution target of Qkθ ( s , a ) and write T̂ in to distinguish the in-distribution target from that of the OOD target , which we introduce in the sequel . We remark that there are two options to penalize the Q-functions through the operator T̂ in . In Eq . ( 4 ) , we penalize the next-Q value Qkθ− ( s ′ , a′ ) with the corresponding uncertainty Uθ− ( s′ , a′ ) . Alternatively , we can also penalize the immediate reward by r̂ ( s , a ) : = r ( s , a ) − Uθ ( s , a ) , and use r̂ ( s , a ) in place of r ( s , a ) for the target . Nevertheless , since the datapoint ( s , a ) ∈ Din lies on rich-data areas , the penalization Uθ ( s , a ) on the immediate reward is usually very small thus having less effect in training . In PBRL , we penalize the uncertainty of ( s′ , a′ ) in the next-Q value . Nevertheless , our empirical findings in §D suggest that solely penalizing the uncertainty for indistribution target is insufficient to control the OOD performance of the fitted Q-functions . To enforce direct regularization over the OOD actions , we incorporate the OOD datapoints directly in training and sample OOD data of the form ( sood , aood ) ∈ Dood . Specifically , we sample OOD states from the in-distribution dataset Din . Correspondingly , we sample OOD actions aood by following the current policy π ( · | sood ) . We highlight that such OOD sampling requires only the offline dataset Din and does not require additional generative models or access to the simulator . It remains to design the target for OOD samples . Since the transition P ( · | sood , aood ) and reward r ( sood , aood ) are unknown , the true target of OOD sample is inapplicable . In PBRL , we propose a novel pseudo-target for the OOD datapoints , T̂ oodQkθ ( sood , aood ) : = Qkθ ( sood , aood ) − βood Uθ ( sood , aood ) , ( 5 ) which introduces an additional uncertainty penalization Uθ ( sood , aood ) to enforce pessimistic Qfunction estimation , and βood is a tuning parameter . For OOD samples that are close to the indistribution data , such penalization is small and the OOD target is close to the Q-function estimation . In contrast , for OOD samples that are distant away from the in-distribution dataset , a larger penalization is incorporated into the OOD target . In our implementation , we introduce an additional truncation to stabilize the early stage training as max { 0 , T oodQkθ ( sood , aood ) } . In addition , we remark that βood is important to the empirical performance . Specifically , • At the beginning of training , both the Q-functions and the corresponding uncertainty quantifications are inaccurate . We use a large βood to enforce a strong regularization on OOD actions . • We then gradually decrease βood in the training process since the value estimation and uncertainty quantification becomes more accurate in training . A decreased βood makes the uncertainty more critical to decide the scale of pessimism in T̂ oodQk . The decay of βood also makes Qkθ ( s ood , aood ) does not diverge . Incorporating both the in-distribution target and OOD target , we conclude the loss function for critic in PBRL as follows , Lcritic = Ê ( s , a , r , s′ ) ∼Din [ ( T̂ inQk −Qk ) 2 ] + Êsood∼Din , aood∼π [ ( T̂ oodQk −Qk ) 2 ] , ( 6 ) where we iteratively minimize the regular TD-error for the offline data and the pseudo TD-error for the OOD data . PBRL obtains a smooth and pessimistic value function by reducing the extrapolation error caused by high-uncertain state-action pairs for both ( s , a ) and ( sood , aood ) . Based on the pessimistic Q-functions , the corresponding policy is obtained by the following maximization problem , πϕ : = max ϕ Ês∼Din , a∼π ( ·|s ) [ min k=1 , ... , K Qk ( s , a ) ] , ( 7 ) where ϕ is the policy parameters . Taking the minimum among ensembleQ-functions is also adopted in previous actor-critic methods like SAC ( Haarnoja et al. , 2018 ) and TD3 ( Fujimoto et al. , 2018 ) . The overall architecture of PBRL is illustrated in Fig . 1 ( b ) .

Summary: This paper proposed an uncertainty-aware offline reinforcement learning algorithm based on a Q-network ensemble. The proposed method penalizes Q-values on OOD actions and performs pessimistic offline Q-learning. The paper interprets the method with LCB framework. The proposed method outperforms existing offline RL algorithms on D4RL Gym and Adroit tasks.