Decentralized Cross-Entropy Method for Model-Based Reinforcement Learning

Cross-Entropy Method ( CEM ) is a popular approach to planning in model-based reinforcement learning . It has so far always taken a centralized approach where the sampling distribution is updated centrally based on the result of a top-k operation applied to all samples . We show that such a centralized approach makes CEM vulnerable to local optima and impairs its sample efficiency , even in a one-dimensional multi-modal optimization task . In this paper , we propose Decentralized CEM ( DecentCEM ) where an ensemble of CEM instances run independently from one another and each performs a local improvement of its own sampling distribution . In the exemplar optimization task , the proposed decentralized approach DecentCEM finds the global optimum much more consistently than the centralized CEM approaches that use either a single or a mixture of Gaussian distributions . Also , we show that DecentCEM is theoretically sound . Further , we extend the decentralized approach to sequential decision-making problems where we show in several continuous control benchmark environments that it provides an effective mechanism to improve the performance of CEM algorithms , under the same sample budget for planning . 1 INTRODUCTION . Model-based reinforcement learning ( MBRL ) uses a model as a proxy of the environment for planning actions in multiple steps . This paper studies planning in MBRL with a specific focus on the Cross-Entropy Method ( CEM ) ( De Boer et al. , 2005 ; Mannor et al. , 2003 ) , which is popular in MBRL due to its ease of use and strong empirical performance ( Chua et al. , 2018 ; Hafner et al. , 2019 ; Wang & Ba , 2020 ; Zhang et al. , 2021 ; Yang et al. , 2020 ) . CEM is a stochastic , derivativefree optimization method . It uses a sampling distribution to generate imaginary trajectories of environment-agent interactions with the model . These trajectories are then ranked based on their returns computed from the rewards given by the model . The sampling distribution is updated to increase the likelihood of producing the top-k trajectories with higher returns . These steps are iterated and eventually yield an improved distribution over the action sequences to guide the action execution in the real environment . Despite the strong empirical performance of CEM for planning , it is prone to two problems : ( 1 ) lower sample efficiency as the dimensionality of solution space increases , and ( 2 ) the Gaussian distribution that is commonly used for sampling may cause the optimization to get stuck in local optima of multi-modal solution spaces commonly seen in real-world problems . Previous works addressing these problems either add gradient-based updates of the samples to optimize the parameters of CEM , or adopt more expressive sampling distributions , such as using Gaussian Mixture Model ( Okada & Taniguchi , 2020 ) or masked auto-regressive neural network ( Hakhamaneshi et al. , 2020 ) . Nevertheless , all CEM implementations to date are limited to a centralized formulation where the ranking step involves all samples . As analyzed below and in Section 3 , such a centralized design makes CEM vulnerable to local optima and impairs its sample efficiency . We propose Decentralized CEM ( DecentCEM ) to address the above problems . Rather than ranking all samples , as in the centralized design , our method distribute the sampling budget across an ensemble of CEM instances . These instances run independently from one another , and each performs a local improvement of its own sampling distribution based on the ranking of its generated samples . The best action is then aggregated by taking an arg max among the solution of the instances . It recovers the conventional CEM when the number of instance is one . We hypothesize that by shifting to this decentralized design , CEM can be less susceptible to premature convergence caused by the centralized ranking step . As illustrated in Fig . 1 , the centralized sampling distribution exhibits a bias toward the sub-optimal solutions near top right , due to the global top-k ranking . This bias would occur regardless of the family of distributions used . In comparison , a decentralized approach could maintain enough diversity thanks to its local top-k ranking in each sampling instance . Through a one-dimensional multi-modal optimization problem in Section 3 , we show that DecentCEM empirically finds the global optimum more consistently than centralized CEM approaches that use either a single Gaussian or a mixture of Gaussian distributions . Also we show that DecentCEM is theoretically sound that it converges almost surely to a local optimum . We further apply DecentCEM to sequential decision making problems and use neural networks to parameterize the sampling distribution in each CEM instance . Empirical results in several continuous control benchmarks suggest that DecentCEM offers an effective mechanism to improve the sample efficiency over the baseline CEM methods under the same sample budget for planning . 2 PRELIMINARIES . We consider a Markov Decision Process ( MDP ) specified by ( S , A , R , P , γ , d0 , T ) . S ⊂ Rds is the state space , A ⊂ Rda is the action space . R : S × A → R is the reward function that maps a state and action pair to a real-valued reward . P ( s′|s , a ) : S × A × S → R+ is the transition probability from a state and action pair s , a to the next state s′ . γ ∈ [ 0 , 1 ] is the discount factor . d0 denotes the distribution of the initial state s0 . At time step t , the agent receives a state st 1 and takes an action at according to a policy π ( ·|s ) that maps the state to a probability distribution over the action space . The environment transitions to the next state st+1 ∼ P ( ·|st , at ) and gives a reward rt = R ( st , at ) to the agent 2 . The return Gt = ∑T i=0 γ irt+i , is the sum of discounted reward within an episode length of T . The agent aims to find a policy π that maximizes the expected return . We denote the learned model in MBRL as fω ( ·|s , a ) , which is parameterized by ω and approximates P ( ·|s , a ) . Planning with the Cross Entropy Method Planning in MBRL is about leveraging the model to find the best action in terms of its return . Model-Predictive-Control ( MPC ) performs online planning at each time step up to a horizon to find the optimal action sequence : πMPC ( st ) = arg max at : t+H−1 E [ ΣH−1i=0 γ ir ( st+i , at+i ) + γ HV ( sH ) ] ( 1 ) whereH is the planning horizon , at : t+H−1 denotes the action sequence from time step t to t+H−1 , and V ( sH ) is the terminal value function at the end of the planning horizon . The first action in this sequence is executed and the rest are discarded . The agent then re-plans at the next time step . The Cross-Entropy Method ( CEM ) is a gradient-free optimization method that can be used for solving Eq . ( 1 ) . The workflow is shown in Fig . 2 . CEM planning starts by generating N samples { τj } Nj=1 = { ( âj,0 , âj,1 , · · · , âj , H−1 ) } Nj=1 from an initial sampling distribution gφ ( τ ) parameterized 1We assume full observability , i.e. , we assume that the agent has access to the state . 2We assume that the agent receives the true reward . This makes the problem easier but is unfair to MBRL methods that do not assume this . Thus , we will limit comparison to methods that also make this assumption . by φ , where each sample τj is an action sequence from the current time step up to the planning horizon H . The domain of gφ ( τ ) has a dimension of dτ = daH . Using a model f , CEM generates imaginary rollouts based on the action sequence { τj } ( in the case of a stochastic model ) and estimate the associated value v ( τj ) = E [ ΣH−1i=0 γir ( sj , i , aj , i ) ] where sj,0 is the current state s and sj , i+1 ∼ f ( sj , i , aj , i ) . The terminal value γHV ( sj , H ) is omitted here following convention in the CEM planning literature but the MPC performance can be further improved if paired with an accurate value predictor ( Bertsekas , 2005 ; Lowrey et al. , 2019 ) . The sampling distribution is then updated by fitting to the current top-k samples in terms of their value estimates v ( τj ) , using the Maximum Likelihood Estimation ( MLE ) which solves : φ′ = arg max φ N∑ j=1 1 ( v ( τj ) ≥ vth ) log gφ ( τj ) ( 2 ) where vth is the threshold equal to the value of the k-th best sample and 1 ( · ) is the indicator function . In practice , the update to the distribution parameters are smoothed by φl+1 = αφ′+ ( 1−α ) φl where α ∈ [ 0 , 1 ] is a smoothing parameter that balances between the solution to Eq . ( 2 ) and the parameter at the current internal iteration l. CEM repeats this process of sampling and distribution update in an inner-loop , until it reaches the stopping condition . In practice , it is stopped when either a maximum number of iterations has been reached or the parameters have not changed for a few iterations . The output of CEM is an action sequence , typically set as the expectation3 of the most recent sampling distribution for uni-modal distributions such as Gaussians µ̂ = E ( gφ ) = ( â0 , â1 , · · · , âH−1 ) . Choices of Sampling Distributions in CEM A common choice of the sampling distribution is a multivariate Gaussian distribution under which Eq . ( 2 ) has a straight-forward analytical solution . But the uni-modal nature of Gaussian makes it inadequate in solving multi-modal optimization that often occur in MBRL . To increase the capacity of the distribution , a Gaussian Mixture Model ( GMM ) can be used ( Okada & Taniguchi , 2020 ) . We denote such an approach as CEM-GMM . Going forward , we use CEM to refer to the vanilla version that uses a Gaussian distribution . Computationally , the major difference between CEM and CEM-GMM is that the distribution update in CEM-GMM involves solving for more parameters . Detailed steps can be found in Okada & Taniguchi ( 2020 ) . 3 DECENTRALIZED CEM . In this section , we first introduce the formulation of the proposed decentralized approach called the Decentralized CEM ( DecentCEM ) . Then we illustrate the intuition behind the proposed approach using a one-dimensional synthetic multi-modal optimization example where we show the issues of the existing CEM methods and how they can be addressed by DecentCEM . Formulation of DecentCEM DecentCEM is composed of an ensemble of multiple CEM instances indexed by i , each having its own sampling distributions gφi . They can be described by a set of distribution parameters Φ = { φi } Mi=1 . Each instance i manages its own sampling and distribution update by the steps described in Section 2 , independently from other instances . Note that the number of samples and elites are evenly split among the M instances . The top- kM sample sets are decentralized and managed by each instance independently whereas the centralized approach only keeps one set of top-k samples regardless of the distribution family used . After the stopping 3Other options are discussed in Appendix A.2 condition is reached for all instances , the final sampling distribution is taken as the best distribution in the set Φ according to ( the arg max uses a deterministic tie-breaking ) : φDecentCEM = arg max φi∈Φ Eφi [ v ( x ) ] ≈ arg max φi∈Φ N M∑ j=1 v ( τi , j ) ( 3 ) where Eφi [ v ( x ) ] denotes the expectation with respect to the distribution gφi , approximated by the sample mean . When M = 1 , it recovers the conventional CEM method . Motivational Example Consider a one dimensional multi-modal optimization problem shown in Fig.3 ( Left ) . There are eight local optima , including one global optimum f ( x∗ ) = −1.9 where x∗ = 5.146 . This objective function mimics the RL value landscape that has many local optima , as shown by Wang & Ba ( 2020 ) . This optimization problem is “ easy ” in the sense that a grid search over the domain can get us a solution close to the global optimum . However , only our proposed DecentCEM method successfully converges to the global optimum consistently under varying population size ( i.e. , number of samples ) and random runs , as shown in Fig.3 ( Right ) 4 . Both CEM-GMM and the proposed DecentCEM are equipped with multiple sampling distributions . The fact that CEM-GMM is outperformed by DecentCEM may appear surprising . To gain some insights , we illustrate in Fig . 4 how the sampling distribution evolves during the iterative update ( more details in Fig . 9 in Appendix ) . CEM updated the unimodal distribution toward a local optimum despite seeing the global optimum . CEM-GMM appears to have a similar issue . During MLE on the top-k samples , it moved most distribution components towards the same local optimum which quickly lead to mode collapse . On the contrary , DecentCEM successfully escaped the local optima thanks to its independent distribution update over decentralized top-k samples and was able to maintain a decent diversity among the distributions . GMM suits density estimation problems like distribution-based clustering where the samples are drawn from a fixed true distribution that can be represented by multi-modal Gaussians . However , 4For a fair comparison , hyperparameter search has been conducted on all three methods for each population size ( Appendix A ) . in CEM for optimization , exploration is coupled with density estimation : the sampling distribution in CEM is not fixed but rather gets updated iteratively toward the top-k samples . And the “ true ” distribution in optimization puts uniform non-zero densities to the global optima and zero densities everywhere else . When there is a unique global optimum , it degenerates into a Dirac measure that assigns the entire density to the optimum . Density estimation of such a distribution only needs one Gaussian but the exploration is challenging . In other words , the conditions for GMM to work well are not necessarily met when used as the sampling distribution in CEM . CEM-GMM is subject to mode collapse during the iterative top-k greedification , causing premature convergence , as observed in Fig 4 . In comparison , our proposed decentralized approach takes care of the exploration aspect by running multiple CEM instances independently , each performing its own local improvement . This is shown to be effective from this optimization example and the benchmark results in Section 6 . CEM-GMM only consistently converge to the global optimum when we increase the population size to the maximum 1,000 which causes expensive computations . Our proposed DecentCEM runs more than 100 times faster than CEM-GMM at this population size , shown in Table A.3 in Appendix . Convergence of DecentCEM We state the convergence result of DecentCEM in Theorem 3.1 . We show that the previous convergence result of CEM ( Hu et al. , 2011 ) applies to DecentCEM under the same sample budget . The key observation is that the convergence property of each CEM instance still holds since the number of samples in each instance is only changed by a constant factor ( the number of instances ) . We leave the detailed proof to appendix H. Theorem 3.1 ( Convergence of DecentCEM ) . If Assumptions 1-5 hold for a CEM instance described in Algorithm 3 and we decentralize it by evenly dividing its sample sizeNk intoM CEM instances in DecentCEM algorithm that satisfies Assumption 6 , then the sequence of iterates { ηi , k } generated by each CEM instance indexed by i converges almost surely to an internally chain recurrent set ( Hirsch et al. , 2001 ) of Equation 7 . Furthermore , the solution of DecentCEM { ηo , k } converges almost surely to the best solution of the individual instances in terms of the expected value of Em−1 ( η ) [ V ( x ) ] .

Summary: This paper studies a novel CEM method for model-based RL. While previous approaches used a centralized method (based on a unimodal Gaussian or Gaussian mixture), they propose a decentralized CEM, where each instance independently tracks its own data and top-k estimates. They first test their new method on a 1D toy task, showing better convergence to the global optimum. They then test their method on several model-based RL tasks, outperforming previous methods in some tasks, while being on par or inferior in others.