Advantage Weighted Regression: Simple and Scalable Off-Policy Reinforcement Learning

1 INTRODUCTION . Model-free reinforcement learning can be a general and effective methodology for training agents to acquire sophisticated behaviors with minimal assumptions on the underlying task ( Mnih et al. , 2015 ; Heess et al. , 2017 ; Pathak et al. , 2017 ) . However , reinforcement learning algorithms can be substantially more complex to implement and tune than standard supervised learning methods . Arguably the simplest reinforcement learning methods are policy gradient algorithms ( Sutton et al. , 2000 ) , which directly differentiate the expected return and perform gradient ascent . Unfortunately , these methods can be notoriously unstable and are typically on-policy ( or nearly on-policy ) , often requiring a substantial number of samples to learn effective behaviors . Our goal is to develop a reinforcement learning algorithm that is simple , easy to implement , and can readily incorporate off-policy experience data . In this work , we propose advantage-weighted regression ( AWR ) , a simple off-policy algorithm for model-free RL . Each iteration of the AWR algorithm simply consists of two supervised regression steps : one for training a value function baseline via regression onto cumulative rewards , and another for training the policy via weighted regression . The complete algorithm is shown in Algorithm 1 . 1Supplementary video : sites.google.com/view/awr-supp/ AWR can accommodate continuous and discrete actions , and can be implemented in just a few lines of code on top of standard supervised learning methods . Despite its simplicity , we find that AWR achieves competitive results when compared to commonly used on-policy and off-policy RL algorithms , and can effectively incorporate fully off-policy data , which has been a challenge for other RL algorithms . Our derivation presents an interpretation of AWR as a constrained policy optimization procedure , and provides a theoretical analysis of the use of off-policy data with experience replay . We first revisit the original formulation of reward-weighted regression , an on-policy RL method that utilizes supervised learning to perform policy updates , and then propose a number of new design decisions that significantly improve performance on a suite of standard continuous control benchmark tasks . We then provide a theoretical analysis of AWR , including the capability to incorporate off-policy data with experience replay . Although the design of AWR involves only a few simple design decisions , we show experimentally that these additions provide for a large improvement over previous methods for regression-based policy search , such as reward-weighted regression ( RWR ) ( Peters & Schaal , 2007 ) , while also being substantially simpler than more modern methods , such as MPO ( Abdolmaleki et al. , 2018 ) . We show that AWR achieves competitive performance when compared to several well-established state-of-the-art on-policy and off-policy algorithms . We further demonstrate our algorithm on challenging control tasks with complex simulated characters . 2 PRELIMINARIES . In reinforcement learning , the objective is to learn a control policy that enables an agent to maximize its expected return for a given task . At each time step t , the agent observes the state of the environment st ∈ S , and samples an action at ∈ A from a policy at ∼ π ( at|st ) . The agent then applies that action , which results in a new state st+1 and a scalar reward rt = r ( st , at ) . The goal is to learn an optimal policy that maximizes the agent ’ s expected discounted return J ( π ) , J ( π ) = Eτ∼pπ ( τ ) [ ∞∑ t=0 γtrt ] = Es∼dπ ( s ) Ea∼π ( a|s ) [ r ( s , a ) ] , ( 1 ) where pπ ( τ ) represents the likelihood of a trajectory τ = { ( s0 , a0 , r0 ) , ( s1 , a1 , r1 ) , ... } under a policy π , and γ ∈ [ 0 , 1 ) is the discount factor . dπ ( s ) = ∑∞ t=0 γ tp ( st = s|π ) represents the unnormalized discounted state distribution induced by the policy π ( Sutton & Barto , 1998 ) , and p ( st = s|π ) is the likelihood of the agent being in state s after following π for t timesteps . A popular class of algorithms for solving this problem is policy gradient ( PG ) methods , which directly estimates the gradient of the expected return with respect to the policy parameters ∇πJ ( π ) , and then updates the policy with gradient ascent . Basic PG algorithms are generally on-policy methods , which require the data to be sampled from the same policy that is being optimized . This can result in poor sample efficiency , but PG algorithms can be modified to utilize off-policy data . An alternative class of RL methods is expectation-maximization algorithms . Instead of estimating the gradient of the expected return , EM algorithms first construct an estimate of the optimal policy using a dataset of experiences ( E-step ) , and then projects this estimate onto the space of parameterized policies ( M-step ) . An early example of an EM-based RL algorithm is reward-weighted regression ( RWR ) ( Peters et al. , 2010 ) . At each iteration , the E-step constructs an estimate of the optimal policy according to π∗ ( a|s ) ∝ πk ( a|s ) exp ( Rs , a/β ) , where πk represents the policy at the kth iteration of the algorithm , Rs , a = ∑∞ t=0 γ trt is the return , and β > 0 is a temperature parameter . Then the M-step projects π∗ onto the space of parameterized policies by solving a supervised regression problem : πk+1 = arg max π Es∼dπk ( s ) Ea∼πk ( a|s ) [ log π ( a|s ) exp ( 1 β Rs , a ) ] . ( 2 ) The RWR update can be interpreted as fitting a new policy πk+1 to samples collected under the current policy πk , where the likelihood of each action is weighted by the exponentiated return received for that action . Since EM algorithms do not directly estimate the gradient of the expected return with respect to the current policy , they can be more amenable to learning from off-policy data . Algorithm 1 Advantage-Weighted Regression 1 : π1 ← random policy 2 : D ← ∅ 3 : for iteration k = 1 , ... , kmax do 4 : add trajectories { τi } sampled via πk to D 5 : V Dk ← arg minV Es , a∼D [ ∣∣∣∣RDs , a − V ( s ) ∣∣∣∣2 ] 6 : πk+1 ← arg maxπ Es , a∼D [ log π ( a|s ) exp ( 1 β ( RDs , a − V Dk ( s ) ) ) ] 7 : end for 3 ADVANTAGE-WEIGHTED REGRESSION . In this work , we present advantage-weighted regression ( AWR ) , a simple off-policy RL algorithm based on reward-weighted regression . We first provide an overview of the complete advantageweighted regression algorithm , and then describe its theoretical motivation and analyze its properties . The complete AWR algorithm is summarized in Algorithm 1 . Each iteration k of AWR consists of the following simple steps . First , the current policy πk ( a|s ) is used to sample a batch of trajectories { τi } that are then stored in the replay buffer D , which is structured as a first-in first-out ( FIFO ) queue , as is common for off-policy reinforcement learning algorithms ( Mnih et al. , 2015 ; Lillicrap et al. , 2016 ) . Then , the entire buffer D is used to fit a value function V Dk ( s ) to the trajectories in the replay buffer , which can be done with simple Monte Carlo return estimates RDs , a = ∑T t=0 γ trt . Finally , the same buffer is used to fit a new policy using advantage-weighted regression , where each state-action pair in the buffer is weighted according to the exponentiated advantage exp ( 1βA D ( s , a ) ) , with the advantage given by AD ( s , a ) = RDs , a − V D ( s ) and β is a hyperparameter . AWR uses only supervised regression as learning subroutines , making the algorithm very simple to implement . In the following subsections , we first motivate the algorithm as an approximation to a constrained policy search problem , and then extend our analysis to incorporate experience replay . 3.1 DERIVATION . In this section , we derive the AWR algorithm as an approximate optimization of a constrained policy search problem . Our goal is to find a policy that maximizes the expected improvement η ( π ) = J ( π ) − J ( µ ) over a sampling policy µ ( a|s ) . We first derive AWR for the setting where the sampling policy is a single Markovian policy . Then , in the next section , we extend our result to the setting where the data is collected from multiple policies , as in the case of experience replay that we use in practice . The expected improvement η ( π ) can be expressed in terms of the advantage Aµ ( s , a ) = Rµs , a − V µ ( s ) with respect to the sampling policy µ ( Kakade & Langford , 2002 ; Schulman et al. , 2015 ) : η ( π ) = Es∼dπ ( s ) Ea∼π ( a|s ) [ A µ ( s , a ) ] = Es∼dπ ( s ) Ea∼π ( a|s ) [ Rµs , a − V µ ( s ) ] , ( 3 ) where Rµs , a denotes the return obtained by performing action a in state s and following µ for the following timesteps , and V µ ( s ) = ∫ a µ ( a|s ) Ras da corresponds to the value function of µ . This objective differs from the ones used in the derivations of related algorithms , such as RWR and REPS ( Peters & Schaal , 2007 ; Peters et al. , 2010 ; Abdolmaleki et al. , 2018 ) , which maximize the expected return J ( π ) instead of the expected improvement . The expected improvement directly gives rise to an objective that involves the advantage . We will see later that this yields weights for the policy update that differ in a subtle but important way from standard reward-weighted regression . As we show in our experiments , this difference results in a large empirical improvement . The objective in Equation 3 can be difficult to optimize due to the dependency between dπ ( s ) and π , as well as the need to collect samples from π . Following Schulman et al . ( 2015 ) , we can instead optimize an approximation η̂ ( π ) of η ( π ) using the state distribution of µ : η̂ ( π ) = Es∼dµ ( s ) Ea∼π ( a|s ) [ Rµs , a − V µ ( s ) ] . ( 4 ) Here , η̂ ( π ) matches η ( π ) to first order ( Kakade & Langford , 2002 ) , and provides a good estimate of η if π and µ are close in terms of the KL-divergence ( Schulman et al. , 2015 ) . Using this objective , we can formulate the following constrained policy search problem : arg max π ∫ s dµ ( s ) ∫ a π ( a|s ) [ Rµs , a − V µ ( s ) ] da ds ( 5 ) s.t . ∫ s dµ ( s ) DKL ( π ( ·|s ) ||µ ( ·|s ) ) ds ≤ . ( 6 ) The constraint in Equation 6 ensures that the new policy π is close to the data distribution of µ , and therefore the surrogate objective η̂ ( π ) remains a reasonable approximation to η ( π ) . We refer the reader to Schulman et al . ( 2015 ) for a detailed derivation and an error bound . We can derive AWR as an approximate solution to this constrained optimization . This derivation follows a similar procedure as Peters et al . ( 2010 ) , and begins by forming the Langrangian of the constrained optimization problem presented above , L ( π , β ) = ∫ s dµ ( s ) ∫ a π ( a|s ) [ Rµs , a − V µ ( s ) ] da ds + β ( − ∫ s dµ ( s ) DKL ( π ( ·|s ) ||µ ( ·|s ) ) ds ) , ( 7 ) where β is a Lagrange multiplier . Differentiating L ( π , β ) with respect to π ( a|s ) and solving for the optimal policy π∗ results in the following expression for the optimal policy π∗ ( a|s ) = 1 Z ( s ) µ ( a|s ) exp ( 1 β ( Rµs , a − V µ ( s ) ) ) , ( 8 ) with Z ( s ) being the partition function . A detailed derivation is available in Appendix A . If π is represented by a function approximator ( e.g. , a neural network ) , a new policy can be obtained by projecting π∗ onto the manifold of parameterized policies , arg min π Es∼D [ DKL ( π∗ ( ·|s ) ||π ( ·|s ) ) ] ( 9 ) = arg max π Es∼dµ ( s ) Ea∼µ ( a|s ) [ log π ( a|s ) exp ( 1 β ( Rµs , a − V µ ( s ) ) ) ] . ( 10 ) While this derivation for AWR largely follows the derivations used in prior work ( Peters et al. , 2010 ; Abdolmaleki et al. , 2018 ) , our expected improvement objective introduces a baseline V µ ( s ) to the policy update , which as we show in our experiments , is a crucial component for an effective algorithm . A similar advantage-weighting scheme has been previously used for fitted Q-iteration ( Neumann & Peters , 2009 ) , where the policy is given by π ( a|s ) = 1Z ( s ) exp ( ( Rµs , a − V µ ( s ) ) /β ) . In this definition , the likelihood of an action does not depend on its likelihood under the sampling distribution , and therefore does not enforce a trust region with respect to the sampling distribution . Next , we extend AWR to incorporate experience replay for off-policy training , where the sampling policy is no longer a single policy , but rather a mixture of policies from past iterations .

Summary: This paper proposes an off-policy reinforcement learning method in which the model parameters have been updated using the regression style loss function. Specifically, this method uses two regression update steps: one update value function and another one update policy using weighted regression. To compare the proposed method with others [main comprasion], 6 MuJoCo tasks are used for continuous control and LunarLander-v2 for discrete space.