Trust-Region Twisted Policy Improvement

We thank the reviewer for their effort and accurate summary of our paper, we appreciate that you recognize our work as a solid advancement to SMC-planners for RL and the positive assessment.

We address your points below.

---

Minor comments:

• We included the derivation for Corollary 2.3 in the appendix A.4, but forgot to reference this. We will update this.
• The reference to Fig.5 vs Fig.3 in the main text and appendix is a mistake, this should be swapped. Good catch!

---

Figure 1:

In Fig.1 we show that running the baseline SMC planner for a longer horizon (depth) causes the distribution over actions $\pi_{SMC}$ at the start to diverge from their optimal distribution $\pi^*$. This is indicated in the y-label with $KL(\pi^* \Vert \pi_{SMC})$. We believe it's important to show the pattern that our method doesn't worsen as a function of depth. The complete details on how we generated this figure are in Appendix B.6, we will add a reference for this in the figure 1 caption.

This pattern is a problem in the baseline because it increases variance of the policy estimation which cascades into exploration and learning (as also noted by Reviewer bnoX). We tested our modification in the ablation-study in figure 4 on the right. We tested this over multiple parameter settings due to hyperparameters interacting with one-another. This follows principles laid out by Patterson et al. (2024). See also our answer below on the ablations.

For clarity, we will copy the ylabel in Fig.1 from the left-bottom also to the plot in the right-bottom.

---

Learning diagram:

We prepared pseudocode for the outer loop to complement our discussion on approximate policy iteration in the background and the inner-loop pseudocode in the appendix (see also our response to reviewer SRnR on how we intend to revise this). The pseudocode can be found here: link to EM-pseudocode), we will put this either in Appendix C or in the main text depending on whether we have space for additional ablation results.

Ablations + Answer to Question:

We tested out all the improvements separately to the best of our abilities and resources, across varying environments and marginalizing over many parameter settings. This is mostly shown in figure 4 and table 1. Because there were so many hyperparameters involved, the details of the ablations were put in the appendix B.3. Furthermore, we simply had to make sacrifices in what to test within this immense combinatorial space of hyperparameters.

We recognize that the details of each ablation study got a bit buried in section 5 when trying to connect this back to section 3, where we described the improvements.

Q: As a minor revision, we will annotate each section in the experiments. Would this help?

For example, the contribution in section 3.2 is tested in the experiments in line 434 (col 1), we propose to 1) reference this part in section 3.2 and 2) cross-reference back to section 3.2 in the paragraph-heading -- e.g., Revived Resampling (§3.2).'

In summary,

• §3.1 improves sampling inside SMC, this is studied in figure 4.
• §3.2 prevents particles from getting trapped in absorbing states, wasting compute resources. This is algorithmically shown in Alg.2 (appendix C). It is studied empirically in table 1, setup in appendix B.3.
• The problem of §3.3 and §3.4 is illustrated in figure 1.
• For §3.3, our 'fix' is shown in figure 1 bottom-right and studied in the right of figure 4.
• §3.4 also improves path-degeneracy, the improvement is studied in table 1, setup in appendix B.3.

Despite the details in the paper, we have included the table below for your convenience that summarizes the performance detriment of turning off each modification for our TRT-SMC shown in Figure 2. Note, that due to limited time and resources within this rebuttal period, we could only show this on RubiksCube for 10m training steps with the same statistical precision (30 seeds 99% intervals). Note that on this environment, and for this combination of hyperparameters, that some choices have less profound effect on the final performance than others. With the twisting having the strongest effect and revived resampling not differing statistically significantly. However, these effects can change as e.g., the depth of the SMC planner increases.

For the camera-ready, we are currently running experiments for this table on all environments.

We thank the reviewer for their time and effort. We divided our rebuttal in two parts, firstly on the experiment design and the second part on answering your comments.

---

Choice of environments:

For the choice of environments, we wanted a mix of discrete and continuous control environments. For this, we tried to closely stick to the choice by Macfarlane et al. (2024) as we believe they achieve this mix well. It also allows us to accurately compare our SMC baseline results. However we replaced Sokoban with Snake to drastically scale down the experiment costs. Based on Bonnet et al. (2024), we believed Snake to be sufficient as it is a sparse reward environment that benefits from planning, and can be learned with a smaller model and fewer environment steps.

