Policy Gradient with Tree Expansion

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 20:29:15

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 decays exponentially with the planning horizon as a function of the chosen tree-expansion policy. Specifically, we show that the closer the induced transitions are to being state-independent, the stronger the 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.

Main paper: 
 introduction: 
 Policy Gradient (PG) methods [Sutton et al., 1999] for Reinforcement Learning (RL) are often the first choice for environments that allow numerous interactions at a fast pace [Schulman et al., 2017] .Their success is attributed to several factors: they are easy to distribute to multiple workers, require no assumptions on the underlying value function, and have both on-policy and off-policy variants.Despite these positive features, PG algorithms are also notoriously unstable due to the high variance of the gradients computed over entire trajectories [Liu et al., 2020 , Xu et al., 2020] . As a result, PG algorithms tend to be highly inefficient in terms of sample complexity. Several solutions have been proposed to mitigate the high variance issue, including baseline subtraction [Greensmith et al., 2004 , Thomas and Brunskill, 2017 , Wu et al., 2018] , anchor-point averaging [Papini et al., 2018] , and other variance reduction techniques [Zhang et al., 2021 , Shen et al., 2019 , Pham et al., 2020] .A second family of algorithms that achieved state-of-the-art results in several domains is based on planning. Planning is exercised primarily in the context of value-based RL and is usually implemented using a Tree Search (TS) [Silver et al., 2016 , Schrittwieser et al., 2020] . In this work, we combine PG with TS by introducing a parameterized differentiable policy that incorporates tree expansion.Namely, our SoftTreeMax policy replaces the standard policy logits of a state and action, with the expected value of trajectories that originate from these state and action. We consider two variants of SoftTreeMax, one for cumulative reward and one for exponentiated reward.Combining TS and PG should be done with care given the biggest downside of PG-its high gradient variance. This raises questions that were ignored until this work: (i) How to design a PG method based on tree-expansion that is stable and performs well in practice? and (ii) How does the tree-expansion Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute. policy affect the PG variance? Here, we analyze SoftTreeMax, and provide a practical methodology to choose the expansion policy to minimize the resulting variance. Our main result shows that a desirable expansion policy is one, under which the induced transition probabilities are similar for each starting state. More generally, we show that the gradient variance of SoftTreeMax decays at a rate of |λ 2 | d , where d is the depth of the tree and λ 2 is the second eigenvalue of the transition matrix induced by the tree expansion policy. This work is the first to prove such a relation between PG variance and tree expansion policy. In addition, we prove that the with an approximate forward model, the bias of the gradient is bounded proportionally to the approximation error of the model.To verify our results, we implemented a practical version of SoftTreeMax that exhaustively searches the entire tree and applies a neural network on its leaves. We test our algorithm on a parallelized Atari GPU simulator [Dalton et al., 2020] . To enable a tractable deep search, up to depth eight, we also introduce a pruning technique that limits the width of the tree. We do so by sampling only the most promising nodes at each level. We integrate our SoftTreeMax GPU implementation into the popular PPO [Schulman et al., 2017] and compare it to the flat distributed variant of PPO. This allows us to demonstrate the potential benefit of utilizing learned models while isolating the fundamental properties of TS without added noise. In all tested Atari games, our results outperform the baseline and obtain up to 5x more reward. We further show in Section 6 that the associated gradient variance is smaller by three orders of magnitude in all games, demonstrating the relation between low gradient variance and high reward.We summarize our key contributions. (i) We show how to combine two families of SoTA approaches: PG and TS by introducing SoftTreeMax: a novel parametric policy that generalizes softmax to planning. Specifically, we propose two variants based on cumulative and exponentiated rewards. (ii) We prove that the gradient variance of SoftTreeMax in its two variants decays exponentially with its tree depth. Our analysis sheds new light on the choice of tree expansion policy. It raises the question of optimality in terms of variance versus the traditional regret; e.g., in UCT [Kocsis and Szepesvári, 2006] . (iii) We prove that with an approximate forward model, the gradient bias is proportional to the approximation error, while retaining the variance decay. This quantifies the accuracy required from a learned forward model. (iv) We implement a differentiable deep version of SoftTreeMax that employs a parallelized GPU tree expansion. We demonstrate how its gradient variance is reduced by three orders of magnitude over PPO while obtaining up to 5x reward.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes]Justification: [NA] Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
preliminaries: 
 Let ∆ U denote simplex over the set U. Throughout, we consider a discounted Markov Decision Process (MDP) M = (S, A, P, r, γ, ν), where S is a finite state space of size S, A is a finite action space of size A, r : S × A → [0, 1] is the reward function, P : S × A → ∆ S is the transition function, γ ∈ (0, 1) is the discount factor, and ν ∈ R S is the initial state distribution. We denote the transition matrix starting from state s by P s ∈ [0, 1] A×S , i.e., [P s ] a,s ′ = P (s ′ |a, s). Similarly, let R s = r(s, •) ∈ R A denote the corresponding reward vector. Separately, let π : S → ∆ A be a stationary policy. Let P π and R π be the induced transition matrix and reward function, respectively, i.e., P π (s ′ |s) = a π(a|s) Pr(s ′ |s, a) and R π (s) = a π(a|s)r(s, a). Denote the stationary distribution of P π by µ π ∈ R S s.t. µ ⊤ π P π = P π , and the discounted state visitation frequency by d π so thatd ⊤ π = (1 -γ) ∞ t=0 γ t ν ⊤ (P π ) t . Also, let V π ∈ R S be the value function of π defined by V π (s) = E π [ ∞ t=0 γ t r (s t , π(s t )) | s 0 = s], and let Q π ∈ R S×A be the Q-function such that Q π (s, a) = E π [r(s, a) + γV π (s ′ )].Our goal is to find an optimal policy π ⋆ such thatV ⋆ (s) ≡ V π ⋆ (s) = max π V π (s), ∀s ∈ S.For the analysis in Section 4, we introduce the following notation. Denote by Θ ∈ R S the vector representation of θ(s) ∀s ∈ S. For a vector u, denote by exp(u) the coordinate-wise exponent of u and by D(u) the diagonal square matrix with u in its diagonal. For a matrix A, denote its i-th eigenvalue by λ i (A). Denote the k-dimensional identity matrix and all-ones vector by I k and 1 k , respectively. Also, denote the trace operator by Tr . Finally, we treat all vectors as column vectors.PG schemes seek to maximize the cumulative reward as a function of the policy π θ (a|s) by performing gradient steps on θ. The celebrated Policy Gradient Theorem [Sutton et al., 1999] states that∂ ∂θ ν ⊤ V π θ = E s∼dπ θ ,a∼π θ (•|s) [∇ θ log π θ (a|s)Q π θ (s, a)] .The variance of the gradient is thusEQUATIONIn the notation above, we denote the variance of a vector random variable X byVar x (X) = Tr E x (X -E x X) ⊤ (X -E x X) ,similarly as in [Greensmith et al., 2004] . From now on, we drop the subscript from Var in (1)for brevity. When the action space is discrete, a commonly used parameterized policy is softmax:π θ (a|s) ∝ exp (θ(s, a)), where θ : S × A → R is a state-action parameterization.3 SoftTreeMax: Exponent of trajectoriesWe introduce a new family of policies called SoftTreeMax, which are a model-based generalization of the popular softmax. We propose two variants: Cumulative (C-SoftTreeMax) and Exponentiated (E-SoftTreeMax). In both variants, we replace the generic softmax logits θ(s, a) with the score of a trajectory of horizon d starting from (s, a), generated by applying a behavior policy π b . In C-SoftTreeMax, we exponentiate the expectation of the logits. In E-SoftTreeMax, we first exponentiate the logits and then only compute their expectation.Logits. We define the SoftTreeMax logit ℓ s,a (d; θ) to be the random variable depicting the score of a trajectory of horizon d starting from (s, a) and following the policy π b :EQUATIONIn the above expression, note that s 0 = s, a 0 = a, a t ∼ π b (•|s t ) ∀t ≥ 1, and r t ≡ r (s t , a t ) .For brevity of the analysis, we let the parametric score θ in (2) be state-based, similarly to a value function. Instead, one could use a state-action input analogous to a Q-function. Thus, SoftTreeMax can be integrated into the two types of implementation of RL algorithms in standard packages. Lastly, the preceding γ -d scales the θ parametrization to correspond to its softmax counerpart.C-SoftTreeMax. Given an inverse temperature parameter β, we let C-SoftTreeMax beEQUATIONC-SoftTreeMax gives higher weight to actions that result in higher expected returns. While standard softmax relies entirely on parametrization θ, C-SoftTreeMax also interpolates a Monte-Carlo portion of the reward.The second operator we propose is E-SoftTreeMax:EQUATIONhere, the expectation is taken outside the exponent. This objective corresponds to the exponentiated reward objective which is often used for risk-sensitive RL [Howard and Matheson, 1972 , Fei et al., 2021 , Noorani and Baras, 2021] . The common risk-sensitive objective is of the form log E[exp(δR)],where δ is the risk parameter and R is the cumulative reward. Similarly to that literature, the exponent in (4) emphasizes the most promising trajectories.SoftTreeMax is a natural model-based generalization of softmax. For d = 0, both variants above coincide since (2) becomes deterministic. In that case, for a state-action parametrization, they reduce to standard softmax. When β → 0, both variants again coincide and sample actions uniformly (exploration). When β → ∞, the policies become deterministic and greedily optimize for the best trajectory (exploitation). For C-SoftTreeMax, the best trajectory is defined in expectation, while for E-SoftTreeMax it is defined in terms of the best sample path.SoftTreeMax convergence. Under regularity conditions, for any parametric policy, PG converges to local optima [Bhatnagar et al., 2009] , and thus also SoftTreeMax. For softmax PG, asymptotic [Agarwal et al., 2021] and rate results [Mei et al., 2020b] were recently obtained, by showing that the gradient is strictly positive everywhere [Mei et al., 2020b, Lemmas 8-9] . We conjecture that SoftTreeMax satisfies the same property, being a generalization of softmax, but formally proving it is subject to future work.SoftTreeMax gradient. The two variants of SoftTreeMax involve an expectation taken over S d many trajectories from the root state s and weighted according to their probability. Thus, during the PG training process, the gradient ∇ θ log π θ is calculated using a weighted sum of gradients over all reachable states starting from s. Our method exploits the exponential number of trajectories to reduce the variance while improving performance. Indeed, in the next section we prove that the gradient variance of SoftTreeMax decays exponentially fast as a function of the behavior policy π b and trajectory length d. In the experiments in Section 6, we also show how the practical version of SoftTreeMax achieves a significant reduction in the noise of the PG process and leads to faster convergence and higher reward.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We included a relevant section at the end of the paper.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
