Policy Gradient with Tree Expansion

Thank you for your detailed review. Let us address your concerns: **Weakness 1: Implementation and Wall-time Details** 1. Network Architecture: We follow the architecture from stable-baselines3 PPO implementation which has been extensively validated by the community. For more details regarding the hardware, software and implementation used see Appendix B.1. 2. Tree Pruning: Beyond depth 3, we limit the tree to a maximum width of 1024 nodes and prune trajectories with low estimated weights. The pruning process reduces computational complexity while still maintaining performance benefits. To better illustrate this, we created a new diagram showing the tree expansion implementation in the GPU: https://ibb.co/5RPRkzF. 3. Comparison to [1]: While TreeQN achieves high fps, it is fundamentally different as it approximates the value, without planning, in the network architecture. In contrast, SoftTreeMax performs tree expansion using the simulator to parametrize the policy. Applying multiple simulator steps, even in the GPU, in every tree expansion is much slower than a forward pass in a single neural network. This likely accounts for the main difference in fps, in addition to other more obvious differences in the setup (hardware, implementation etc.). **Weakness 2: Source of Improvement - Variance Reduction vs Computation** This is an important point that deserves clarification. While we cannot definitively attribute all improvements to variance reduction, several observations suggest it plays a significant role: 1. Figure 3 demonstrates a strong correlation between lower variance and higher reward across different depths, though causation remains to be proven. 2. When controlling for computation time (Figure 4 in Appendix B.2), SoftTreeMax still outperforms PPO with equivalent computation budget - PPO uses 256 GPU workers while SoftTreeMax uses equivalent compute for tree expansion. 3. To isolate the effect of increased computation, we ran experiments with larger networks for PPO (2x and 4x parameter count) and found no significant improvement, suggesting the gains are not simply from increased model capacity. These results suggest that variance reduction may be a key factor in the improved performance, though further investigation would be needed to establish a causal relationship. **Weakness 3: Model-based Baseline & Related Work** Following this feedback, we've added two improvements: 1. Added EfficientZero [Ye et al., 2021] as a model-based baseline. Initial results (due to the time constraints) here: https://ibb.co/fNzG5cR show SoftTreeMax outperforming it in three out of four games. We will add full experiments in the final version. 2. We will expand the related work section to include: - Value Iteration Networks [Tamar et al., 2016] - TreeQN [Farquhar et al., 2018] Could you also please add a reference to the A3Tree algorithm? We could not find it in an online search. **Q1: Why does SoftTreeMax start with lower scores?** Different games benefit differently from long-term vs short-term planning. Games requiring longer-term planning need better value estimates, which take longer to obtain with larger depths. When applying the network to states far from those in the buffer, if the model hasn't learned a good value function and there's no guiding dense reward, the policy can be initially noisier. **Q2: Performance decrease on Krull** The apparent decrease is due to our fixed one-week training time limit. Deeper trees mean fewer steps completed in the same time period. When comparing the same number of steps (Figure 5), performance improves monotonically with depth. **Q3: Relationship to [1]** While both works use tree-structured computation, TreeQN approximates the value, without planning, in the network architecture. In contrast, SoftTreeMax performs tree expansion using the simulator to parametrize the policy. Our approach directly reduces policy gradient variance through tree expansion rather than learning the model through a differentiable tree. **Q4: Formula for logits with full tree expansion** For C-SoftTreeMax, the logits are computed as: $\log \pi_{d,\theta}(a|s) = \log(\sum_{s_1,...,s_d} P(s_1|s,a)...P(s_d|s_{d-1},a_{d-1})[\gamma^d W_\theta(s_d) + \sum_{t=0}^{d-1} \gamma^t r(s_t,a_t)])$. We will add this explicit formula to the paper. **Q5: Extension to stochastic environments** You raise a valid concern. While the theory extends naturally, a practical implementation in highly stochastic environments would require careful consideration of how to handle the branching factor growth. We will clarify this limitation in the paper. We thank you for helping us improve the paper and will incorporate these additions in the revised version.

