AWAC: Accelerating Online Reinforcement Learning with Offline Datasets

1 INTRODUCTION . Learning models that generalize effectively to complex open-world settings , from image recognition ( Krizhevsky et al. , 2012 ) to natural language processing ( Devlin et al. , 2019 ) , relies on large , high-capacity models and large , diverse , and representative datasets . Leveraging this recipe for reinforcement learning ( RL ) has the potential to yield real-world generalization for control applications such as robotics . However , while deep RL algorithms enable the use of large models , the use of large datasets for real-world RL has proven challenging . Most RL algorithms collect new data online every time a new policy is learned , which limits the size and diversity of the datasets for RL . In the same way that powerful models in computer vision and NLP are often pre-trained on large , general-purpose datasets and then fine-tuned on task-specific data , RL policies that generalize effectively to open-world settings will need to be able to incorporate large amounts of prior data effectively into the learning process , while still collecting additional data online for the task at hand . For data-driven reinforcement learning , offline datasets consist of trajectories of states , actions and associated rewards . This data can potentially come from demonstrations for the desired task ( Schaal , 1997 ; Atkeson & Schaal , 1997 ) , suboptimal policies ( Gao et al. , 2018 ) , demonstrations for related tasks ( Zhou et al. , 2019 ) , or even just random exploration in the environment . Depending on the quality of the data that is provided , useful knowledge can be extracted about the dynamics of the world , about the task being solved , or both . Effective data-driven methods for deep reinforcement learning should be able to use this data to pre-train offline while improving with online fine-tuning . Since this prior data can come from a variety of sources , we would like to design an algorithm that does not utilize different types of data in any privileged way . For example , prior methods that incorporate demonstrations into RL directly aim to mimic these demonstrations ( Nair et al. , 2018 ) , which is desirable when the demonstrations are known to be optimal , but imposes strict requirements on the type of offline data , and can cause undesirable bias when the prior data is not optimal . While prior methods for fully offline RL provide a mechanism for utilizing offline data ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ) , as we will show in our experiments , such methods generally are not effective for fine-tuning with online data as they are often too conservative . In effect , prior methods require us to choose : Do we assume prior data is optimal or not ? Do we use only offline data , or only online data ? To make it feasible to learn policies for open-world settings , we need algorithms that learn successfully in any of these cases . In this work , we study how to build RL algorithms that are effective for pre-training from offpolicy datasets , but also well suited to continuous improvement with online data collection . We systematically analyze the challenges with using standard off-policy RL algorithms ( Haarnoja et al. , 2018 ; Kumar et al. , 2019 ; Abdolmaleki et al. , 2018 ) for this problem , and introduce a simple actor critic algorithm that elegantly bridges data-driven pre-training from offline data and improvement with online data collection . Our method , which uses dynamic programming to train a critic but a supervised learning style update to train a constrained actor , combines the best of supervised learning and actor-critic algorithms . Dynamic programming can leverage off-policy data and enable sample-efficient learning . The simple supervised actor update implicitly enforces a constraint that mitigates the effects of distribution shift when learning from offline data ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ) , while avoiding overly conservative updates . We evaluate our algorithm on a wide variety of robotic control and benchmark tasks across three simulated domains : dexterous manipulation , tabletop manipulation , and MuJoCo control tasks . Our algorithm , Advantage Weighted Actor Critic ( AWAC ) , is able to quickly learn successful policies on difficult tasks with high action dimension and binary sparse rewards , significantly better than prior methods for off-policy and offline reinforcement learning . Moreover , AWAC can utilize different types of prior data without any algorithmic changes : demonstrations , suboptimal data , or random exploration data . The contribution of this work is not just another RL algorithm , but a systematic study of what makes offline pre-training with online fine-tuning unique compared to the standard RL paradigm , which then directly motivates a simple algorithm , AWAC , to address these challenges . 2 PRELIMINARIES . We consider the standard reinforcement learning notation , with states s , actions a , policy π ( a|s ) , rewards r ( s , a ) , and dynamics p ( s′|s , a ) . The discounted return is defined as Rt = ∑T i=t γ ir ( si , ai ) , for a discount factor γ and horizon T which may be infinite . The objective of an RL agent is to maximize the expected discounted return J ( π ) = Epπ ( τ ) [ R0 ] under the distribution induced by the policy . The optimal policy can be learned directly by policy gradient , estimating ∇J ( π ) ( Williams , 1992 ) , but this is often ineffective due to high variance of the estimator . Many algorithms attempt to reduce this variance by making use of the value function V π ( s ) = Epπ ( τ ) [ Rt|s ] , action-value function Qπ ( s , a ) = Epπ ( τ ) [ Rt|s , a ] , or advantage Aπ ( s , a ) = Qπ ( s , a ) − V π ( s ) . The action-value function for a policy can be written recursively via the Bellman equation : Qπ ( s , a ) = r ( s , a ) + γEp ( s′|s , a ) [ V π ( s′ ) ] = r ( s , a ) + γEp ( s′|s , a ) [ Eπ ( a′|s′ ) [ Qπ ( s′ , a′ ) ] ] . ( 1 ) Instead of estimating policy gradients directly , actor-critic algorithms maximize returns by alternating between two phases ( Konda & Tsitsiklis , 2000 ) : policy evaluation and policy improvement . During the policy evaluation phase , the critic Qπ ( s , a ) is estimated for the current policy π . This can be accomplished by repeatedly applying the Bellman operator B , corresponding to the right-hand side of Equation 1 , as defined below : BπQ ( s , a ) = r ( s , a ) + γEp ( s′|s , a ) [ Eπ ( a′|s′ ) [ Qπ ( s′ , a′ ) ] ] . ( 2 ) By iterating according to Qk+1 = BπQk , Qk converges to Qπ ( Sutton & Barto , 1998 ) . With function approximation , we can not apply the Bellman operator exactly , and instead minimize the Bellman error with respect to Q-function parameters φk : φk = argmin φ ED [ ( Qφ ( s , a ) − y ) 2 ] , y = r ( s , a ) + γEs′ , a′ [ Qφk−1 ( s′ , a′ ) ] . ( 3 ) During policy improvement , the actor π is typically updated based on the current estimate of Qπ . A commonly used technique ( Lillicrap et al. , 2016 ; Fujimoto et al. , 2018 ; Haarnoja et al. , 2018 ) is to update the actor πθk ( a|s ) via likelihood ratio or pathwise derivatives to optimize the following objective , such that the expected value of the Q-function Qπ is maximized : θk = argmax θ Es∼D [ Eπθ ( a|s ) [ Qφk ( s , a ) ] ] ( 4 ) Actor-critic algorithms are widely used in deep RL ( Mnih et al. , 2016 ; Lillicrap et al. , 2016 ; Haarnoja et al. , 2018 ; Fujimoto et al. , 2018 ) . With a Q-function estimator , they can in principle utilize off-policy data when used with a replay buffer for storing prior transition tuples , which we will denote β , to sample previous transitions , although we show that this by itself is insufficient for our problem setting . 3 CHALLENGES IN OFFLINE RL WITH ONLINE FINE-TUNING . In this section , we study the unique challenges that exist when pre-training using offline data , followed by fine-tuning with online data collection . We first describe the problem , and then analyze what makes this problem difficult for prior methods . Problem definition . A static dataset of transitions , D = { ( s , a , s′ , r ) j } , is provided to the algorithm at the beginning of training . This dataset can be sampled from an arbitrary policy or mixture of policies , and may even be collected by a human expert . This definition is general and encompasses many scenarios , such as learning from demonstrations , random data , prior RL experiments , or even from multi-task data . Given the dataset D , our goal is to leverage D for pre-training and use some online interaction to learn the optimal policy π∗ ( a|s ) , with as few interactions with the environment as possible ( depicted in Fig 1 ) . This setting is representative of many real-world RL settings , where prior data is available and the aim is to learn new skills efficiently . We first study existing algorithms empirically in this setting on the HalfCheetah-v2 Gym environment1 . The prior dataset consists of 15 demonstrations from an expert policy and 100 suboptimal trajectories sampled from a behavioral clone of these demonstrations . All methods for the remainder of this paper incorporate the prior dataset , unless explicitly labeled “ scratch ” . 3.1 ) Data Efficiency . One of the simplest ways to utilize prior data such as demonstrations for RL is to pre-train a policy with imitation learning , and fine-tune with on-policy RL ( Gupta et al. , 2019 ; Rajeswaran et al. , 2018 ) . This approach has two drawbacks : ( 1 ) prior data may not be optimal ; ( 2 ) on-policy fine-tuning is data inefficient as it does not reuse the prior data in the RL stage . In our setting , data efficiency is vital . To this end , we require algorithms that are able to reuse arbitrary offpolicy data during online RL for data-efficient fine-tuning . We find that algorithms that use on-policy fine-tuning ( Rajeswaran et al. , 2018 ; Gupta et al. , 2019 ) , or Monte-Carlo return estimation ( Peters & Schaal , 2007 ; Wang et al. , 2018 ; Peng et al. , 2019 ) are generally much less efficient than off-policy actor-critic algorithms , which iterate between improving π and estimating Qπ via Bellman backups . This can be seen from the results in Figure 2 plot 1 , where on-policy methods like DAPG ( Rajeswaran et al. , 2018 ) and Monte-Carlo return methods like AWR ( Peng et al. , 2019 ) and MARWIL ( Wang et al. , 2018 ) are an order of magnitude slower than off-policy actor-critic methods . Actor-critic methods , shown in Figure 2 plot 2 , can in principle use off-policy data . However , as we will discuss next , naïvely applying these algorithms to our problem suffers from a different set of challenges . 3.2 ) Bootstrap Error in Offline Learning with Actor-Critic Methods . When standard off-policy actor-critic methods are applied to this problem setting , they perform poorly , as shown in the second plot in Figure 2 : despite having a prior dataset in the replay buffer , these algorithms do not benefit significantly from offline training . We evaluate soft actor critic ( Haarnoja et al. , 2018 ) , a state-of-theart actor-critic algorithm for continuous control . Note that “ SAC-scratch , ” which does not receive the prior data , performs similarly to “ SACfD-prior , ” which does have access to the prior data , indicating that the off-policy RL algorithm is not actually able to make use of the off-policy data for pre-training . Moreover , even if the SAC is policy is pre-trained by behavior cloning , labeled “ SACfD-pretrain ” , we still observe an initial decrease in performance , and performance similar to learning from scratch . This challenge can be attributed to off-policy bootstrapping error accumulation , as observed in several prior works ( Sutton & Barto , 1998 ; Kumar et al. , 2019 ; Wu et al. , 2020 ; Levine et al. , 2020 ; 1We use this environment for analysis because it helps understand and accentuate the differences between different algorithms . More challenging environments like the ones shown in Fig 3 are too hard to solve to analyze variants of different methods . Fujimoto et al. , 2019 ) . In actor-critic algorithms , the target value Q ( s′ , a′ ) , with a′ ∼ π , is used to update Q ( s , a ) . When a′ is outside of the data distribution , Q ( s′ , a′ ) will be inaccurate , leading to accumulation of error on static datasets . Offline RL algorithms ( Fujimoto et al. , 2019 ; Kumar et al. , 2019 ; Wu et al. , 2020 ) propose to address this issue by explicitly adding constraints on the policy improvement update ( Equation 4 ) to avoid bootstrapping on out-of-distribution actions , leading to a policy update of this form : argmax θ Es∼D [ Eπθ ( a|s ) [ Qφk ( s , a ) ] ] s.t . D ( πθ , πβ ) ≤ . ( 5 ) Here , πθ is the actor being updated , and πβ ( a|s ) represents the ( potentially unknown ) distribution from which all of the data seen so far ( both offline data and online data ) was generated . In the case of a replay buffer , πβ corresponds to a mixture distribution over all past policies . Typically , πβ is not known , especially for offline data , and must be estimated from the data itself . Many offline RL algorithms ( Kumar et al. , 2019 ; Fujimoto et al. , 2019 ; Siegel et al. , 2020 ) explicitly fit a parametric model to samples for the distribution πβ via maximum likelihood estimation , where samples from πβ are obtained simply by sampling uniformly from the data seen thus far : π̂β = maxπ̂β Es , a∼πβ [ log π̂β ( a|s ) ] . After estimating π̂β , prior methods implement the constraint given in Equation 5 in various ways , including penalties on the policy update ( Kumar et al. , 2019 ; Wu et al. , 2020 ) or architecture choices for sampling actions for policy training ( Fujimoto et al. , 2019 ; Siegel et al. , 2020 ) . As we will see next , the requirement for accurate estimation of π̂β makes these methods difficult to use with online fine-tuning . 3.3 ) Excessively Conservative Online Learning . While offline RL algorithms with constraints ( Kumar et al. , 2019 ; Fujimoto et al. , 2019 ; Wu et al. , 2020 ) perform well offline , they struggle to improve with fine-tuning , as shown in the third plot in Figure 2 . We see that the purely offline RL performance ( at “ 0K ” in Fig . 2 ) is much better than the standard off-policy methods shown in Section 3.2 . However , with additional iterations of online fine-tuning , the performance increases very slowly ( as seen from the slope of the BEAR curve in Fig 2 ) . What causes this phenomenon ? This can be attributed to challenges in fitting an accurate behavior model as data is collected online during fine-tuning . In the offline setting , behavior models must only be trained once via maximum likelihood , but in the online setting , the behavior model must be updated online to track incoming data . Training density models online ( in the “ streaming ” setting ) is a challenging research problem ( Ramapuram et al. , 2017 ) , made more difficult by a potentially complex multi-modal behavior distribution induced by the mixture of online and offline data . To understand this , we plot the log likelihood of learned behavior models on the dataset during online and offline training for the HalfCheetah task . As we can see in the plot , the accuracy of the behavior models ( log πβ on the y-axis ) reduces during online fine-tuning , indicating that it is not fitting the new data well during online training . When the behavior models are inaccurate or unable to model new data well , constrained optimization becomes too conservative , resulting in limited improvement with fine-tuning . This analysis suggests that , in order to address our problem setting , we require an off-policy RL algorithm that constrains the policy to prevent offline instability and error accumulation , but not so conservatively that it prevents online fine-tuning due to imperfect behavior modeling . Our proposed algorithm , which we discuss in the next section , accomplishes this by employing an implicit constraint , which does not require any explicit modeling of the behavior policy .

Summary: This paper studies shorting coming of existing off-policy methods when it comes to prior data and fine-tuning and shows that those existing methods can't effectively utilize previously collected data with online updates. To address this problem, they propose to constraint policy updates with respect to behavioral policy. Their proposed method is built mainly on the top of AWR [1].