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 . However , RL 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 , often requiring a substantial number of samples to learn effective behaviors . Our goal is to develop an RL algorithm that is simple , easy to implement , and can readily incorporate off-policy 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 . 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 ( RWR ) ( Peters & Schaal , 2007 ) , 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 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 RWR , 1Supplementary video : sites.google.com/view/awr-supp/ while also being substantially simpler than more modern methods , such as MPO ( Abdolmaleki et al. , 2018b ) . We show that AWR achieves competitive performance when compared to several well-established state-of-the-art on-policy and off-policy algorithms . 2 PRELIMINARIES . In reinforcement learning , the objective is to learn a policy that maximizes an agent ’ s expected return . At each time step t , the agent observes the state of the environment st , and samples an action 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 a policy that maximizes the expected return J ( π ) , J ( π ) = Eτ∼pπ ( τ ) [ ∞∑ t=0 γtrt ] = Es∼dπ ( s ) , a∼π ( 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 . Our proposed AWR algorithm builds on ideas from reward-weighted regression ( RWR ) ( Peters & Schaal , 2007 ) , a policy search algorithm based on an expectation-maximization framework . 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 , 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 from the current policy πk , where the likelihood of each action is weighted by the exponentiated return for that action . 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 AWR algorithm , and then describe its theoretical motivation and analyze its properties . The 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 ( Mnih et al. , 2015 ) . Then , a value function V Dk ( s ) is fitted to all trajectories in the replay buffer D , which can be done with simple Monte Carlo return estimatesRDs , 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 . In the following subsections , we first motivate AWR as 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 data from multiple policies , as in the case of experience replay . The expected improvement η ( π ) can be expressed in terms of the advantage Aµ ( s , a ) = Rµs , a − V µ ( s ) with respect to µ ( Kakade & Langford , 2002 ; Schulman et al. , 2015 ) : η ( π ) = 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 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 REPS ( Peters & Schaal , 2007 ; Peters et al. , 2010 ; Abdolmaleki et al. , 2018b ) , 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 a policy update that differ in a subtle but important way from standard RWR . 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 Lagrangian of the 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. , 2018b ) , 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 the sampling distribution , and therefore does not enforce a trust region with respect to µ .

Summary: This paper focuses on developing an RL learning algorithm that is simple and can significantly improve performance over existing algorithms. The paper presents the algorithm AWR, which is an extension of the algorithm reward weighted regression. The primary extensions of RWR are using the advantage function instead of the q function and the ability to use experience replay. Experiments on common environments are conducted to evaluate the performance of AWR and compare it to other algorithms. There are ablation experiments to justify the choice of some of the extensions.