SPO: Sequential Monte Carlo Policy Optimisation

We thank the reviewer for their detailed review which will be used to clarify details of the paper

1. We aim to sample from the target policy over sequences $\tau= (s_0, a_0, s_1, a_1, ..., s_t, a_t, s_{t+1})$ similar to [60]. In the following notation we switch to using $\tau$ from $x$ and give definitions for the density $p$ and proposal $\beta$ at policy iteration $i$.

• $p_i(\tau_{1:t}) \propto \mu(s_0) \prod_{j=0}^t \mathcal{T}(s_{j+1}|s_j,a_j) \pi_i(a_j|s_j, \theta_i) \exp\left(\frac{\bar{A_i}(a_j,s_j)}{\eta^*_i}\right)$

• $p_i(\tau_t|\tau_{1:t-1}) \propto \mathcal{T}(s_{t+1}|s_{t},a_{t}) \pi_i(a_t|s_t, \theta_i) \exp\left(\frac{\bar{A_i}(a_t, s_t)}{\eta_{i}^*}\right)$

• $\beta(\tau_{1:t}) \propto \mu(s_0) \prod_{j=0}^t \mathcal{T_{\text{model}}}(s_{j+1}|s_j,a_j) \pi_i(a_j|s_j, \theta_i)$

• $\beta(\tau_{t}|\tau_{1:t-1}) \propto \mathcal{T_{\text{model}}}(s_{j+1}|s_j,a_j) \pi_i(a_j|s_j, \theta_i)$

The weight update:

$w_t(\tau_{1:t}) \propto w_{t-1}(\tau_{1:t-1}) \cdot \frac{p(\tau_t|\tau_{1:t-1})}{\beta(\tau_t|\tau_{1:t-1})} = w_{t-1}(\tau_{1:t-1}) \cdot \left(\frac{\mathcal{T}(s_{t+1}|s_{t}, a_{t})}{\mathcal{T_{model}}(s_{t+1}|s_{t}, a_{t})}\right) \cdot \frac{\exp(\bar{A_i}(a_t, s_t)/\eta_{i}^*) \cdot \pi_i(a_t|s_t, \theta_i)}{\pi_i(a_t|s_t, \theta_i)}$

2. You are correct, the target distribution over sequences utilises the sum of advantages and not explicitly per step rewards. This does account for the sum of per step rewards, as we calculate $A^{\pi}(a_t,s_t) = E_{\pi}[r(a_t,s_t) + \gamma V(s_{t+1})] - V(s_t)$ therefore for $\gamma = 1$, over sequences, we have the following estimate $\sum_{t=0}^h A^{\pi}(a_t,s_t) = \sum_{t=0}^h r(a_t,s_t) + V^{\pi}(s_{h+1}) - V^\pi(s_{0})$, the sum of per step rewards plus the value of the last state, minus the value of the first state. If the weight only uses the per step rewards, future rewards beyond the planning horizon are not accounted for. This destroys performance, see new results fig 2. Without a baseline, the weight update would be $\prod_{j=0}^t \exp(Q^{\pi}(a_t,s_t)/\eta^*)$, see line 660 for ablation.

3. In optimising 4, the first term $\mathbb{E}_{q(a|s)} \left[Q^q(s,a)\right]$ is dependent on the scale of reward in the environment. This can make it hard to choose $\alpha$, as the first and second terms are on arbitrary scales. As explored in previous work [1,69] we choose to change this from a soft constraint to a hard constraint (line 158). Instead of choosing $\alpha$ we choose $\epsilon$ which is the maximum KL divergence between $q$ and $\pi$ which is practically less sensitive to reward scale.

4. Eq 9 is a supervised learning step, we have choice over the parameterisation of our policy and prior for regularisation. RL algorithms often constrain policies from moving too far from the current policy, in this work we utilise a Gaussian prior around $\theta_i$ optimised from the previous iteration. Therefore, $\theta \sim \mathcal{N}(\mu, \Sigma)$ where $\mu = \theta_i$ and $\Sigma^{-1} = \lambda F(\theta_i)$, where $F$ is the Fisher information matrix. The prior term in (9) can written as $\log p(\theta) = -\lambda(\theta - \theta_i)^T F(\theta_i)^{-1}(\theta - \theta_i) + c$. Note $c$ is a term that do not depend on $\theta$. The first term is the quadratic approximation of the KL divergence [2*], so we can rewrite equation (9) as:

$$\max_{\pi} \mathbb{E}\_{\mu_q(s)} \left[ \mathbb{E}_{q(a|s)} \left[ \log \pi(a|s, \boldsymbol{\theta}) \right] - \lambda \text{KL} \left( \pi(a|s, \boldsymbol{\theta_i}), \pi(a|s, \boldsymbol{\theta}) \right) \right]$$

