Paper: SoftTreeMax: Exponential Variance Reduction in Policy Gradient via Tree Expansion (ICLR-2024)

Reviewer comment (Experiments):
One weakness of the paper is in its experimental evaluation section. While the paper presents promising results for SoftTreeMax in Atari, some of the claims made are not well-supported. For example, the paper does not include enough baselines to make a fair comparison with SoftTreeMax. This makes it difficult to determine the extent of SoftTreeMax's improvement over existing methods.

Author response:
**W1: Add baselines … to determine the extent of SoftTreeMax's improvement over existing methods?** Following this important comment, we now added a SOTA baseline: EfficientZero [Ye et al., 2021]. We chose EfficientZero because it is a highly sample-efficient version of MuZero -- one of the best known RL algorithms today -- and is also open source. Due to the time limitation of the rebuttal, we present the experiments on four games, for a runtime of 96 hours. The results are given here: https://ibb.co/6Zd24yk. As seen in the plots, SoftTreeMax surpasses EfficientZero in all games but one. We will add full experiments in the final version. 

Experiment details: We used the same hardware to run all experiments, as described in Appendix B.1.  For “SoftTreeMax”, we present the best depth out of those in Appendix B.2, Figure 4. 

We thank the reviewer for their useful suggestion, which shows that there are indeed better baselines than PPO to compare to. With that said, we also recall that the PPO baseline we run here has “enhanced” sample efficiency due to its multiple GPU actors.

Reviewer comment (Experiments):
The paper lacks in-depth comparison with other related methods. While the paper compares SoftTreeMax with distributed PPO, it does not provide a comprehensive comparison with other state-of-the-art methods in the field. This makes it difficult to determine the generalizability of SoftTreeMax and its performance in comparison to other methods.

Author response:
**W1: Add baselines … to determine the extent of SoftTreeMax's improvement over existing methods?** Following this important comment, we now added a SOTA baseline: EfficientZero [Ye et al., 2021]. We chose EfficientZero because it is a highly sample-efficient version of MuZero -- one of the best known RL algorithms today -- and is also open source. Due to the time limitation of the rebuttal, we present the experiments on four games, for a runtime of 96 hours. The results are given here: https://ibb.co/6Zd24yk. As seen in the plots, SoftTreeMax surpasses EfficientZero in all games but one. We will add full experiments in the final version. 

Experiment details: We used the same hardware to run all experiments, as described in Appendix B.1.  For “SoftTreeMax”, we present the best depth out of those in Appendix B.2, Figure 4. 

We thank the reviewer for their useful suggestion, which shows that there are indeed better baselines than PPO to compare to. With that said, we also recall that the PPO baseline we run here has “enhanced” sample efficiency due to its multiple GPU actors.

Reviewer comment (Experiments):
In the experimental evaluation part, only PPO is used as the baseline algorithm. There have been many efforts to reduce the variance and improve sample complexity for policy gradient methods (e.g., including a baseline). It is therefore better to also compare state-of-the-art approaches that also solve the same problem.

Author response:
**Q1: There have been efforts to reduce the variance and improve sample complexity (e.g., including a baseline). Compare to SoTA that solves that problem?** 
Thank you for bringing this critical point. Our implementations of SoftTreeMax and SoftMax PPO already employ a baseline subtraction technique. This means that on top of our technique to reduce “direct” variance (i.e. SoftTreeMax), we also run a SOTA technique to reduce “estimation” variance (i.e. baseline subtraction). Importantly, SoftTreeMax puts forward a completely new category of variance reduction techniques because it shrinks the *inherent variance of the gradient* itself as defined in Section 2.1 below Eq. (1), instead of the variance of the *gradient estimator*. As such, it can be applied orthogonally in addition to previous variance reduction techniques. 

We will clarify this point in the revised version.

Reviewer comment (Evaluation):
The method involves sub-sampling only the most promising branches at each level, which limits the width and reduces runtime and GPU memory costs with only a small performance sacrifice. Does this also introduce additional variance by sub-sampling and pruning?

Author response:
**Q2: What is the effect of sub-sampling and pruning on the variance?** We perform sub-sampling as part of the pruning process, beyond depth 3. Its effect on variance is complex but small, and may vary with the numbers of pruned nodes. It is small because scores in the leaves are exponentiated to form the policy; thus, most pruned nodes have very low probability and the effect is insignificant compared to the computational gain. Nonetheless, in ablation studies, we saw that extreme pruning leads to low performance and slow convergence.

Reviewer comment (Experiments):
For depths \( d \geq 3 \), the method limits the tree to a maximum width of 1024 nodes and prunes trajectories with low estimated weights. Does it suffer from a limitation when used with a larger value of \( d \), which may lead to a more significant limitation of the allowed maximum width considering costs?

