SoftTreeMax: Exponential Variance Reduction in Policy Gradient via Tree Expansion

Thank you for your response. Please find our answers in the following.

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.

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.

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.

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.

Q4: Is there a better way to compare sample efficiency instead of run-time? To answer the question, we first define “sample efficiency” more carefully. We distinguish “online environment interactions” from “planning interactions” (i.e. tree search predictions). If we compare only online interactions, we will overestimate the sample efficiency of SoftTreeMax because it will hide the extra planning iterations during policy inference. However, if we also naively add in the planning interactions, we will underestimate the sample efficiency of SoftTreeMax. The reason being that the computational complexity of planning is lower than online interaction, thanks to the GPU batching mechanism that our tree search takes advantage of.

For a fair comparison, we chose the metric of the overall wallclock time that accounts for both outer and inner iteration complexity. Appendix B.2, Figure 4, shows a sweet spot that balances online vs. planning iterations. This occurs at different depths, depending on the Atari game, with optimal data efficiency usually obtained at $d>0$. Appendix B.3, Figure 5, complements Figure 4 and highlights the need for fewer outer (online) interactions as the tree deepens. For more details, see the discussions in Sections B.2 and B.3.

Thank you for your response. Please find our answers in the following.

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.

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

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.

Thank you for your response. Please find our answers in the following.

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.

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.

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

Q1a : Do the claimed exponential variance reduction still hold with the approximate model? Thank you for raising this important question. Yes, the variance decays exponentially at a rate dictated by $\hat{P}^{\pi_b}$ instead of $P^{\pi_b}.$ As the two terms get closer when the approximation error decays, the rate naturally also becomes closer due to eigenvalue continuity. To study the consequence of using an approximate model, we analyzed the bias in Theorem 4.8.

To formally see why the variance decays exponentially in the approximate case, recall the three key steps in our proof of Theorem 4.4:

1. Eigendecomposition of $P^{\pi_b}$ with the eigenvalue 1 not repeating and noting that the leading eigenvector is $1_A$ (the all ones vector of size $A$).

2. Noting that $P_s 1_S = 1_A.$

3. Noting that $[I_A - 1_A \pi^C_{d, \theta}] 1_A = 0.$

The first two steps also hold for the approximate model since these steps only rely on row-stochasticity (leaving aside the pathological case where the eigenvalue $1$ repeats in $\hat{P}^{\pi_b}$, which will not occur when the latter is close to $P^{\pi_b}.$). Finally, the third step is valid as long as $\hat{\pi}_{d, \theta}^C$ is a distribution, which is trivially true.

We will add this new discussion to the revised version.

Q1b: Is it useful to make PG, a model-free method, to a model-based method by combining it with tree search? We believe that PG as a model-based method has great potential to improve upon its model-free version. To demonstrate this beyond what was shown in the submitted version regarding PPO, 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 model-based algorithms today -- and is also open source. 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.

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.