theory assumptions and proofs: 
 Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [Yes]Answer: [Yes] Justification: We attached the repository, together with a docker environment, as supplementary material.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
theoretical analysis: 
 In this section, we first bound the variance of PG when using the SoftTreeMax policy. Later, we discuss how the gradient bias resulting due to approximate forward models diminishes as a function of the approximation error, while retaining the same variance decay.We show that the variance decreases exponentially with the tree depth, and the rate is determined by the second eigenvalue of the transition kernel induced by π b . Specifically, we bound the same expression for variance as appears in [Greensmith et al., 2004, Sec. 3.5 ] and [Wu et al., 2018, Sec. A, Eq. (21) ]. Other types of analysis could instead have focused on the estimation aspect in the context of sampling [Zhang et al., 2021 , Shen et al., 2019 , Pham et al., 2020] . Indeed, in our implementation in Section 5, we manage to avoid sampling and directly compute the expectations in Eqs.(3) and(4). As we show later, we do so by leveraging efficient parallel simulation on the GPU in feasible run-time. In our application, due to the nature of the finite action space and quasi-deterministic Atari dynamics [Bellemare et al., 2013] , our expectation estimator is noiseless. We encourage future work to account for the finite-sample variance component. We defer all the proofs to Appendix A.We begin with a general variance bound that holds for any parametric policy.Lemma 4.1 (Bound on the policy gradient variance). Let ∇ θ log π θ (•|s) ∈ R A×dim(θ) be a matrix whose a-th row is ∇ θ log π θ (a|s) ⊤ . For any parametric policy π θ and functionQ π θ : S × A → R, Var (∇ θ log π θ (a|s)Q π θ (s, a)) ≤ max s,a [Q π θ (s, a)] 2 max s ∥∇ θ log π θ (•|s)∥ 2 F .Hence, to bound (1), it is sufficient to bound the Frobenius norm ∥∇ θ log π θ (•|s)∥ F for any s.Note that SoftTreeMax does not reduce the gradient uniformly, which would have been equivalentto a trivial change in the learning rate. While the gradient norm shrinks, the gradient itself scales differently along the different coordinates. This scaling occurs along different eigenvectors, as a function of problem parameters (P , θ) and our choice of behavior policy (π b ), as can be seen in the proof of the upcoming Theorem 4.4. This allows SoftTreeMax to learn a good "shrinkage" that, while reducing the overall gradient, still updates the policy quickly enough. This reduction in norm and variance resembles the idea of gradient clipping Zhang et al. [2019] , where the gradient is scaled to reduce its variance, thus increasing stability and improving overall performance.A common assumption in the RL literature [Szepesvári, 2010] that we adopt for the remainder of the section is that the transition matrix P π b , induced by the behavior policy π b , is irreducible and aperiodic. Consequently, its second highest eigenvalue satisfies |λ 2 (P π b )| < 1.From now on, we divide the variance results for the two variants of SoftTreeMax into two subsections. For C-SoftTreeMax, the analysis is simpler and we provide an exact bound. The case of E-SoftTreeMax is more involved and we provide for it a more general result. In both cases, we show that the variance decays exponentially with the planning horizon.We express C-SoftTreeMax in vector form as follows.Lemma 4.2 (Vector form of C-SoftTreeMax). For d ≥ 1, (3) is given byEQUATIONwhereC s,d = γ -d R s + P s d-1 h=1 γ h-d (P π b ) h-1 R π b .The vector C s,d ∈ R A represents the cumulative discounted reward in expectation along the trajectory of horizon d. This trajectory starts at state s, involves an initial reward dictated by R s and an initial transition as per P s . Thereafter, it involves rewards and transitions specified by R π b and P π b , respectively. Once the trajectory reaches depth d, the score function θ(s d ) is applied,.Lemma 4.3 (Gradient of C-SoftTreeMax). The C-SoftTreeMax gradient is given by∇ θ log π C d,θ = β I A -1 A (π C d,θ ) ⊤ P s (P π b ) d-1 , in R A×S, where for brevity, we drop the s index in the policy above, i.e., π C d,θ ≡ π C d,θ (•|s).We are now ready to present our first main result:Theorem 4.4 (Variance decay of C-SoftTreeMax). For every Q : S × A → R, the C-SoftTreeMax policy gradient variance is bounded by d-1) .Var ∇ θ log π C d,θ (a|s)Q(s, a) ≤ 2 A 2 S 2 β 2 (1 -γ) 2 |λ 2 (P π b )| 2(We provide the full proof in Appendix A.4, and briefly outline its essence here.Proof outline. Lemma 4.1 allows us to bound the variance using a direct bound on the gradient norm. The gradient is given in Lemma 4.3 as a product of three matrices, which we now study from right to left. The matrix P π b is a row-stochastic matrix. Because the associated Markov chain is irreducible and aperiodic, it has a unique stationary distribution. This implies that P π b has one and only one eigenvalue equal to 1; all others have magnitude strictly less than 1. Let us suppose that all these other eigenvalues have multiplicity 1 (the general case with repeated eigenvalues can be handled via Jordan decompositions as in [Pelletier, 1998, Lemma1] ). Then, P π b has the spectraldecomposition P π b = 1 S µ ⊤ π b + S i=2 λ i v i u ⊤ i, where λ i is the i-th eigenvalue of P π b (ordered in descending order according to their magnitude) and u i and v i are the corresponding left and right eigenvectors, respectively, and therefore(P π b ) d-1 = 1 S µ ⊤ π b + S i=2 λ d-1 i v i u ⊤ i .The second matrix in the gradient relation in Lemma 4.3, P s , is a rectangular transition matrix that translates the vector of all ones from dimension S to A : P s 1 S = 1 A . Lastly, the first matrixI A -1 A (π C d,θ) ⊤ is a projection whose null-space includes the vector 1 A , i.e.,I A -1 A (π C d,θ ) ⊤ 1 A = 0.Combining the three properties above when multiplying the three matrices of the gradient, it is easy to see that the first term in the expression for (P π b ) Theorem 4.4 guarantees that the variance of the gradient decays exponentially with d. It also provides a novel insight for choosing the behavior policy π b as the policy that minimizes the absolute second eigenvalue of the P π b . Indeed, the second eigenvalue of a Markov chain relates to its connectivity and its rate of convergence to the stationary distribution [Levin and Peres, 2017] .Optimal variance decay. For the strongest reduction in variance, the behavior policy π b should be chosen to achieve an induced Markov chain whose transitions are state-independent. In that case, P π b is a rank one matrix of the form 1 S µ ⊤ π b , and λ 2 (P π b ) = 0. Then, Var (∇ θ log π θ (a|s)Q(s, a)) = 0.Naturally, this can only be done for pathological MDPs; see Appendix C.1 for a more detailed discussion. Nevertheless, as we show in Section 5, we choose our tree expansion policy to reduce the variance as best as possible.Worst-case variance decay. In contrast, and somewhat surprisingly, when π b is chosen so that the dynamics is deterministic, there is no guarantee that it will decay exponentially fast. For example, if P π b is a permutation matrix, then λ 2 (P π b ) = 1, and advancing the tree amounts to only updating the gradient of one state for every action, as in the basic softmax.The proof of the variance bound for E-SoftTreeMax is similar to that of C-SoftTreeMax, but more involved. It also requires the assumption that the reward depends only on the state, i.e. r(s, a) ≡ r(s).This is indeed the case in most standard RL environments such as Atari and Mujoco.Lemma 4.5 (Vector form of E-SoftTreeMax). For d ≥ 1, (4) is given byEQUATIONwhereE s,d = P s d-1 h=1 D exp(βγ h-d R) P π b .The vector R above is the S-dimensional vector whose s-th coordinate is r(s).The matrix E s,d ∈ R A×S has a similar role to C s,d from ( 5), but it represents the exponentiated cumulative discounted reward. Accordingly, it is a product of d matrices as opposed to a sum. It captures the expected reward sequence starting from s and then iteratively following P π b . After d steps, we apply the score function on the last state as in (6).Lemma 4.6 (Gradient of E-SoftTreeMax). The E-SoftTreeMax gradient is given by∇ θ log π E d,θ = β I A -1 A (π E d,θ ) ⊤ × D π E d,θ -1 E s,d D(exp(βΘ)) 1 ⊤ A E s,d exp(βΘ) ∈ R A×S ,where for brevity, we drop the s index in the policy above, i.e.,π E d,θ ≡ π E d,θ (•|s).This gradient structure is harder to handle than that of C-SoftTreeMax in Lemma 4.3, but here we also can bound the decay of the variance nonetheless.Theorem 4.7 (Variance decay of E-SoftTreeMax). There exists α ∈ (0, 1) such that,Var ∇ θ log π E d,θ (a|s)Q(s, a) ∈ O β 2 α 2d , for every Q. Further, if P π b is reversible or if the reward is constant, then α = |λ 2 (P π b )|.Theory versus Practice. We demonstrate the above result in simulation. We draw a random finite MDP, parameter vector Θ ∈ R S + , and behavior policy π b . We then empirically compute the PG variance of E-SoftTreeMax as given in (1) and compare it to |λ 2 (P π b )| d . We repeat this experiment three times for different P π b : (i) close to uniform, (ii) drawn randomly, and (iii) close to a permutation matrix. As seen in Figure 1 , the empirical variance and our bound match almost identically. This also suggests that α = |λ 2 (P π b )| in the general case and not only when P π b is reversible or when the reward is constant.The definition of the two SoftTreeMax variants involves the knowledge of the underlying environment, in particular the value of P and r. However, in practice, we often can only learn approximations of the dynamics from interactions, e.g., using NNs [Ha and Schmidhuber, 2018, Schrittwieser et al., 2020] . Let P and r denote the approximate kernel and reward functions, respectively. In this section, we study the consequences of the approximation error on the C-SoftTreeMax gradient. 𝑊(𝑆 !"# $,$ ) logits for logits for logits for𝑎 !"$ ($) Policy network 𝑎 !"$ (() 𝑎 !"$ ($) 𝑎 !"$ (() 𝑎 !"$ ($) 𝑎 !"$ (() 𝑎 !") ($) 𝑎 !") (#) 𝑎 !") (() 𝑎 !"# (%) 𝑎 !"# (') 𝑎 !"# (() 𝑆 !") 𝑆 !"$ ($) 𝑆 !"$ (#) 𝑆 !"$ (() 𝑆 !"# ($,$) 𝑆 !"# (#,$) 𝑆 !"# ($,() 𝑆 !"# (#,() 𝑆 !"# ((,$) 𝑆 !"# ((,() 𝑊(𝑆 !"# $,( ) 𝑊(𝑆 !"# #,$ ) 𝑊(𝑆 !"# #,( ) 𝑊(𝑆 !"# (,$ ) 𝑊(𝑆 !"# (,( )Figure 2 : SoftTreeMax policy. Our exhaustive parallel tree expansion iterates on all actions at each state up to depth d (= 2 here).The leaf state of every trajectory is used as input to the policy network. The output is then added to the trajectory's cumulative reward as described in (2). I.e., instead of the standard softmax logits, we add the cumulative discounted reward to the policy network output. This policy is differentiable and can be easily integrated into any PG algorithm. In this work, we build on PPO and use its loss function to train the policy network.Let πC d,θ be the C-SoftTreeMax policy defined given the approximate forward model introduced above. That is, let πC d,θ be defined exactly as in ( 5), but using Rs , Ps , Rπ b and P π b , instead of their unperturbed counterparts from Section 2. Then, the variance of the corresponding gradient again decays exponentially with a decay rate of λ 2 ( P π b ). However, a gradient bias is introduced. In the following, we bound this bias in terms of the approximation error and other problem parameters. The proof is provided in Appendix A.9.Theorem 4.8. Let ϵ be the maximal model mis-specification, i.e., let max{∥P -P ∥, ∥r -r∥} = ϵ.Then the policy gradient bias due to πC d,θ satisfiesEQUATIONTo the best of our knowledge, Theorem 4.8 is the first result that bounds the bias of the gradient of a parametric policy due to an approximate model. It states that if the learned model is accurate enough, we expect similar convergence properties for C-SoftTreeMax as we would have obtained with the true dynamics. It also suggests that higher temperature (lower β) reduces the bias. In this case, the logits get less weight, with the extreme of β = 0 corresponding to a uniform policy that has no bias. Lastly, the error scales linearly with d : the policy suffers from cumulative error as it relies on further-looking states in the approximate model. 
softtreemax: deep parallel implementation: 
 Following impressive successes of deep RL [Mnih et al., 2015 , Silver et al., 2016] , using deep NNs in RL is standard practice. Depending on the RL algorithm, a loss function is defined and gradients on the network weights can be calculated. In PG methods, the scoring function used in the softmax is commonly replaced by a neural network W θ : π θ (a|s) ∝ exp (W θ (s, a)) . Similarly, we implement SoftTreeMax by replacing θ(s) in ( 2) with a neural network W θ (s). Although both variants of SoftTreeMax from Section 3 involve computing an expectation, this can be hard in general. One approach to handle it is with sampling, though these introduce estimation variance into the process.We leave the question of sample-based theory and algorithmic implementations for future work.Instead, in finite action space environments such as Atari, we compute the exact expectation in SoftTreeMax with an exhaustive TS of depth d. Despite the exponential computational cost of spanning the entire tree, recent advancements in parallel GPU-based simulation allow efficient expansion of all nodes at the same depth simultaneously [Dalal et al., 2021 , Rosenberg et al., 2022] . This is possible when a simulator is implemented on GPU [Dalton et al., 2020 , Makoviychuk et al., 2021 , Freeman et al., 2021] , or when a forward model is learned [Kim et al., 2020, Ha and Schmidhuber, 2018] . To reduce the complexity to be linear in depth, we apply tree pruning to a limited width in all levels. We do so by sub-sampling only the most promising branches at each level.Limiting the width drastically improves runtime, and enables respecting GPU memory limits, with only a small sacrifice in performance.To summarize, in the practical SoftTreeMax algorithm we perform an exhaustive tree expansion with pruning to obtain trajectories up to depth d. We expand the tree with equal weight to all actions, which corresponds to a uniform tree expansion policy π b . We apply a neural network on the leaf states, and accumulate the result with the rewards along each trajectory to obtain the logits in (2). Finally, we aggregate the results using C-SoftTreeMax. We leave experiments E-SoftTreeMax for future work on risk-averse RL. During training, the gradient propagates to the NN weights of W θ . When the gradient ∇ θ log π d,θ is calculated at each time step, it updates W θ for all leaf states, similarly toSiamese networks [Bertinetto et al., 2016 ]. An illustration of the policy is given in Figure 2 . 
experiments: 
 We conduct our experiments on multiple games from the Atari simulation suite [Bellemare et al., 2013] . As a baseline, we train a PPO [Schulman et al., 2017] agent with 256 GPU workers in parallel [Dalton et al., 2020] . For the tree expansion, we employ a GPU breadth-first as in [Dalal et al., 2021] .We then train C-SoftTreeMax1 for depths d = 1 . . . 8, with a single worker. For depths d ≥ 3, we limited the tree to a maximum width of 1024 nodes and pruned trajectories with low estimated weights. Since the distributed PPO baseline advances significantly faster in terms of environment steps, for a fair comparison, we ran all experiments for one week on the same machine. For more details see Appendix B.In Figure 3 , we plot the reward and variance of SoftTreeMax for each game, as a function of depth.The dashed lines are the results for PPO. Each value is taken after convergence, i.e., the average over the last 20% of the run. The numbers represent the average over five seeds per game. The plot conveys three intriguing conclusions. First, in all games, SoftTreeMax achieves significantly higher reward than PPO. Its gradient variance is also orders of magnitude lower than that of PPO. Second, the reward and variance are negatively correlated and mirror each other in almost all games. This phenomenon demonstrates the necessity of reducing the variance of PG for improving performance.Lastly, each game has a different sweet spot in terms of optimal tree depth. Recall that we limit the run-time in all experiments to one week The deeper the tree, the slower each step and the run consists of less steps. This explains the non-monotone behavior as a function of depth. For a more thorough discussion on the sweet spot of different games, see Appendix B.3.Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: [NA] Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
related work: 
 Softmax Operator. The softmax policy became a canonical part of PG to the point where theoretical results of PG focus specifically on it [Zhang et al., 2021 , Mei et al., 2020b , Li et al., 2021 , Ding et al., 2022] . Even though we focus on a tree extension to the softmax policy, our methodology is general and can be easily applied to other discrete or continuous parameterized policies as in [Mei et al., 2020a , Miahi et al., 2021 , Silva et al., 2019] . Tree Search. One famous TS algorithm is Monte-Carlo TS (MCTS; [Browne et al., 2012] ) used in AlphaGo [Silver et al., 2016] and MuZero [Schrittwieser et al., 2020] . Other algorithms such as Value Iteration, Policy Iteration and DQN were also shown to give an improved performance with a tree search extensions [Efroni et al., 2019 , Dalal et al., 2021] .Parallel Environments. In this work we used accurate parallel models that are becoming more common with the increasing popularity of GPU-based simulation [Makoviychuk et al., 2021 , Dalton et al., 2020 , Freeman et al., 2021] . Alternatively, in relation to Theorem 4.8, one can rely on recent works that learn the underlying model [Ha and Schmidhuber, 2018, Schrittwieser et al., 2020] and use an approximation of the true dynamics. Risk Aversion. Previous work considered exponential utility functions for risk aversion [Chen et al., 2007 , Garcıa and Fernández, 2015 , Fei et al., 2021] .This utility function is the same as E-SoftTreeMax formulation from ( 4), but we have it directly in the policy instead of the objective. Reward-free RL. We showed that the gradient variance is minimized when the transitions induced by the behavior policy π b are uniform. This is expressed by the second eigenvalue of the transition matrix P π b . This notion of uniform exploration is common to the reward-free RL setup [Jin et al., 2020] . Several such works also considered the second eigenvalue in their analysis [Liu and Brunskill, 2018, Tarbouriech and Lazaric, 2019] .Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: [NA] Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
discussion: 
 In this work, we introduced for the first time a differentiable parametric policy that combines TS with PG. We proved that SoftTreeMax is essentially a variance reduction technique and explained how to choose the expansion policy to minimize the gradient variance.It is an open question whether optimal variance reduction corresponds to the appealing regret properties the were put forward by UCT [Kocsis and Szepesvári, 2006] . We believe that this can be answered by analyzing the convergence rate of SoftTreeMax, relying on the bias and variance results we obtained here.As the learning process continues, the norm of the gradient and the variance both become smaller.On the face of it, one can ask if the gradient becomes small as fast as the variance or even faster can there be any meaningful learning? As we showed in the experiments, learning happens because the variance reduces fast enough (a variance of 0 represents deterministic learning, which is fastest).Finally, our work can be extended to infinite action spaces. The analysis can be extended to infinitedimension kernels that retain the same key properties used in our proofs. In the implementation, the tree of continuous actions can be expanded by maintaining a parametric distribution over actions that depend on θ. This approach can be seen as a tree adaptation of MPPI [Williams et al., 2017] . 
reproducibility and limitations: 
 In this submission, we include the code as part of the supplementary material. We also include a docker file for setting up the environment and a README file with instructions on how to run both training and evaluation. The environment engine is an extension of Atari-CuLE [Dalton et al., 2020] ,a CUDA-based Atari emulator that runs on GPU. Our usage of a GPU environment is both a novelty and a current limitation of our work.Var (∇ θ log π θ (a|s)Q(a, s)) ≤E s∼dπ θ ,a∼π θ (•|s) ∇ θ log π θ (a|s) ⊤ ∇ θ log π θ (a|s)Q(s, a) 2 = s d π θ (s) a π θ (a|s)∇ θ log π θ (a|s) ⊤ ∇ θ log π θ (a|s)Q(s, a) 2 ≤ max s,a [Q(s, a)] 2 π θ (a|s) s d π θ (s) a ∇ θ log π θ (a|s) ⊤ ∇ θ log π θ (a|s) ≤ max s,a [Q(s, a)] 2 max s a ∇ θ log π θ (a|s) ⊤ ∇ θ log π θ (a|s) = max s,a [Q(s, a)] 2 max s ∥∇ θ log π θ (•|s)∥ 2 F .A.2 Proof of Lemma 4.2 -Vector form of C-SoftTreeMaxIn vector form, (3) is given byEQUATIONwhereEQUATIONProof. Consider the vector ℓ s,• ∈ R |A| . Its expectation satisfiesE π b ℓ s,• (d; θ) = E π b d-1 t=0 γ t-d r t + θ(s d ) = γ -d R s + d-1 t=1 γ t-d P s (P π b ) t-1 R π b + P s (P π b ) d-1 Θ.As required.The C-SoftTreeMax gradient of dimension A × S is given by∇ θ log π C d,θ = β I A -1 A (π C d,θ ) ⊤ P s (P π b ) d-1 ,where for brevity, we drop the s index in the policy above, i.e.,π C d,θ ≡ π C d,θ (•|s). Proof. The (j, k)-th entry of ∇ θ log π C d,θ satisifes [∇ θ log π C d,θ ] j,k = ∂ log(π C d,θ (a j |s)) ∂θ(s k ) = β[P s (P π b ) d-1 ] j,k - a exp β C s,d + P s (P π b ) d-1 Θ a β P s (P π b ) d-1 a,k 1 ⊤ A exp β C s,d + P s (P π b ) d-1 Θ = β[P s (P π b ) d-1 ] j,k -β a π C d,θ (a|s) P s (P π b ) d-1 a,k = β[P s (P π b ) d-1 ] j,k -β (π C d,θ ) ⊤ P s (P π b ) d-1 k = β[P s (P π b ) d-1 ] j,k -β 1 A (π C d,θ ) ⊤ P s (P π b ) d-1 j,k. Moving back to matrix form, we obtain the stated result.The C-SoftTreeMax policy gradient is bounded byVar ∇ θ log π C d,θ (a|s)Q(s, a) ≤ 2 A 2 S 2 β 2 (1 -γ) 2 |λ 2 (P π b )| 2(d-1) .Proof. We use Lemma 4.1 directly. First of all, it is know that when the reward is bounded in [0, 1],the maximal value of the Q-function is 1 1-γ as the sum as infinite discounted rewards. Next, we bound the Frobenius norm of the term achieved in Lemma 4.3, by applying the eigen-decomposition on P π b :EQUATION)where µ is the stationary distribution of P π b , and u i and v i are left and right eigenvectors correspondingly.∥β I A,A -1 A π ⊤ P s (P π b ) d-1 ∥ F = β∥ I A,A -1 A π ⊤ P s 1 S µ ⊤ + S i=2 λ d-1 i u i v ⊤ i ∥ F (P s is stochastic) = β∥ I A,A -1 A π ⊤ 1 A µ ⊤ + S i=2 λ d-1 i P s u i v ⊤ i ∥ F (projection nullifies 1 A µ ⊤ ) = β∥ I A,A -1 A π ⊤ S i=2 λ d-1 i P s u i v ⊤ i ∥ F (triangle inequality) ≤ β S i=2 ∥ I A,A -1 A π ⊤ λ d-1 i P s u i v ⊤ i ∥ F (matrix norm sub-multiplicativity) ≤ β|λ d-1 2 | S i=2 ∥I A,A -1 A π ⊤ ∥ F ∥P s ∥ F ∥u i v ⊤ i ∥ F = β|λ d-1 2 |(S -1)∥I A,A -1 A π ⊤ ∥ F ∥P s ∥ F .EQUATION)From the Cauchy-Schwartz inequality,∥P s ∥ 2 F = a s [[P s ] a,s ] 2 = a ∥[P s ] a,• ∥ 2 2 ≤ a ∥[P s ] a,• ∥ 1 ∥[P s ] a,• ∥ ∞ ≤ A.So,Var ∇ θ log π C d,θ (a|s)Q(s, a) ≤ max s,a [Q(s, a)] 2 max s ∥∇ θ log π C d,θ (•|s)∥ 2 F ≤ 1 (1 -γ) 2 ∥β I A,A -1 A π ⊤ P s (P π b ) d-1 ∥ 2 F ≤ 1 (1 -γ) 2 β 2 |λ 2 (P π b )| 2(d-1) S 2 (2A 2 ),which obtains the desired bound.For completeness we also supply a lower bound on the Frobenius norm of the gradient. Note that this result does not translate to the a lower bound on the variance since we have no lower bound equivalence of Lemma 4.1.Lemma A.1. The Frobenius norm on the gradient of the policy is lower-bounded by:EQUATION)Proof. We begin by moving to the induced l 2 norm by norm-equivalence:∥β I A,A -1 A π ⊤ P s (P π b ) d-1 ∥ F ≥ ∥β I A,A -1 A π ⊤ P s (P π b ) d-1 ∥ 2 .Now, taking the vector u to be the eigenvector of the second eigenvalue of P π b :∥β I A,A -1 A π ⊤ P s (P π b ) d-1 ∥ 2 ≥ ∥β I A,A -1 A π ⊤ P s (P π b ) d-1 u∥ 2 = β∥ I A,A -1 A π ⊤ P s u∥ 2 = β|λ 2 (P π b )| (d-1) ∥ I A,A -1 A π ⊤ P s u∥ 2 .Note that even though P s u can be 0, that is not the common case since we can freely change π b (and therefore the eigenvectors of P π b ).A.6 Proof of Lemma 4.5 -Vector form of E-SoftTreeMaxEQUATIONwhereEQUATIONwith R being the |S|-dimensional vector whose s-th coordinate is r(s).Proof. Recall thatEQUATIONand, hence,EQUATIONTherefore,EQUATIONBy repeatedly using iterative conditioning as above, the desired result follows. Note that exp(βγ -d r(s)) does not depend on the action and is therefore cancelled out with the denominator.The E-SoftTreeMax gradient of dimension A × S is given by∇ θ log π E d,θ = β I A -1 A (π E d,θ ) ⊤ D π E d,θ -1 E s,d D(exp(βΘ)) 1 ⊤ A E s,d exp(βΘ),where for brevity, we drop the s index in the policy above, i.e., π E d,θ ≡ π E d,θ (•|s).Proof. The (j, k)-th entry of ∇ θ log π E d,θ satisfies[∇ θ log π E d,θ ] j,k = ∂ log(π E d,θ (a j |s)) ∂θ(s k ) = ∂ ∂θ(s k ) log[(E s,d ) ⊤ j exp(βΘ)] -log[1 ⊤ A E s,d exp(βΘ)] = β(E s,d ) j,k exp(βθ(s k )) (E s,d ) ⊤ j exp(βΘ) - β1 ⊤ A E s,d e k exp(βθ(s k )) 1 ⊤ A E s,d exp(βΘ) = β(E s,d e k exp(βθ(s k ))) j (E s,d ) ⊤ j exp(βΘ) - β1 ⊤ A E s,d e k exp(βθ(s k )) 1 ⊤ A E s,d exp(βΘ) = β e ⊤ j e ⊤ j E s,d exp(βΘ) - 1 ⊤ A 1 ⊤ A E s,d exp(βΘ) E s,d e k exp(βθ(s k )).Hence,[∇ θ log π E d,θ ] •,k = β D(E s,d exp(βΘ)) -1 -(1 ⊤ A E s,d exp(βΘ)) -1 1 A 1 ⊤ A E s,d e k exp(βθ(s k ))From this, it follows thatEQUATION)The desired result is now easy to see.There exists α ∈ (0, 1) such that, for any functionQ : S × A → R, Var ∇ θ log π E d,θ (a|s)Q(s, a) ∈ O β 2 α 2d . If all rewards are equal (r ≡ const), then α = |λ 2 (P π b )|.Proof outline. Recall that thanks to Lemma 4.1, we can bound the PG variance using a direct bound on the gradient norm. The definition of the induced norm is∥∇ θ log π E d,θ ∥ = max z:∥z∥=1 ∥∇ θ log π E d,θ z∥,with ∇ θ log π E d,θ given in Lemma 4.6. Let z ∈ R S be an arbitrary vector such that ∥z∥ = 1. Then, z = S i=1 c i z i , where c i are scalar coefficients and z i are vectors spanning the S-dimensional space.In the full proof, we show our specific choice of z i and prove they are linearly independent given that choice. We do note that z 1 = 1 S .The first part of the proof relies on the fact that (∇ θ log π E d,θ )z 1 = 0. This is easy to verify using Lemma 4.6 together with (6), and becauseI A -1 A (π E d,θ) ⊤ is a projection matrix whose null-space is spanned by 1 S . Thus,∇ θ log π E d,θ z = ∇ θ log π E d,θ S i=2 c i z i .In the second part of the proof, we focus on E s,d from ( 6), which appears within∇ θ log π E d,θ . Notice that E s,d consists of the product d-1 h=1 D exp(βγ h-d R P π b .Even though the elements in this product are not stochastic matrices, in the full proof we show how to normalize each of them to a stochastic matrix B h . We thus obtain thatE s,d = P s D(M 1 ) d-1 h=1 B h ,where M 1 ∈ R S is some strictly positive vector. Then, we can apply a result by Mathkar and Borkar [2016] , which itself builds on [Chatterjee and Seneta, 1977] . The result states that the product of stochastic matrices d-1 h=1 B h of our particular form converges exponentially fast to a matrix of theform 1 S µ ⊤ s.t. ∥1 S µ ⊤ - d-1 h=1 B h ∥ ≤ Cα d for some constant C.Lastly, 1 S µ ⊤ π b gets canceled due to our choice of z i , i = 2, . . . , S. This observation along with the above fact that the remainder decays then shows that∇ θ log π E d,θ S i=2 z i = O(α d), which gives the desired result.EQUATIONand that R refers to the S-dimensional vector whose s-th coordinate is r(s). DefineEQUATION)and the vectorEQUATION)where • denotes the element-wise product. Then,EQUATION)It is easy to see that each B i is a row-stochastic matrix, i.e., all entries are non-negative andB i 1 S = 1 S .Next, we prove that all non-zeros entries of B i are bounded away from 0 by a constant. This is necessary to apply the next result from Chatterjee and Seneta [1977] . The j-th coordinate of M i satisfiesEQUATION)Separately, observe that∥M d-1 ∥ ∞ ≤ ∥ exp(βγ -1 R)∥ ∞ . Plugging these relations in (26) gives ∥M 1 ∥ ∞ ≤ d-1 h=1 ∥ exp[βγ h-d R]∥ ∞ = d-1 h=1 ∥ exp[βγ -d R]∥ γ h ∞ = ∥ exp[βγ -d R]∥ d-1 h=1 γ h ∞ ≤ ∥ exp[βγ -d R]∥ 1 1-γ ∞ . (29)Similarly, for every 1 ≤ i ≤ d -1, we have thatEQUATION)The jk-th entry ofEQUATION)Hence, for non-zero P π b jk , the entries are bounded away from zero by the same. We can now proceed with applying the following result. Now, by [Chatterjee and Seneta, 1977, Theorem 5] (see also ( 14) in [Mathkar and Borkar, 2016] ),lim d→∞ d-1i=1 B i exists and is of the form 1 S µ ⊤ for some probability vector µ. Furthermore, there is some α ∈ (0, 1) such that ε(dEQUATION)Pick linearly independent vectors w 2 , . . . , w S such thatEQUATIONSince S i=2 α i w i is perpendicular to µ for any α 2 , . . . α S and because µ ⊤ exp(βΘ) > 0, there exists no choice of α 2 , . . . , α S such that S i=2 α i w i = exp(βΘ). Hence, if we let z 1 = 1 S and z i = D(exp(βΘ)) -1 w i for i = 2, . . . , S, then it follows that {z 1 , . . . , z S } is linearly independent.In particular, it implies that {z 1 , . . . , z S } spans R S .EQUATIONwhere (34) follows from the fact that ∇ θ log π E d,θ z 1 = ∇ θ log π E d,θ 1 S = 0, (35) follows from Lemma 4.6, (36) holds since z i = D(exp(βΘ)) -1 w i , (38) because µ is perpendicular w i for each i, while (39) follows by reusing z i = D(exp(βΘ)) -1 w i relation along with the fact that z 1 = 1 S .From (39), it follows thatEQUATIONLastly, we prove that ∥z -c 1 1 S ∥ is bounded independently of d. First, denote by c = (c 1 , . . . , c S ) ⊤ and c = (0, c 2 , . . . , c S ) ⊤ . Also, denote by Z the matrix with z i as its i-th column. Now,EQUATIONwhere the last relation is due to z being a unit vector. All matrix norms here are l 2 -induced norms.Next, denote by W the matrix with w i in its i-th column. Recall that in (33) we only defined w 2 , . . . , w S . We now set w 1 = exp(βΘ). Note that w 1 is linearly independent of {w 2 , . . . , w S } because of (33) together with the fact that µ ⊤ w 1 > 0. We can now express the relation between Z and W by Z = D -1 (exp(βΘ))W. Substituting this in (53), we haveEQUATION)It further holds thatEQUATIONwhere the last relation equals 1 if θ(s) < 0 for all s. Similarly,EQUATIONFurthermore, by the properties of the l 2 -induced norm,EQUATIONLastly,EQUATIONEQUATION)The determinant of W is a sum of products involving its entries. To upper bound (65) independently of d, we lower bound its denominator by upper and lower bounds on the entries [W ] i,1 that are independent of d, depending on their sign:EQUATIONUsing this, together with ( 53), ( 55), ( 56), (57), and (61), we showed that ∥z -c 1 1 S ∥ is upper bounded by a constant independent of d. This concludes the proof.A.9 Bias Estimates Lemma A.2. For any matrix A and Â,Âk -A k = k h=1 Âh-1 ( Â -A)A k-h .Proof. The proof follows from first principles: EQUATIONEQUATIONProof. The proof follows easily from the fact that the differences above are convex combinations of P -P and r -r.Lemma A.5. Let π θ be as in (5), and let πθ also be defined as in (5), but with R s , P s , P π b replaced by their perturbed counterparts Rs , Ps , P π b throughout. Then,EQUATIONProof. To prove the desired result, we work with (5) to bound the error between R s , P s , P π b , R π band their perturbed versions.First, we apply Lemma A.2 together with Lemma A.4 to obtain that∥(P π b ) k -( P π b ) k ∥ = O(kϵ).Next, denote by M the argument in the exponent in (5), i.e.M := β[C s,d + P s (P π b ) d-1 Θ].Similarly, let M be the corresponding perturbed sum that relies on P and r. Combining the bounds from Lemma A.4, and using the triangle inequality, we have that ∥ M -M ∥ = O(βdϵ).Eq. ( 5) states that the C-SoftTreeMax policy in the true environment isπ θ = exp(M )/(1 ⊤ exp(M )).Similarly define πθ using M for the approximate model. Then,πθ = (π θ • exp(M -M ))1 ⊤ exp(M )/(1 ⊤ exp( M )),where • denotes element-wise multiplication. Using the above relation, we have that ∥π θ -π θ ∥ = ∥π θ ∥∥ exp(M -M )1 ⊤ exp(M ) 1 ⊤ exp( M )-1∥. Using the relation |e x -1| = O(x) as x → 0, the desired result follows.Theorem A.6. Let ϵ be as in Definition A.3. Further let πC d,θ being the corresponding approximate policy as given in Lemma 4.2. Then, the policy gradient bias is bounded byEQUATIONWe first provide a proof outline for conciseness, and only after it the complete proof.Proof outline. First, we prove that max{∥Rs -Rs ∥, ∥P s -Ps ∥, ∥R π b -Rπ b ∥, ∥P π b -P π b ∥} = O(ϵ).This follows from the fact that the differences above are suitable convex combinations of either the rows of P -P or r -r. We use the above observation along with the definitions of π C d,θ and πCd,θ given in (5) to show that ∥π C d,θ -πC d,θ ∥ = O(βdϵ).The proof for the latter builds upon two key facts:(a) ∥(P π b ) k -( P π b ) k ∥ ≤ k h=1 ∥ P π b ∥ h-1 ∥ P π b -P π b ∥∥p π b ∥ k-h = O(kϵ) for any k ≥ 0, and (b) |e x -1| = O(x) as x → 0. Next, we decompose the LHS of (7) to get s 4 i=1 X i (s) - 4 i=1 Xi (s) = s 4 i=1 X1 (s) • • • Xi-1 (s) X i (s) -Xi (s) ×X i+1 (s) • • • X 4 (s), where X 1 (s) = d π C d,θ (s) ∈ R, X 2 (s) = (∇ θ log π C d,θ (•|s)) ⊤ ∈ R S×A , X 3 (s) = D(π C d,θ (•|s)) ∈ R A×A , X 4 (s) = Q π C d,θ (s, •) ∈ R A×A, and X1 (s), . . . , X4 (s) are similarly defined with π C d,θ replaced by πC d,θ . Then, we show that, for i = 1, . . . , 4, (i) ∥X i (s) -Xi (s)∥ = O(ϵ) and (ii) max{∥X i ∥, ∥ Xi ∥} is bounded by problem parameters. From this, the desired result follows.Proof. We haveEQUATION-EQUATIONEQUATIONwhereX 1 (s) = d π θ (s) ∈ R, X 2 (s) = (∇ θ log π θ (•|s)) ⊤ ∈ R S×A , X 3 (s) = D(π θ (•|s)) ∈ R A×A , X 4 (s) = Q π θ (s, •) ∈ R A×A, and X1 (s), . . . , X4 (s) are similarly defined with π θ replaced by πθ .Therefore,EQUATIONwhere EQUATIONSeparately, using Lemma 4.3, we haveEQUATIONSince all rows of the above matrices have non-negative entries that add up to 1, we getEQUATIONIn the rest of the proof, we bound each of ∥X 1 -X1 ∥, . . . , ∥X 4 -X4 ∥.Finally,EQUATION)Similarly, the same bounds hold for X1 , X2 , X3 and X4 .From, we haveEQUATION)The last relation follows from the fact that (1 -γ) -1 = ∞ t=0 γ t , which in turn impliesEQUATIONFrom Lemma A.5, it follows thatEQUATIONNext, recall that from Lemma 4.3 thatX 2 (s, •) = β I A -1 A (π θ ) ⊤ P s (P π b ) d-1 .Then,EQUATION)EQUATION)Following the same argument as in ( 85) and applying Lemma A.2, we have that ( 93) is O(βdϵ).Similarly, from the argument of (85), Eq. ( 94) is O(βϵ). Lastly, ( 95) is O(βdϵ) due to Lemma A.5.Putting the above three terms together, we have thatEQUATIONSince the state-action value function satisfies the Bellman equation, we haveEQUATIONandEQUATIONwhich finally shows that∥X 4 -X4 ∥ = ∥Q π θ -Q πθ ∥ ≤ ϵ (1 -γ) 2 .(102)The environment engine is the highly efficient Atari-CuLE [Dalton et al., 2020] , a CUDA-based version of Atari that runs on GPU. Similarly, we use Atari-CuLE for the GPU-based breadth-first TS as done in Dalal et al. [2021] : In every tree expansion, the state S t is duplicated and concatenated with all possible actions. The resulting tensor is fed into the GPU forward model to generate the tensor of next states (S 0 t+1 , . . . , S A-1 t+1 ). The next-state tensor is then duplicated and concatenated again with all possible actions, fed into the forward model, etc. This procedure is repeated until the final depth is reached, for which W θ (s) is applied per state.We train SoftTreeMax for depths d = 1 . . . 8, with a single worker. We use five seeds for each experiment.For the implementation, we extend Stable-Baselines3 [Raffin et al., 2019] with all parameters taken as default from the original PPO paper [Schulman et al., 2017] . For depths d ≥ 3, we limited the tree to a maximum width of 1024 nodes and pruned non-promising trajectories in terms of estimated weights. Since the distributed PPO baseline advances significantly faster in terms of environment steps, for a fair comparison, we ran all experiments for one week on the same machine and use the wall-clock time as the x-axis. We use Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz equipped with one NVIDIA Tesla V100 32GB.We provide the training curves in Figure 4 . For brevity, we exclude a few of the depths from the plots.As seen, there is 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 although it observes much less episodes, for the same running time. In Figure 5 we also provide the same convergence plots where the x-axis is now the number of online interactions with the environment, thus excluding the tree expansion complexity. As seen, due to the complexity of the tree expansion, less steps are conducted during training (limited to one week) as the depth increases. In this plot, the monotone improvement of the reward with increasing tree depth is noticeable in most games. The training curves correspond to the evaluation runs in Figure 3 .We note that not for all games we see monotonicity. Our explanation for this phenomenon relates to how immediate reward contributes to performance compared to the value. Different games benefit differently from long-term as opposed to short-term planning. Games that require longer-term planning need a better value estimate. A good value estimate takes longer to obtain with larger depths, in which we apply the network to states that are very different from the ones observed so far in the buffer (recall that as in any deep RL algorithm, we train the model only on states in the buffer). If the model hasn't learned a good enough value function yet, and there is no guiding dense reward along the trajectory, the policy becomes noisier, and can take more steps to converge -even more than those we run in our week-long experiment.For a concrete example, let us compare Breakout to Gopher. Inspecting Fig. 5 , we observe that Breakout quickly (and monotonically) gains from large depths since it relies on the short term goal of simply keeping the paddle below the moving ball. In Gopher, however, for large depths (>=5), learning barely started even by the end of the training run. Presumably, this is because the task in Gopher involves multiple considerations and steps: the agent needs to move to the right spot and then hit the mallet the right amount of times, while balancing different locations. This task requires long-term planning and thus depends more strongly on the accuracy of the value function estimate.In that case, for depth 5 or more, we would require more train steps for the value to "kick in" and become beneficial beyond the gain from the reward in the tree.The figures above convey two key observations that occur for at least some non-zero depth: (1) The final performance with the tree is better than PPO (Fig. 3 ); and (2) the intermediate step-based results with the tree are better than PPO (Fig. 5 ). This leads to our main takeaway from this work -there is no reason to believe that the vanilla policy gradient algorithm should be better than a multi-step variant. Indeed, we show that this is not the case.C.1 The case of λ 2 (P π b ) = 0When P π b is rank one, it is not only its variance that becomes 0, but also the norm of the gradient itself (similarly to the case of d → ∞). Note that such a situation will happen rarely, in degenerate MDPs. This is a local minimum for SoftTreeMax and it would cause the PG iteration to get stuck, and to the optimum in the (desired but impractical) case where π b is the optimal policy. However, a similar phenomenon was also discovered in the standard softmax with deterministic policies:θ(s, a) → ∞ for one a per s. PG with softmax would suffer very slow convergence near these local equilibria, as observed in Mei et al. [2020a] . To see this, note that the softmax gradient is∇ θ log π θ (a|s) = e a -π θ (•|s), where e a ∈ [0, 1] A is the vector with 0 everywhere except for the a-th coordinate. I.e., it will be zero for a deterministic policy. SoftTreeMax avoids these local optima by integrating the reward into the policy itself (but may get stuck in another, as discussed above).NeurIPS Paper Checklist 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes]Justification: [NA] Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.We also experimented with E-SoftTreeMax and the results were almost identical. This is due to the quasideterministic nature of Atari, which causes the trajectory logits (2) to have almost no variability. We encourage future work on E-SoftTreeMax using probabilistic environments that are risk-sensitive. 
 
 Appendix:  