Author response:
**Q3: Can you use higher depths than d=3 without limitation?** Not with the current hardware. Beyond a depth of 3, the GPU memory becomes a bottleneck and without pruning searches in such depths become impossible. With Pruning, the complexity of applying the policy grows linearly with the depth. Supplemental Figure 5 In Appendix B.3 compares the number of online steps taken during one week for several depths up to 8. Indeed, for higher depths, the complexity of the policy inference becomes a limiting factor for training. We do note that this cost depends heavily on the simulator itself, and for RL problems with faster simulators higher depths are more tractable.

Reviewer comment (Writing):
What do you mean by "reward and variance are negatively correlated" on page 8?

Author response:
**Q1: What do you mean by "reward and variance are negatively correlated" on page 8?** We refer to the empirical finding in Figure 3. Its data shows that low variance is correlated with high reward.

Reviewer comment (Writing):
The definition of Var_x(X) seems to have a typo.

Author response:
**Q2: Can you fix the typo in Var$_x(X)$ definition?** Typo fixed, thank you.

Reviewer comment (Theory):
How would you expect the sampling variance to impact the policy gradient if the expectations cannot be computed exactly?

Author response:
**Q3: How would sampling impact the policy gradients if expectations cannot be computed exactly?** We perform sub-sampling as part of the pruning process, beyond depth 3. In ablation studies, we saw that extreme pruning leads to low performance and slow convergence. This is because sub-sampling increases the bias of the gradient, as it is no longer an exact expectation over trajectories. Sub-sampling also affects the variance, but the effect is complex and may vary for different choices of pruned nodes.

Reviewer comment (Theory):
Looking at Lemma 4.1 and Lemma 4.3, it seems the exponential decay of variance is from $\nabla_\theta \log{ \pi_{\theta}(\cdot | s) } $. If $\pi_{\theta}(\cdot | s)$ has softmax parameterization then this basically means the policy is nearly deterministic?

Author response:
**Q1: When the policy is nearly deterministic, how do softmax and SoftTreeMax behave?** In short: Indeed, softmax may become ``stuck” in deterministic policies, but this does not happen in SoftTreeMax. This question highlights the advantage of SoftTreeMax over softmax. 

In more detail, the softmax gradient is indeed almost zero near deterministic policies, as also observed in [Mei et al., 2020a]. In contrast, SoftTreeMax avoids these local optima by integrating the reward into the policy itself. More formally, for softmax, $\nabla_\theta\log \pi_\theta(a|s) = e_{a} - \pi_\theta(\cdot|s),$ where $e_{a} \in [0,1]^A$ is the vector with 0 everywhere except for the $a$-th coordinate. Thus, for a deterministic policy, we get $\theta(s, a) \rightarrow \infty$ for one $a$ per $s$; this expression will be zero for a deterministic policy. In contrast, from Lemmas 4.3 and 4.6 it follows that the SoftTreeMax gradient will not be zero in the respective cumulative or exponential variants.

Reviewer comment (Theory):
If this is true, then this also means the policy gradient has to be close to zero (softmax policy has almost zero gradient near deterministic policies), which is expected to slow down the convergence. Could you explain why reducing both gradient and variance to exponentially small would help learning?

Author response:
**Weakness & Q2: Why would reducing both gradient and variance to exponentially small help learning?** Thank you for a very important question. The key observation is that SoftTreemax does not reduce the gradient uniformly (which would have been equivalent to a change in learning rate). While the gradient norm shrinks, the gradient itself scales differently along the different eigenvectors, as a function of problem parameters ($P$, $\theta$) and our choice of behavior policy ($pi_b$). This allows SoftTreeMax to learn a good “shrinkage” that, while reducing the overall gradient, still updates the policy quickly enough. We will clarify this in the revised paper. 

Figure 3 in the paper presents empirical evidence that the theory holds in our case. It depicts the gradient variance in each (environment, depth),  showing a strong correlation between lower variance and higher reward. In some games, performance is non-monotonic, because all experiments used the same computational budget (same HW for the same amount of limited time). So, in some cases the variance was very small, each training iteration took very long to compute, and fewer iterations were completed by the end of the run.

Reviewer comment (Evaluation):
Although authors mention that formally proving the conjectured global convergence with fast rate as in (Mei et al 2020b) is subject to future work, it is hard to demonstrate its advantage over the traditional SoftMax policy gradient, or more generally traditional policy gradient methods, without comparing the sample complexity between SoftTreeMax policy gradient and traditional SoftMax policy gradient. The main missing piece is that the reduction in variance does not necessarily imply the faster convergence of smaller sample complexity if bringing such variance reduction needs to use a form of policy gradient that sacrifices performance in the deterministic setting (for example, there is uncertainty whether the proposed SoftTreeMax policy gradient will even converge in the deterministic setting).

