Policy Gradient with Tree Expansion

Thank you for your detailed and thoughtful review.

Dependency of bound on $S$ and lower bound

RL analysis on tabular MDPs often include $S$ terms, usually stemming from the triangle inequality applied on summation over states. These dependencies can sometimes be replaced by structural assumptions on the transition matrix (like branching). This is a common and important theoretical-practical gap. While the theory suggests meaningless bounds for large $S$ in practice we see exponential variance reduction with $d$ for the huge state space in Atari.

On a related note, a lower bound to the variance is given in the appendix A.5 (it does not include $S$).

Bound on E-SoftTreeMax does not contain $S$

There is a dependence on $S,A$ also in the case of E-SoftTreeMax. However, notice that the bound in Thm 4.7 is presented in O-notation while the C-SoftTreeMax bound is expressed explicitly. The critical theoretical contribution in both cases is the identification of the exponential decay rate with respect to tree depth.

Regarding the missing numerical variance measurements for C-SoftTreeMax: we conducted similar variance analyses for both SoftTreeMax variants, but included only the E-SoftTreeMax results in the paper to avoid redundancy. The numerical results for C-SoftTreeMax show similar behavior and correspondence to the theory. We will clearly state this in the revision.

The primary reason we included the detailed variance measurements for E-SoftTreeMax was to verify our theoretical hypothesis from Thm 4.7 that the parameter $\alpha$ is related to the $\lambda_2$. Our numerical experiments confirmed this in the general case by showing that variance decays at a rate determined by the mixing properties of the baseline policy.

We will add a brief discussion of C-SoftTreeMax's empirical variance characteristics in the revised manuscript to ensure completeness.

Comparison to MCTS

Following your comment, we added comparisons with a strong baseline to our experiments: EfficientZero [Ye et al., 2021]. We chose it because it is a highly sample-efficient version of MuZero that is also open source. MuZero is one of the best known RL algorithms to date. The results are given here: https://ibb.co/zHG3KP8m, showing that SoftTreeMax surpasses EfficientZero in all the games we tested except one. We will add full experiments in the final version.

Question 1: How can one handle a stochastic environment with an infinite number of states?

For stochastic environments with infinite state spaces, we propose three complementary approaches:

1. We can use Monte Carlo sampling to expand the tree. This approach maintains a parametric distribution over actions (potentially dependent on $\theta$) and samples from it to build the tree. This method can be viewed as a tree adaptation of Model Predictive Path Integral control (MPPI) with a value function.

2. Function approximation for leaf states: For infinite state spaces, we already use neural networks to approximate the value function at leaf states. This approach naturally extends to continuous state spaces.

3. Theoretical extension: The key concepts in our analysis can be adapted to infinite state spaces by:

   3a. Replacing transition matrices with kernels in the continuous case.

   3b. Exploiting that $\pi_b$'s transition kernel is a non-expansive operator.

   3c. Leveraging that the eigenvalue 1 gets canceled for the policy gradient.

These properties can be shown to hold for decision models with infinite state and action spaces, though we leave the complete theoretical development for future work.

Question 2: Does it make sense to change the behavioral policy over time? Specifically, does it make sense to use the current policy as the tree-expansion policy?

Yes, adapting the behavioral policy over time is a promising direction. As training progresses and the policy improves, using it as the tree expansion policy could lead to exploring more relevant parts of the state space. This creates a bootstrapping effect similar to what's done in MCTS algorithms like AlphaZero, where the current policy guides the search.

Still, care must be taken as trained policies often become deterministic leading to unfavorable variance properties for semi-deterministic environments (when $P^{\pi_b}$ approaches a permutation matrix, $\lambda_2(P^{\pi_b})$ could potentially approach 1, eliminating the variance reduction benefit). Also we can’t use the SoftTreeMax policy directly to open the tree as it requires expanding a tree while expanding a tree.

A reasonable compromise might be to learn a proxy policy as a mixture of the current policy and a uniform policy, with the mixing weight potentially decreasing over time to balance exploration and exploitation while maintaining favorable variance properties.

Thank you for raising this excellent point, which provides fertile ground for future work. We will address this in the revised paper.

Thank you for your detailed and thoughtful review. We appreciate the careful reading of our work and the constructive feedback. Below, we address your concerns:

Additional MCTS baselines

