V-MPO: On-Policy Maximum a Posteriori Policy Optimization for Discrete and Continuous Control

1 INTRODUCTION . Deep reinforcement learning ( RL ) with neural network function approximators has achieved superhuman performance in several challenging domains ( Mnih et al. , 2015 ; Silver et al. , 2016 ; 2018 ) . Some of the most successful recent applications of deep RL to difficult environments such as Dota 2 ( OpenAI , 2018a ) , Capture the Flag ( Jaderberg et al. , 2019 ) , Starcraft II ( Vinyals et al. , 2019 ) , and dexterous object manipulation ( OpenAI , 2018b ) have used policy gradient-based methods such as Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ) and the Importance-Weighted Actor-Learner Architecture ( IMPALA ) ( Espeholt et al. , 2018 ) , both in the approximately on-policy setting . Policy gradients , however , can suffer from large variance that may limit performance , especially for high-dimensional action spaces ( Wu et al. , 2018 ) . In practice , moreover , policy gradient methods typically employ carefully tuned entropy regularization in order to prevent policy collapse . As an alternative to policy gradient-based algorithms , in this work we introduce an approximate policy iteration algorithm that adapts Maximum a Posteriori Policy Optimization ( MPO ) ( Abdolmaleki et al. , 2018a ; b ) to the on-policy setting . The modified algorithm , V-MPO , relies on a learned state-value function V ( s ) instead of the state-action value function used in MPO . Like MPO , rather than directly updating the parameters in the direction of the policy gradient , V-MPO first constructs a target distribution for the policy update subject to a sample-based KL constraint , then calculates the gradient that partially moves the parameters toward that target , again subject to a KL constraint . ∗Equal contribution As we are particularly interested in scalable RL algorithms that can be applied to multi-task settings where a single agent must perform a wide variety of tasks , we show for the case of discrete actions that the proposed algorithm surpasses previously reported performance in the multi-task setting for both the Atari-57 ( Bellemare et al. , 2012 ) and DMLab-30 ( Beattie et al. , 2016 ) benchmark suites , and does so reliably without population-based tuning of hyperparameters ( Jaderberg et al. , 2017a ) . For a few individual levels in DMLab and Atari we also show that V-MPO can achieve scores that are substantially higher than has previously been reported in the single-task setting , especially in the challenging Ms. Pacman . V-MPO is also applicable to problems with high-dimensional , continuous action spaces . We demonstrate this in the context of learning to control both a 22-dimensional simulated humanoid from full state observations—where V-MPO reliably achieves higher asymptotic performance than previous algorithms—and a 56-dimensional simulated humanoid from pixel observations ( Tassa et al. , 2018 ; Merel et al. , 2019 ) . In addition , for several OpenAI Gym tasks ( Brockman et al. , 2016 ) we show that V-MPO achieves higher asymptotic performance than has previously been reported . 2 BACKGROUND AND SETTING . We consider the discounted RL setting , where we seek to optimize a policy π for a Markov Decision Process described by states s , actions a , initial state distribution ρenv0 ( s0 ) , transition probabilities Penv ( st+1|st , at ) , reward function r ( st , at ) , and discount factor γ ∈ ( 0 , 1 ) . In deep RL , the policy πθ ( at|st ) , which specifies the probability that the agent takes action at in state st at time t , is described by a neural network with parameters θ . We consider problems where both the states s and actions a may be discrete or continuous . Two functions play a central role in RL : the state-value function V π ( st ) = Eat , st+1 , at+1 , ... [ ∑∞ k=0 γ kr ( st+k , at+k ) ] and the state-action value function Qπ ( st , at ) = Est+1 , at+1 , ... [ ∑∞ k=0 γ kr ( st+k , at+k ) ] = r ( st , at ) + γEst+1 [ V π ( st+1 ) ] , where s0 ∼ ρenv0 ( s0 ) , at ∼ π ( at|st ) , and st+1 ∼ Penv ( st+1|st , at ) . In the usual formulation of the RL problem , the goal is to find a policy π that maximizes the expected return given by J ( π ) = Es0 , a0 , s1 , a1 , ... [ ∑∞ t=0 γ tr ( st , at ) ] . In policy gradient algorithms ( Williams , 1992 ; Sutton et al. , 2000 ; Mnih et al. , 2016 ) , for example , this objective is directly optimized by estimating the gradient of the expected return . An alternative approach to finding optimal policies derives from research that treats RL as a problem in probabilistic inference , including Maximum a Posteriori Policy Optimization ( MPO ) ( Levine , 2018 ; Abdolmaleki et al. , 2018a ; b ) . Here our objective is subtly different , namely , given a suitable criterion for what are good actions to take in a certain state , how do we find a policy that achieves this goal ? As was the case for the original MPO algorithm , the following derivation is valid for any such criterion . However , the policy improvement theorem ( Sutton & Barto , 1998 ) tells us that a policy update performed by exact policy iteration , π ( s ) = arg maxa [ Qπ ( s , a ) − V π ( s ) ] , can improve the policy if there is at least one state-action pair with a positive advantage and nonzero probability of visiting the state . Motivated by this classic result , in this work we specifically choose an exponential function of the advantages Aπ ( s , a ) = Qπ ( s , a ) − V π ( s ) . Notation . In the following we use ∑ s , a to indicate both discrete and continuous sums ( i.e. , integrals ) over states s and actions a depending on the setting . A sum with indices only , such as ∑ s , a , denotes a sum over all possible states and actions , while ∑ s , a∼D , for example , denotes a sum over sample states and actions from a batch of trajectories ( the “ dataset ” ) D . 3 RELATED WORK . V-MPO shares many similarities , and thus relevant related work , with the original MPO algorithm ( Abdolmaleki et al. , 2018a ; b ) . In particular , the general idea of using KL constraints to limit the size of policy updates is present in both Trust Region Policy Optimization ( TRPO ; Schulman et al. , 2015 ) and Proximal Policy Optimization ( PPO ) ( Schulman et al. , 2017 ) ; we note , however , that this corresponds to the E-step constraint in V-MPO . It is worth noting here the following main differences with MPO , which is conceptually quite similar to V-MPO . MPO is primarily designed to be a sample-efficient off-policy algorithm in which the E-step constructs a conditional target distribution q ( a|s ) , which requires a state-action value function Q ( s , a ) that can evaluate multiple sampled actions for a given state . In contrast , V-MPO is primarily ( though not exclusively ) designed to be an on-policy algorithm in which the E-step constructs a joint distribution ψ ( s , a ) , and in the absence of a learned Q-function only one action per state is used . In this regard V-MPO can also be compared to Fitted Q-iteration by Advantage Weighted Regression ( Neumann & Peters , 2009 ) , which learns a Q-function but uses only one action per state . V-MPO can also be related to Relative Entropy Policy Search ( REPS ) ( Peters et al. , 2008 ) . Two distinguishing features of V-MPO from REPS are the introduction of the M-step KL constraint and the use of top-k advantages . Moreover , in REPS the value function is a linear function of a learned feature representation whose parameters are trained by matching the feature distributions under the policy ’ s stationary state distribution . In V-MPO , the nonlinear neural network value function is instead learned directly from n-step returns . Interestingly , previous attempts to use REPS with neural network function approximators reported very poor performance , being particularly prone to local optima ( Duan et al. , 2016 ) . In contrast , we find that the principles of EM-style policy optimization , when combined with this learned value function and appropriate constraints , can reliably train powerful neural networks , including transformers , for RL tasks . Like V-MPO , Supervised Policy Update ( SPU ) ( Vuong et al. , 2019 ) seeks to exactly solve an optimization problem and fit the parametric policy to this solution . As we argue in Appendix D , however , SPU uses this nonparametric distribution quite differently from V-MPO ; as a result , the final algorithm is closer to a policy gradient algorithm such as PPO . 4 METHOD . V-MPO is an approximate policy iteration ( Sutton & Barto , 1998 ) algorithm with a specific prescription for the policy improvement step . In general , policy iteration uses the fact that the true state-value function V π corresponding to policy π can be used to obtain an improved policy π′ . Thus we can 1 . Generate trajectories τ from an old policy πθold ( a|s ) whose parameters θold are fixed . To control the amount of data generated by a particular policy , we use a target network which is fixed for Ttarget learning steps ( Fig . 1a ) . 2 . Evaluate the policy πθold ( a|s ) by learning the value function V πθold ( s ) from empirical returns and estimating the corresponding advantages Aπθold ( s , a ) for the actions that were taken ( Section 4.1 ) . 3 . Based onAπθold ( s , a ) , estimate an improved policy πθ ( a|s ) which we call the “ online ” policy to distinguish it from the fixed target network ( Section 4.2 ) . The first two steps are standard , and describing V-MPO ’ s approach to step 3 is the essential contribution of this work . At a high level , our strategy is to first construct a nonparametric target distribution for the policy update , then partially move the parametric policy towards this distribution subject to a KL constraint . We first review policy evaluation ( step 2 ) in Section 4.1 , then derive the V-MPO policy improvement ( step 3 ) in Section 4.2 . Ultimately , we use gradient descent to optimize a single , relatively simple loss , which is given in Eq . 10 following the derivation . A summary of the full algorithm is also presented in Algorithm 1 . 4.1 POLICY EVALUATION . In the present setting , policy evaluation means learning an approximate state-value function V π ( s ) given a policy π ( a|s ) , which we keep fixed for Ttarget learning steps ( i.e. , batches of trajectories ) . We note that the value function corresponding to the target policy is instantiated in the “ online ” network receiving gradient updates ; bootstrapping uses the online value function , as it is the best available estimate of the value function for the target policy . Thus in this section π refers to πθold , while the value function update is performed on the current φ , which may share parameters with the current θ . We fit a parametric value function V πφ ( s ) with parameters φ by minimizing the squared loss LV ( φ ) = 1 2|D| ∑ st∼D ( V πφ ( st ) −G ( n ) t ) 2 , ( 1 ) where G ( n ) t is the standard n-step target for the value function at state st at time t ( Sutton & Barto , 1998 ) . This return uses the actual rewards in the trajectory and bootstraps from the value function for the rest : for each ` = t , . . . , t + n − 1 in an unroll , G ( n ) ` = ∑t+n−1 k= ` γ k− ` rk + γ t+n− ` V πφ ( st+n ) . The advantages , which are the key quantity of interest for the policy improvement step in V-MPO , are then given by Aπ ( st , at ) = G ( n ) t − V πφ ( st ) for each st , at in the batch of trajectories . PopArt normalization . As we are interested in the multi-task setting where a single agent must learn a large number of tasks with differing reward scales , we used PopArt ( van Hasselt et al. , 2016 ; Hessel et al. , 2018 ) for the value function , even when training on a single task . We observed benefits in using PopArt even in the single-task setting , partly due to the fact that we do not tune the relative weighting of the policy evaluation and policy improvement losses despite sharing most parameters for the policy and value networks . Specifically , the value function outputs a separate value for each task in normalized space , which is converted to actual returns by a shift and scaling operation , the statistics of which are learned during training . We used a scale lower bound of 10−2 , scale upper bound of 106 , and learning rate of 10−4 for the statistics . The lower bound guards against numerical issues when rewards are extremely sparse . Importance-weighting for off-policy data . It is possible to importance-weight the samples using V-trace to correct for off-policy data ( Espeholt et al. , 2018 ) , for example when data is taken from a replay buffer . For simplicity , however , no importance-weighting was used for the experiments presented in this work , which were mostly on-policy .

Summary: The paper proposes an online variant of MPO, V-MPO. Compared to MPO, the main difference seems to be in the E-step. Instead of optimizing the non-parametric distribution towards a parameterized Q-function, V-MPO learns the V-function and updates the non-parametric distribution towards the advantages, which can be estimated on the samples of the last roll-outs based on the empirical returns and the learned V-function. Of course, updating towards (exponentiated and normalized) Q- or A-function does not make a difference. There are also some minor changes (which might still be crucial) such as an entropy constraint in the M-step and using only the top-k advantages during the E-Step (an option that was discussed in Abdolmaleki, et al. 2018a).