Probabilistic Mixture-of-Experts for Efficient Deep Reinforcement Learning

1 INTRODUCTION . The mixture-of-experts method ( MOE ) ( Jacobs et al. , 1991a ) is testified to be capable of improving the generalisation ability of reinforcement learning ( RL ) agents ( Hausknecht & Stone , 2016a ; Peng et al. , 2016 ; Neumann et al. ) . Among these methods , the Gaussian Mixture Models ( GMM ) are promising to model multimodal policy in RL ( Peng et al. , 2019 ; Akrour et al. , 2020 ) , in which distinguishable experts or so-called primitives are learned . The distinguishable experts can propose several solutions for a task and have a larger range of exploration space , which can potentially lead to better task performance and sample efficiency compared to its unimodal counterpart ( Bishop , 2007 ) . The multimodal policy can be learned by various methods , such as a two-stage training approach ( Peng et al. , 2019 ) , specific clustering method ( Akrour et al. , 2020 ) , or especially parameterised actions design ( Hausknecht & Stone , 2016b ) . However , these methods are limited , neither applicable to complicated scenarios such as high-dimensional continuous control tasks nor the training algorithms are too complex to deal with general utility . To the best of our knowledge , the present DRL algorithms for general utility do not deploy MOE to model the multimodal policy mainly due to the lack of differentiability , or without explicit probabilistic representation . Therefore , in the policy gradientbased algorithms ( Sutton et al. , 1999a ) , the gradient of the performance concerning the policy parameters is undifferentiated . The undifferentiability problem also remains to learn a deep neural network policy thus making the combinations of MOE and DRL not trivial . In this paper , we propose a probabilistic framework to tackle the undifferentiated problem by holding the mixture distribution assumption . We will still use the GMM to model the multimodal policies . Once the undifferentiated problem is solved , our training methods can be combined with the policy gradient algorithms by simply setting the number of experts ( mixtures ) greater than one . Hereafter , the contribution can be summarised as follows : • We analyse the undifferentiability problem of approximating policy as the GMM in DRL and its associated drawbacks . • We propose an end-to-end training method to obtain the primitives with probability in a frequentist manner to solve the undifferentiability problem . • Our experiments show the proposed method can achieve better task performance and sample efficiency by exploring larger behaviours space , especially in complicated continuous control tasks , compared with unimodal RL algorithms and three different MOE methods or option frameworks . 2 RELATED WORK . Hierarchical Policies There are two main related hierarchical policy structures . The feudal schema ( Dayan & Hinton , 1992 ) has two types of agents : managers and workers . The managers first make high-level decisions , then the workers make low-level actions according to these high-level decisions . The options framework ( Sutton et al. , 1999b ) has an upper-level agent ( policy-over-options ) , which decides whether the lower level agent ( sub-policy ) should start or terminate . In the early years , it ’ s the subject of research to discover temporal abstractions autonomously often in discrete actions and the state space ( McGovern & Barto , 2001 ; Menache et al. , 2002 ; Simsek & Barto , 2008 ; Silver & Ciosek , 2012 ) . Recently , ( Mankowitz et al. , 2016 ) proposes a method that assumes the initiation sets and termination functions have particular structures . ( Kulkarni et al. , 2016 ) uses internal and extrinsic rewards to learn sub-policies and policy-over-options . ( Bacon et al. , 2017 ) trains sub-policies and policy-over-options in end-to-end fusion with a deep termination function . ( Vezhnevets et al. , 2017 ) generalises the feudal schema into continuous action space and uses an embedding operation to solve the indifferentiable problem . ( Peng et al. , 2016 ) introduces a mixture of actor-critic experts approaches to learn terrain-adaptive dynamic locomotion skills . ( Peng et al. , 2019 ) changes the mixture-of-experts distribution addition expression into the multiplication expression . Mixture-of-Experts and Ensemble Methods To speed up the learning and improve the generalisation ability on different scenarios , Jacobs et al . ( 1991a ) proposed to use several different expert networks instead of a single one . To partition the data space and assign different kernels for different spaces , Lima et al . ( 2007 ) ; Yao et al . ( 2009 ) combines MOE with SVM . To break the dependency among training outputs and speed up the convergence , Gaussian process ( GP ) is generalised similarly to MOE ( Tresp , 2000 ; Yuan & Neubauer , 2008 ; Luo & Sun , 2017 ) . MOE can be also combined with RL ( Doya et al. , 2002 ; Neumann et al . ; Peng et al. , 2016 ; Hausknecht & Stone , 2016a ; Peng et al. , 2019 ) , in which the policies are modelled as probabilistic mixture models and each expert aim to learn distinguishable policies . Policy-based RL Policy-based RL aims to find the optimal policy to maximise the expected return through gradient updates . Among various algorithms , Actor-critic is often employed ( Barto et al. , 1983 ; Sutton & Barto , 1998 ) . Off-policy algorithms ( O ’ Donoghue et al. , 2016 ; Lillicrap et al. , 2016 ; Gu et al. , 2017 ; Tuomas et al. , 2018 ) are more sample efficient than on-policy ones ( Peters & Schaal , 2008 ; Schulman et al. , 2017 ; Mnih et al. , 2016 ; Gruslys et al. , 2017 ) . However , the learned policies are still unimodal . 3 METHOD . 3.1 NOTATION . The model-free RL problem can be formulated by Markov Decision Process ( MDP ) , denoted as a tuple ( S , A , P , r ) , where S and A are continuous state and action space , respectively . The agent observes state st ∈ S and takes an action at ∈ A at time step t. The environment emits a reward r : S ×A → [ rmin , rmax ] and transitions to a new state st+1 according to the transition probabilities P : S × S × A → [ 0 , ∞ ) . In deep reinforcement learning algorithms , we always use the Q-value functionQ ( st , at ) to describe the expected return after taking an action at in the state st . The Q-value can be iteratively computed by applying the Bellman backup given by : Q ( st , at ) , Est+1∼P [ r ( st , at ) + γEat+1∼π [ Q ( st+1 , at+1 ) ] ] . ( 1 ) Our goal is to maximise the expected return : πΘ∗ ( at|st ) = arg max πΘ ( at|st ) Eat∼πΘ ( at|st ) [ Q ( st , at ) ] , ( 2 ) where Θ denotes the parameters of the policy network π . With Q-value network ( critic ) Qφ parameterised by φ , Stochastic gradient descent ( SGD ) based approaches are usually used to update the policy network : Θ = Θ +∇ΘEa∼πΘ ( at|st ) [ Qφ ( st , at ) ] . ( 3 ) 3.2 PROBABILISTIC MIXTURE-OF-EXPERTS ( PMOE ) . The proposed PMOE method decomposes a stochastic policy π as a mixture of low-level policies while retaining the probabilistic properties of the stochastic policy as a probability distribution , with the following formula : π { θ , ψ } ( at|st ) = K∑ i=1 wθi ( st ) πψi ( at|st ) , s.t . K∑ i=1 wθi = 1 , wθi > 0 , ( 4 ) where each πψi denotes the action distribution within each low-level policy , i.e . a primitive , and K denotes the number of primitives . wψi is the weight that specifies the probability of the activating primitive πψi , which is called the routing function . θi and ψi are parameters of wθi and πψi , respectively . After the policy decomposition with PMOE method , we can rewrite the update rule in Eq . 3 as : θ = θ +∇θEat∼π { θ , ψ } ( at|st ) [ Qφ ( st , at ) ] , ψ = ψ +∇ψEat∼π { θ , ψ } ( at|st ) [ Qφ ( st , at ) ] . ( 5 ) In practice , we usually apply a Gaussian distribution for either a unimodal policy or the low-level policies here in PMOE , making the overall stochastic policy with PMOE to be a GMM . However , sampling from the distributions of primitives usually embeds a sampling process from a categorical distribution w , which makes the differential calculation of policy gradients commonly applied in DRL hard to achieve . We provide a theoretically guaranteed solution for approximating the gradients in the sampling process of PMOE and successfully optimising the PMOE policy model within DRL , which will be described in details . 3.3 LEARNING THE ROUTING . To optimise the routing function w , which involves a sampling process from a categorical distribution , we propose a frequency loss function to approximate the gradients , which is theoretically proved to approximate w as the probability of the corresponding primitive being the optimal one w.r.t . the Q-value function . Specifically , given a state st , we sample one action ait from each primitive πψi , to get a total of K actions { ait ; i = 1 , 2 , · · · , K } , and compute K Q-value estimations { Qφ ( st , ait ) ; i = 1 , 2 , · · · , K } for each of the actions . Then we select an “ optimal ” primitive index as j = arg maxiQφ ( st , a i t ) . Then we encode a one-hot code vector v = [ v1 , v2 , · · · , vK ] with : vj = { 1 , if j = arg max i Qφ ( st , a i t ) ; 0 , otherwise . ( 6 ) Here we define a frequency loss function as : Lfreq = ( v − w ) ( v − w ) T , w = [ wθ1 , wθ2 , · · · , wθK ] . ( 7 ) We use the proposed frequency loss Lfreq as a smooth and differentiable function to update the routing function parameters θ , which is guaranteed to approximate wθi as the probability of the i-th primitive being the optimal primitive for current state . Detailed proof is provided in Appendix B . 3.4 LEARNING THE PRIMITIVE . To update the ψi within each primitive , we provide two approaches of optimising the primitives : back-propagation-all and back-propagation-max manners . For the back-propagation-all approach , we update all the primitive : Lbpapri = − K∑ i Qφ ( st , a i t ) , a i t ∼ πψi ( at|st ) . ( 8 ) For the back-propagation-max approach , we use the highest Q-value estimation as the primitive loss : Lbpmpri = −max i { Qφ ( st , ait ; i = 1 , 2 , · · · , K ) } , ait ∼ πψi ( at|st ) . ( 9 ) With either approach , we have the stochastic policy gradients as following : ∇ψiLpri =−∇ψjEπψi [ Qφ ( st , at ) ] =Eπψi [ −Qφ ( st , at ) ∇ψj log πψj ( at|st ) ] ( 10 ) Ideally , both approaches are feasible for learning a PMOE model . However , in practice , we find that the back-propagation-all approach will tend to learn primitives that are close to each other , while the back-propagation-max approach is capable of keeping primitives distinguishable . The phenomenon is demonstrated in our experimental analysis . Therefore , we adopt the back-propagation-max approach as the default setting of PMOE model without additional clarification .

Summary: The paper studies the problem of differentiating through the policy return when the policy is a Gaussian mixture model. The main contribution of the paper is a heuristic approach for computing this gradient. Having defined the policy update, the authors integrate it to two RL algorithms: PPO and SAC. The experiments show that the algorithms with GMMs behave roughly the same as with a single Gaussian save for one (SAC) or two (PPO) environments. On the other side the authors show that their algorithm can learn multiple solutions on a reaching task and that exploration is better behaved on this task too.