Policy Gradient with Tree Expansion

Thank you for your response. Please find our answers below.

Question 1 - Do you learn the forward model and if not how can you use learned models?

In the experiments, the forward model is not learned but given. In cases where access to a forward model is not given, it should be learned and replaced with the exact model in the tree expansion. Note that learned forward models are typically represented using a neural network which allows efficient tree expansion with Breadth-First Search (BFS), as we do in this work with the GPU simulator. In the last few years several prominent works have shown a forward model can be learned successfully such as MuZero (Schrittwieser et al., 2020), EfficientZero (Ye et al., 2021), and Dreamer (Hafner et al., 2020). We discuss this alternative in Section 4.3 and provide a corresponding result in Theorem 4.8 for an approximate forward model.

Question 2 - Is “NameThisGame” the name of an Atari game or a typo?

Indeed an unusual name, but it is an Atari game. see https://en.wikipedia.org/wiki/Name_This_Game.

Question 3 - Is tree expansion worth the computation?

Evidently yes. As we show in Figure 3, the SoftTreeMax policy achieves substantially higher reward than those obtained by vanilla softmax. To see a comparison of time-based runs we refer the reviewer to Figure 4 in the Appendix.

Question 4 - What is the motivation of observing the variance only at the end of training?

In the initial stages of training, the return (cumulative reward) fluctuates wildly, since the policy is still almost random and exploration is intense (Anschel et al, 2016). This directly injects noise to the critic and gradient variance estimator, as they build on the return. Thus, we consider the gradient variance after convergence for a fair comparison of the estimator operator itself. More formally, the theoretical analysis relates to the expected values obtained in the limit of the stationary distribution over states.

Question 5.1 - Add a schematic of the forward pass and gradient accumulation.

We tried capturing the forward-pass mechanism (lines 649-658 that you referred to) in the following diagram: https://ibb.co/5RPRkzF .

Question 5.2 - Add an algorithm description with pseudo code.

Per your request, we wrote the following pseudo code https://ibb.co/s2CRMp8, which we will add to the final version.

Question 6 - Change figure location in the paper.

Thank you for this suggestion. We tried re-arranging the figure locations and (i) see that this hurts the readability and flow more than it helps, in addition to (ii) violating the strict 9-page limit.

Question 7 & 8 - Future work and related work sections:

We have prepared new sections addressing future work and related work. Due to space constraints, we cannot include the full content here, but only include their outline. We would be happy to provide these expanded sections as additional comments later if the reviewer wishes to see them.

The future work section focuses on:

1. Efficient tree expansion using GPU-based simulators like Atari-CuLE, IsaacGym, and Brax.

2. Learning forward models when simulators are unavailable, building on works like MuZero.

3. Extending to continuous control by sampling from parametric action distributions, adapting MPPI with a value function.

4. Theoretical analysis of non-expansive operators and eigenvalue cancellation for infinite action spaces.

The related work section covers variance reduction techniques, including:

1. Baseline subtraction (state and action-dependent)

2. Importance sampling

3. Sub-sampling (e.g., SVRPG)

4. Control variates

5. Natural policy gradients

6. Experience replay

Thank you for the rebuttal. I am still not convinced that the computations of tree expansions are worth it because:

• SoftTreeMax does not consistently imrpove on PPO.
• SoftTreeMax perf sometimes decreases with planning horizon $d$ which is weird
• I do not agree with what you answered to reviewer m94g : "If instead we used SAC with n-step TD as a baseline, then the comparison will be with SAC with n-step TD and SoftTreeMax as a policy." I think it would be fair to compare PPO/SAC with n-step TD against PPO/SAC with SoftTreeMax only. The point being that it would be a fair way to compare the performance gain over PPO/SAC with two different types of added information that are shown to reduce variance respectively in the Critic and in the Actor.

Can you provide more insight as to why the performances of SoftTreeMax can decrease with $d$ as opposed to not increase? Why not compare an Actor-Critic baseline with n-step return estimation of the Critic to an Actor-Crtic baseline with tree expansion in the Actor estimation? This would answer the question of where is it best to spend compute and use more information.

Respectfully

We appreciate your feedback on the method's soundness and reproducibility. We hope the following clarifications and new content address your concerns.

Weakness 1 - Difference between SoftTreeMax and n-step return