Thank you for your detailed review. Let us address your main concerns: **Weakness 1 & Q1: Limited novelty - Is SoftTreeMax just a Q-value function variation of AlphaZero/MuZero?** While Eq. (2) may appear similar to a Q-value function, SoftTreeMax differs fundamentally from AlphaZero/MuZero: 1. Novel Policy Parametrization: *SoftTreeMax is a new type of policy* that directly incorporates planning into policy parametrization. While the logits in Eq. (2) share some similarities with Q-values, they serve a different purpose - they define a policy rather than estimate state-action values. They also don’t explicitly minimize any Bellman-type loss function. This is a key distinction that enables variance reduction in policy gradient methods. 2. Architecture vs Algorithm: *AlphaZero/MuZero are value-based algorithms* that use tree search for action selection. In contrast, SoftTreeMax introduces a new policy architecture that can be integrated into any policy gradient method. This enables combining the benefits of tree expansion with policy-based approaches - a direction previously unexplored in RL. **Weakness 2 & Q2: Limited comparison - Only compared against PPO rather than model-based methods** Following this important feedback, we now add comparisons with EfficientZero [Ye et al., 2021], a highly sample-efficient version of MuZero. We chose it because: - It is a state-of-the-art model-based algorithm - It represents one of the best known RL algorithms today - It is open source, enabling fair comparison The results that were completed during the rebuttal are given here: https://ibb.co/fNzG5cR, showing that SoftTreeMax surpasses EfficientZero in three out of the four games we tested. 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 you for helping us improve the paper. We will incorporate these additions in the revised version.

You raise an important point about model access. Indeed, EfficientZero must learn its model while we use the true dynamics. However: 1. Our comparison aims to demonstrate the benefits of integrating tree expansion with PG when a model is available, not to claim superiority over model-learning approaches. 2. Our Theorem 4.8 actually analyzes how SoftTreeMax performs with approximate models, showing the gradient bias diminishes with approximation error. This suggests SoftTreeMax could work well with learned models too, though proving this empirically is left for future work. 3. As noted in our limitations section, our use of Atari-CuLE (a GPU-based simulator) is "both a novelty and a current limitation." Our results demonstrate what's achievable with a good model ("upper bound" on performance), while empirical validation with learned models remains future work. We will clarify this better in the appropriate locations of the revised version.

Thank you for your thoughtful review. Let us address your concerns: **Weakness 1 & Q2: Limitation to discrete action spaces and softmax parameterization** The method can be extended to continuous action spaces, though it requires some modifications: 1. For continuous control, one can maintain a parametric distribution over actions which depends on $\theta$. This method can be seen as a tree adaptation of MPPI [Williams et al., 2017] with a value function. 2. On the theory side, our key concepts can extend to continuous spaces: - The transition kernel in the continuous case remains a non-expansive operator. - The eigenvalue cancellation in policy gradient still holds. - These properties can be shown to hold for decision models with infinite actions. While this extension requires significant research effort, the finite action space setup studied here is an important one that has been extensively studied in RL literature. **Weakness 2: Model dynamics and offline setting** Indeed, in this work we used accurate dynamics models. There are two possible approaches for settings without exact models: 1. Use GPU-based simulators, which are becoming increasingly common [Dalton et al., 2020; Makoviychuk et al., 2021; Freeman et al., 2021]. 2. Learn the forward model, as successfully demonstrated in recent works like MuZero [Schrittwieser et al., 2020] and Dreamer [Hafner et al., 2020]. Our analysis in Theorem 4.8 provides bounds on the gradient bias when using approximate models, showing the bias diminishes with approximation error. We will expand the discussion of these alternatives in the revised version. **Weakness 3: Computational costs of tree expansion** You raise an important point about computational complexity. We handle this in several ways: 1. For depths ≥3, we limit the tree width to 1024 nodes and prune trajectories with low estimated weights. 2. Our GPU implementation allows efficient parallel expansion - advancing 1,024 environments simultaneously takes substantially less compute time than advancing a single environment 1,024·A steps sequentially. 3. We provide detailed runtime analysis in Appendix B.2 (Figure 4), showing the tradeoff between: - Number of policy updates (reduced with depth) - Complexity per iteration (increased with depth) Each game has a different optimal depth balancing this tradeoff. **Q1: Tree expansion beyond softmax parameterization** The tree expansion technique could theoretically be adapted to other policy parameterizations. One such alternative is to replace the exponentiated ratio with a direct normalization (Escort transform: Mei et al 2020a, Eq. 3). Nonetheless, we (as most other researchers in the field) chose softmax for its attractive characteristics: 1. Its analytical form allows clear theoretical analysis of variance reduction 2. It naturally handles the normalization of trajectory scores 3. It provides a clear connection to temperature-based exploration We will discuss these aspects and the potential for other parameterizations in the revised version. **Q2: discrete action space** Please see response to Weakness 1.