Author response:
**Q2a: Direct comparison of sample complexity is missing.** For a fair comparison, we estimate sample complexity through wall-clock time; We explain in detail the considerations for this decision in our reply to Q4, R eGu2. 

For wall clock time, in Appendix B.2, Figure 4, we provide the convergence plots as a function of wallclock time. It shows a clear benefit for SoftTreeMax over distributed PPO with the standard softmax policy. In most games, PPO with the SoftTreeMax policy shows very high sample efficiency: it achieves higher episodic reward even though it observes far fewer episodes, for the same running time. In addition, in Appendix B.3, Figure 5, we provide convergence plots of SoftTreeMax where the x-axis is the number of online interactions with the environment, thus excluding the tree expansion complexity. The plot demonstrates a monotone improvement as a function of the tree depth.

**Q2b: Does a reduction in variance necessarily imply faster convergence?** 
There are several theoretical results that support this relation. First, the general fact that variance reduction translates to faster convergence was established in stochastic approximation theory. In the context of PG, several recent works proved improved convergence rates thanks to variance reduction techniques such as SVRG [Papini et al., 2018] (see also [Liu et al., 2020]). 

For SoftTreemax discussed here, note that SoftTreemax does not reduce the gradient uniformly (which would have been equivalent to a change in learning rate). While the gradient norm shrinks, the gradient itself scales differently along the different eigenvectors, as a function of problem parameters ($P$, $\theta$) and our choice of behavior policy ($pi_b$). This allows SoftTreeMax to learn a good “shrinkage” that, while reducing the overall gradient, still updates the policy quickly enough.  

Figure 3 in the paper presents empirical evidence that supports the theory in our case. It depicts the gradient variance in each (environment, depth),  showing a strong correlation between lower variance and higher reward. In some games, performance is non-monotonic, because all experiments used the same computational budget (same HW for the same amount of limited time). So, in some cases, the variance was very small, each training iteration took very long to compute, and fewer iterations were completed by the end of the run.

**Q2c: I am not sure whether the proposed SoftTreeMax policy gradient will even converge in the deterministic setting.** 
SoftTreeMax does converge in a deterministic setting. To test this, we conducted experiments in both the deterministic and the quasi-deterministic variants of the Atari environment. In both cases, we obtained convergence, and the results are almost identical. The quasi-deterministic setting is a popular variation that adds to the deterministic emulator a small stochastic component via “sticky actions” [Machado et al., 2017]. We will discuss this in the revised version.

Reviewer comment (Novelty):
There are some relevant papers that address related problems that authors may need to add to the related work.

Author response:
**Q3a: How does your work relate to “Interpretable Differentiable Decision Trees” from [Silva et al., 2020]?** Although the terminology appears closely related, the work by Silva et al is only loosely connected to our work. There is a delicate distinction to be made here, and one should be aware of the overload of the term “tree”. The differentiable tree-policy that we propose relates to a tree of futures, in the context of planning. The DDT paper examines a decision tree as a function approximator (instead of, e.g., a neural network). The two are orthogonal and can be used in conjunction. We added this reference to the related work section in the revised version.

Reviewer comment (Novelty):
It is important to compare with "Optimization Methods for Interpretable Differentiable Decision Trees in Reinforcement Learning," Andrew Silva, et al., 2020 AISTATS, to evaluate whether this paper is still the first work on proposing a differentiable parametric policy that combines tree expansion with PG.

Author response:
**Q3a: How does your work relate to “Interpretable Differentiable Decision Trees” from [Silva et al., 2020]?** Although the terminology appears closely related, the work by Silva et al is only loosely connected to our work. There is a delicate distinction to be made here, and one should be aware of the overload of the term “tree”. The differentiable tree-policy that we propose relates to a tree of futures, in the context of planning. The DDT paper examines a decision tree as a function approximator (instead of, e.g., a neural network). The two are orthogonal and can be used in conjunction. We added this reference to the related work section in the revised version.

Reviewer comment (Novelty):
"On the Global Optimum Convergence of Momentum-based Policy Gradient," Yuhao Ding, et al., 2022 AISTATS, also studies the convergence and the variance reduction for softmax PG and should be considered in the related work.

Author response:
**Q3b: Consider [Ding et al., 2022] as a reference on convergence and variance reduction for softmax PG.** Thank you, we added it in the revised version.