1. Thank you for this insightful point. SoftTreeMax may seem similar to n-step TD, but it differs in several important aspects: In n-step TD, one uses observed n-long trajectories from the interaction to estimate the value. In our (model-based) approach, we expand the tree of all future trajectories using the forward model instead. In fact, averaging over multiple future scenarios is the engine behind variance reduction.

2. In policy gradient (PG) methods, n-step returns merely serve as a Monte Carlo estimation technique for the advantage function $A$ in the gradient estimator $E[(\nabla \log \pi)A]$, but do not directly influence the policy itself. In contrast, in SoftTreeMax the multi-step reward is part of the policy itself, and it affects $\pi$ in $E[(\nabla \log \pi)A]$, rather than affecting $A$.

3. In our SoftTreeMax, $\theta$ (defined in Section 2) appears in the softmax similarly to a Q-function. However, it is not necessarily the Q-function itself and is not explicitly minimizing any Bellman-type loss function. While n-step TD with value bootstrapping is common in value-based methods, its application to policy gradient algorithms remains unclear. Our work aims to address an important question: “how should planning be incorporated into PG methods?” As far as we know, the community hasn't fully explored this path.

Since SoftTreeMax and n-step TD operate on different terms of the gradient, both can be used at the same time by choosing SoftTreeMax as the policy.

Interestingly, Reviewer mmxv independently highlighted a key distinction of our approach. They noted the originality in using look-ahead information directly in policy parametrization, contrasting it with more common advantage estimation using n-step returns. This observation aligns with our view on the novel aspects of SoftTreeMax compared to traditional n-step TD approaches.

We will add a discussion on this topic to avoid any misunderstanding.

Weakness 2 - E-SoftTreeMax experiments

We actually repeated all experiments with E-SoftTreeMax, following the same protocols as C-SoftTreeMax (see footnote 1 in page 8). Results for E-SoftTreeMax are almost identical to C-SoftTreeMax, so we left the figures out. The reason why the results of the two variants are almost identical, is that for quasi-deterministic Atari, the trajectory logits (Eq. (2) in the paper) have almost no variability. This makes the two SoftTreeMax variants (Eqs. (3) and (4)) nearly identical. We believe there is room for future work and a deeper investigation of E-SoftTreeMax using probabilistic environments that have inherent risk associated with them.

Weakness 3 - SAC baseline with n-step TD

The reason we used PPO as our baseline is that we built SoftTreeMax on top of it, as a model-based variant of PPO. This gives a fair comparison with similar compute power, using 256 GPU workers in the original PPO, considered as “depth 0” in the paper. SAC can be extended in similar fashion, and in that case we believe the “depth 0 SAC” would be a relevant baseline. Again, to highlight the point of Weakness 1 - If instead we used SAC with n-step TD as a baseline, then the comparison will be with SAC with n-step TD and SoftTreeMax as a policy.

Weakness 4 - Five seeds is not enough.

We followed the community standard for this parameter. Using 5 seeds is very common; for example, this is what is used in the seminal papers introducing DQN, PPO, SAC, TRPO, IMPALA, TD3, MuZero, EfficientZero, Dreamer, and many more. Figure 4 shows that while there is some variance between the experiments, the differences in the plots is considerably higher, demonstrating the efficacy of our proposed method.

Weakness 5 - Continuous control

Extending to continuous control is doable but requires significant research effort. To expand the tree, one can maintain a parametric distribution over actions which can depend on $\theta$. This method can be seen as a tree adaptation of MPPI (Williams et al., 2017) with a value function. Theoretically, we can exploit the non-expansive operator property of $P^{\pi_b}$ and the policy gradient's eigenvalue cancellation. These properties can be shown to hold for decision models with infinite actions. Nonetheless, we wish to highlight that a finite action space is an important setup in RL that was studied abundantly.

Weakness 6 - How to pick $d$?

The choice of d depends on the game's planning requirements. Games needing long-term planning require better value estimates, which take longer to obtain at larger depths. For example, Breakout benefits quickly from large depths due to its short-term goal, while Gopher's complex task requires more training time for larger depths to be beneficial. We suggest three approaches for choosing d: 1) Increase until results stop improving, 2) Use a "natural" time horizon if applicable, or 3) Learn d online (though our preliminary tests found this unstable). We'll discuss these in the final version.

Question 1 - Is $W_\theta$ the value function?