Thank you for your thorough and supportive review. Let us address your concerns: **Weakness 1: Unbalanced presentation of E-SoftTreeMax and C-SoftTreeMax** We acknowledge this presentation imbalance. To clarify: 1. The results in Section 4.2 for E-SoftTreeMax mirror the results in Section 4.1 for C-SoftTreeMax – each includes the vector representation for the operator, its gradient, and the resulting variance bound. The additional bias result indeed only holds for C-SoftTreeMax, and is more intricate to obtain for E-SoftTreeMax; thus we leave it for future work. 2. We actually repeated all experiments with E-SoftTreeMax, following the same protocols as C-SoftTreeMax (see footnote 1 in page 8). The results were almost identical, which is why we left them out. The reason for this similarity is that in quasi-deterministic Atari environments, the trajectory logits (Eq. 2) have almost no variability, making the two variants (Eqs. 3 and 4) nearly identical. 3. The reason we chose E-SoftTreeMax for the experiment in Figure 1 is to demonstrate that the $\alpha$ parameter from Theorem 4.7 is indeed $\lambda_2$ as we conjecture. For C-SoftTreeMax we managed to prove this for any MDP, while for E-SoftTreeMax we could only show it exists and equals $\lambda_2$ in certain cases. Repeating the same experiment for C-SoftTreeMax gives the same results. We will clarify this better in the revised version. 4. Choice between variants: - C-SoftTreeMax corresponds to Boltzmann exploration when the NN approximates the Q-function. - E-SoftTreeMax is more suitable for risk-averse objectives since it exponentiates rewards, giving high-cost trajectories more influence. Overall, we included both C-SoftTreeMax and E-SoftTreeMax as natural extensions of softmax to the multi-step setting. Their theoretical properties, while derived differently, lead to similar variance reduction guarantees. **Weakness 2: Lack of finite-sample variance components** We may have misunderstood your concern about finite-sample variance components. If you're referring to sampling-based implementation, we do employ pruning beyond depth 3, effectively implementing a form of sampling. If your concern is about a different aspect of finite-sample analysis, we would appreciate clarification so we can address it properly. In our current analysis, we focus on the inherent variance of the gradient itself rather than estimation variance. This distinction is important because: 1. Our technique reduces the "direct" variance of the gradient as defined in Section 2.1 2. It can be applied orthogonally to traditional variance reduction techniques that tackle estimation variance (e.g., baseline subtraction) We will clarify this distinction and its implications in the revised version. **Q1: Choice and realism of behavior policy** The reviewer raises excellent points about the behavior policy. Our choice of uniform policy is motivated by theoretical considerations: 1. As our theory (supported by Figure 1) shows, optimal variance decay is achieved when the induced transition matrix is close to being rank-one (i.e., its rows are similar). Without further assumptions on the MDP, a uniform policy that smoothens transition probabilities is a good proxy. 2. Regarding irreducibility and aperiodicity: These are standard assumptions in the analysis of Markov chains and reinforcement learning [Puterman, 1994; Bertsekas and Tsitsiklis, 1996]. They ensure the existence of a unique stationary distribution and are commonly used [Bhandari et al., 2018; Khodadadian et al., 2022]. Note that it is sufficient for even one deterministic policy to exist for which the resulting Markov chain is irreducible and aperiodic, for the uniform policy to be also irreducible and aperiodic. 3. We agree that uniform expansion may not be optimal for exploration. This opens an interesting direction for future work on adaptive expansion policies. **Q2: Linear scaling with d in Theorem 4.8** The linear dependence on d appears because of error accumulation along the trajectories. More specifically: 1. When using an approximate model, errors compound at each step of tree expansion. 2. Each level of the tree introduces a new source of error through both transition and reward approximations. 3. These errors accumulate additively due to the structure of the gradient computation, leading to the O(d) term. However, this linear scaling with depth is balanced by the exponential decay term $\gamma^d$ in the same theorem, suggesting that moderate depths can still be effective even with approximate models. We will add this discussion to clarify the interplay between approximation error and tree depth.