Following your comment, we added comparisons with a strong baseline to our experiments: EfficientZero [Ye et al., 2021]. We chose it because it is a highly sample-efficient version of MuZero that is also open source. MuZero is one of the best known RL algorithms to date. The results that were completed during the rebuttal are given here: https://ibb.co/zHG3KP8m, showing that SoftTreeMax surpasses EfficientZero in all the games we tested except one. We will add full experiments in the final version.

Question 1: Does the number of online interactions include the samples generated during the tree search?

No, the number of online interactions does not include the samples generated during tree search. We distinguish between two types of interactions:

1. Online environment interactions: These are actual interactions with the environment during training (e.g., when collecting trajectories for PPO updates).

2. Planning interactions (tree search): These are simulated forward passes using our GPU-based simulator to expand the tree.

As we explained in our experiments, the computational complexity of planning interactions is substantially lower than online interactions, especially with our GPU batching mechanism for tree expansion. This is why we track them separately, and our sample complexity plots show only the online interactions.

For a complete view of computational efficiency, we also provide wallclock time comparisons in Figure 4, which accounts for the total computation including tree expansion. As seen there, despite the additional computation for tree search, SoftTreeMax still demonstrates favorable performance per unit time for appropriate choices of depth.

Question 2: Why can increasing d to 1 or 2 harm performance in some experiments?

This is an excellent observation. While theoretically more information should help, this illustrates the basic bias-variance tradeoff in the SoftTreeMax policy:

At small depths (d=1,2), we introduce a structural bias by incorporating tree search, but the variance reduction benefit hasn't fully kicked in yet. The theoretical variance reduction is exponential in d, so the effect becomes more pronounced at larger depths.

We will address this valuable insight in the revised paper.

Thank you for your thoughtful review. We appreciate the careful reading of our proofs and the constructive feedback. We address your concerns below.

Proof issues in Lemma A.5 and Theorem A.6

You're correct that $\gamma^{-d}$ from $C_{s,d}$ should appear in this bound. This oversight makes the proper bound $O(\beta d \gamma^{-d} \epsilon)$. When the approximation error $\epsilon$ is small enough and $\beta$ properly controlled, this term still approaches 0 as required for the subsequent proof. We will correct this in the revised version.

Variance bounds discrepancy

E-SoftTreeMax also depends on $S,A$, but Theorem 4.7 uses O-notation while C-SoftTreeMax shows explicit constants. Both establish the crucial exponential decay rate with tree depth.This exponential decay explains why deeper trees provide more stable policy gradient estimates and better overall performance, regardless of specific constants or dependencies. The key contribution is establishing this relationship between tree depth and variance reduction through the eigenstructure of the MDP's transition dynamics.

We'll highlight this insight more clearly.

Assumption that r(s,a) equals r(s)

This simplification was made for theoretical analysis of E-SoftTreeMax. Our practical implementation handles general reward structures correctly in all experiments.

Definition of $\theta$ and $\Theta$

Thank you for highlighting this inconsistency. In Section 2.1, we use $\theta$ conventionally for policy parameters, but later use $\Theta \in \mathbb{R}^S$ and $\theta(s)$ as a scalar state function. We'll replace the notation in Section 2.1 with a different symbol, reserving $\theta$ and $\Theta$ exclusively for our SoftTreeMax formulation to ensure consistency throughout the paper.

Verification for C-SoftTreeMax

We analyzed both variants but included only E-SoftTreeMax results (Figure 1) to avoid redundancy. C-SoftTreeMax shows similar behavior and correspondence to theory. We focused on E-SoftTreeMax to verify that parameter $\alpha$ relates to the transition matrix's second eigenvalue in the general case, which our experiments confirmed. We'll add a brief discussion of C-SoftTreeMax's empirical variance characteristics for completeness.

Behavior policy choice

Our theory identifies properties of ideal behavior policies but doesn't provide concrete procedures for complex domains. In Atari experiments, we used uniform policies for simplicity. Preliminary experiments with other policies showed similar trends but weren't comprehensive enough to include.

We'll clarify our claims and add discussion about the practical impact of behavior policy choice.

Minor comments and title suggestion

We appreciate your writing suggestions and will consider revising the title to better reflect our focus.

Question 1 regarding loose variance bounds