Not necessarily. When making that assumption, we get very interesting connections to SAC and n-step TD as pointed out in your first comment and by the other reviewers. But we do not constrain or provide any loss that links $W_\theta$ directly to the value or Q-function. Our work is a natural extension from vanilla PPO which does not hold any requirements from the weights per action either.

Question 2 - Could you explain the empirical setting in more detail?

Yes. We added a new diagram of the tree expansion implementation in the GPU https://ibb.co/5RPRkzF, and the pseudo code for SoftTreeMax-PPO https://ibb.co/s2CRMp8. We will add them to their appropriate location in the final version.

Thank you for supporting the paper and for your detailed and insightful response. We are delighted that you found the concept intriguing and innovative. Please find our answers below.

Weakness 1 - “Not a weakness”: I expect to see improvement to any policy gradient method

Indeed, our SoftTreeMax method can be used with any other PG algorithm since it provides a parametric policy that is an alternative to softmax. We agree that it is expected to improve other algorithms besides PPO studied here. We focused the experimental efforts on PPO because it is commonly used, has good performance overall and can be reduced and extended to multiple other variants in the literature (such as REINFORCE).

Weakness 2 - Highlight the difference between "Tree search" and "tree expansion"

Good point; we will emphasize this when introducing the expansion mechanism in the third paragraph of the introduction in the final version.

Weakness 3,4 - You can cite Efroni [2018] and Barakalt [2024], Yuan [2022]

Thank you for the references; we will add them to the appropriate parts of the paper. Specifically intriguing is Efroni et al., 2018. There, they show that multi-step policies improve convergence rate at the expense of more complex iterations for policy iteration. We believe that a similar property can be shown for such policies used in policy gradient, and specifically for SoftTreeMax, with higher depth leading to better convergence rate. This behavior is also expected because the policy then gains access to more information (reward).

Weakness 1 - Implementation details

We've added pseudo-code for SoftTreeMax-PPO (https://ibb.co/s2CRMp8), a new GPU tree expansion diagram (https://ibb.co/5RPRkzF), and open-source code (supplementary). We used hyper-parameters from the original PPO paper (Appendix B.1) and defaults from stable-baselines3 for the other details requested.

Weakness 2 - Continuous action spaces

For continuous control, we propose sampling from a parametric distribution over actions depending on $\theta$, similar to MPPI (Williams et al., 2017) with a value function. Theoretically, we can exploit the non-expansive operator property of $P^{\pi_b}$ and the policy gradient's eigenvalue cancellation.

Weakness 3.1 - Exponentially slow convergence

Unlike Softmax, SoftTreeMax's convergence is less likely to depend exponentially on state space size. The lookahead and added reward prevent gradient collapse except in degenerate cases where $P^{\pi_b}$ becomes rank-one, enabling more stable learning across various state space sizes.

Weakness 3.2 – Choice of behavior policy

We use a uniform tree expansion policy, assigning equal weight to all actions. This approximates the optimal variance decay, achieved when the induced transition matrix is near rank-one. Such a policy smoothens transition probabilities, minimizing the associated Markov chain's mixing time without requiring specific MDP assumptions.

Weakness 4 - SoftTreeMax's nature and learned theta values

SoftTreeMax is an alternative parametric policy to Softmax, incorporating reward and forward model. It differs from MCTS in framework, extensions, and loss functions. Learning significantly improves performance, especially at larger depths in complex games. Figure 4 (Appendix B.2) shows performance variations across games with and without learning. The importance of learned theta values increases with depth, particularly in games requiring long-term planning. We can provide more concrete examples and explanations if needed.

Weakness 5 - Sampling from learned policy

Correct. Inference requires tree expansion using the same behavior policy as training. Potential mitigations like policy distillation need further research. We'll add this as a limitation in the paper.

Weakness 6 - Results

Our results show significant improvement over PPO in all 8 games, including Breakout and VideoPinball. In Breakout, we reach a reward over 800 for depths 4 and 5, while PPO doesn't. In VideoPinball, we obtain 800K for depth 3 while vanilla PPO reaches less than 200K.

Question 1 - Convergence and variance reduction

We believe convergence is sped up by variance reduction. In most games, variance monotonically reduces as depth increases, while reward acts oppositely. Figure 1 shows SoftTreeMax's gradient variance is consistently three orders of magnitude lower than PPO's, correlating with tree depth.

Question 2 - Figure 4 interpretation