Concerning the ablations, the environments in the ablations were chosen to 1) spread out to not overfit results on one single environment, 2) to scale down experiment cost where possible, and 3) to verify that our modification actually solves an issue of interest.

This third point is important as on some environments the effect of e.g., $\alpha$ is more profound than others (see extra results below). For instance, sparse reward environments can benefit more from value-guidance in search. Just like how exploration papers often test their method specifically on Montezuma's revenge, we also want to test our modifications in settings where they are most needed. Showing isolated effects of our modifications in ablations while still having good performance across the board keeps the discussion self-contained.

Choice of $\alpha$

The specific choice of $\alpha=0.1$ we used in the main results (figures 1 and 2) seemed to be robust across environments (not tuned per environment). This value also mirrors common trust-region parameters of $\epsilon = 0.1$ used in TRPO, MPO, and SPO.

Extra results: We did not have complete results for $\alpha$ ablations at the time of submission (only a few seeds) and have since ran experiments for $\alpha$ to a greater extent. These results can be found here for: on discrete envs, on continuous envs, and aggregated across envs. The figures report final min-max normalized average returns on all environments (30 seeds). The different lines show different planning budgets of $N=K \cdot m$ where $K$ (particles) and $m$ (depth) with $K=m$, and small $\alpha$ means sticking to the prior $\pi_\theta$ and large $\alpha$ means a greedy policy $\hat{\pi}^*$ over the predicted $Q$ values.

Similar to e.g., the clipping-ratio in PPO, different environments have different optimal $\alpha$ ($0.1$ for discrete, and $0.01$ for continuous). Overall $\alpha=0.01$ seems to work even better than what was used in our submission $\alpha =0.1$ (which we picked from preliminary testing). However, both of these values are robust across environments.

Click here for supplementary results folder

A2C:

A2C is indeed a weak baseline that could be removed. But, we included it to contrast the difference in implementations so that readers can compare against results by the official jumanji code.

---

Answers to Questions:

1. This means that the SMC and MCTS agents use more environment samples (planning budget) per 'real' environment step. We do planning with the real environment model. This mirrors having a highly accurate simulator. See also our response to Reviewer SRnR on Figure 2.

2. PPO is slower in figure 3 due to the bottleneck of neural network training and relative cheapness of the environment model. This means we can use SMC planning to cheaply draw more environment samples which can then 1) extract a more informative learning target and 2) generate more diverse data (Hamrick et al., 2021).

Outer-loop pseudocode:

We prepared pseudocode for the outer loop to extend our discussion on approximate policy iteration in the background (see our response to reviewer SRnR on how we intend to revise this).

On $q^\star$

Sadly, the convention of optimal state-action values coincides with variational inference for which $q^*$ is standard to the optimal variational distribution. So we opted to write value functions with capitals $Q$ (with a hat $\hat{Q}$ for approximations) and use small $q$ for variational distributions. It is a sacrifice.

---

We hope our response addresses your concerns, and we are happy to discuss on any further concerns or questions, and explain the extra results in more detail.

We thank the reviewer for their time, effort and, accurate summary of our paper.

Regarding limited novelty, we understand the concern but would also like to point to a slightly different perspective. With the benefit of hindsight it is easy to dismiss contributions by being reductive. For instance, the SMC-planner by Piche (2018) combined soft-actor critic with a policy extracted from a bootstrapped particle-filter, without significant changes to this filter. An overly reductive summary would say that they simply 'copy-paste' a sequential importance resampler for policy inference into the RL-loop. A similar thing can be stated about SPO by Macfarlane (2024) which replaced the MCTS-planner in AlphaZero with SMC, and then merged this into the MPO framework. There are always components being rehashed, but just like these prior works, it requires deep understanding to identify which parts of an algorithm work or don't work well, and how to deal with this.

We directly point out non-trivial issues concerning SMC-planning for RL that were not yet addressed, or overlooked, in the literature. Their identification requires deep technical understanding on integrating RL with SMC. They are non-trivial since applying a particle-filter for policy inference is fundamentally different to how the method should be used for conventional state-estimation, this introduces challenges that only emerge when integrating them. These problems relate to path-degeneracy, inefficient trajectory sampling, and wasteful data generation. We also directly provide solutions for these issues that work well and have sound theoretical grounding. We also believe that further study of this topic holds significant value to the broader RL community.

