Keep Doing What Worked: Behavior Modelling Priors for Offline Reinforcement Learning

Off-policy reinforcement learning algorithms promise to be applicable in settings where only a fixed data-set ( batch ) of environment interactions is available and no new experience can be acquired . This property makes these algorithms appealing for real world problems such as robot control . In practice , however , standard off-policy algorithms fail in the batch setting for continuous control . In this paper , we propose a simple solution to this problem . It admits the use of data generated by arbitrary behavior policies and uses a learned prior – the advantage-weighted behavior model ( ABM ) – to bias the RL policy towards actions that have previously been executed and are likely to be successful on the new task . Our method can be seen as an extension of recent work on batch-RL that enables stable learning from conflicting data-sources . We find improvements on competitive baselines in a variety of RL tasks – including standard continuous control benchmarks and multitask learning for simulated and real-world robots . Videos are available at https : //sites.google.com/view/behavior-modelling-priors . 1 INTRODUCTION . Batch reinforcement learning ( RL ) ( Ernst et al. , 2005 ; Lange et al. , 2011 ) is the problem of learning a policy from a fixed , previously recorded , dataset without the opportunity to collect new data through interaction with the environment . This is in contrast to the typical RL setting which alternates between policy improvement and environment interaction ( to acquire data for policy evaluation ) . In many real world domains collecting new data is laborious and costly , both in terms of experimentation time and hardware availability but also in terms of the human labour involved in supervising experiments . This is especially evident in robotics applications ( see e.g . Riedmiller et al . 2018 ; Haarnoja et al . 2018b ; Kalashnikov et al . 2018 for recent examples learning on robots ) . In these settings where gathering new data is expensive compared to the cost of learning , batch RL promises to be a powerful solution . There exist a wide class of off-policy algorithms for reinforcement learning designed to handle data generated by a behavior policy µ which might differ from π , the policy that we are interested in learning ( see e.g . Sutton & Barto ( 2018 ) for an introduction ) . One might thus expect solving batch RL to be a straightforward application of these algorithms . Surprisingly , for batch RL in continuous control domains , however , Fujimoto et al . ( 2018 ) found that policies obtained via the naïve application of off-policy methods perform dramatically worse than the policy that was used to generate the data . This result highlights the key challenge in batch RL : we need to exhaustively exploit the information that is in the data but avoid drawing conclusions for which there is no evidence ( i.e . we need to avoid over-valuing state-action sequences not present in the training data ) . As we will show in this paper , the problems with existing methods in the batch learning setting are further exacerbated when the provided data contains behavioral trajectories from different policies µ1 , . . . , µN which solve different tasks , or the same task in different ways ( and thus potentially execute conflicting actions ) that are not necessarily aligned with the target task that π should accomplish . We empirically show that previously suggested adaptations for off-policy learning ( Fujimoto et al. , 2018 ; Kumar et al. , 2019 ) can be led astray by behavioral patterns in the data that are consistent ( i.e . policies that try to accomplish a different task or a subset of the goals for the target task ) but not relevant for the task at hand . This situation is more damaging than learning from noisy or random data where the behavior policy is sub-optimal but is not predictable , i.e . the randomness is not a correlated signal that will be picked up by the learning algorithm . We propose to solve this problem by restricting our solutions to ‘ stay close to the relevant data ’ . This is done by : 1 ) learning a prior that gives information about which candidate policies are potentially supported by the data ( while ensuring that the prior focuses on relevant trajectories ) , 2 ) enforcing the policy improvement step to stay close to the learned prior policy . We propose a policy iteration algorithm in which the prior is learned to form an advantage-weighted model of the behavior data . This prior biases the RL policy towards previously experienced actions that also have a high chance of being successful in the current task . Our method enables stable learning from conflicting data sources and we show improvements on competitive baselines in a variety of RL tasks – including standard continuous control benchmarks and multi-task learning for simulated and real-world robots . We also find that utilizing an appropriate prior is sufficient to stabilize learning ; demonstrating that the policy evaluation step is implicitly stabilized when a policy iteration algorithm is used – as long as care is taken to faithfully evaluate the value function within temporal difference calculations . This results in a simpler algorithm than in previous work ( Fujimoto et al. , 2018 ; Kumar et al. , 2019 ) . 2 BACKGROUND AND NOTATION . In the following we consider the problem of reinforcement learning , modeling the environment as a markov decision process ( MDP ) consisting of the continuous states s ∈ S , actions a ∈ A , and transition probability distribution p ( st+1|st , at ) – describing the evolution of the system dynamics over time ( e.g . probability of reaching st+1 from state st when executing action at ) – together with the state-visitation distribution p ( s ) . The goal of reinforcement learning is to find a policy π ( a|s ) that maximizes the cumulative discounted return J ( π ) = Eπ [ ∑∞ t=1 γ tr ( st , at ) |at ∼ π ( ·|st ) , st+1 ∼ p ( ·|st , at ) , s1 ∼ p ( · ) ] , for the reward function r ( s , a ) ∈ R. We also define the state-action value function for taking action at in state st , and thereafter following π : Qπ ( st , at ) = Eπ [ ∑∞ i=t γ ir ( si , ai ) |ai ∼ π , si+1 ∼ p ( ·|si , ai ) ] , which we can relate to the objective J via J ( π ) = Es∼p ( · ) [ Ea∼π ( ·|s ) [ Qπ ∗ ( s , a ) ] ] , where π∗ is the optimal policy . We parameterize the policy πθ ( a|s ) by θ but we will omit this dependency where unambiguous . In some of the experiments we will also consider a setting where we learn about multiple tasks k ∈ { 1 , . . .K } , each with their own reward function rk ( s , a ) . We condition the policy and Q-function on the task index k ( i.e . Qπk and π ( a|s , k ) ) , changing the objective to maximize the sum of returns across all tasks . For the batch RL setting we assume that we are given a dataset Dµ containing trajectory snippets ( i.e . sub-trajectories of length N ) τ = { ( s0 , a0 ) , · · · , ( sT , aT ) } , with τ ∈ Dµ . We assume access to the reward function r for the task of interest and can evaluate it for all transitions in Dµ ( for example , r may be some function of s ) . We further assume Dµ was filled , prior to training , by following a set of arbitrary N behavior policies µ1 , · · · , µN . Note that these behavior policies may try to accomplish the task we are interested in ; or might indeed generate trajectories unrelated to the task at hand . 3 A LEARNED PRIOR FOR OFFLINE OFF-POLICY RL FROM IMPERFECT DATA . To stabilize off-policy RL from batch data , we want to restrict the learned policy to those parts of the state-action space supported by the batch . In practice this means that we need to approximately restrict the policy to the support of the empirical state-conditional action distribution . This prevents the policy from taking actions for which the Q-function can not be trained and for which it might thus give erroneous , overly optimistic values ( Fujimoto et al. , 2018 ; Kumar et al. , 2019 ) . In this paper we achieve this by adopting a policy iteration procedure – in which the policy is constrained in the improvement step . As in standard policy iteration ( Sutton & Barto , 2018 ) , the procedure consists of two alternating steps . First , starting with a given policy πi = πθi in iteration i ( with πθ0 corresponding to a randomly initialized policy distribution ) , we find an approximate action-value function Qπi ( s , a ) ≈ Q̂ ( s , a ; φi ) , with parameters φ ( Section 3.1 ) ( as with the policy we will drop the dependence on φi and write Q̂πi ( s , a ) where unambiguous ) . Second , we optimize for πi+1 with respect to Q̂πi subject to a constraint that ensures closeness to the empirical state-conditional action distribution of the batch ( Section 3.2 ) . Iterating these steps , overall , optimizes J ( π ) . We realize both policy evaluation and improvement via a fixed number of gradient descent steps – holding π and Q̂πi fixed via the use of target networks ( Mnih et al. , 2015 ) . We refer to Algorithm 1 for details . 3.1 POLICY EVALUATION . To learn the task action-value function in each iteration we minimize the squared temporal difference error for a given reward – note that when performing offline RL from a batch of data the reward r might be computed post-hoc and does not necessarily correspond to the reward optimized by the behavior policies µ1 , · · · , µN . The result after iteration i is given as φi = arg min φ E τ∼Dµ [ ( r ( st , at ) + γV̂ πi ( st+1 ) − Q̂ ( st , at ; φ ) ) 2∣∣∣ ( st , at , st+1 ) ∼ τ ] , with V̂ πi ( s ) = Ea∼πi ( ·|s ) [ Q̂ ( s , a ; φi−1 ) ] ( 1 ) We approximate the expectation required to calculate V̂ πi ( s ) with M samples from πi , i.e . V̂ πi ( s ) ≈ 1 M [ ∑M j=1 Q̂ ( s , aj ; φi−1 ) | aj ∼ πi ( ·|s ) ] . As further discussed in the related work Section 4 , the use of policy evaluation is different from the Q-learning approach pursued in Fujimoto et al . ( 2018 ) ; Kumar et al . ( 2019 ) , which requires a maximum over actions and may be more susceptible to overestimation of Q-values . We find that when enough samples are taken ( we use M = 20 ) and the policy is appropriately regularized ( see Section 3.2 ) learning is stable without additional modifications . 3.2 PRIOR LEARNING AND POLICY IMPROVEMENT . In the policy improvement step we solve the following constrained optimization problem πi+1 = arg max π E τ∼Dµ [ Ea∼π ( ·|s ) [ Q̂πi ( s , a ) ] |s ∼ τ ] s.t . E τ∼Dµ [ KL [ π ( ·|s ) ‖πprior ( ·|s ) ] |s ∼ τ ] ≤ , ( 2 ) where Dµ is the behavior data , π the policy being learned , and πprior is the prior policy . This is similar to the policy improvement step in Abdolmaleki et al . ( 2018 ) but instead of enforcing closeness to the previous policy here the constraint is with respect to a separately learned “ prior ” policy , the behavior model . The role of πprior in Equation 2 is to keep the policy close to the regime of the actions found in Dµ . We consider two different ways to express this idea by learning a prior alongside the policy optimization . For learning the prior , we first consider simply modeling the raw behavior data . This is similar to the approach of BCQ and BEAR-QL ( Fujimoto et al. , 2018 ; Kumar et al. , 2019 ) , but we use a parametric behavior model and measure distance by KL ; we refer to the related work for a discussion . The behavior model can be learned by maximizing the log likelihood of the observed data θbm = arg max θbm E τ∼Dµ |τ |∑ t=1 log πθbm ( at|st ) , ( 3 ) where θbm are the parameters of the behavior model prior . Regularizing towards the behavior model can help to prevent the use of unobserved actions , but it may also prevent the policy from improving over the behavior in Dµ . In effect , the simple behavior prior in Equation 3 regularizes the new policy towards the empirical state-conditional action distribution in Dµ . This may be acceptable for datasets dominated by successful trajectories for the task of interest or when the unsuccessful trajectories are not predictable ( i.e . they correspond to random behaviour ) . However , we here are interested in the case where Dµ is collected from imperfect data and from multiple tasks . In this case , Dµ will contain a diverse set of trajectories – both ( partially ) successful and actively harmful for the target task . With this in mind , we consider a second learned prior , the advantage-weighted behavior model , πabm , with which we can bias the RL policy to choose actions that are both supported by Dµ and also good for the current task ( i.e . keep doing actions that work ) . We can formulate this as maximizing the following objective : θabm = arg max θabm E τ∼Dµ |τ |∑ t=1 log πθabm ( at|st ) f ( R ( τt : N ) − V̂ πi ( s ) ) , with R ( τt : N ) = γN−tV̂ πi ( sN ) + N−1∑ j=t γj−tr ( sj , aj ) , ( 4 ) where f is an increasing , non-negative function , and the difference R ( τt : N ) − V̂ πi is akin to an n-step advantage function , but here calculated off-policy representing the “ advantage ” of the behavior snippet over the policy πi . This objective still tries to maximize the log likelihood of observed actions , and avoids taking actions not supported by data . However , by “ advantage weighting ” we focus the model on “ good ” actions while ignoring poor actions . We let f = 1+ ( the unit step function with f ( x ) = 1 for x ≥ 0 and 0 otherwise ) both for simplicty – to keep the number of hyperparameters to a minimum while keeping the prior broad – and because it has an intuitive interpretation : such a prior will start by covering the full data and , over time , filter out trajectories that would lead to worse performance than the current policy , until it eventually converges to the best trajectory snippets contained in the data . We note that Equation 4 is similar to a policy gradient , though samples here stem from the buffer and it will thus not necessarily converge to the optimal policy in itself ; πθabm will instead only cover the best trajectories in the data due to no importance weighting being performed for off-policy data . This bias is in fact desirable in a batch-RL setting ; we want a broad prior that only considers actions present in the data . We also note that we tried several different functions for f including exponentiation , e.g . f ( x ) = exp ( x ) , but found that choice of function did not make a significant difference in our experiments . Using either πθbm or πθabm as πprior , Equation 2 can be solved with a variety of optimization schemes . We experimented with an EM-style optimization following the derivations for the MPO algorithm ( Abdolmaleki et al. , 2018 ) , as well as directly using the stochastic value gradient of Q̂πi wrt . policy parameters ( Heess et al. , 2015 ) . It should be noted that , if the prior itself is already good enough to solve the task we can learn Qπprior – e.g . if the data stems from an expert or has sufficiently high quality . In this case learning both Qπprior and πprior becomes independent of the RL policy improvement step ; if we then set = 0 , skipping the policy improvement step , we obtain a further simplified algorithm consisting only of learning πθabm and Q πθabm ( see Figure 9 for an ablation ) . EM-style optimization We can optimize the objective from Equation 2 using a two-step procedure . Following ( Abdolmaleki et al. , 2018 ) , we first notice that the optimal π for Equation 2 can be expressed as π̂ ( a|s ) ∝ πprior ( a|s ) exp ( Q̂πi ( s , a ) /η ) , where η is a temperature that depends on the used for the KL constraint and can be found automatically by a convex optimization ( Appendix B.1 ) . Conveniently , we can sample from this distribution by querying Q̂πi using samples from πprior . These samples can then be used to learn the parametric policy by minimizing the divergence KL ( π̂‖πθi+1 ) , which is equivalent to maximizing the weighted log likelihood θi+1 = arg max θ E τ∼Dµ [ Ea∼πprior ( ·|s ) [ exp ( Q̂ πi ( s , a ) /η ) log πθ ( a|s ) |s ∼ τ ] ] , ( 5 ) which we optimize via gradient descent subject to an additional trust-region constraint on πθ given as KL ( πθi‖πθ ) < trust to ensure conservative updates ( Appendix B.1 ) . Stochastic value gradient optimization Alternatively , we can use Langrangian relaxation to turn Equation 2 into an objective amenable to gradient descent . Inserting πθ for π and relaxing results in ( θi+1 , η ) = arg max θ min η E τ∼Dµ [ Ea∼πθ ( ·|s ) [ Q̂πi ( s , a ) ] + η ( − KL [ πθ ( ·|s ) ‖πprior ( ·|s ) ] ) ] , ( 6 ) for η > 0 and which we can optimize by alternating gradient descent steps on θ and η respectively , taking the stochastic gradient of the Q-value ( Heess et al. , 2015 ) through the sampling of a ∼ πθ ( · , s ) via re-parameterization . See Appendix B.2 for a derivation of this gradient . Full pseudocode for our approach can be found in Appendix A , Algorithm 1 .

Summary: This paper proposes a method for “offline RL” (a.k.a “batch RL”), i.e. reinforcement learning from a given static dataset, with no option to perform on-policy data collection. Contrary to prior work (Fujimoto et al, Kumar et al, Agarwal et al) in this area which focuses on making Q-learning robust in the offline RL setting, the authors in this paper instead propose making the policy update step robust to the batch RL setting. They achieve this by constraining the policy that is being learned to stay close to a learned “prior policy” (using KL divergence as a distance metric). The authors provide two different ways of learning this prior policy  - the first one is a simple behavior model policy (which involves fitting a model using maximum likelihood on the entire dataset - effectively performing behavior cloning on the entire dataset), while the second is a more sophisticated “advantage behavior model”, which fits a model only on the “good” data (i.e. environment transitions with a positive advantage, where the advantage is estimated using the current policy). The authors extend the MPO and SVG algorithms to work with this constrained update step.