Correct. 'Varying number of steps' means algorithms run for the same wallclock time but different environment interactions. This balances online vs. planning iterations, accounting for both outer and inner iteration complexity. Appendix B.2, Figure 4, shows a sweet spot balancing these iterations at different depths depending on the Atari game.

Question 3 - PG algorithms stability

We'll revise the statement about PG algorithms' instability. Actor-critic methods mitigate variance from entire trajectories, but instability persists for other reasons. We'll cite TRPO and PPO papers for relevance.

Question 4 - Reward exponentiation

E-SoftTreeMax (Equation 4) can be viewed as exponentiating rewards within logits before expectation, leading to a product of exponents of the reward sequence. We'll clarify this phrasing.

Question 5 - General comments

We appreciate your thorough suggestions on wording, figure locations, notations, and preliminary section. We'll incorporate these in the final version and fix any typos and inaccurate phrasing.

Question 6 - Gradient positivity

We meant the norm of the gradient is strictly positive everywhere. We'll clarify this and elaborate on the relevant lemmas from Mei et al., 2020b. Lemma 8 provides a lower bound on the gradient norm, while Lemma 9 shows this bound is strictly positive.

Question 7 - Variance in permutation matrix case

In this case, variance isn't zero. We have a single leaf per action selection, with rewards and logits exponentiated in the softmax operation. The zero-variance case occurs with zero second eigenvalue, detailed in Section C.1.

Question 8 - Norm in Theorem 4.8

We use the induced infinity norm, mentioned in line 597. The bound's independence from behavior policy and transition dynamics stems from approximating dynamics and reward up to $\epsilon$ error. We can bound relevant matrices by a factor of $\epsilon$ due to their convex combinations of transition and reward. For further details please refer to the proof in Lemma A.4 and line 624 in the appendix.

Question 9 - Search process

For depths up to 3, we search all possible state-actions due to the quasi-deterministic nature of Atari environments. For larger depths, we prune low-weight trajectories. We'll clarify this in the paper.

Limitations

We've added pseudo-code and implementation details as requested. For continuous action spaces, we've proposed approaches but acknowledge further research is needed. We'll expand the limitations section, incorporating feedback from reviewers. The observation regarding linear dependence on $S$ is a great one. in our bias bound may be an artifact of our proof, potentially improvable in future work, with a tighter bound in Eq. (81). Pruning introduces bias through sub-sampling, affecting performance and convergence. The behavior policy controls the bias-variance tradeoff, balancing between proximity to $\pi_{t,d}$ (lower bias) and similarity of $\pi_b$ rows (lower variance).

I have read the rebuttal and I thank the author for their clarifications; they helped a lot. I also read the other reviews and the corresponding rebuttals. I think all of them raise relevant points. I think if the authors were to incorporate these suggestions, the quality of the paper would increase a lot. I retain my score of 7; it seems to be a fair evaluation of the work.

Some additional / minor comments:

Related to my review:

Lines 12 and 13 in https://ibb.co/s2CRMp8 can't really be implemented (since the objective is highly non-convex in terms of the policy / value function weights). So maybe, put "take a gradient step" or something similar there instead of "maximizing".

The diagram https://ibb.co/5RPRkzF looks very nice, but because it is somewhat complicated, it needs a proper caption. I hope that will be added too.

"Our results show significant improvement over PPO in all 8 games, including Breakout and VideoPinball. In Breakout, we reach a reward over 800 for depths 4 and 5, while PPO doesn't. In VideoPinball, we obtain 800K for depth 3 while vanilla PPO reaches less than 200K."

I still think that the results are very similar for these problems. I will suggest not writing "obtain up to 5x more reward"; it is not well supported by the experiments and seems like an oversell. Instead, just saying "competitive" or "outperforms" is enough.

Related to other reviewer's comments:

It seems like making the distinction between this work and n-step TD based actor-critic methods would be very useful. In particular, the nice and short analysis by m94G could be included and clarified.

The comment on five seeds not being enough is correct, and I don't agree with the author's responses. Now I'm not necessarily asking you to run more seeds, but at the very least, acknowledging that five seeds is probably not enough in the paper would be great! (Also see Figure 5 of https://arxiv.org/pdf/1709.06560 and Figure 15 of https://arxiv.org/pdf/2304.01315).

May I request the authors to include the experiments they already ran on the E-Softtreemax?

More importantly, is it possible to run the experiments for a learned model? I think having even just one experiment demonstrating that SoftTreeMax works with a learned model could be very encouraging for the readers.