---

On figure 2:

We understand your suggestion to adjust the $x$-axis of figure 2 to account for the extra environment samples due to the planning budget. In this case, PPO and A2C would indeed perform much better in comparison to the other agents! For completeness, we added this figure here. However, this sidesteps the research question as we wrote at the start of section 5.

We really want to study which method implements a stronger local policy improvement operator, for which we compare against baseline SMC and gumbel MCTS. These all use the exact same planning budget and number of environment samples. This setting assumes we have a highly accurate simulator, which is also standard in prior work (AlphaZero, SMC for RL, SPO).

We will include the above 2 sentences to the figure 2 caption and emphasize it more in the main text of section 5 as to not mislead the reader. We will include the attached figure to the appendix.

Click here for supplementary results folder

---

On the background:

In line with reviewers bnoX and ub84 we will add pseudocode of the outer-loop training (see link to EM-pseudocode) either in Appendix C or in the main text depending on whether we have space for additional ablation results. We can 1) add a paragraph heading in the last paragraph of section 2.1 with 'Outer-loop' to emphasize/ structure the approximate policy iteration part, and 2) extend this text on how the estimated posterior policies $\hat{q}_t$ are used within this Expectation-Maximization loop in combination with neural network training. Section 2.2 and section 3 then continue in making improvements to the E-step approximation of $\hat{q}_t$.

If needed, we can shorten the background, however we do need the definitions of the posterior-policy and value functions. Furthermore, the property of theorem 2.2 leads into how we motivate the exponentially twisted proposals. We can definitely shorten the formulas and move part of this derivation to the appendix (e.g., equations 2 and 3 can be shortened to 1 line) so that the proposed pseudocode and guiding text fit.

We thank the reviewer for their effort and accurate summary of our paper. We address your points below.

---

Theoretical results:

Our paper does not introduce new theoretical results, nor do we claim so. It is correct that theorem 2.2 is known, it also goes quite farther back than Levine (2018), however this paper has a nice presentation on the theory. We restated this property since, as you also write, to motivate how we use the sequential Monte-Carlo planner and our proposed changes.

The theorem shows that 1) posterior-policy estimation implies a regularized policy improvement, and 2) that learning on samples of this posterior-policy should maximize this lower-bound. We used the former insight to impose trust-regions on the proposal distribution in the SMC-planner and the latter to perform neural network training in the outer loop.

We believe that motivating our new method with known prior results should not merit rejection but instead strengthen our approach. But, could you elaborate whether the critique is that we put too much focus on the background? If so, we will shorten this part.

---

Response to Technical difficulty:

It is unclear why the result is novel, however, as there does not seem to be a technical difficulty underpinning the result

It is unclear whether you refer to the theory, our empirical results, or our method.

As mentioned above, the theory is there to support our modifications and show how SMC planning works (e.g., the importance-sampling weights). We don't understand why supporting a newly proposed method with known previous results should be a basis for rejection.

Regarding our method, we discovered and described non-trivial issues in current approaches for SMC based planning in RL, and present concrete approaches on how to deal with this. Prior work overlooked these issues, and their identification requires deep technical understanding on integrating RL with SMC.

Also note that simplicity (if that is the concern) of a new method does not detract from its novelty, instead it should improve its clarity and broader applicability.

Perhaps you could clarify your position?

---

Minor comments:

Our method is technically not a parallel TRPO (cf. Fickinger et al. 2021). The main distinctions are:

1. The KL-divergence constraint: the arguments to this constraint would be swapped for TRPO vs. our approach (MPO-based; theorem 2.2).
2. TRPO can already be parallelized, simply by increasing the number of asynchronous environments.
3. There is a fundamental difference in how data is collected and how policy learning is performed -- relating back to our previous comment, we can extend the outer-loop discussion in the background to make this distinction more clear.

---

References

Fickinger, A., Hu, H., Amos, B., Russell, S., & Brown, N. (2021). Scalable online planning via reinforcement learning fine-tuning. Advances in Neural Information Processing Systems, 34, 16951-16963.