Decision: **Summary**

This paper proposes a combination of tree expansion and policy gradient. The core idea is to replace the softmax parameter with Eq. (2), using cumulative reward plus $\gamma^d$ discounted parameter. Two methods are derived based on this idea, C-SoftTreeMax and E-SoftTreeMax, with the difference of taking expectations (over trajectories) inside or outside exponential in the softmax.

The major claim is that by doing parameterization this way. The variance of policy gradient can be largely reduced, decaying exponentially on the expansion depth $d$. The authors conducted experiments on Atari to show that the proposed methods achieve better performances as well as lower variance than PPO.

**Strengths**

The idea of combining policy gradient with tree expansion is novel to me (as also noted by reviewers). I believe this idea is reasonable and merits more investigations. The paper is also well written and the results are clearly presented.

**Weaknesses**

After reading, I considered that the major insight the authors wanted to tell the audience in this paper, i.e., "reducing variance exponentially", is not convincing, and is misleading.

Eq. (2) uses $\gamma^d$ as a scaling factor to the parameter $\theta$. The issue with that is the gradient w.r.t. $\theta$ itself is reduced exponentially by $\gamma^d$ as well. This makes the "variance reduction" claim not valid, since we wanted to preserve the scale of policy gradient and then reduce variances, rather than to make the whole gradient small.