This is indeed an artifact of our analysis approach. Our bound becomes loose in this regime because we use worst-case inequalities at multiple steps in the proof that don't capture the special structure that emerges in deterministic settings. The $\beta^2$ term appears because we bound the gradient norm without accounting for how the policy structure changes as $\beta$ increases. For practical parameter ranges used in training, our bound provides useful insights about variance reduction with tree depth. In the revised paper, we'll clarify this limitation and discuss how the bound could be refined to better handle the deterministic limit case.

We'll clarify this limitation and discuss potential refinements.

Question 2 on the potential issue with the proof

Please see our response above regarding the missing $\gamma^{-d}$ term.

Question 3 on variance bounds intuition

C-SoftTreeMax has a more decomposable gradient structure (Lemma 4.3): $\nabla_\theta\log \pi_{d,\theta}^C = \beta\left[I_{A} - 1_A (\pi_{d,\theta}^C)^\top \right]P_s \left(P^{\pi_b}\right)^{d-1}.$ This allows direct spectral decomposition. When the projection matrix interacts with the decomposed $(P^{\pi_b})^{d-1}$, the stationary distribution term gets canceled, leaving terms scaled by $\lambda_i^{d-1}$, with $|\lambda_2|$ dominating.

E-SoftTreeMax's analysis is more complex as its gradient involves matrix products and ratios with exponentiated rewards, requiring sophisticated techniques to establish the convergence rate $\alpha$.

This explains why we provide an exact variance bound for C-SoftTreeMax, while for E-SoftTreeMax we give an asymptotic bound where $\alpha = |\lambda_2(P^{\pi_b})|$ only under certain conditions.

Thank you for your thoughtful review and for recognizing the value of our theoretical and empirical contributions. We appreciate your feedback and address your questions below.

MCTS vs. our search approach

Thank you for this important point of clarification. SoftTreeMax and MCTS represent fundamentally different approaches with distinct goals. While both involve planning, our work specifically focuses on how to directly integrate forward modeling into the policy gradient framework itself. The key innovation of SoftTreeMax is that it incorporates planning directly into policy parameterization, creating a differentiable policy that naturally reduces gradient variance. Unlike MCTS, which operates as a separate planning algorithm that produces improved rollouts, our approach transforms the underlying architecture of policy gradient methods. This enables us to maintain the theoretical elegance of policy gradient approaches while addressing their fundamental variance issues. Our analytical results show how this integration affects gradient properties in ways that would not be directly applicable to MCTS-based methods. We believe this approach opens new avenues for research at the intersection of planning and policy-based methods.

That said, we acknowledge the value of comparing against and potentially combining with MCTS-based approaches in future work.

Additional MCTS baselines

Following your comment, we added comparisons with a strong baseline to our experiments: EfficientZero [Ye et al., 2021]. We chose it because it is a highly sample-efficient version of MuZero that is also open source. MuZero is one of the best known RL algorithms to date. The results that were completed during the rebuttal are given here: https://ibb.co/zHG3KP8m, showing that SoftTreeMax surpasses EfficientZero in all the games we tested except one. We will add full experiments in the final version.

Variance scaling with eigenvalue assumptions

You've identified an important theoretical-practical gap. While our theory depends on properties of transition matrices that may not always hold in practice, our empirical results consistently show variance reduction across various environments. In cases where the eigenvalue assumptions are violated, we would expect diminished (but still positive) variance reduction benefits rather than variance increases. Figure 3 in our paper demonstrates this empirically - the variance consistently decreases with depth across different games, even though the theoretical assumptions likely vary between environments.

Bias-variance tradeoff with approximate models

The bias-variance tradeoff we observed follows this pattern:

• The closer $\pi_b$ is to $\pi_{d,t}$ at time t, the lower the bias

• The closer $\pi_b$ is to having similar rows in $P^{\pi_b}$, the lower the variance

For practical solutions to this tradeoff, we recommend:

1. Adaptive depth selection based on model certainty
2. Progressive increase in depth as model accuracy improves during training
3. Ensemble methods to better estimate model uncertainty

Our Theorem 4.8 shows that the gradient bias diminishes with approximation error while retaining similar variance reduction properties. Even with approximation errors, the variance continues to decay exponentially, just at a rate dictated by $\hat{P}^{\pi_b}$ instead of $P^{\pi_b}$.

We will expand this discussion in the revised paper to provide clearer guidance on managing this tradeoff in practical implementations.