Policy Gradient with Tree Expansion

Authors: Gal Dalal, Assaf Hallak, Gugan Thoppe, Shie Mannor, Gal Chechik
Primary area: reinforcement learning
Keywords: Reinforcement Learning, Policy Gradient, Tree Expansion, Softmax

Abstract:
Policy gradient methods are notorious for having a large variance and high sample complexity. To mitigate this, we introduce SoftTreeMax---a generalization of softmax that employs planning. In SoftTreeMax, we extend the traditional logits with the multi-step discounted cumulative reward, topped with the logits of future states. We analyze SoftTreeMax and explain how tree expansion helps to reduce its gradient variance. We prove that the variance depends on the chosen tree-expansion policy. Specifically, we show that the closer the induced transitions are to being state-independent, the stronger the variance decay. With approximate forward models, we prove that the resulting gradient bias diminishes with the approximation error while retaining the same variance reduction. Ours is the first result to bound the gradient bias for an approximate model. In a practical implementation of SoftTreeMax we utilize a parallel GPU-based simulator for fast and efficient tree expansion. Using this implementation in Atari, we show that SoftTreeMax reduces the gradient variance by three orders of magnitude. This leads to better sample complexity and improved performance compared to distributed PPO.

Decision: reject
Reviewer scores: 8, 6, 3, 5

Reviews and discussion:
<thread>
creation date:
1729776209132

thread_type:
official_review

<thread_object>
creator: reviewer_qpzq

<content>
summary:
The paper proposes SoftTreeMax, a novel policy gradient method that integrates tree search (TS) and policy gradient (PG) techniques to reduce the variance of policy gradients, a common issue in traditional PG methods. The SoftTreeMax approach uses a differentiable policy that extends standard softmax by incorporating planning, leading to improved sample efficiency and performance.  Theoretical analysis shows that the variance reduction depends on the choice of the tree expansion policy, with deeper tree expansions leading to lower variance. Furthermore, the authors provide bounds on the gradient bias when using approximate models.

strengths:
1. An Interesting and important question.
2. Comprehensive conceptual and theoretical analysis.

weaknesses:
I am much concerned about the empirical evaluations.
1. The wall time is much affected by implementational details. These should be, at the very least, appropriately discussed.
	1. For example, [1] reports very high fps, exceeding what is obtained in the paper.
2. I do not see details of the network used. 
3. It is not clear what the pruning is and how it affects training
4. **important**: It is not clear if the advantage of the proposed method stems from the variance reduction (as claimed) or increased computations per action. E.g., what would happen if the baseline had a bigger network? 
5. The algorithm is a model-based one; thus, having a baseline from this realm is also desirable.

The literature on embedding search as an inductive bias is quite broad and is completely missing in the related work section. Eg. value iteration networks, A3Tree etc.

questions:
1. In many cases (Sec B.3), the SoftTree max policies start with a lower score than PPO; why is that so?
2. Why did the performance on Krull decrease (PPO)?
3. How is your work related to [1]
4. Could you provide an explicit formula for calculating the logits with the full tree expansion?
5. Following footnote 1 on page 8. I am unsure if extending the proposed algorithm to stochastic environments is that easy.

[1] https://arxiv.org/abs/1710.11417
</content>
</thread_object>

</thread>
<thread>
creation date:
1729967771278

thread_type:
official_review

<thread_object>
creator: reviewer_ki8s

<content>
summary:
The paper aims to integrate a novel tree search methodology with policy gradient with the goal of reducing the variance and improving the performance of PG agents.

strengths:
1. The experiments show significant improvements over the chosen baseline.
2. The methodology is supported with theoretical results.
3. The paper implements a parallel version of BFS using GPUs that seems to be scalable.

weaknesses:
1. The novelty of the methodology seems quite limited. Eq. (2) is supposed to be the logit of the policy but I cannot see any difference from the definition of a q-value function, which would entail that the paper is proposing a slight variation of AlphaZero.
2. The authors only compare against PPO. Given that the proposed algorithm is a tree search model-based algorithm, it would be more appropriate to compare against model-based methods, ideally MuZero.

questions:
1. Can you further explain how the proposed method is different from the AlphaZero/MuZero?
2. Why was PPO the only baseline?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730721680404

thread_type:
official_review

<thread_object>
creator: reviewer_tvwf

<content>
summary:
The paper introduces a novel approach that replaces logits by the score of a trajectory starting from that state-action. Authors provide solid theoretical results showing their methods reduces the variance of policy gradient approach and validate it by conducting experiments on Atari games.

