Addressing Distribution Shift in Online Reinforcement Learning with Offline Datasets

1 INTRODUCTION . Offline reinforcement learning ( RL ) , the task of training a sequential decision-making agent with a static offline dataset , holds the promise of a data-driven approach to reinforcement learning , thereby bypassing the laborious and often dangerous process of sample collection ( Levine et al. , 2020 ) . Accordingly , various offline RL methods have been developed , some of which are often capable of training agents that are more performant than the behavior policy ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Wu et al. , 2019 ; Siegel et al. , 2020 ; Agarwal et al. , 2020 ; Kidambi et al. , 2020 ; Yu et al. , 2020 ; Kumar et al. , 2020 ) . However , agents trained via offline RL methods may be suboptimal , for ( a ) the dataset they were trained on may only contain suboptimal data ; and ( b ) environment in which they are deployed may be different from the environment in which dataset was generated . This necessitates an online fine-tuning procedure , where the agent improves by interacting with the environment and gathering additional information . Fine-tuning an offline RL agent , however , poses certain challenges . For example , Nair et al . ( 2020 ) pointed out that offline RL algorithms based on modeling the behavior policy ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Wu et al. , 2019 ; Siegel et al. , 2020 ) are not amenable to fine-tuning . This is because such methods require sampling actions from the modeled behavior policy for updating the agent , and fine-tuning such a generative model online for reliable sample generation is a challenging task . On the other hand , a more recent state-of-the-art offline RL algorithm , conservative Q-learning ( CQL ; Kumar et al. , 2020 ) , does not require explicit behavior modeling , and one might expect CQL to be amenable to fine-tuning . However , we make an observation that fine-tuning a CQL agent is a non-trivial task , due to the so-called distribution shift problem – the agent encounters out-ofdistribution samples , and in turn loses its good initial policy from offline RL training . This can be attributed to the bootstrapping error , i.e. , error introduced when Q-function is updated with an inaccurate target value evaluated at unfamiliar states and actions . Such initial training instability is a severe limitation , given the appeal of offline RL lies in safe deployment at test time , and losing such safety guarantees directly conflicts with the goal of offline RL . Contribution . In this paper , we first demonstrate that fine-tuning a CQL agent may lead to unstable training due to distribution shift ( see Section 3 for more details ) . To handle this issue , we introduce a simple yet effective framework , which incorporates a balanced replay scheme and an ensemble distillation scheme . Specifically , we propose to maintain two separate replay buffers for offline and online samples , respectively . Then we modulate the sampling ratio between the two , in order to balance the effects of ( a ) widening the data distribution the agent sees ( offline data ) , and ( b ) exploiting the environment feedback ( online data ) . Furthermore , we propose an ensemble distillation scheme : first , we learn an ensemble of independent CQL agents , then distill the multiple policies into a single policy . During fine-tuning , we improve the policy using the mean of Q-functions , so that policy updates are more robust to error in each individual Q-function . In our experiments , we demonstrate the strength of our method based on MuJoCo ( Todorov et al. , 2012 ) datasets from the D4RL ( Fu et al. , 2020 ) benchmark suite . Our goal is to achieve both ( a ) strong initial performance as well as maintaining it during the initial training phase , and ( b ) better sample-efficiency . For evaluation , we measure the final performance and sample-efficiency of RL agents throughout the fine-tuning procedure . We demonstrate that our method achieves stable training , while significantly outperforming all baseline methods considered , including BCQ ( Fujimoto et al. , 2019 ) and AWAC ( Nair et al. , 2020 ) , both in terms of final performance and sample-efficiency . 2 BACKGROUND . Reinforcement learning . We consider the standard RL framework , where an agent interacts with the environment so as to maximize the expected total return . More formally , at each timestep t , the agent observes a state st , and performs an action at according to its policy π . The environment rewards the agent with rt , then transitions to the next state st+1 . The agent ’ s objective is to maximize the expected return Eπ [ ∑∞ k=0 γ krk ] , where γ ∈ [ 0 , 1 ) is the discount factor . In this work , we mainly consider off-policy RL algorithms , a class of algorithms that can , in principle , train an agent with samples generated by any behavior policy . These algorithms are well-suited for fine-tuning a pretrained RL agent , for they can leverage both offline and online samples . Offline RL algorithms are off-policy RL algorithms that only utilize static datasets for training . Here , we introduce an off-policy RL algorithm and an offline RL algorithm we build on in this work . Soft Actor-Critic . SAC ( Haarnoja et al. , 2018 ) is an off-policy actor-critic algorithm that learns a soft Q-function Qθ ( s , a ) parameterized by θ and a stochastic policy πφ modeled as a Gaussian with its parameters φ . To update parameters θ and φ , SAC alternates between a soft policy evaluation and a soft policy improvement . During soft policy evaluation , soft Q-function parameter θ is updated to minimize the soft Bellman residual : LSACcritic ( θ ) = Eτt∼B [ LQ ( τt , θ ) ] , ( 1 ) LSACQ ( τt , θ ) = ( Qθ ( st , at ) − ( rt + γEat+1∼πφ [ Qθ̄ ( st+1 , at+1 ) − α log πφ ( at+1|st+1 ) ] ) ) 2 , ( 2 ) where τt = ( st , at , rt , st+1 ) is a transition , B is the replay buffer , θ̄ is the moving average of soft Q-function parameter θ , and α is the temperature parameter . During soft policy improvement , policy parameter φ is updated to minimize the following objective : LSACactor ( φ ) = Est∼B [ Lπ ( st , φ ) ] , where Lπ ( st , φ ) = Eat∼πφ [ α log πφ ( at|st ) −Qθ ( st , at ) ] . ( 3 ) Conservative Q-Learning . CQL ( Kumar et al. , 2020 ) is an offline RL algorithm that learns a lower bound of the Q-function Qθ ( s , a ) , in order to prevent extrapolation error – value overestimation caused by bootstrapping from out-of-distribution actions . To this end , CQL ( H ) , a variant of CQL , imposes a regularization that minimizes the expected Q-value at unseen actions , and maximizes the expected Q-value at seen actions by minimizing the following objective : LCQLcritic ( θ ) = 1 2 · LCQLQ ( θ ) + α CQL · LCQLreg ( θ ) , ( 4 ) LCQLQ ( θ ) = Eτt∼D [ ( Qθ ( st , at ) − ( rt + γEat+1∼πφ [ Qθ̄ ( st+1 , at+1 ) ] ) ) 2 ] , ( 5 ) LCQLreg ( θ ) = Est∼D [ log ∑ a exp ( Qθ ( st , a ) − Eat∼π̂β [ Qθ ( st , at ) ] ] , ( 6 ) where αCQL is the tradeoff parameter , D the offline dataset , and π̂β ( at|st ) = ∑ ( s , a ) ∈D 1 { s=st , a=at } ∑ s∈D 1 { s=st } the empirical behavior policy . 3 CHALLENGE : DISTRIBUTION SHIFT . In the context of fine-tuning an offline RL agent , there may exist a distribution shift between offline and online data distribution : an offline RL agent encounters data distributed away from the offline data , as the agent starts interacting with the environment ( see Figure 1a for a visualization ) . More formally , let pπ ( s , a ) = dπ ( s ) π ( a|s ) denote the discounted stationary state-action distribution of the policy , where dπ denotes the discounted marginal state distribution , and π denotes the policy . Then distribution shift refers to the difference between pβ ( s , a ) and pπθ ( s , a ) , where β refers to the behavior policy , and πθ refers to the current policy . To demonstrate the existence of distribution shift empirically , we report the negative log-likelihiood estimates of offline and online samples in Figure 1b . Specifically , we first split the walker2d-medium offline dataset into training and validation datasets , then train a variational autoencoder ( VAE ; Kingma & Welling , 2014 ) on the training dataset . Then , we estimate the log-likelihood of online samples and offline samples from the validation dataset , using the importance weight estimator ( Burda et al. , 2016 ) . We can observe a difference between offline and online data distributions , i.e. , they follow different distributions . This shift in turn affects fine-tuning in a very complicated manner , for it involves an interplay between actor and critic updates with newly collected out-of-distribution samples . To elaborate in more detail , we train an agent offline via state-of-the-art CQL algorithm in ( 4 ) , then fine-tune it via SAC algorithm in ( 1 ) and ( 3 ) 1 , either using samples drawn uniformly at random ( Uniform replay ) as in Nair et al . ( 2020 ) , or using online samples exclusively ( Online only ) . Figure 1c shows how fine-tuning may suffer from instability in both cases . Essentially , this instability occurs due to the shift between offline and online data distribution . In the case of using both offline and online data , the chance of agent seeing online samples for update becomes too low , especially when the offline dataset contains massive amount of data . This prevents timely updates at unfamiliar states encountered online , and may lead to performance degradation as seen in Figure 1c . On the other hand , when using online samples exclusively , the agent is exposed to unseen samples only , for which Q function does not provide a reliable value estimate . This may lead to bootstrapping error , and hence a dip in performance as seen in Figure 1c . This points to leveraging a mixture of offline and online samples for fine-tuning , so that the agent can balance the trade-off between them . 1Fine-tuning with CQL updates ( 4 ) leads to slow ( if at all ) improvement in most setups we consider ( see Appendix B for more results ) . We also provide additional experimental results when using different RL algorithms than SAC for fine-tuning a pre-trained offline CQL agent in Appendix C . 4 BRED : BALANCED REPLAY WITH ENSEMBLE DISTILLATION . In this section , we present BRED : Balanced Replay with Ensemble Distillation , in order to address the distribution shift between offline and online data distributions . First , we introduce separate offline and online replay buffers in order to select a balanced mix of samples from them for better Q-function updates . This has the dual advantage of updating Q-function with a wide distribution of samples , while making sure that Q-values are updated at novel , unseen states from online interaction . Furthermore , we train multiple actor-critic offline RL agents , and distill the policies into a single policy . Then , during fine-tuning , we improve the distilled policy via Q-ensemble .

Summary: This paper considers the problem of policy learning in Markov Decision Process (MDP) from the combination of online and offline samples. The offline samples are generated by a behavior policy in the same MDP model, i.e., the behavior agent and the learning agent share the same state-action space. The learning procedure goes as follows. One first trains a MDP policy from the offline data; the online samples are then used to fine-tune the learned policy.