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: The paper looks at the offline RL problem and claims that existing techniques penalize out-of-distribution actions too aggressively and so perform poorly at generalization and extrapolation. To remedy this, the paper proposes an offline RL algorithm in which an ensemble of critics is trained with an objective composed of (1) a TD error based on actions seen in the dataset with target value penalized by standard deviation over the ensemble of the next Q-values, and (2) a squared error on the Q-values of OOD actions regressing to those same Q-values penalized by the standard deviation over the ensemble. The proposed algorithm is paired with a theoretical analysis connecting it to recent theoretical offline RL papers. The algorithm is evaluated on the D4RL benchmarks and shows favorable performance compared to baselines.