strengths:
1. The paper presents a novel method that combines policy gradient techniques with tree search planning, offering a fresh perspective on reducing gradient variance in reinforcement learning.
2. The authors provide solid theoretical analysis, including variance bounds and proofs that show how the tree expansion policy affects variance reduction. They also address the gradient bias introduced by approximate models.
3. The experimental evaluation on Atari games shows significant reductions in gradient variance and improvements in performance, validating the effectiveness of the proposed methods.

weaknesses:
1. The paper primarily considers softmax parameterization for policies, which is restrictive, especially when considering environments with continuous action spaces.
2. The proposed methods rely on the model dynamcis, which need to be infer from the interactions. When employed on offline setting, this may introduce extra bias.
3. Tree expansion, especially with deeper trees, can be computationally intensive. The paper could benefit from a more detailed analysis of the computational costs and how they impact practical deployment.

questions:
1. Is this tree expansion limited to softmax parameterization?
2. Is the method designed only for discrete action space?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730911966032

thread_type:
official_review

<thread_object>
creator: reviewer_u4vp

<content>
summary:
The paper introduces a novel method to reduce variance in Policy Gradient (PG) methods by integrating Tree Search (TS) with a differentiable parametric policy. By incorporating planning into the PG procedure, the authors propose "SoftTreeMax," presented in two variants: C (exp-Expectation) and E (Expectation-exp). The paper offers a theoretical analysis demonstrating how this method reduces variance concerning tree expansion policy and environment parameters. It also analyzes the soundness of using an estimated model for planning, showing that the gradient bias correlates with model approximation error. The practical utility of this approach is demonstrated through experiments on Atari, employing a GPU-based simulator that efficiently utilizes distributed nature for tree expansion.

strengths:
- The work introduces an innovative approach to integrating planning with policy search methods, supported by substantial theoretical analysis that highlights the benefits of this approach.
- The empirical results on Atari, using a GPU-based simulator, are highly promising.

weaknesses:
- The presentation of the two versions of SoftTreeMax (E and C) is unbalanced. While it is beneficial to discuss both, the paper does not adequately represent each. For instance, the Atari experiments focus on C-SoftTreeMax, the empirical demonstration of variance reduction (Figure 1) pertains to the E-version, the theoretical analysis in Section 4.2 is solely for the E version, and the discussion involving the forward model in Section 4.3 is exclusively for C-SoftTreeMax. This uneven representation can make the paper difficult to follow and raises questions about the rationale for using different versions. It also gives the impression that the work is incomplete. Additionally, guidance on which version to use is lacking; although the authors provide a motivation related to risk-aversion/safety for one version, it is not explored in detail.

- In the experiments, exact expectations are considered instead of sample-based methods. This seems limiting, given the appeal of PG methods in sampling-based scenarios. The absence of finite-sample variance components is a significant omission.

questions:
- The behavior/expansion policy $\pi_b$ is crucial in the approach. The paper sets the expansion policy to uniform, but this may not be realistic. Even with pruning, uniform expansion seems suboptimal for exploration. Moreover, how realistic is it to have a behavior policy that induces an irreducible and aperiodic transition matrix?

- Theorem 4.8: Additional insights on why the error scales linearly with $d$ would be valuable.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732182386478

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
To all reviewers

comment:
We thank the reviewers for their dedicated and comprehensive efforts. We are encouraged that the reviewers found our work to "present a novel method that combines policy gradient techniques with tree search planning" [tVWf], supported by "solid theoretical analysis" [tVWf] and "comprehensive conceptual and theoretical analysis" [QPZQ]. The reviewers noted that we address "an interesting and important question" [QPZQ] with "highly promising" empirical results [u4vP]. Specifically, they found that our "methodology is supported with theoretical results" and our "parallel version of BFS using GPUs seems to be scalable" [Ki8S]. Beyond these encouraging feedbacks, the reviewers also made valuable suggestions that we will incorporate in the revised version and comments that we now address.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposed and analyzed a softmax-linear policy parameterization with an interesting feature design that depends on the behavioral policy that generate the data and a tree search approach. Two methods called C-Soft-TreeMax and E-Soft-TreeMax have been proposed. Though the paper has provided interesting preliminary experimental results, the theoretical results are not convincing enough. In particular, as the main motivation of the new PG method is to reduce the variance, the $S^2A^2$ dependence in Theorem 4.4 and the $S$ dependence in 4.8 do not seem to support this motivation. In fact, even the variance bound for naive REINFORCE estimator does not have such dependence. Given the empirical bounds in the experiments, this suggests that the theoretical analysis of variance bound for Soft-TreeMax should be improved. 

Though the paper has a good potential after fixing the theoretical issues, we decide to reject this paper for this year's ICLR.