SPO: Sequential Monte Carlo Policy Optimisation

Authors: matthew v macfarlane
edan  toledo
donal  byrne
paul  duckworth
alexandre  laterre
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-08 10:13:14

Abstract: Leveraging planning during learning and decision-making is central to the long-term development of intelligent agents. Recent works have successfully combined tree-based search methods and self-play learning mechanisms to this end. However, these methods typically face scaling challenges due to the sequential nature of their search. While practical engineering solutions can partly overcome this, they often result in a negative impact on performance. In this paper, we introduce SPO: Sequential Monte Carlo Policy Optimisation, a model-based reinforcement learning algorithm grounded within the Expectation Maximisation (EM) framework. We show that SPO provides robust policy improvement and efficient scaling properties. The sample-based search makes it directly applicable to both discrete and continuous action spaces without modifications. We demonstrate statistically significant improvements in performance relative to model-free and model-based baselines across both continuous and discrete environments. Furthermore, the parallel nature of SPO’s search enables effective utilisation of hardware accelerators, yielding favourable scaling laws.

Main paper: 
 introduction: 
 The integration of reinforcement learning (RL) and neural-guided planning methods has recently achieved considerable success. Such methods effectively leverage planning during training via iterative imitation learning [7, 68, 18, 8] . Applying additional computation via planning generates improved policies which is then amortised [4] into a neural network policy. Using this policy as part of planning itself creates a powerful self-improvement cycle. They have demonstrated state of the art performance in applications ranging from chess [75] and matrix multiplication [24] , to language modelling [61] . However, one of the most commonly-used search-based policy improvement operators, MCTS [16] : i) performs poorly on low budgets [33] , ii) is inherently sequential limiting its scalability [51, 73] , and iii) requires modifications to adapt to large or continuous action spaces [57, 41] . These limitations underscore the need for more scalable, efficient and generally applicable algorithms.In this work, we introduce SPO: Sequential Monte Carlo Policy Optimisation, a model-based RL algorithm that utilises scalable sampled-based Sequential Monte Carlo planning for policy improvement. We formalise SPO as an approximate policy iteration algorithm, and show that it is formally grounded within the Expectation Maximisation (EM) framework. 
related work: 
 There is an extensive history of prior works that frame control and RL as an inference problem [82, 45] . Expectation Maximization [19, 28, 60] is a popular approach to directly optimising the evidence lower bound (ELBO) of the probability of optimality of a given policy. A number of model free methods implement EM, including RWR [63] , REPS [64] , MPO [1] , V-MPO [76] , and AWAC [59] each with different approaches to the E-step and M-step (summarised in Furuta et al. [29] ). Model based algorithms such as MD-GPS [58] and AlphaZero [74] can also be viewed under the EM framework using LQR [5] and MCTS [9, 16] during the E-step respectively. Grill et al. [33] show MCTS is a form of regularised policy optimisation equivalent to optimising the ELBO. Such model-based algorithms, that iterate between policy improvement using search, and projecting this improvement to the space of parameteriseable policies, are also often referred to as Expert iteration (ExIt) [74, 7] or Dual Policy Iteration [77] . Such approaches have also been applied to active inference [27] where MCTS has been used to estimate expected free energy, with this computation then amortised into a neural policy [25] . Our work can also be framed as ExIt, but differs by using Neural SMC for planning with an explicit KL regularisation on the improvement. Sequential Monte Carlo (SMC) [31, 43, 52] , also referred to as particle filters [66, 6] , is an inference method often employed to sample from intractable distributions. Gu et al. [34] explore the parameterisation of the proposal in SMC using LSTMs [40] , showing that posterior estimates from SMC can be used to train a parameterised proposal in non-control based tasks. Li et al. [49] demonstrate the same concept but using an MCMC sampler [30] . SMC can be applied to sample from distributions over trajectories in RL. For example, SMC-Learning [44] updates SMC weights using a Boltzmann exploration strategy, but unlike SPO it does not control policy improvement with KL regularisation or leverage neural approximators. Piché et al. [65] derive an SMC weight update to directly estimate the posterior over optimal trajectories. However this update requires the optimal value function, which they approximate using a SAC [37] value function. This approach requires high sample counts and does not leverage SMC posterior estimates for policy improvement, instead utilising a Gaussian distribution or SAC to train the proposal. CriticSMC [50] estimates the same posterior but utilises a soft action-value function to score particles [36] , reducing the number of steps performed in the environment, enabling more efficient exploration and improved performance on lower SMC budgets compared with Piché et al. [65] . However their method is slower in wall clock time and uses a static proposal, therefore not performing iterative policy improvement. A useful properly of SMC search methods is that it is inherently parallelisable. Parallelising MCTS with a virtual loss has been explored, however this often leads to performance degradation [51] , increasing exploration and leading to out-of-distribution states that are difficult to evaluate Dalal et al. [17] . 