I believe the reason for why the proposed methods work is something else, rather than reducing variance this way. From the paper, the authors seemed to realize that (the gradient norm is also reduced exponentially), while still attributed the improvements to variance reduction. This in total delivered a misleading message of "reducing the variance and improving policy gradient is the major contribution of this work".

**Comments**

That being said, I think the proposed methods work for some other deeper reasons (but not the variance reduction as explained above), and they are worth investigating. It could be the use of tree expansion guided faster exploration (to learn more important states), or it is the situation that the gradient here makes the parameter easier to be learned than the usual policy parameterizations (like in PPO). However, those are all my guesses and they require more efforts from the authors to confirm or disprove.

Reviews: review: paper_summary: paper_summary: This paper proposes a model based policy gradient (PG) algorithm that tries to decrease the variance of the gradient updates in PG methods and thus (hopefully) improving their sample complexity. The proposed approach works by modifying the softmax policy to work not just not with simple score functions (mapping states to real numbers), but rather with the truncated expected discounted return from a state. The resulting policy is called "SoftTreeMax". As such, this paper seems a little bit similar to decision time planning (or MCTS type) algorithms.

The paper proposes two variants of SoftTreeMax and analyzes their variance properties; in particular, the results show that as the size of truncation increases, the variance decreases exponentially. The paper also has a theorem which shows that SoftTreeMax can estimate the gradients well even with an approximate model. There are also some experiments on certain Atari games which show the efficacy of SoftTreeMax against the softmax policy using the popular PPO algorithm.
strength_weakness: strength_weakness: - The major weaknesses of this paper are not enough discussion on the exact mechanisms for implementing SoftTreeMax. So SoftTreeMax is akin to a policy parameterization (such as softmax). As such, when implementing an RL agent, we need a large number of other details (some details being how is the exploration policy chosen, what optimizer is used to update the gradients, how are the function approximators initialized, is there any processing of the observation features, etc.). These implementation details are important since they can affect the performance of the RL algorithms more than the core idea itself (for instance, recall RainbowDQN https://arxiv.org/abs/1710.02298 and PPO https://arxiv.org/abs/2005.12729).

- Can you guess how SoftTreeMax run on a continuous action space MDP, such as those from Mujoco?

- Does SoftTreeMax solve the problems of Softmax policies? I know it helps with variance, but what about softmax's dependence on the initialization? Or the possibly slow convergence of sotmax PG in MDPs due to an exponential dependence on the state space size (https://arxiv.org/abs/2102.11270)? Do you have any intuition? From the nice discussion in Appendix C.1, it seems that the answer is no. In particular, choosing the behavior policy while generating the tree could be quite important. Maybe some comments on how this policy is chosen in the implementation could be helpful (I saw the high-level comments about controlling for the eigenvalues of transition matrix; but those seem high-level ideas).

- What does SoftTreeMax really do? Is it just an MCTS like decision time planning algorithm that is helped a tiny bit by the theta values in Eq. 2? For instance, what if I do not update the theta values at all (i.e. set learning rate to zero). In that case, would the performance decrease drastically? Related to that, do you have any comments on how it relates with MCTS? I remember Bertsekas (http://web.mit.edu/dimitrib/www/LessonsfromAlphazero.pdf) said that as the depth increases, the utility / importance of the learned values (in this case theta values) in decision making decreases. Would that apply here as well?

- Once the policy is learned, how do you sample from it? I guess, you still need a model to sample from this policy. This seems like a clear disadvantage (although, this is true for all tree search based / decision time planning methods). Maybe mention this as a limitation.

line 53: "In all tested Atari games, our results outperform the baseline and obtain up to 5x more reward." --> this is not true? For instance, results on VideoPinball or Breakout show that PPO is pretty much as good. Also, the paper doesn't specify how the hyperparameters were chosen. So maybe PPO with a better choice of hyper-params might outperform SoftTreeMax?
questions: questions: The method seems to be highly computationally expensive. Thus, a major question is whether the convergence is sped up by the variance reduction? The paper claim that SoftTreeMax algorithm would result in a faster convergence than traditional policy gradient methods.

Now, it looks like Figure 4 (page 24) is talking about exactly that. (Consider moving it into the main paper? This seems like a very important experiment.)
Also, in Figure 4 caption "The x-axis is the wall-clock time. The runs ended after one week with varying number of time-steps.", what does 'varying number of steps' mean? Does this mean that the algorithms run for the same amount of time but different number of (non-simulated) environment interactions?


Minor comments (feel free to ignore)

line 22: "PG algorithms are also notoriously unstable due to the high variance of the gradients computed over entire trajectories" --> I don't think actor-critic type methods do that.

line 34: "for cumulative reward and one for exponentiated reward." --> reward is not exponentiated, right? The SoftTreeMax logits (Eq. 2) are exponentiated.

line 55: "demonstrating the relation between low gradient variance and high reward" --> seems like a strong statement? We don't know if low gradient is necessarily good, right? As Appendix C.1 says, this could also be sometimes harmful.

line 69-81: Nice preliminaries section! Also consider specifying the agent-MDP interaction: $s_0 \sim \nu, a_t \sim( \cdot | s_t), s_{t+1} \sim p( \cdot | s_t, a_t), r_{t+1} = r(s_t, r_t)$, etc..

line 76: define $P^{\pi'}(s' | s)$ as $P^{\pi'}(s, s')$ instead?

line 82: consider writing bold $\Theta \in \mathcal{R}^S$ as bold $\theta \in \mathcal{R}^S$ --- capital letters are often reserved for matrices.

line 84: consider replacing $D(u)$ with $\textrm{diag}(u)$ throughout the paper -- midway, I got confused by $D$.

line 127: "by showing that the gradient is strictly positive everywhere [Mei et al., 2020b, Lemmas 8-9]." --> Gradient being positive is unclear, since it is a vector. Also, I'm not sure if Lemmas 8 and 9 actually show this.

line 165: "where the gradient is scaled to reduce its variance," --> do you mean that gradient is clipped to reduce variance?

line 181: "the C-SoftTreeMax gradient is given" --> try using 'Jacobian' instead of 'gradient'

line 214-217: would the variance zero in this case?

Figure 2: Consider writing "I.e., instead of the" as --> "That is, instead of the"

Figure 254: what norm is being used?

line 254, Theorem 4.8: It is quite interesting that bound is unaffected by the the behavior policy and the tranistion dynamics $\mathcal{P}^\pi$. Any comments?

line 273: "SoftTreeMax with an exhaustive TS of depth d" --> Is the search really going through all possible state-actions? If not, consider dropping the work exhaustive.

line 282: maybe figure 2 should be moved somewhare closer.

lines 311-326: maybe put a colon, instead of a period, after the bolded letters. So do "true dynamics. Aversion Risk. Previous wor"
limitations: limitations: The major limitation seems to be that the paper doesn't provide

1.  a written algorithm - while I understand that this is almost straightforward, just putting a clear algorithm (specifying order of planning/simulation steps, how to take real world actions with SoftTreeMax policy, computing gradients, and updating the policy) would be highly useful.

2. The paper doesn't run experiments on continuous action space environments, which is arguably the more relevant domain for PG methods. (For discrete environments like Atari, we already have a plethora of value function + tree search based methods - like AlphaZero or MuZero.) Further, there are no hyper-parameter settings specified; it would be difficult to reproduce these results without looking at the codebase.


Further, a limitations section is missing
- if the PG is biased, would that affect convergence? In particular, note that Theorem 4.8 shows that the bias would scale with the state space size S (which is generally huge). Maybe setting d = \log(S) / \log(1 / gamma) helps with this and reduces this error to \log(S). But the computational complexity is probably exponential in d, so each gradient update becomes O(S), which is also bad..

Maybe a discussion between these factors would be helpful
---- how does the increase in bias affect the convergence speed (and the quality of the final policy obtained). What bias does the pruning approach introduce?
---- how does the decrease in the variance of gradient update affect the convergence speed (and the quality of the final policy obtained). Does the additional variance (in computing Lemma 4.6) introduced by pruning remove counteract the variance reduction (as given in Theorem 4.4)?
---- how does changing the bias / variance affect d and thus the computational complexity (edited)
paper_summary: paper_summary: This paper proposes a new policy parameterization called SoftTreeMax, which can reduce the variance of the stochastic policy gradient. The authors consider finite state and action spaces throughout the paper. They start by taking a softmax tabular policy and replacing the logit $\theta(s,a)$ with a score from a trajectory of horizon $d$ sampled from a behavior policy. They then address the normalization accordingly. Because of the randomness of the trajectory of size $d$, the authors take an expectation with respect to the randomness of the trajectory of size $d$ to well-define the policy parameterization. By taking the expectation before or after the exponent operator, the authors propose two policy parameterizations: C-SoftTreeMax and E-SoftTreeMax. It turns out that under either C-SoftTreeMax or E-SoftTreeMax parameterization, the variance of the stochastic gradient of the value function decays exponentially with the planning horizon $d$. The theoretical findings are well supported by simulations. Finally, the authors propose a parallel GPU-based simulator for the practical implementation of SoftTreeMax. By applying the C-SoftTreeMax parameterization to PPO and adapting $\theta(s)$ with a neural network, as well as using the parallel GPU-based implementation, leads to better performance compared to distributed PPO.
strength_weakness: strength_weakness: - Not really a weakness. Since the SoftTreeMax can be beneficial to different policy gradient methods, I expect to see its improvement applied on some (original) policy gradient methods for the experiments, such as REINFORCE and actor-critic, to demonstrate its generality, which is another strength of the paper.
- Since SoftTreeMax is inspired from tree search, the authors can highlight the difference between the "tree search" literature and the "tree expansion" used in this paper in order to avoid confusion. The tree search method uses a max to find the best policy, while here the authors use a fixed behavior policy.
- For the tree search literature, the authors can cite Efroni et al. [2018] and Protopapas et Barakat [2024]. In particular, Protopapas et Barakat [2024] combine a variant of PG, the policy mirror descent, with tree search.
- For the theoretical PG literature in Line 312, the authors can cite Yuan et al. [2022], where they provide a general analysis of PG, including softmax tabular parameterization as a special case.

Efroni, Y., Dalal, G., Scherrer, B., and Mannor, S. Beyond the one-step greedy approach in reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1387–1396. PMLR, 10–15 Jul 2018.

Kimon Protopapas, Anas Barakat (2024). Policy Mirror Descent with Lookahead.

Rui Yuan, Robert M. Gower, and Alessandro Lazaric. A general sample complexity analysis of vanilla policy gradient. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pages 3332–3380. PMLR, 28–30 Mar 2022.
questions: questions: No specific questions.
limitations: limitations: Yes.
paper_summary: paper_summary: The paper introduces a model-based online planning method called SoftTreeMax. The method acts as an extension of the softmax, by replacing the logit in softmax with a n-step return. Based on SoftTreeMax, the paper proposes two learning algorithms, C-SoftTreeMax and E-SoftTreeMax. The work includes a mathematical analysis of both algorithms to prove the policy gradient variance is bounded and exponentially decays when the planning horizon becomes longer. Then, the paper provides an empirical test on C-SoftTreeMax to compare the learned policy and the variance with PPO.
strength_weakness: strength_weakness: * The main concern I have is the difference between the proposed method and the n-step return. 

    In Section 3, the paper defines the SoftTreeMax logit as 

    $$l_{s,a}(d; \theta) = \gamma^{-d} \left[ \sum^{d-1}_{t=0} \gamma^t r_t + \gamma^d \theta(s_d) \right] .$$

    In the Actor-Critic method, it is normal to use $Q(s,a)$ for the scoring function $\theta$, and learn a policy $\pi(a|s) \propto exp(Q(s,a))$, for example, SAC, a commonly known online learning algorithm, with state-of-the-art performance. In this case, $l_{s,a}(d; \theta)$ can be rewritten to 

    $$ \gamma^{-d} \left[ \sum^{d-1}_{t=0} \gamma^{t-1} r_t + \gamma^d E_a Q(s_d, a_d) \right] .$$
 
    This equation is exactly the same as the equation in n-step TD, except the normalizer $\gamma^{-d}. However, the normalizer can be written as part of the temperature in softmax. 

    In C-SoftTreeMax, the learned policy is written as $\pi^C_{d,\theta} (a|s) \propto \exp [\beta E^{\pi_b} l_{s,a}(d; \theta)] $ in formula (3). 

    Replacing $l_{s,a}$ with the equation in (2), we have $\pi^C_{d,\theta} (a|s) \propto \exp [\beta E^{\pi_b} [\gamma^{-d}  \sum^{d-1}_{t=0} \gamma^t r_t + \gamma^d \theta(s_d) ]] $. 

     As the expectation is taken on $\pi_b$, the normalizer $\gamma^{-d}$ does not depend on $\pi_b$, thus could be moved out of the expectation. Then (3) is changed to $$ \pi^C_{d,\theta} (a|s) \propto \exp [(\beta\gamma^{-d}) \mathbb{E}^{\pi_b} [ \sum^{d-1}_{t=0} \gamma^t r_t + \gamma^d \theta(s_d)) ]]  .$$ 

     Now the temperature becomes $\beta\gamma^{-d}$, which is a tunable parameter, and the expectation term becomes n-step TD when using $Q(s,a)$ for $\theta$.
     Similarly, when using the state value $V(s) = \theta$, the normalizer can be written as part of the temperature, and the expectation term becomes n-step TD.
     In this case, using the proposed method seems to be no different from simply changing the value update from TD(0) to n-step TD. 

* The empirical test seems to be incomplete. The paper proposes two algorithms, C-SoftTreeMax and E-SoftTreeMax. But only C-SoftTreeMax was empirically tested. The paper explains the empirical test for E-SoftTreeMax was left for future work on risk-averse RL. But it does not make sense to me that the empirical test is completely omitted while the method is listed as a new algorithm in the paper. 

* There could be a stronger baseline. Given the similarity between SAC policy ($\pi(a|s) \propto exp(Q(s,a))$) and the policy learned by C-SoftTreeMax ($\pi(a|s) \propto exp(W_\theta(s,a))$), it may worth testing SAC, both 1-step TD and n-step TD in critic update.

* The empirical test could include more seeds. Currently there are only five seeds. The variance calculated based on 5 seeds could be highly inaccurate.

* The empirical test is limited to discrete action space. It is acceptable and understandable to leave continuous control experiments to future work, but still, adding continuous control results could make the paper more convincing.

* The new method introduces a new parameter, the planning length $d$. According to the empirical results (Figures 3, 4, and 5), different $d$ gave very different performances. The best setting is different across tasks. In 3 out of 8 tasks, (NameThisGame, Phoenix, and VideoPinball), longer planning hurts the performance. It is acceptable to say the performance is sensitive to the value of $d$ and leave this parameter as a tunable one, but it could be better to give an indication of how to pick this value when the reader would like to apply this method to a new task.


**Typo:**
* Line 307, “one week The” -> “one week. The”
questions: questions: * In Section 5, the paper says $\theta$ is replaced by a neural network $W_\theta$. It would be good if the authors could further explain what the target is when updating $W_\theta$, i.e. is $W_\theta$ a value function estimation or not?


* It would be better to explain the empirical setting in more detail. In Section 6, the paper explains the tree is pruned according to an estimated weight. It would be more clear if the paper could include a description of how the weights are updated.
limitations: limitations: Yes
paper_summary: paper_summary: Disclaimer: I do not have the mathematical background required to check all proofs. It is also my first time reviewing for NeurIPS. However, I have contributions in classical (deep) RL and finishing my PhD.

In this paper, authors use a novel policy parametrization for RL. Namely, the SoftTreeMax policy that replaces the traditional logit values by small horizon trajectories rewards values. The authors claim and prove that SoftTreeMax policy gradient has less variance than traditional policy gradient. They do experiments with PPO on Atari. 

I summarize my review as follows. The idea and work in this paper are novel, strong and well-motivated. However, it is in my opinion poorly presented, and poorly situated compared to existing related work. 

I very slightly lean towards accepting this paper as is but give feedbacks for the authors to increase their score in the following sections.
strength_weakness: strength_weakness: In the following comments, I assume that the techincal results are correct.
 
First weakness: clarity.

The proposed algorithm (PPO with SoftTreeMaxPolicy) makes heavy use of three major branches of reinforcement learning namely policy gradient methods (core of the paper), model-based RL (use of a forward model to compute small trajectories), efficient implementations of tree search (control the exponential cost of look-ahead search). Those components are all central to your work but the paper lacks a summary of how those components work together.  I recommend to summarize those in a schematic (see Questions section).  

Second weakness: baselines and related work.

The related work mentionned in this work does not clearly help to understand the significance of your work. There are well-studied tools reducing the variance in the gradient estimates (n-steps returns and gae-lambda), and some work combining TS and PG. It would have been nice to see comparisons with those work. I recommend the authors to rewrite and extend their related work section (see Questions section). 

REFERENCES.

TS + PG references:
[An Actor-Critic Algorithm Using a Binary Tree Action Selector, Action Guidance with MCTS for Deep Reinforcement Learning, Policy Gradient Algorithms with Monte-Carlo Tree Search for Non-Markov Decision Processes, Policy Gradient Search: Online Planning
and Expert Iteration without Search Trees]

Better sample efficiency and variance reduction:
[Schulman et al., 2015b; Mnih et al., 2016; Munos
et al., 2016; Schulman et al., 2017; Hessel et al., 2018;
Schrittwieser et al., 2020; Wurman et al., 2022; Chebotar
et al., 2023; Schwarzer et al., 2023, Brett Daley, Martha White, 2024]
questions: questions: Questions: 

- How do you learn the forward model? If it is not learned, how do you plan to use your algorithm for environments for which you have to learn the model?

- Is there an Atari game named "NameThisGame' or is this a typo in your plots ?

Suggestions:

I provide some suggestions to increase the score of the paper. I believe increasing my score by 1 for every three implemented suggestions seems fair.

- Could you please briefly conclude on if it is worth spending all this compute on tree expansion with respect to how much variance is reduced ? 

- Could you motivate looking at the variance only at the end of training of PPO (line 301)? 

- Could you please add a schematic of the forward pass and gradient accumulation in the fashion of fig 1 from [Fully-Convolutional Siamese Networks for Object Tracking] ? Basically combining your figure 2 with the text in lines 282-285 and lines 649-658 (appendix B1) in a summary figure OR in an algorithm description with pseudo code ?

- Could you please move to appendix results on one of C-SoftTreeMax or E-SoftTreeMax and move in the main paper the walltime plots and sample complexity plots (B2, B3)?

- Could please add a dedicated future work section with dedicated paragraphs on efficient tree expansion implems (better pruning, parallelism, gpu, ...) and on forward model learning ?

- Could you please have a dedicated related work section (you could again move some theory to the appendix) with dedicated paragraphs  on other variance reduction method in PG RL (you could have a look at the introduction of Averaging n-step Returns Reduces Variance in Reinforcement Learning 2024 ), on combining TS and PG and an other paragraph on GPU-based environment that have differentiable models ? 

Advice:

If I may, having also worked with expensive tree structures, I recommend you to limit your experiments to e.g. depths 2, 4, 6,8 rather than 2,3,4,5,6,7,8. 

Thank you
limitations: limitations: No limitations.

Comments: title: Thank you
Convinced and making sure the paper is accepted
Rebuttal follow-up
Rebuttal follow-up
Rebuttal follow-up
Reply
Sharing concern with reviewer m94G
comment: Thank you for your response. I have read it and I appreciate it.
Thank you so much for the kind words and insightful feedback! Your thorough reviews have been invaluable in improving our work, and we're committed to incorporating all these additions to enhance the paper's clarity and reach for the camera-ready version.
I believe this is great work. Really. So please keep your word on clarifying the discussed point in the camera-ready version of the paper. Make the presentation as clear as possible to reach as many readers as possible! I raised my score to 7. Great work.
Thank you for your reply and for raising you score in light of the discussion!

**The statistical significance of 5 seeds**

Thank you for raising this concern. This claim coincides with the two references shared by Reviewer d2zq and indeed highlights an important point regarding deep RL research. While using 5 seeds is a common practice in the field, we acknowledge the growing concerns about the statistical reliability of this approach. We will explicitly discuss this limitation in the paper, noting both the standard practice and the recent critiques. We will make an effort to include more seeds for the final version.
Thank you for your reply. 

**Question 1: Can you provide more insight as to why the performances of SoftTreeMax can decrease with d?**

This is a good observation. The deterioration is not inherent to the method, but due to experiments having a fixed training time of one week. The deeper the tree expansion, the slower each step and less steps are finished by the end of the run. Consequently, performance may decline beyond a certain depth d in these games, as the number of value updates becomes insufficient for proper training. Figure 5 illustrates this phenomenon: as depth increases, the number of online interactions (corresponding to value updates) decreases. However, when comparing performance at the same number of interactions (x-axis value), we observe a monotonic improvement with increasing d in almost all cases. This explanation clarifies that the method's effectiveness is not diminishing; rather, the fixed time constraint is limiting the learning process for deeper tree expansions in some games.

**Question 2: Why not compare an Actor-Critic baseline with n-step return estimation of the Critic to an Actor-Critic baseline with tree expansion in the Actor estimation? This would answer the question of where it is best to spend computation and use more information.**

We wish to first clarify that we in fact use a variant n-step TD in the critic: the PPO implementation from stable-baselines3 that we extend uses, by default, the GAE estimator in the critic with lambda=0.95, which is very similar n-step TD. We apologize for not making this clear earlier and will elaborate on this in the experiments section.

However, it's crucial to understand that SoftTreeMax is a new type of policy (actor) that can be combined with various types of critics (pure Monte-Carlo, TD(0), TD(N), GAE, etc). Our intention in the rebuttal was that one needs to compare actor to actor and critic to critic separately. Still, your question seems to approach this comparison from a computational angle.

In that aspect, the comparison between n-step return and SoftTreeMax is not straightforward, as they fundamentally differ in their approach and computational requirements. While n-step return samples one possible outcome, SoftTreeMax averages over all possible outcomes through tree expansion, providing more comprehensive information at a higher computational cost. A fair comparison might involve using tree expansion for value function estimation (like the concept of back-propagation in MCTS) or sampling multiple future trajectories for both methods. However, such a comparison, while interesting, is beyond the scope of our current work, which focuses on introducing and analyzing SoftTreeMax as a novel policy approach. We will clarify these points in the final version.
Thank you for your reply.

**Comment 1: Maximize -> take a gradient step**

Thank you for the suggestion, see the revised version here: https://ibb.co/L9Jfw9G

**Comment 2: Complicated diagram requires captioning**

We agree, here is the captioning we will add:

“A diagram of the tree expansion used by SoftTreeMax. In every step, the states in the current level of the tree are duplicated and concatenated with each possible action. The resulting state-action pairs are then fed as a batch to the GPU simulator to generate the next level of states. Finally, the states of the last level $d$ are inserted into the neural network $W_\theta$ and the logits are computed using the corresponding rewards along each trajectory.”

**Comment 3: Presentation of the results**

We understand your concern and will rephrase accordingly.

**Comment 4: Connection to n-step TD is useful**

Thank you, we will add a summary of the discussion with the reviewers in the rebuttal; please also see the follow up conversation with Reviewer mmxv.

**Comment 5: The statistical significance of 5 seeds**

Thank you for the references; they indeed highlight an important point regarding deep RL research. While using 5 seeds is a common practice in the field, we acknowledge the growing concerns about the statistical reliability of this approach. We will explicitly discuss this limitation in the paper, noting both the standard practice and the recent critiques. We will make an effort to include more seeds for the final version.


**Comment 6: Please add the results on E-SoftTreeMax to the paper as well**

Certainly, we will include these plots in the final revision.

**Comment 7: Is it possible to run the experiments for a learned model?**

We appreciate the suggestion to include experiments with a learned model. Given the time constraints before the final submission, implementing this would be challenging. However, we recognize its importance and will make our best effort to include preliminary results if time allows. If not feasible for this submission, we will explicitly address this as an important direction for future work in our paper.
I appreciate the detailed reply from the authors. The reply addressed my main concern about the novelty. Now, the difference between the proposed method and the n-step TD is clear to me. 

For the experiments, I am aware that many papers use 5 seeds, but usually, only a sample size of no less than 30 is considered sufficient for the central limit theorem to hold. Compared with 30, 5 is not a large number. Therefore, I still believe that 5 seeds are not good enough to make a conclusion on variance. Increasing the number of seeds is definitely making the conclusion more convincing. However, considering the limited time, I would not insist on asking authors to add more seeds to this paper.

After reading other reviews and the rebuttal, I agree with what reviewers d2zq, dZVs, and mmxv said about the novelty. The paper does provide a new aspect of policy gradient learning. Given the good novelty and concern about experiments (number of random seeds and baseline), I have increased my score accordingly.
I have read the rebuttal and I thank the author for their clarifications; they helped a lot. I also read the other reviews and the corresponding rebuttals. I think all of them raise relevant points. I think if the authors were to incorporate these suggestions, the quality of the paper would increase a lot. I retain my score of 7; it seems to be a fair evaluation of the work. 

Some additional / minor comments:

Related to my review:
------------------------------
Lines 12 and 13 in https://ibb.co/s2CRMp8 can't really be implemented (since the objective is highly non-convex in terms of the policy / value function weights). So maybe, put "take a gradient step" or something similar there instead of "maximizing".

The diagram https://ibb.co/5RPRkzF looks very nice, but because it is somewhat complicated, it needs a proper caption. I hope that will be added too.

> "Our results show significant improvement over PPO in all 8 games, including Breakout and VideoPinball. In Breakout, we reach a reward over 800 for depths 4 and 5, while PPO doesn't. In VideoPinball, we obtain 800K for depth 3 while vanilla PPO reaches less than 200K."

I still think that the results are very similar for these problems. I will suggest not writing "obtain up to 5x more reward"; it is not well supported by the experiments and seems like an oversell. Instead, just saying "competitive" or "outperforms" is enough.


Related to other reviewer's comments:
----------------------------------------------------
It seems like making the distinction between this work and n-step TD based actor-critic methods would be very useful. In particular, the nice and short analysis by m94G could be included and clarified.

The comment on five seeds not being enough is correct, and I don't agree with the author's responses. Now I'm not necessarily asking you to run more seeds, but at the very least, acknowledging that five seeds is probably not enough in the paper would be great! **(Also see Figure 5 of https://arxiv.org/pdf/1709.06560 and Figure 15 of https://arxiv.org/pdf/2304.01315).**

May I request the authors to include the experiments they already ran on the E-Softtreemax?

More importantly, is it possible to run the experiments for a learned model? I think having even just one experiment demonstrating that SoftTreeMax works with a learned model could be very encouraging for the readers.
Thank you for the rebuttal. I am still not convinced that the computations of tree expansions are worth it because:
- SoftTreeMax does not consistently imrpove on PPO.
- SoftTreeMax perf sometimes decreases with planning horizon $d$ which is weird
- I do not agree with what you answered to reviewer m94g :  "If instead we used SAC with n-step TD as a baseline, then the comparison will be with SAC with n-step TD and SoftTreeMax as a policy." I think it would be fair to compare PPO/SAC with n-step TD against PPO/SAC with SoftTreeMax only. The point being that it would be a fair way to compare the performance gain over PPO/SAC with two different types of added information that are shown to reduce variance respectively in the Critic and in the Actor.

Can you provide more insight as to why the performances of SoftTreeMax can decrease with $d$ as opposed to not increase?
Why not compare an Actor-Critic baseline with n-step return estimation of the Critic to an Actor-Crtic baseline with tree expansion in the Actor estimation? This would answer the question of where is it best to spend compute and use more information.

Respectfully