Like the E-step, choosing $\lambda$ can be non-trivial so we convert it into a hard constraint optimisation (eq 10). $\epsilon$ in Eq (10) is different from (5). The constrained M-step (with a Gaussian prior) is not required for SPO, however it adds stability to training.

5. SPO does not require a Q-function and only learns a value function since we leverage observed rewards and state-values from planning to calculate advantages. We learn the value function like PPO, i.e using GAE values. GAE combines multiple N-step td-errors estimates, $A_\text{GAE}(s_t) = \sum_{l=0}^{N} (\gamma \lambda)^l \delta_{t+l}$ where $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$. Such that $V_{target}(s_t) = V(s_t, \theta) + A_\text{GAE}(s_t)$. These value targets are not constructed during planning but on real rollouts. It is a practical choice as to what method is used to generate targets, we opt for GAE to balance bias and variance.

6. In Eq 4 when optimising with respect to q there is a moving target, each update to q results in a new critic $Q^q$. Our approach aligns with most actor-critic approaches [65] that fix a critic within an iteration of policy optimisation. For the suggestion of optimising the $Q^q$ term and fixing the acting policy $q$, this would be a highly non-trivial optimisation. Since $Q^q$ depends on the transition dynamics and reward function of the environment unless we knew these we would need to perform off-policy evaluation over the space of possible policies $q$, with such evaluation known to be extremely challenging [1*], as behaviour of policies is difficult to estimate using different distributions. We are not aware of RL algorithms that approach optimisation from this perspective.

Minor

1. $\gamma$ is introduced in the background as the discount factor however we will add to 3) for clarity

2. $p(x_t| x_{1:t-1})$ is unnormalised, we define p(x) before eq 2) as an arbitrary target distribution

3. See major q4

4. $\mu$ is the distribution over states. Any policy acting in the MDP will induce a different state distribution. We therefore denote $\mu_q$ in Eq 11 as this relies on the policy $q$. We will add this clarification to the manuscipt.

References

[1*] Dudık et al "Doubly Robust Policy Evaluation and Learning" ICML 2011

[2*] Schulman et al "Trust region policy optimization" International conference on machine learning PMLR 2015

On Novelty: Re: "limited novelty'', we respectfully disagree with the reviewer. We highlight three contributions (all included in the manuscript lines 27, 39, 38):

1. This is the first instance of Sequential Monte Carlo being used as a policy improvement operator for RL. To date, MCTS has been the most widely adopted search-based policy improvement operator. The paper presents a robust search algorithm, with a theoretical foundation, resulting in an operator that outperforms competing methods on a range of benchmarks.
2. The introduction of an inherently parallelisable operator is also a strong contribution. We are not aware of any successful search based policy improvement operations that are able to leverage parallelism over search to speed up inference during training. In fact most works note that parallelisable versions of MCTS strongly reduce performance [1,2] and have not successfully implemented parallelisable versions of AlphaZero for policy improvement[3].
3. Introducing a method that requires no algorithmic modifications between continuous and discrete environments while demonstrating strong performance in both settings, should not be considered as limited. The generality of popular RL algorithms is often limited with algorithms such as AlphaZero requiring considerable modifications just to be applied to complex continuous domains. This is an important barrier for the applicability of search based policy

On Quality/Clarity: The manuscript highlights deficiencies with the current SOTA methods (lines 22-25). After introducing SPO (lines 26-36), we explain how SPO tackles each of the highlighted problems. To clarify, we aim to solve the scalability problem in current search based policy improvement methods, such that search can be executed in parallel without performance deterioration observed in previous works [1,2]. This allows for search methods to leverage hardware accelerators to perform search faster reducing train times (lines 38-39). We explain that SPO is widely applicable to discrete and continuous environments without modifications (lines 37-39). Lastly we explain SPO is theoretically backed (line 29) and that the method scales for different levels of planning (line 41-42).

On Ablations: "Ablation studies are missing'' This is an incorrect statement from the reviewer. In the original manuscript we provide two important ablations. We ablate the difference between sampling the target policy in (6) and the alternative formulation without a baseline (line 660), important as previous works differ on these objectives [1,69]. We also provide a second ablation that demonstrates how the appropriate temperature for SMC varies significantly between environments (line 649). This backs up our approach of learning the temperature variable, and demonstrates this is an important contribution to reduce hyperparamter tuning. We add a 3rd ablation to elicit insight into how SPO enforces the KL constraint in policy improvement (Figs 1). We show two charts with varying runs for different maximum KL constraints. We see when viewed on the iteration level that SPO is able to accurately target the KL of choice (fig 1 a). Fig 1 b) shows the impact different KL constraints have on training curves.