background: 
 Sequential decision-making can be formalised under the Markov Decision Process (MDP) framework [67] . An MDP is a tuple (S, A, T , r, γ, µ) where, S is the set of possible states, A is the set of actions, T : S × A → P(S) is the state transition probability function, r : S × A → R is the reward function, γ ∈ [0, 1] is the discount factor, and µ is the initial state distribution. An agent interacts with the MDP using a policy π : S → P(A), that associates to every state a distribution over actions. We quantify the quality of a policy by the expected discounted return, that the agent seeks to maximise π * ∈ arg max π∈Π E π [ ∞ t=0 γ t r t ] , where r t = r(s t , a t ) is the reward received at time t and Π is the set of all realisable policies. The value function V π (s t ) = E π [ ∞ t=t γ t r t | s t ] maps a state s t to the expected discounted sum of future rewards when acting according to π. Similarly, the state-action value function Q π (s t , a t ) = E π [ ∞ t=t γ t r t | s t , a t ] maps a state s t to the expected discounted return, when taking the initial action a t and following π thereafter.Control as inference formulates the RL objective as an inference problem within a probabilistic graphical model [82, 42, 45] . For a horizon T , the distribution over trajectories τ = (s 0 , a 0 , s 1 , a 1 , ..., s T , a T ) is given by p(τ ) = µ(s 0 ) T t=0 p(a t )T (s t+1 |s t , a t ), which is a function of the initial state distribution, the transition dynamics, and an action prior. Note that this distribution is insufficient for solving control problems, because it has no notion of rewards. We therefore have to introduce an additional optimality variable into this model, which we will denote O t , that is defined such that p(O t = 1|τ ) ∝ exp(r t ). Therefore, a high reward at time t means a high probability of having taken the optimal action at that point. We are then concerned with the target distribution p(τ |O 1:T ), which is the distribution of trajectories given optimality at every step. We denote O 1:T as O going forward. The RL objective is then formulated as finding a policy to maximise log p π (O = 1) = log π(τ )p(O = 1|τ )dτ , which intuitively can be thought of as maximising the distribution of optimality at all timesteps, given actions sampled according to π. To optimise this challenging objective, we can derive the evidence lower bound (ELBO) using an auxiliary distribution q [60]:log p π (O = 1) = log π(τ )p(O = 1|τ )dτ = log q(τ ) π(τ )p(O = 1|τ ) q(τ ) dτ ≥ q(τ ) log p(O = 1|τ ) + log π(τ ) q(τ ) dτ = E q t r t α -KL(q(τ )∥π(τ )).(1) Since p(O = 1|τ ) ∝ exp( t r t ) and α is a normalising constant, ensuring a valid probability distribution.Expectation Maximisation (EM) [20] has been widely applied to solve the optimisation problem max π log p π (O = 1) [19, 63, 64, 1, 76] . After deriving the lower bound J (q, π) on the objective in eq. ( 1) using an auxiliary non-parametric distribution q, EM then performs a coordinate ascent, iterating between optimizing the bound with respect to q (E-step) and with respect to the parametric policy π (M-step). This generates a sequence of policy pairs {(π 0 , q 0 ), (π 1 , q 1 ), . . . , (π n , q n )} such that in each step i in the EM sequence, J (q i+1 , π i+1 ) ≥ J (q i , π i ). Viewed through a traditional policy improvement lens, the E-step corresponds to a policy evaluation phase where we perform rollouts, generating states with their associated estimates for q and value estimates. The M-step corresponds to a policy improvement phase where π and V are updated. This step can also be thought of as amortising the probabilistic inference computation performed in the E-step, into a neural network forward pass operation. 2Sequential Monte Carlo (SMC) methods [31] are designed to sample from an intractable distribution p referred to as the target distribution by using a tractable proposal distribution β (we use β to prevent overloading of q). Importance Sampling does this by sampling from β and weighting samples by p(x)/β(x). The estimate is performed using a set of N particles {x (n) } N n=1 , where x (n) represents a sample from the support that the target and proposal distributions are defined over, along with the associated importance weights {w (n) } N n=1 . Each particle uses a sample from the proposal distribution to improve the estimation of the target, increasing N naturally improves the estimation of p [10] . Once importance weights are calculated, the target is estimated asN n=1 w(n) δ x (n) (x),where w is the normalised importance sample weight and δ x (n) is the dirac measure located at x (n) . Sequential Importance Sampling generalises this for sequential problems, where x = (x 1 , . . . , x T ). In this case importance weights can be calculated iteratively according to:EQUATIONParticle filtering methods can be affected by weight degeneracy [56] . This occurs when a few particles dominate the normalised importance weights, rendering the remaining particles negligible. As the variance of particle weights is guaranteed to increase over sequential updates [21] , this phenomenon is unavoidable. When the majority of particles contribute little to the estimation of the target distribution, computational resources are wasted. Sequential Importance Resampling (SIR) [53] mitigates this problem by periodically resampling particles according to their current weights, subsequently resetting these weights. 
spo method: 
 In this section, we present a novel method that combines Sequential Monte Carlo (SMC) sampling with the Expectation Maximisation (EM) framework for policy iteration. We begin by formulating the objective function that we aim to optimise. We then outline our iterative approach to maximising this function, alternating between an expectation step (E-step) and a maximisation step (M-step).Within the E-step, we derive the analytical solution for optimising the objective with respect to the auxiliary distribution q. We then demonstrate how SMC can be employed to effectively estimate this target distribution. The M-step can be viewed as a projection of the non-parametric policy obtained in the E-step back onto the space of feasible policies. A comprehensive algorithmic outline of our proposed approach is provided in appendix D.2.We add an additional assumption to the lower bound objective defined in eq. ( 1) and assume that the auxiliary distribution q, defined over trajectories τ , can be decomposed into individual state dependent distributions, i.e. q(τ ) = µ(s 0 ) t≥0 q(a t |s t )T (s t+1 |s t , a t ). We parameterise π using θ which decomposes in the same way. This enables the objective to be written with respect to individual states instead of full trajectories. Multiplying by α, the objective can then be written as follows:EQUATIONPrevious works have explored this formulation where p(θ) is a prior over the parameters of π [38, 71, 1].Within the expectation step of EM, we maximise eq. ( 3) with respect to q. As in previous work, instead of optimising for the rewards, we consider an objective written according to Q-values [1, 59, 54] .EQUATIONThis aims to maximise the expected Q-value of a policy q, with the constraint that it doesn't move too far away from π i . Framing optimisation with respect to Q-values is useful as it enables the use of function approximators to estimate Q.Maximising this objective is difficult due to the dependence of both the expectation terms and Qvalues on q. Following previous works, we fix the Q-values with respect to a fixed policy π, resulting in a partial E-step optimisation [74, 1, 76] . We use the most recent estimate of q as the fixed policy we perform maximisation with respect to, similar to AlphaZero's approach of acting according to the expert policy. The initial state distribution µ(s) is likewise fixed to be distributed according to µ π . In practice, we use a FIFO replay buffer and sample from it to determine µ π .Equation (4) consists of balancing two objectives. To optimize this, we frame this as a constrained maximization problem, to avoid scaling problems, see appendix G.3 for further information. This objective limits the KL between q and π from exceeding a certain threshold: max q E µπ(s) E q(a|s) Q π (s, a) s.t. E µπ(s) [KL (q(a|s)∥π(a|s, θ i ))] < ϵ.(The following analytic solution to the constrained optimisation problem can then be derived using the Lagrangian multipliers method outlined in appendix G.1. Likewise η * is obtained by minimising the convex dual function eq. ( 7), and intuitively can be thought of as enforcing the KL constraint on q, preventing q i from moving to far from π i :EQUATIONwhere V π (s) is an action independent baseline. Q(s, a) -V (s) is referred to as the advantage, where optimising for Q-values vs advantages is equivalent [59] . Optimising with advantages however, has demonstrated to practically outperform optimising Q-values directly [62, 76, 59] . This is also a natural update to perform as the policy improvement theorem [78] outlines that an update by policy iteration can be performed if at least one state-action pair has positive advantage and a non-zero probability of reaching such a state. Although we have a closed form solution for q, we do not have either the value function or action-value function in practice. In the next section we outline our approach to estimating this distribution.Regularised Policy Optimisation: In our closed form solution to the optimisation eq. ( 6), we are required to solve for η * which ensures that the KL is constrained. This can be calculated by minimising the following objective [84] , see appendix G.2 for derivation:EQUATIONPractically, we estimate eq. ( 7) using a sample-based estimator according to the distribution of states from the replay buffer and stored values for π(a|s, θ i ) and A π (s, a).Building upon previous work using Sequential Monte Carlo for trajectory distribution estimation [44, 65] , we propose a novel approach that integrates SMC within the Expectation Maximisation . . , x n , each of which represents an SMC trajectory sample, are performed in parallel according to π i (left to right). At each environment step, the weights of the particles are adjusted (indicated in the diagram by circle size). We show two resampling regions where particles are resampled, favouring those with higher weights, and their weights are reset. The target distribution is estimated from the initial actions of the surviving particles (rightmost particles). This target estimate, q i , is then used to update π in the M-step.framework to estimate the target distribution defined in Equation 6. Our SMC-based method enables sampling from q i over trajectories of length h, incorporating multiple predicted future states and rewards for a more accurate distribution estimation. A key property of estimating the target using SMC is that it uses both breadth (we initialise multiple particles to calculate importance weights) but also depth (leveraging future states and rewards) to estimate the distribution. In contrast, MPO evaluates several actions per state, focusing on breadth without depth, while V-MPO calculates advantages using n-step returns from trajectory sequences, leveraging depth but only for a single action sample. Our method combines both aspects, enhancing the accuracy of target distribution estimation, crucial for both action selection and effective policy improvement updates in the M-step. An outline of the SMC algorithm is provided in algorithm 1, with a visual representation in fig. 1 .SMC estimates the target by maintaining a set of particles {x (n) } N n=1 each of which maintains an importance weight for a particular sample. Given calculated importance weights SMC, estimates the target distribution according to qi (τ ) = N n=1 w(n) δ x (n) (τ ) where w(n) is the normalised importance sample weight for a particular sample. We next outline the sequential importance sampling weight update needed to estimate eq. ( 6). Since we can sample from π(a t , s t , θ i ), we leverage this as our proposal distribution β. Given the target distribution q can be decomposed into individual state dependent distributions, we can define p i (τ t |τ 1:t-1 ) and β(τ t |τ 1:t-1 ) as:p i (τ t |τ 1:t-1 ) ∝ T (s t+1 |s t , a t )π i (a t |s t , θ i ) expĀi (a t , s t )EQUATIONleading to the following convenient SMC weight update according to eq. ( 2):EQUATIONwhere Q π (a t , s t ) -V π (s t ) is simplified to A π (a t , s t ), τ 1:t is a sequence of state, action pairs from timestep 1 to t, and T model is the environment transition function of the planning model. Note that our work assumes the availability of a model that accurately represents the transition dynamics of the environment T , and therefore simplify the update to w(τ 1:t ) ∝ w(τ 1:t-1 ) • exp (A π (a t , s t )/η * ).Algorithm 1 outlines the method for estimating eq. ( 6) over h planning steps for N particles (line 3) using the advantage based weight update (line 7) in eq. (10) . Once h steps in the environment model have been performed in parallel, we marginalise all but the first actions as a sample based estimate of q i (line 14).Algorithm 1 SMC q target estimation (timestep t)EQUATIONi , .., w(N ) i ) 10: {w (n) i = 1} N n=1 11: end if 12: end for 13: {a (n) t } as the set of first actions of {x (n) t:t+h } N n=1 14: q(a|st) = N n=1 w(n) δ a (n) t (a)The advantage function A π is typically unknown, so we compute a 1-step estimate at each iteration using the value function and observed environment rewards Â(s t , a) = r t + V π (s t+1 ) -V π (s t ). Practically, we parameterise the value function V using a neural network and train it using GAE [70] on true environment rollouts collected. After h steps, the importance weights leverage h-steps of observed states and rewards during planning and corresponding advantage estimates. Compared to tree-based methods such as MCTS, SMC does not require maintaining the full tree in memory. Instead, after each planning step, it only needs to retain the initial action, current state, and current particle weight. It also does not require spending computation sending backward messages to update statistics of previous nodes everytime a new node is added to the tree.Resampling Adaptive Search: To mitigate the issue of weight degeneracy [56] , SPO conducts periodic resampling (lines [8] [9] [10] [11] . This involves generating a new set of particles from the existing set by duplicating some trajectories and removing others, based on the current particle weights. This process enhances computational efficiency in estimating the target distribution. By resampling, we avoid updating weights for trajectories with low likelihood under the target, thereby reallocating computational resources to particles with high importance weights [22] .After completing the E-step (which generates q i ), we proceed with the M-step, which optimises eq. ( 3) with respect to π, parametrised by θ. By eliminating terms that are independent of π, we optimise the following objective, corresponding to a maximum a posteriori estimation with respect to the distribution q i : maxEQUATIONThis optimisation can be viewed as projecting the non-parametric policy q i back to the space of parametrisable policies Π θ , as performed in expert iteration style methods such as AlphaZero. p(θ) represents a prior over the parameter θ. Previous works find that utilising a prior for θ to be close to the estimate from the previous iteration θ i leads to stable training [1, 76] . Therefore we assume a gaussian prior over the current policy parameters, see appendix G.5 where we show utilising such a prior leads to the following constrained objective:EQUATIONPrevious approaches to using SMC for RL either do not perform iterative policy improvement using SMC [65] , or lack policy improvement constraints needed for iterative improvement [44] .Assuming that SMC provides perfect estimates of the analytic target distribution at each step i, Expectation Maximisation algorithm will guarantee that successive iterations will result in monotonic improvements in our lower bound objective, with details outlined in appendix G.4 [54] . Proposition 1. Given a non-parametric variational distribution q i and a parametric policy π θi . Given q i+1 , the analytical solution to E-step optimisation eq. (3) , and π θi+1 , the solution to maximisation problem in the M-step eq. ( 12) then the ELBO J is guaranteed to be monotonically increasing: J (q i+1 , π θi+1 ) ≥ J (q i , π θi ).In practice we are unlikely to generate perfect estimates of the target through sample based inference and leave derivations regarding the impact of this estimation on overall convergence for future work.We also draw the connection between our EM optimisation method and Mirror Descent Guided Policy Search (MD-GPS) [58] . Our objective can be viewed as a specific instance of MD-GPS (see appendix G.6). Depending on whether dynamics are linear or not, optimising the EM objective can be viewed either as exact or approximate mirror descent [11] . Monotonic improvement guarantees in MD-GPS follow from those of mirror descent.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: [NA] Guidelines:• The answer NA means that the paper does not include experiments.• If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example , with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
experiments: 
 In this section, we focus on three main areas of analysis. First, we demonstrate the improved performance of SPO in terms of episode returns, relative to both model-free and model-based algorithms. We conduct evaluations across a suite of common environments for both continuous control and discrete action spaces. Secondly, we examine the scaling behaviour of SPO during training, showing that asymptotic performance scales with particle count and depth. Finally, we explore the performance-to-speed trade-off at test time by comparing SPO directly to AlphaZero as the search budget increases.In order to ensure the robustness of our conclusions, we follow the evaluation methodology proposed by Agarwal et al. [3] , see appendix H for further details. This evaluation methodology groups performance across tasks within an environment suite, enabling clearer conclusions over the significance of learning algorithms as a whole. We include individual results in appendix C, along with additional analysis measuring the statistical significance of our results.Environments: For continuous control we evaluate on the Brax [26] benchmark environments of: Ant, HalfCheetah, and Humanoid. For discrete environments, we evaluate on Boxoban [35] (a specific instance of Sokoban), commonly used to assess planning methods, and Rubik's Cube, a sparse reward environment with a large combinatorial state-action space. See appendix A for further details regarding environments.Baselines: Our model-free baselines include PPO [72] , MPO [1] and V-MPO [76] for both continuous and discrete environments. For model-based algorithms, we compare performance to AlphaZero (including search improvements from MuZero [68] ) and an SMC method introduced by Piché et al. [65] 3 (which we refer to as SMC-ENT due to its use of maximum entropy RL to train the proposal and value function used within SMC). Our AlphaZero benchmark follows the official open-source implementation 4 . For continuous environments we baseline our results to Sampled MuZero [41] , a modern extension to MuZero for large and/or continuous action spaces. We utilise a true environment model, aligning our implementation of Sampled MuZero more closely with AlphaZero. Our core experiments configure SPO with 16 particles and a horizon of 4, and AlphaZero with 64 simulations to equalise search budgets. For remaining parameters, see appendix E.1 and appendix D.3. Each environment and algorithm were evaluated with five random seeds. 5Policy Improvement: In fig. 2 we show empirical evidence that SPO outperforms all baseline methods across both discrete (left) and continuous (right) benchmarks for policy improvement. This conclusion also holds for the per-environment results reported in appendix C. SMC-ENT [65] is a strong SMC based algorithm, however, SPO outperforms it for both discrete and continuous environments providing strong evidence of the direct benefit of using SMC posterior estimates for policy improvement. We also highlight the importance regularising policy optimisation, which is practically achieved by solving for η * eq. ( 7), and is re-estimated at each iteration. In appendix B we ablate the impact of this adaptive method by comparing varying fixed temperature schedules. While a good temperature can be found through expensive grid search for each new environment we find the adaptive temperature ensures stable updates leading to strong learning performance across all environments, with minimal overhead to compute η * .We find that SPO performance is statistically significant compared to AlphaZero across the evaluated environments, see appendix C for detailed results. The substantial variation in AlphaZero performance across these environments underscores the difficulty in tuning it for diverse settings. For instance, while AlphaZero performs well on Boxoban and HalfCheetah, its performance drops significantly on other discrete and continuous problems. This inconsistency poses a major challenge for its applicability to real-world problems. In contrast, SPO performs consistently well across all environments (both discrete and continuous), highlighting its general applicability and robustness to various problem settings.In fig. 3 (left) we investigate how SPO performance is impacted by scaling particle counts N and horizon length h during training, both of which we find improve the estimation of the target distribution. Our results show that scaling both the particle count and the horizon leads to improvements in the asymptotic performance of SPO. It also suggests that both variables should be scaled together for maximum benefit as having a long horizon with low particle counts can actually negatively impact performance. Secondly, we highlight that while for our main results we use a total search budget (particles × depth) of 64 for policy improvement, our scaling results show that competitive results can be achieved by reducing this budget by a factor of four, which demonstrates competitive performance at low compute budgets.We can scale particle counts and horizon during training to improve target distribution estimates for the M-step. However, this scaling can also be performed after training when π θi is fixed, enhancing performance, since actions are sampled in the environment according to the improved policy q. This equates to additional E-step optimization, which can enhance performance due to the representational constraints of the space of parameterised policies. In fig. 3 (right) we demonstrate how the performance of both SPO and AlphaZero search scales with additional search budget at test time using the same high performing checkpoint. We plot the time taken for a single actor step (measured on a TPUv3-8) against solve rate for the Rubik's Cube problem with time on a logarithmic axis. Specifically, we evaluate on 1280 episodes for cubes ten scrambles away from solved. While we recognise that such analysis can be difficult to perform due to implementation discrepancies, we used the official JAX implementation of AlphaZero (MCTX) within our codebase and compare this to SPO. Additionally, we exclusively measure inference thus no training implementation details affect our measurements. This provides evidence that AlphaZero has worse scaling when compared to SPO on horizons four and eight, noting that for horizon of two, SPO performance converges early, as low depths can act as a bottleneck for further performance improvements. This highlights the benefits of the SPO parallelism. This chart also shows how for different compute preferences at test time, a different horizon length is preferable. For limited compute, low horizon lengths and relatively higher particle counts provide the best performance, but as compute availability increases, the gains from increasing particle counts decrease and instead horizon length h should be increased.Lastly, the plot illustrates the increase in available search budget, by using SPO, given a time restriction rather than a compute restriction. For example, SPO can use 4 times more search budget, when compared to AlphaZero, given a requirement of 0.1 second per step.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [No] Justification: Inference code to be released at publication Guidelines:• The answer NA means that paper does not include experiments requiring code. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion: 
 Planning-based policy improvement operators have proven to be powerful methods to enhance the learning of policies for complex environments when compared to model-free methods. Despite the success of these methods, they are still limited in their usefulness, due to the requirement of high planning budgets and the need for algorithmic modifications to adapt to large or continuous action spaces. In this paper, we introduce a modern implementation of Sequential Monte Carlo planning as a policy improvement operator, within the Expectation Maximisation framework. This results in a general training methodology with the ability to scale in both discrete and continuous environments.Our work provides several key contributions. First, we show that SPO is a powerful policy improvement operator, outperforming our model-based and model-free baselines. Additionally, SPO is shown to be competitive across both continuous and discrete domains, without additional environmentspecific alterations. This illustrates the effectiveness of SPO as a generic and robust approach, in contrast to prior methods that require domain-specific enhancements. Furthermore, the parallelisable nature of SPO results in efficient scaling behaviour of search. This allows SPO to achieve a significant boost in performance at inference time by scaling the search budget. Finally, we demonstrate that scaling SPO results in faster wall-clock inference time compared to previous work utilising tree-based search methods. The presented work culminates in a versatile, neural-guided, sample-based planning method that demonstrates superior performance and scalability over baselines.Limitations & Future Work: This work demonstrates the efficacy of combining probabilistic inference with reinforcement learning and amortisation. While our work considers a relatively standard form of Sequential Monte Carlo, future work could investigate making improvements to this inference method, in order to further improve the estimate of the target distribution, while maintaining the scalability benefits demonstrated. We also draw the connection to Active Inference [27] which has been explored in the context of agent behaviour in complex environments, where Monte Carlo Tree Search has been used to scale previous methods along with deep learning [25] . Leveraging Sequential Monte Carlo in such methods along with amortisation could be a promising direction of research to further scale such methods. Our work exclusively uses exact world models of the environment in order to isolate the effects of various planning methods on performance. Extending our research to include a learned world model would broaden the applicability of SPO to more complex problems that may not have a perfect simulator and are therefore unsuitable for direct planning. Additionally, we apply Sequential Monte Carlo (SMC) only in deterministic settings. Adapting our approach for stochastic environments presents challenges. Specifically, the importance weight update in SMC can lead to the selection of dynamics most advantageous to the agent, potentially fostering risk-seeking behaviour in stochastic settings. However, mitigation strategies exist, as discussed in Levine [45] . We use the specific instance of Sokoban outlined in Guez et al. [35] illustrated in Figure 4 with the codebase available at https://github.com/instadeepai/jumanji. The datasets employed in this study are publicly accessible at https://github.com/google-deepmind/boxoban-levels. These datasets are split into different levels of difficulty, which are categorised in Table 1 .Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: [NA] Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: We include error bars for the main results of the paper in addition to statistical significance tests in the appendix. For additional scaling plots we do not include error bars. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?Answer: [Yes]• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 
crowdsourcing and research with human subjects: 
 Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 
institutional review board (irb) approvals or equivalent for research with human subjects: 
 Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?Answer: [NA]Justification: [NA] Guidelines:• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: In this research, we always train on the Unfiltered-Train dataset. Evaluations are performed on the Hard dataset, important for differentiating the strongest algorithms. Additionally, the architecture includes two output heads (policy and value), each comprising two layers of size 128 with ReLU activations. The output heads share the same torso network. We use the same architecture for all algorithms.A.2 Rubik's CubeFor our Rubik's Cube experiments, we utilize the implementation in Bonnet et al. [12] available at https: //github.com/instadeepai/jumanji. Rubik's cube problems can be be made progressively difficult by scaling the number of random actions performed in a solution state to generate problem instances. We always perform training on a uniform distribution sampled from the range [3, 7] of scrambled states, followed by exclusively evaluating on 7 scrambles. The observation is represented using an array of size (6, 3, 3) . The action space is represented using an array of size (6,3) corresponding to each face and the depth at which it can be rotated.We utilise an embedding layer to first embed the face representations to size (6, 3, 3, 4) . The embedding is flattened and we use a torso layer consisting of a two layer residual network with layer size 512 and ReLU activations. Additionally, the architecture includes two output heads (policy and value), each comprising of a single layer of size 512 with ReLU activations. We use the same architecture for all algorithms.For experimentation on continuous settings, we make use of Brax [26] . Brax is a library for rigid body simulation much like the popular MuJoCo setting [79] simulator. However, Brax is written entirely in JAX [14] to fully exploit the parallel processing capabilities of accelerators like GPUs and TPUs.It is important to note, that at the time of writing, there are 4 different physics back-ends that can be used. These back-ends have differing levels of computational complexity and the results between them are not comparable.The results generated in this paper utilise the Spring backend. Table 2 contains the observation and action specifications of the scenarios used for our experiments. We specify the dimension size of the observation and action vectors.In practice, we found the smaller networks used in the original Brax publication to limit overall performance. We instead use a 4 layer feed forward network 6 to represent both the value and policy network. Aside from the output layers, all layers are of size 256 and use SiLU non-linearities. We use the same architecture for all algorithms. Unlike Rubik's Cube and Sokoban, we do not use a shared torso to learn embeddings.The KL constraint for the E-step within the Expectation Maximisation framework is an important component of ensuring policy improvement, due to the KL term in the ELBO. We investigate the impact of the adaptive temperature for both discrete and continuous environments. We compare SPO with an adaptive temperature updated every iteration to a variety of fixed temperatures. This ablation shows that SPO with an adaptive temperature is among the top performing hyperparameter settings across all environments. However we also note that it is possible to tune a temperature that works well when considering a wide range of temperatures. This is consistent with previous results in Peng et al. [62] that also find practically for specific problems a fixed temperature can be used. Of course in practice having an algorithm that can learn this parameter itself is practically beneficial, removing the need for costly hyperparameter tuning, since the appropriate temperature is likely problem dependant. Subsequently, we evaluated whether the partial optimisation of the temperature parameter η effectively maintained the desired KL divergence constraint and how different values of this constraint affected performance. As shown in Figure 7 (a), the KL divergence constraint is tightly maintained throughout training for different values of ϵ. Figure 7 (b) demonstrates that, for the Ant task, allowing a larger KL divergence between the prior and target policies (q and π, respectively) leads to improved performance.For our analytic solution to the optimisation problem in the E-step we choose to add a value baseline such that our target distribution re-weights with respect to Advantages instead of Q-values. In practice there are no restrictions on utilising Q-values instead for the importance weight update of SMC eq. ( 10). Below we show results on the continuous environments comparing SPO using advantages to SPO using Q-values.We see in Figure 8 , that the use of Q values produces a reduction in performance that is relatively constant throughout training. It is possible that this reduction is due to different hyperparameters being required to effectively utilise the Q-value distribution but in order to accurately judge the effect, we kept hyperparameters constant Estimating the target distribution accurately is important for policy improvement in SPO. Therefore, we aim to validate the ability of SMC search to better approximate the true target distribution than a simple 1-step function evaluation as is done in algorithms such as MPO [1] . To do this, while the true target distribution is not directly accessible, we approximate it using an unbiased Monte Carlo oracle. We use a large computational budget of 1280 rollouts, to the end of the episode, for every state and every action. We evaluate the impact of varying planning horizons (depth) and particle counts on the KL divergence between the SMC-estimated target policy and the Monte Carlo oracle in the Sokoban environment. As shown in Figure 9 , increasing the number of particles and the planning depth reduces the KL divergence to the oracle, indicating improved estimation of the target distribution. This aligns with SMC theory [22] , which suggests that more particles and deeper planning lead to better approximations. These results highlight the importance of leveraging both breadth (more particles) and depth (longer planning horizons) in SMC for accurate target estimation. In Figure 12 , we present the performance profiles which visually illustrates the full distribution of scores across all tasks and seeds for each algorithm. We see that SPO has a higher lower bound on performance, in addition to upper bound, indicating lower variance across tasks and seeds. Additionally, we see that SPO is strictly above all other algorithms indicating that SPO is stochastically dominant7 to all baselines. Lastly, we present the probability of improvement plots in Figure 13 . We see that SPO has a high probability of improvement compared to all baselines. Additionally, all probabilities are larger than 0.5 and have their CIs above 0.5 thus indicating statistical significance as specified by Agarwal et al. [3] . Furthermore, we see all CIs have upper bounds greater than 0.75, thereby indicating that these results are statistically meaningful as specified by Bouthillier et al. [13] .Training was performed using a mixture of Google v4-8 and v3-8 TPUs. Each experiment was run using a single TPU and only v3-8 TPUs were used to compare wall-clock time.All individual task results are presented in fig. 14 . Specifically, we present the IQM of returns achieved during the evaluation intervals throughout training. Below we outline the practical policy-based losses used within SPO,Lπ(θ) = - s,a∼D qi(a|s) log π θ i (a|s) g(η) = ηε + η µ(s) log π(a|s, θi) exp A π (s, a) η da ds. Lα(θ, α) = α ϵα -E s∼p(s) [sg [DKL(π θ old ∥ π θ )]] LKL(θ, α) = sg [α] E s∼p(s) [DKL(π θ old ∥ π θ )]The temperature loss is calculated using advantages collected during SPO search with rollouts performed according to policy π. If resampling is utilised, after a resample the distribution over trajectories shifts towards the target distribution away from π. Therefore in order to calculate the loss g(η), we only utilise the advantages up to the first resampling period with rollouts performed according to π.Below we outline the overall SPO algorithm with the main loop split into the E-step and M-step.Algorithm 2 SPO Algorithm 1: B ← ∅ 2: Initialize the policy, value function, temperature, and alpha:3:Initialize θ 0 , ϕ 0 , η 0 , α 0 4: for iteration k = 0, 1, 2, . . . , num_iterations do 5:Step (E-step):for agent m = 1, 2, . . . , M in parallel do ā = {a (n) } N n=1 ∼ qSMC (s i , η k , θ k , ϕ k ) ▷ sample-based estimate of q i 10: a ∼ {a (n) } N n=1▷ Sample an action to execute in the environment Maximization Step (M-step):17:for batch b = 1, 2, . . . , num_batches do 18:Sample batch from replay buffer:19:Sample batch of size N from replay buffer (s, r, ā) ∼ BUpdate value function using GAE targets:21:ϕ k+1 ← arg min ϕ E s∼B (V GAE (s, ϕ ′ ) -V (s, ϕ)) 2 22:Update parameterized policy using sample-based estimate of q:23:θ k+1 , α k+1 ← arg min θ E s∼B [L α (θ, α) + L KL (θ, α)] 24:Update dual temperature:25:η k+1 ← arg min η E s∼B [g(η)] 26: ϕ ′ ← polyak(ϕ ′ , ϕ k+1 ) 27:end for 28: end for D.3 Hyperparameters The temperature parameter η plays a crucial role in balancing exploration and exploitation during the search process in SPO. In this work, it is derived via first choosing the desired KL divergence ϵ between the target policy qi generated by search and the current policy πi. If the temperature is too low (i.e., η is small), the exponential weighting exp A π (s, a)/η becomes sharply peaked around the highest-advantage actions. This causes the importance weights to concentrate on a few particles, and during resampling, the particles can collapse onto a single root action. Such collapse reduces diversity in the search and renders the rest of the planning process ineffective. Conversely, if the temperature is too high (i.e., η is large), the weighting becomes flatter, leading to excessive exploration. While exploration is necessary, too much can prevent the search from effectively focusing on promising paths. This work shows that deriving this value via a target KL leads to stable and strong performance across different environments.For the baseline implementations, we expanded and adapted the existing implementations from Stoix8 [80] . All implementations were conducted using JAX [14] , with off-policy algorithms leveraging Flashbax [81] . EM algorithms differ on a small number of dimensions through which the algorithm can be understood. The first dimension, G, is the optimization objective that is maximized in the E-step. This includes whether advantages, Q-values, or rewards are used, and with respect to which policy. We then consider whether a trust region constraint is used both in the E-step and M-step.With a defined optimization objective, various methods can be used to estimate it. Most methods derive an analytic solution to the optimization problem and then estimate this distribution using techniques such as TD(0) or function approximation. For example, AlphaZero does not explicitly derive the analytic solution but estimates the solution to this optimization objective through Monte Carlo Tree Search (MCTS).Table 10 provides a summary of some common EM methods, highlighting their core differences.Method G E-step TR M-step TR G estimate Depth Breadth MPO Q πp Yes Yes Analytic + Function Approximation 0 M V-MPO A πp Yes Yes Analytic + n-step TD + top-k Adv T 1 AWR A πp No No Analytic + n step TD T 1 AWAC A πp No No Analytic + Function Approximation 0 1 PoWER η log Q πp Yes No Analytic + n step TD T 1 RWR η log r No No - 1 1 REPS A πp Yes No Analytic + TD(0) 1 1 AlphaZero Q πq Yes No MCTS > 0 > 0 SPO A πq Yes Yes Analytic + SMC D MNote that T refers to an episode length. In additon we add two important dimensions of the G estimate (breadth and depth). Methods like MPO estimate the analytic solution directly using a Q-function. This can be leveraged to estimate the target distribution for a selection of M actions, but without leveraging depth, so rewards and future states are not used to improve the estimates. In contrast, V-MPO leverages depth to form an n-step TD estimate, but only for one of the actions, leading to a relatively poor estimate of the analytic solution.We outline the analytic solution to eq. ( 5), by writing a Lagrangian equation and solving for q. We can first represent our constrained optimisation exactly with the following optimisation problem and associated 2 constraints. We add a state dependent baseline V (s) to the optimisation objective and notate Q q (s, a) -V q (s) as A q (s, a). max q µq(s) q(a|s) [Q q (s, a) -V q (s)] da ds s.t. µq(s) [KL(q(a|s)∥π(a|s, θi))] ds < ϵ, µq(s) q(a|s)da ds = 1.(EQUATIONThe following Lagrangian L can be constructed, L(q, η, γ) = µq(s) q(a|s)A q (s, a) da ds + η ϵ -µq(s) q(a|s) log q(a|s) π(a|s, θi) da ds + γ 1 -µq(s) q(a|s) da ds .(EQUATIONWe can then solve for the value of q that maximises this expression by taking the derivative with respect to q and setting it to 0.∂L(q, η, γ) ∂q = A q (s, a) -η log q(a|s) + η log π(a|s, θi)EQUATIONThe analytic form for the optimal distribution q can then be calculated as:EQUATIONWe now derive the dual function to be minimised in order to obtain η within the analytic solution for q, which is crucial for enforcing the KL constraint.The final term -η+γ η acts as a normalising constant since it is independent of q and therefore we can construct the following equality.EQUATIONNow that we have expressions for -η+γ η and q we can form the dual function g(η) by substituting these terms back into the original Lagrangian. After simplifying we recover the following expression. g(η) = ηϵ + η µq(s) log π(a|s, θi) exp A q (a, s) η da dsThe optimal dual variable can be calculated as followsEQUATIONIn optimising eq. ( 4), the first term E q(a|s) [Q q (s, a)] is dependent on the scale of reward in the environment. This can make it hard to choose α, as the first and second terms are on arbitrary scales. Practically this can mean that for each new environment a new α should be used requiring costly hyperparameter tuning to find. As explored in previous work [1, 76] we choose to enforce a hard constraint, as opposed to a soft constraint. Instead of choosing α we choose ϵ which is the maximum KL divergence between q and π which is practically less sensitive to reward scale. We found that choosing an ϵ to be far easier, resulting in a value that generalised across all environments explored. This is important as there has been a trend in recent years to add hyperparameters to Reinforcement Learning algorithms [2] resulting in the need for costly hyperparameter sweeps for algorithms to work on each new problem they are applied to.The optimality of policy π θ , log pπ θ (O = 1), is lower bounded by the following ELBO objective:J (q, π θ ) = Eτ∼q ∞ t=0 γ t rt -αDKL(q(•|st)∥π(•|st, θ)) + log p(θ).We improve π θ by optimizing the ELBO alternatively via EM. We will outline the proof of the monotonic improvement guarantee of the ELBO using the EM procedure, at the i-th iteration of training. We note that this holds under the assumption that we calculate the true value of the closed form solution qi in the E-step, which is in practice unlikely to be the case.By the definition of E-step, we improve the ELBO with respect to q. We show the E-step update will increase ELBO: Therefore, we have: J (qi+1, π θ i+1 ) ≥ J (qi+1, π θ i ). Combining these two results we have that after successive applications of the E-step and M-step,qi+1 = arg max q Eq E a∼q(•|s) [A q i (s, a)] -αDKL[q(•|s)∥π(•|s, θi)] = arg max q Eτ∼q ∞ t=0 γ t rt -αDKL(q(•|st)∥π(•|st, θi)) = arg max q J (q, θi) ⇒ J (qi+1, π θ i ) ≥ J (qi, π θ i ).J (qi+1, π θ i+1 ) ≥ J (qi, π θ i ).In Reinforcement Learning it can be beneficial to constrain policies from moving too far from the current policy, often leading to increased stability or performance of an algorithm [69, 72] . In this work we utilise a Gaussian prior around θi, our current value of θ, optimised from the previous iteration.Therefore, θ ∼ N (µ, Σ) where µ = θi and Σ -1 = λF (θi), where F is the Fisher information matrix. The prior term in eq. ( 11) can then be written as log p(θ) = -λ(θ -θi) T F (θi) -1 (θ -θi) + c. The first term in this expression is the quadratic approximation of the KL divergence [69] while the second term c is a term that does not depend on θ and therefore when optimising eq. ( 11) with respect to θ, can be dropped. Using this approximation we rewrite eq. ( 11) as: Like the E-step, choosing λ can be non-trivial so we convert it into a hard constraint optimisation eq. ( 12). Note that ϵ in eq. ( 12) is different from eq. ( 5). We note that considering a Gaussian prior in the M-step is not strictly required for SPO performance, however practically it adds stability to training.We highlight the connection between the use of Expectation Maximisation to improve the evidence lower bound and Mirror Descent Guided Policy Search (MD-GPS). Mirror Descent Guided Policy Search builds upon previous guided policy search work [46] . Rather than only enforcing the constraint between q and π at convergence (referred to a the local policy:pi and global policy: π θ respectively), they constrain qi against the current policy πi at every iteration. This results in a very similar algorithm to EM optimisation. Below we provide the outline of MD-GPS algorithm, clearly showing the equivalence to EM [58] : C-step:p i ← arg min pi E pi(τ ) T t=1 ℓ(x t , u t ) such that D KL (p i (τ ) ∥ π θ (τ )) ≤ ϵ 3: S-step: π θ ← arg min θ i D KL (p i (τ ) ∥ π θ (τ )) 4: end forComparing MD-GPS to EM, the C-step is equivalent to the E-step and S-step equivalent to the M-step, with the loss ℓ being the negative of the expected discounted sum of returns over trajectories. MD-GPS assumes that the S-step can perfectly minimise the objective in the case of linear dynamics and quadratic costs [48] , resulting in exact mirror descent. In practice, most applications of MD-GPS will not satisfy such constraints and so are akin to approximate mirror descent. However, as long as a constraint between the local and global policy can be enforced in terms of KL, various bounds on cost of the global policy can be constructed, see Montgomery and Levine [58] for further details.• Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced.• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.Many related RL algorithms can be understood by the different approaches to either the E-step or M-step,see appendix F.1 for a summary of EM approaches.We use MDQN[83], and SAC[37] for discrete and continuous environments, respectively. Note that in the discrete case, SMC-ENT using SAC was very unstable, so we chose MDQN as a better alternative.Open-source implementation available at mctx.Inference code and checkpoints are available at https://github.com/instadeepai/spoexcept for SMC-ENTR,for which we use 3 to stabilise performance.A random variable X is termed stochastically dominant over another random variable Y if P (X > τ ) ≥ P (Y > τ ) for all τ , and for some τ , P (X > τ ) > P (Y > τ ).Available at: https://github.com/EdanToledo/Stoix 


Decision: The importance of Monte Carlo planning in RL is well known (see AlphaGO), and this paper proposes extensions to Sequential MC. The proposed methods are interesting and justified mathematically. The scores given by the reviewers were mixed, with an average on the positive side. To emphasize the value of the paper, one could mention that the mathematical justification of this paper establishes a close relationship with the use of Monte Carlo methods in Active Inference, which tries to solve reinforcement learning problems without the use of rewards. In particular, the following past NeurIPS paper is very related: “Fountas, Zafeirios, Noor Sajid, Pedro Mediano, and Karl Friston. Deep active inference agents using Monte-Carlo methods. Advances in neural information processing systems 33 (2020): 11662-11675.”. We believe that the current paper could be seen as a step to narrow the gap between traditional reinforcement learning and active inference (which as seen through the work of Fountas et. al is still problematic). Given this parallel interest in MC methods in neuroscience, we believe that the current paper on pushing forward RL and SMC methods is useful, and it may eventually help to bring the two communities together.

Reviews: review: paper_summary: paper_summary: This paper introduces a novel iterative, particle-based approach to sequential Monte Carlo planning by combining expectation maximization with importance sampling. Their approach uses a model to sample real state-action trajectories in a problem domain then, after importance-weighting said transitions, computes a target policy distribution for the agent. This target policy is used for policy improvement. The paper supplies a theoretical derivation of their approach and provides experimental results against baseline algorithms.
strength_weakness: strength_weakness: I think that an evaluation with an imperfect model should have been included rather than left as future work.
questions: questions: Line 204: How exactly is $\hat{A}$ computed? It isn't very clear on that.

How much does removing samples based on weight affect exploration? Is this only helpful when you have a perfect model?
limitations: limitations: They broach the limitations of their work sufficiently.
paper_summary: paper_summary: This paper propose SPO:  Sequential Monte Carlo Policy Optimisation, a RL algorithm with the Expectation Maximisation (EM) framework for MDP. Experiments on both discrete and continuous environments show that SPO outperforms baselines in terms of sample efficiency and scalability.
strength_weakness: strength_weakness: Originality: The paper propose SPO, which is the combination of Sequential Monte Carlo (SMC) sampling with the EM framework. The novelty is limited.
Quality: The paper does not discuss what problem the paper aims to tackle. The paper uses SMC to better estimate the target distribution. But the paper does not give any experimental results showing better estimation results. Ablation studies are missing. Theoretical insights of proposition 1 are missing.
Clarity: The paper is hard to follow as it is full of formulas.
questions: questions: See the weakness.
limitations: limitations: Yes.
paper_summary: paper_summary: This paper maximizes the expected reward wrt the policy by rewriting it as a log marginal likelihood in a probabilistic model where latents comprise of states, actions, and “optimality” observations.

The prior over states is the world model and the prior over actions is the policy which being optimized. The likelihood is the exponentiated reward which can be normalized to make it a distribution over auxiliary “optimality” variables.

The marginal likelihood can be optimized via EM where in the E step, we maximize the evidence lower bound (ELBO) in (3) with respect to the variational distribution q and in the M step, we fix q and maximize the objective with respect to the policy π.

In the E-step (Sec. 4.2), the paper derives an analytic expression for q via the Lagrange multipliers method (eq. 6).

In the M-step (Sec. 4.3), the paper uses this q to maximize the objective with respect to π. The key difficulty is that the q derived in the E-step is only known up to a normalizing constant and so the paper proposes to get approximate samples from q using SMC (Sec. 4.2.1).

The experiments compare the proposed method to baselines on several tasks with discrete and continuous action spaces, and show wins. There are ablations showing the tradeoff between using more particles vs unrolling for longer.

In addition to performance, the paper argues that SPO is preferable to the popular MCTS method because it parallelizes more easily and performs well across more domains.
strength_weakness: strength_weakness: -
questions: questions: Major
- SMC (Sec. 4.2.1): From (6) it seems like there is just one distribution we want to sample from instead of a sequence of them. What’s the sequence of distributions that SMC targets? In other words, in the background section about SMC, it’s unclear what is $p(x)$ and what is $p(x_t | x_{1:t - 1})$; how do terms in (8) relate to this $p$?
- If I reverse-engineer a sequence of targets distribution from (8), it should be $p(a_{1:t}, s_{1:t}) = \prod_{j = 1}^t \mathcal T(s_t | …) \pi(a_t | …) \exp(A^{\bar \pi}(a_t, s_t) / \eta^*)$. That is, the final target distribution is the transition times the policy times the exponent of the *sum* of advantage functions. I think this is incorrect since it doesn’t correspond to the sum of per-step rewards. Something that would make more sense to me is to have the incremental weights contain the exponentiated rewards for all time steps instead of the last one, where it’s the exponentiated advantage function. This is obviously a completely different algorithm.
- I don’t understand why (5) follows from (4). We have an unconstrained optimization of a regularized objective in (4) which is turned into a constrained optimization problem in (5). I also don’t understand how $\alpha$ turns into $\epsilon$.
- Why does (10) rewrite (9) as a constrained objective where it seems like in practice, we’re actually just optimizing (9).
- How is the value function and the Q-function learned? These seem critical for the method but I couldn’t find a description.
- E-step: I agree that it is difficult to optimize with respect to $q$ since it’s the distribution of the expectation and the term inside the expectation, $Q^q$ depends on it (eq. 4). The paper chooses to fix $Q^q$. I would have expected that it is the $q$ in $\mathbb E_q$ that makes optimization difficult since typically we’d have to resort to the reparameterization trick or REINFORCE. If we fix the $q$ in $\mathbb E_q$, optimizing the expectation is simple: just sample from $q$ and directly optimize $Q^q$.

Minor
- $\gamma$ in (3) is not introduced before.
- Should there be a likelihood term in the incremental weight in (2)? Or is $p(x_t | x_{1:t - 1})$ an unnormalized distribution? $p$ is not defined in this paragraph.
- Swapping the prior term $\log p(\theta)$ in (9) to the KL term in (10) is a little odd. Would it simply not work otherwise?
- In App. G.1, why is $\mu_q$ dependent on $q$?

Other relevant work
- [AESMC](https://arxiv.org/abs/1705.10306), [FIVO](https://arxiv.org/abs/1705.09279), [VSMC](https://arxiv.org/abs/1705.11140)
- [DVRL](https://arxiv.org/abs/1806.02426)
- [SIXO](https://arxiv.org/abs/2206.05952), [NAS-X](https://arxiv.org/abs/2308.14864)
limitations: limitations: -

Comments: title: Acknowledgement of the rebuttal
Reviewer Engagement
comment: I acknowledge that I have read the rebuttal and other reviews and I maintain my score.
Dear AC/Reviewers:

We kindly highlight the feedback window is coming to a close, and would appreciate feedback from the reviewers in response to our clarifications of their reviews.

We feel we addressed the raised concerns. Specifically, Reviewer X2nr, we clarified their theoretical concerns and engaged constructively regarding the method description. We had hoped for a further interaction, as per their comment "I’m giving this paper a reject in its current form but I am willing to raise the score if the method or its description is improved.".

Thank you for your time.