Re: "the paper does not give experimental results showing better estimation results". In practice we do not have access to the target distribution, therefore it is non-trivial to construct an ablation that directly measures how good methods are at estimating it. However, we add new ablation that leverages Monte Carlo samples to generate an unbiased estimate (given large compute budget). It compares the KL divergence of our SMC policy to a Monte Carlo oracle (Fig 3). We investigate the impact of depth and particle number on the KL divergence to the oracle for SMC on Sokoban. This concludes that scaling particles is particularly important for improving the target estimation, aligning with SMC theory [4], also that leveraging depth improves estimation. Methods like MPO that limit particles to the number of actions and depth 0 are severely limited in their estimation of the target, likewise V-MPO only uses 1 sample albeit to a depth of N. Lastly comparing MPO to SPO there is a performance gap across all environments. Both methods have similar theoretical foundations the difference being how samples are obtained from the target. Therefore any performance differences can be attributed to changes in target estimation.

On Theoretical insights: Similarly, the following statement made by the reviewer is untrue "Theoretical insights of proposition 1 are missing" We provide full theoretical insights for proposition 1 in Appendix section G.3 (line 749). This is referenced in the main manuscript on line 238.

On Formulas: Many of the formulas are either existing work, or modifications to existing work. We believe that the provided theoretical justification is in line with the standards of rigour expected at NeurIPS. It is reasonable to expect that experts in the field would feel comfortable with work that provides such level of detail in order to clearly understand the theoretical foundations and methodology. If you think certain formulas could be re-positioned in the paper, with additional insights provided, we welcome specific feedback in order to increase readability.

References

[1] Liu, Anji, et al. "Watch the unobserved: A simple approach to parallelizing monte carlo tree search" arXiv preprint (2018)

[2] Segal Richard B. "On the scalability of parallel UCT." International Conference on Computers and Games. Berlin, Heidelberg: Springer Berlin Heidelberg 2010

[3] Seify Arta, and Michael Buro. "Single-agent optimization through policy iteration using monte-carlo tree search." arXiv preprint (2020)

[4] Del Moral et al Branching and interacting particle systems approximations of Feynman-Kac formulae with applications to non-linear filtering Springer Berlin Heidelberg 2000

We thank the reviewer for their positive review.

Re: evaluation of SPO with an imperfect model. Our primary goal was to focus on the novel planning aspects of SPO and its use within training. We believe that the results demonstrate that SMC is a valuable planning method in its own right demonstrably outperforming other commonly used planning methods like AlphaZero when the model is known.

However, we acknowledge that incorporating a learned model is a crucial next step for demonstrating practical applicability. Recent advancements in model-based methods, such as DreamerV3 [2] and EfficientZeroV2 [3], offer advanced architectures for modelling the underlying MDP. These methods are policy-learning agnostic and could easily integrate with the SPO algorithm. As long as the world model accurately predicts rewards, SPO should perform effectively. Additionally, Piché et al. (2018) [1] demonstrated the feasibility of using a learned model with a SMC-based search for maximum entropy RL. This prior work supports the potential integration of learned models in our approach, which we plan to explore in future work.

Questions

Q: Line 204: How exactly is $\hat{A}$ computed? It isn't very clear on that.

A: Practically, at each step in the sequence we estimate the advantage of a single state action pair using a 1-step estimate $\hat{A}(a_t,s_t) = r(a_t,s_t) + \gamma V(s_{t+1}, \theta_i) - V(s_t, \theta_i)$. As we perform search this we accumulate advantages at each step to calculate the overall importance weights of the sequence according to Eq 8. However, if desired, any advantage estimation technique that can be constructed during the search from a sequence of rewards and values. For example, GAE to allow a trade-off between bias and variance.

Q: How much does removing samples based on weight affect exploration? Is this only helpful when you have a perfect model?

A: This is a good question. The most important factor is actually the temperature variable which is controlled by the desired KL divergence $\epsilon$ between the target policy $q_i$ generated by search and the current policy $\pi_i$. If the temperature is too low, then during resampling, the particles can collapse onto a single root action which would make the rest of the search pointless. But when the temperature is too high, the search can be too exploratory and not effectively travel down promising paths. This does need to be considered when scaling the search method depth-wise and when choosing the resample period. This wouldn't only be helpful when using a perfect model as ultimately as long as a learned model has accurate reward prediction and dynamics prediction, the exploration/exploitation acts the same. We will amend our manuscript to include this explanation in order to provide greater intuition on the exploration of the SPO search.

Thank you once again for your valuable feedback.

References

[1] Piché, Alexandre, et al. "Probabilistic planning with sequential monte carlo methods." International Conference on Learning Representations. 2018.

[2] Hafner, Danijar, et al. "Mastering diverse domains through world models." arXiv preprint arXiv:2301.04104 (2023).

[3] Wang, Shengjie, et al. "EfficientZero V2: Mastering Discrete and Continuous Control with Limited Data." arXiv preprint arXiv:2403.00564 (2024).