Policy Gradient with Tree Expansion
Authors: Gal Dalal
Assaf Hallak
Gugan Thoppe
Shie Mannor
Gal Chechik
Venue: ICML 2025 poster (2025)
Keywords: Reinforcement Learning
Policy Gradient
Tree Expansion
Softmax
Variance Reduction
Primary area: reinforcement_learning
TL;DR: SoftTreeMax: A novel parametric policy based on tree expansion that can be optimized using policy gradient methods

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

Policy Gradient with Tree Expansion
Gal Dalal Assaf Hallak" Gugan Thoppe
Shie Mannor
13 Gal Chechik
Abstract
Despite these positive features, PG algorithms are also no- toriously unstable due to the high variance of the gradients computed over entire trajectories (Liu et al., 2020; Xu et al., 2020). As result, PG algorithms tend to be highly ineffi- cient in terms of sample complexity. Several solutions have been proposed to mitigate the high variance issue, including baseline subtraction (Greensmith et al., 2004; Thomas Brunskill, 2017: Wu et al.. 2018), anchor-point averaging (Papini al., 2018), and other variance reduction techniques (Zhang et al., 2021; Shen et al., 2019; Pham et al., 2020).
Policy gradient methods are notorious for hav- ing large variance and high sample complexity. To mitigate this, introduce SoftTreeMax generalization of softmax that employs planning. In SoferreeMax, we extend the traditional logits with the multi-step discounted cumulative reward, topped with the logits of future states. We an- alyze Mof and explain how tree expan- sion helps to reduce its gradient variance. We prove that the variance depends the chosen tree-expansion policy. Specifically, we show that the closer the induced transitions are to being state-independent, the stronger the variance decay. With upproximate forward models, we prove that the resulting gradient bias diminishes with the approximation error while retaining the same vari- ance reduction. Ours is the first result to bound the gradient bias for approximate model. In practical implementation of SoffTreeMax, we utilize a parallel GPU-based simulator for fast and efficient tree expansion. Using this implementa- tion in Atari, we show that SoftTreeMax reduces the gradient variance by three orders of magni- tude. This leads better sample complexity and improved performance compared to distributed PPO.
A second family of algorithms that achieved statc-of-the-art results in several domains is based on planning. Planning is exercised primarily the context of value-based RL and is usually implemented using a Tree Search (TS) (Silver et al., 2016; Schrittwieser al., 2020). In this work, we combine PG with TS by introducing parameterized differ- entiable policy that incorporates tree expansion. Namely, our SoffTreeMax 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 rex, one for cumulative reward and one for exponentiated reward.
Unlike approaches that incorporate multi-step returns for value function estimation (e.g., n-step TD methods), our work explicitly integrates planning into the policy parame- terization itself. This distinction is crucial-wil turns serve as a Monte Carlo estimation technique for the ad- vantage function A in the gradient estimator E log x Al, our Soae affects the policy directly. This enables to obtain fundamentally different variance reduction prop- erties while keeping the policy gradient framework intact.
n-step re-
1. Introduction
Policy Gradient (PG) methods (Sutton et al., 1999) for Re- inforcement Learning (RL) are often the first choice for environments that allow numerous interactions at fast pace (Schulman al., 2017). Their success is attributed to sev- eral 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.
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 PG method based on exeaassnnnoon that is sta- ble and performs well in practice? and (ii) How does the tree-expansion policy affect the PG variance? Here, we analyze SoftTreeMax, and provide 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 probabili- ties are similar for each starting state. More generally, we show that the gradient variance of SoftTreeMax decays at a
'Equal contribution 'NVIDIA Research IIdiia Institute of Science 'Technion University "Bar-Ilan University. Correspon- dence to: Gal dalal @ Assaf Hallak <ahal- lak @ nom
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Policy Gradient with Tree Expansion
rate of A21l, where d is the depth of the tree and X is the second eigenvalue of the transition matrix induced by the tree expansion policy. This work is the first to prove such relation between PG variance and tree expansion policy. In addition, we prove that with an approximate forward model, the bias of the gradient is bounded proportionally to the approximation error of the model.
state distribution. We denote the transition matrix starting from state by P. E [0. 1|xxs i.e., [Pa]a.x' P(s' (iss|aa,1 s). Similarly, let R. = r(s,-) E EA denote the correspond- ing reward vector. Separately, let т : S + AA be sta- tionary policy. Let PM and RG be the induced transition matrix and reward function, respectively, i.e., PT (s'(s) = Ea T[aas) Pr(s'|s, a) and R=(s)) - En <(als)r(s, a). De- note the stationary distribution of p by Ae E R s.t. m p = P", and the discounted state visitation fre- quency by dx so that de = (1 - 7) Ei=o I'v TPT.. Also, let V E RR be the value function of defined by (T(s) = E= [_=== Y"r (st, (st)) so =s], and let QT E RSXA be the Q-function such that Q" a) = E [r(s, a) + +*V]())) Our goal is to find an optimal pol- icy x- such that V*(s) V (s) max, (T)),, Vs S.
To verify our results, we implemented a practical version of SoferreeMxx 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 SoffTreeMax GPU im- plementation 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.
For the analysis in Section 4, we introduce the following notation. Denote by Q E R the vector representation of 0(s) Vs E S. For a vector 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 (A)) Denote the k-dimensional iden- tity matrix and all-ones vector by L7 and 1g. respectively, Also, denote the trace operator by Tr. Finally, we treat all vectors as column vectors.
2.1. Policy Gradient
We summarize our key contributions. (i) We show how to combine two familics of SoTA approaches: PG and TS by introducing Max: 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 Saax in its two variants decays 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 & Szepesvári, 2006). (iii) We prove that with an approximate forward model, the gradient bias is proportional to the approx- imation 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 expan- sion. We demonstrate how its gradient variance is reduced by three orders of magnitude over PPO while obtaining up to 5x reward.
PG schemes seek to maximize the cumulative reward as function of the policy no (a(s)) by performing gradient steps on A. The celebrated Policy Gradient Theorem (Sutton et al., 1999) states that
2 V V = Esndd a log (" a)) 00
where and d are as defined above. The variance of the gradient is thus
Var Var aag (1s)) (Vo log (salss) (1)
In the notation above, we denote the variance of vector random variable X by
Var, (X) = Tr E. (xx) (x (X- (X-,x) Exx]]]
similarly as in (Greensmith et al., 2004). From now on, we drop the subseript from Var in (1) for brevity. When the action space is discrete, a commonly used parameterized policy is softmax: (a(ass) or exp (0(s.a)) where @ S x A -> R is state-action parameterization.
2. Preliminaries
Let AU denote simplex over the set U. Throughout, we con- sider a discounted Markov Decision Process (MDP) M = finite state space of size S,
3. Moxt Exponent of trajectories
(S, A, P.r. 7,:7,,) where A is finite action space of size A, r.::: S x A - [0, 1] is the reward function, P: S A + As is the transition function, y E (0, is the discount factor, and E g is the initial
We introduce new family of policies called SoftTreeMax, which are a model-based generalization of the popu- lar softmax. We propose two variants: Cumulative
2

---

Policy Gradient with Tree Expansion
(C=)of)TreeMax) and Exponentiated ME-SoffTreeMax). In both variants, we replace the generie softmax logits 0(s, a) with the score of trajectory of horizon d starting from (s, a), generated by applying a behavior policy Tb- In C-SoffTreeMax, we exponentiate the expectation of the log- its. In E-SoftrreeMax, we first exponentiate the logits and then only compute their expectation.
the best trajectory is defined in expectation, while for E-SoftTreeMax it is defined in terms of the best sample path.
SoeerreeMax behavior policy selection. The choice of behavior policy xb plays a crucial role in the performance of SoftTreeMax. As we show in Section 4, the gradient vari- ance minimized when the transitions induced by ne result in similar distributions across states, which is achieved when the second eigenvalue of p is small. Without specific as- sumptions on the MDP, a uniform policy that smoothens transition probabilities across all actions serves as a reason- able approximation to this goal. In practice, this leads to better mixing properties in the associated Markov chain.
Logits. We define the SofxTreeMax logit (e,a(d; 0) to be the random variable depicting the score of trajectory of horizon d starting from (s. a) and following the policy Ta:
Fd-1
0) = =
Σ Et + +70(()))
SoftTreeMax convergence. Under regularity conditions, for any parametric policy, PG converges local optima (Bhatnagar al., 2009), and thus also SofTTreeMax. For softmax PG, asymptotic (Agarwal et al., 2021) and rate re- sults (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 sat- isfies the same property, being generalization of softmax, but formally proving it is subject to future work.
8, au a. ae In the above expression, note that To(-19)) vI. 2 1, and rt = r (st, For brevity of the analysis, we let the parametric score 9 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 yd scales the 0 parametrization to correspond to its softmax counerpart.
TreeMax gradient. The two variants of TreeMrx involve an expectation taken over Std many trajectories from the root state 8 and weighted according to their probability. Thus, during the PG training process, the gradient Vo log Ta is calculated using a weighted sum of gradients over all reachable states starting from s. Our method exploits the exponential number of trajectories 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 n. 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 faster convergence and higher reward.
Tre, Given an inverse temperature parameter 3. we let C-SoftTreeMax be
Tg,o(a|s) x exp (s, 0).
(3)
C-SoffTreeMax gives higher weight to actions that result in higher expected returns. While standard softmax relies
entirely on parametrization 6, C-SoftTreeMax also interpo- lates a Monte-Carlo portion of the reward.
Tee... The second operator we propose is E-SoftTreeMux:
(als) x E exp ((3(s,a(d:))
(4)
here, the expectation is taken outside the exponent. This objective corresponds to the exponentiated reward objec- tive which is often used for risk-sensitive RL (Howard & Matheson, 1972; Fei et al., 2021; Noorani & Baras, 2021). The common risk-sensitive objective is of the form log E ())), where 5 is the risk parameter and R is the cumulative reward. Similarly to that literature, the exponent in (4) emphasizes the most promising trajectories.
4. Theoretical Analysis
In this section, we first bound the variance of PG when using the SoffTreeMax 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 with the tree depth, and the rate is determined by the second eigenvalue of the transition kemel induced by me 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
MoxxTreeMax properties. SoftTreeMax is natural model- based generalization of softmax. For d O, both variants above coincide since (2) becomes deterministic. In that case, for a state-action parametrization, they reduce to standard softmax. When 3 + O, both variants again coincide and sample actions uniformly (exploration). When 3 -> oo, the policies become deterministic and greedily optimize for the best trajectory (exploitation). For C-SoftTreeMax,
3

---

Policy Gradient with Tree Expansion
our implementation in Section 5, we manage avoid sam- pling 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 noiseless. We encourage future work to account for the finite-sample variance component. We defer all the proofs to Appendix A.
Lemma 4.2 (Vector form of C-SoffTreeMax). For d 2 1, (3) is given by
exp B (Cod + (P (0)
(5)
1 exp ( (Cad Ps (pto)a- =)]
where
We begin with general variance bound that holds for any parametric policy.
Ca.d= Ce.==-ARR + Ps
Σ bnd (pmoya-)
R
Lemma 4.1 (Bound on the policy gradient variance). Let
Na log (-Is) E (xxdim(0) be matrix whose a-th row is No log (als)" T. For any parametric policy and func- tion Q S x A A-R R,
The vector C..d E IR represents the cumulative discounted reward in expectation along the trajectory of horizon d. This trajectory starts at state 8; involves an initial reward dictated by R, and an initial transition as per Pa. Thereafter, it involves rewards and transitions specified by Rs, and pm respectively. Once the trajectory reaches depth d, the score function ((sa) is applied,.
Var (Vo log No n (als)) (s,a)) S
max [QT (s, ( a))² max Vo log no (1
Lemma 4.3 (Gradient of C-Sof(TreeMax). The C-SoftTreeMax gradient is given by
Hence, to bound (1), it is sufficient to bound the Frobenius norm || V log mo (-(+)s]] P for any 8.
Vo log C = =3 [B [T Ps (pro)d-1,
Note that SoftTreeMax does not reduce the gradient uni- formly, which would have been equivalent to a trivial change in the learning rate. While the gradient norm shrinks, the gradient itself scales differently along the different coor- dinates. This scaling occurs along different eigenvectors, as function of problem parameters (P, 0) and our choice of behavior policy (T&); as can be seen in the proof of the upcoming Theorem 4.4. This allows SoffTreeMax 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.
in AAXSS where for brevity, we drop the index in the policy above, ie., n0 = =
We are now ready to present our first main result:
Theorem 4.4 (Variance decay of C-SoftTreeMax). For ev- ery Q : S x A R, the C-SoftTreeMax policy gradient variance is bounded by
Var (Vo log (. ((. (s)
2 S²
(1 (1- 7)2 1a (pms) (2(d)))
common assumption in the RL literature (Szepesvári, 2010) that we adopt for the remainder of the section is that the transition matrix pms, induced by the behavior policy To, is irreducible and aperiodic. Consequently, its second highest eigenvalue satisfies (>2(P")) < 1.
We provide the full proof in Appendix A.4, and briefly outline its essence here.
Proof outline. Lemma 4. allows to bound the variance using direct bound on the gradient norm. The gradient is given in Lemma 4.3 as product of three matrices, which we now study from right to left. The matrix pm, is a row- stochastic matrix. Because the associated Markov chain irreducible and aperiodic, has a unique stationary distri- bution. This implies that dM has one and only one eigen- value 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, Lemmal)). Then, pm has the spectral decomposi- tion pd 1syT, + D Aavzu, where X the i-th
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-SoffTreeMMx more involved and we provide for it more general result. In both cases, we show that the variance decays exponentially with the planning horizon.
4.1. Variance of C-SoffTreeMax
We express C-SoffTreeMax in vector form follows.
4

---

Policy Gradient with Tree Expansion
eigenvalue of pad (ordered in descending order accord- ing to their magnitude) and u; and is are the correspond- ing left and right eigenvectors, respectively, and therefore (P"o)a-1= = Is + Ei=g " The second matrix in the gradient relation in Lemma 4.3, Ps, is rectangu- lar transition matrix that translates the vector of all ones from dimension 5 to : PaIs = 1a- Lastly, the first matrix [A - (=( is projection whose null-space
4.2. Variance of E-SoftTreeMax
The proof of the variance bound for E-SoftTreeMax is sim- ilar that of C-SoffTreeMax, but more involved. It also requires the assumption that the reward depends only 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 2 1, (4) is given by
includes the vector LA, i..., LA - (T LA 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*s) gets canceled, and we are left with bounded summands scaled by (e(Pr))d-1 Recalling that [A((P=)) < 1 and that [A2] 2 AAsl 2 i ,,...,.. S, we obtain the desired result.
x=)0(11s== = E exp(3)) 1) Ea,d exp(30))
(6)
where
d-1
E.d P.
II hmm1
( (exp(8>> pro).
The vector R above is the S-dimensional vector whose s-th coordinate is r(s)-
It's important to note that Sofx does not reduce the gradient uniformly, which would be equivalent to simply decreasing the learning rate. Rather, the gradient is scaled differently along different cigenvectors, with scaling fac- tors that depend on the MDP structure and behavior policy T6- This non-uniform scaling allows the policy to continue learning effectively while reducing harmful variance. The empirical results in Section 6 demonstrate that this variance reduction leads to faster, more stable convergence rather than slowing it down.
The matrix Es,d E RASS has similar role to C..d from (5), but it represents the exponentiated cumulative discounted reward. Accordingly, is product of d matrices as opposed to a sum. It captures the expected reward sequence starting from and then iteratively following p=s After d steps, we apply the score function on the last state as in (6).
Lemma 4.6 (Gradient of E-SoftTreeMax). The E-SoffTreeMax gradient is given by
Theorem 4.4 guarantees that the variance of the gradient decays with d. More importantly, it also provides a novel insight for choosing the behavior policy ne as the policy that minimizes the absolute second eigenvalue of the pr. Indeed, the second eigenvalue of Markov chain relates to its connectivity and its rate of convergence to the stationary distribution (Levin& & Peres, 2017).
V9 log Tē.0 3 -1 (ccx((())(33) Td.0) 1 exp((e)
=1111 - 1A(
E AXS
where for brevity, drop the index in the policy above, i.e., n = (1)))
Optimal variance decay. For the strongest reduction in variance, the behavior policy Tr should be chosen to achieve an induced Markov chain whose transitions state-independent. In that case, pa is a rank one matrix of the form IsHT, and (p(P)) = 0. Then, Var (V log (e (ss a)) O. Naturally, this can only be done for pathological MDPs; see Appendix C.1 for a more detailed discussion. Nevertheless, as we show in Sec- tion 5, we choose our tree expansion policy to reduce the
This gradient structure is harder to handle than that of C-SofTrreeax in Lemma 4.3, but here we also can bound the decay of the variance nonetheless.
Theorem 4.7 (Variance decay of E-SoffTreeMax). There exists a E (0, 1) such that,
Var (Vo log v (alaaaa)))ss)))) E E (Baa2))
for every Q. Further, if p is reversible or if the reward is constant, then = (P
variance as best as possible.
Worst-case variance decay. In contrast, and somewhat surprisingly, when nú is chosen so that the dynamics is deterministic, there is no guarantee that it will decay expo- nentially fast. For example, if pa is permutation matrix, then 2₂ (P*)) = 1, and advancing the tree amounts to only updating the gradient of one state for every action, as in the basic softmax.
Theory versus Practice. We demonstrate the above result in simulation. We draw random finite MDP, parameter vector θ E R and behavior policy Tb. We then empirically compute the PG variance of E-SoftTreeMax as given in (1) and compare it to (A2(P**) We repeat this experiment three times for different ps, : (i) close to uniform, (ii) drawn
5

---

Policy Gradient with Tree Expansion
gradient bias due T0.0 satisfies
10a
1 00
(
(
V (- )
10-4
00
SoftTreeMax Gradient variance
10-10
7)- 722 ST8d
(1 1
1
10-16
Permutation: Empirical variance
10-22
Permutation: Variance bound
Proof outline. First, prove that max R2 -R,11, R2,l, IP, Pall: RT, Rr, ll, || - pm ||} = (()) This follows from the fact that the differences above are suitable convex combinations of either the rows of P - P or r - f. We use the above observation along with the definitions of n2.0 and #L given in (5) to show that ||= - * - = ((8dc). The proof for the latter builds upon two key facts: (a) [1(PTo) )
Random: Empirical variance
10-28
Random: Variance bound
10-34
Uniform: Empirical variance
Uniform: Variance bound
2
10
Depth
Figure/ 1. comparison of the empirical PG variance and our bound for E-SoftTreeMai on randomly drawn MDPs. We present three cases for pT (i) close to uniform, (ii) drawn randomly, and (iii) close to a permutation matrix. This exper- iment verifies the optimal and worse-case rate decay cases. The variance bounds here are taken from Theorem 4.7 where we substitute a = ()*))))) To account for the constants, we match the values for the first point in d == 1.
(P=o)* I Σm || | 1 |m - p "1lp" || b
(ke) for any & 2 O, and (b) |ee - 1| = O(r) as x - O.
Next, we decompose the LHS of (7) to get
Σ
II
(s)
II
{:(())
i=1
4
Σ X1(s)
X1-1(s) (Xi(s) - {((())
X1+1(8) ... ({()),))())))
randomly, and (iii) close to permutation matrix. As seen in Figure 1, the empirical variance and our bound match almost identically. This also suggests that a = ( [ in the general case and not only when pm is reversible or when the reward is constant.
where X1(s) drg, (0) E R, X2(0)
(Vo log T (-(s)) RSXA Xz(s) = D(= (T..0s.))
RAXA X1(s)
RAXA and
(1(s), ................................ Xi(s) are similarly defined with Tg. re- placed by &5 Then, we show that, for i = I.... 4, (i)
4.3. Bias with an Approximate Forward Model
[X((s) X,) = O(e) and (ii) max |
is
The definition of the two SoftTreeMax variants involves the knowledge of the underlying environment, in particular the value of P and r. However, practice, we often can only learn approximations of the dynamics from interactions, e.g., using NNs (Ha & Schmidhuber, 2018; Schrittwieser et al., 2020). Let P and f denote the approximate ker- nel and reward functions, respectively. In this section, we study the consequences of the approximation error on the C-SoftTreeMax gradient.
bounded by problem parameters. From this, the desired result follows.
To the best of our knowledge, Theorem 4.8 is the first re- sult that bounds the bias of the gradient of a parametric policy due to an approximate model. This thcorem reveals an intriguing trade-off in SoftTreeMax: while the variance decays exponentially with tree depth d as ((())(P) ((-)) according to Theorem 4.4, the gradient bias with an ap- proximate model grows O(dy-4) Since 7 grows exponentially with d, there exists an optimal tree depth that balances these opposing effects.
Let #G be the C-SoftTreeMax policy defined given the approximate forward model introduced above. That is, let T0 be defined exactly in (5), but using R., Ps, Ras and pT,, instead of their unperturbed counterparts from Section 2. Then, the variance of the corresponding gradient again decays exponentially with a decay rate of X² (mm). However, a gradient bias 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.
The bias also depends on the temperature parameter B, with higher temperature (lower (3) reducing bias at the expense of more exploratory policies. In the extreme case of 33 O. the policy becomes uniform with no bias. Additionally, the error scales linearly with d because the policy suffers from cumulative error as it relies on further-looking states in the approximate model.
Theorem 4.8. Let e be the maximal model mis- speciffifii,,,nn Le., let max P|], ||r i|l} 6. Then the policy
These results suggest that if the learned model is accurate enough, we can expect similar convergence properties for
6

---

Policy Gradient with Tree Expansion
MaSoffTreeMax as we would obtain with the true dynamics. In practice, one should adjust both d and /3 based on the estimated accuracy of the forward model to achieve the best performance. This is particularly important for practitioners implementing SofeTreeMxx with learned forward models.
During training, the gradient propagates to the NN weights of Wo- When the gradient Vo log 74,0 is calculated each time step, it updates We for all leaf states, similarly to Siamese networks (Bertinetto et al., 2016). An illustration of the policy is given in Figure 2.
5. Maxx Deep Parallel Implementation
6. Experiments
We conduct our experiments on multiple games from the Atari simulation suite (Bellemare et al., 2013). As 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-firs as in (Dalal et al., 2021). We then train C-SofCTreeMax for depths d = 1...8 8, with a single worker. For depths d 2 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 fair comparison, we ran all experiments for one week on the same machine. For more details see Appendix B.
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, loss function 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 Wo: (a(ss) x exp (Wo(s,)) a)). Similarly, we implement SoftTreeMax by replacing ((s) in (2) with a neural network (s()). Although both variants of SoffTreeMax from Sec- tion 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.
In Figure 3, we plot the reward and variance of Max for each game, as function of depth. The dashed lines are the results for PPO. Each value is taken after conver- gence, 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 function of depth. For more thorough discussion on the sweet spot of different games, see Appendix B.5.
Instead, in finite action space environments such as Atari, we compute the exact expectation in Soft with an exhaustive TS of depth d. Despite the exponential computa- tional cost of spanning the entire tree, recent advancements in parallel GPU-based simulation allow efficient expansion of all nodes at the same depth simultancously (Dalal al., 2021; Rosenberg et al., 2022). This is possible when a simu- lator is implemented on GPU (Dalton et al., 2020; Makoviy- chuk et al., 2021; Freeman et al., 2021 1), or when a forward model is learned (Kim et al., 2020; Ha & Schmidhuber, 2018). To reduce the complexity to be linear in depth, we apply tree pruning to 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 small sacrifice in performance.
To summarize, in the practical SoftTreeMax algorithm we perform an exhaustive tree expansion with pruming to obtain trajectories up depth d. We expand the tree with equal weight to all actions, which corresponds to uniform tree expansion policy T&. 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 T-eeft We leave experiments E-SoftTreeMax for future work on risk-averse RL.
7. Related Work
Reducing variance in PG estimates is essential for improving efficiency and stability. Approaches include baseline sub- traction (Sutton et al., 1999; Weaver & Tao, 2001), action- dependent baselines (Wu et al., 2018), sub-sampling tech- niques like SVRPG (Papini et al., 2018), and natural policy
We also experimented with TreeMax and the results were almost identical. This is due to the quasi-determinittic nature of Atari, which causes the trajectory logits (2) to have almost no variability. We encourage future work on E-SofTreeeMa using probabilistic environments that are risk-
For detailed illustration of our GPU-based tree expansion implementation, see Appendix B.2. In addition, we provide the psudeocode for C-SefxTreeMax in Algorithm 1 and its integration with PPO in Algorithm 2 in Appendix B.3.
7

---

Policy Gradient with Tree Expansion
C1.13 2-2
Palicy
(S
(1))
(4)
logits for di
4:-0
[=]
4
(LA)
((
12.13
(
2-2
S
2-1
(2)
(A)
lngits for non
12,4
((A)
(A,1)
t=2
W(S
logits for 46m0
(A)
dr=1
(A.A)
Figure 2. SofxTreeMax policy. Our exhaustive parallel expansion iterates all actions each up depth (= 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 's cumulative reward as described in (2). I.c., instead of the standard softmax logits, add the cumulative discounted reward 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.
Asteroids
Breakout
Gopher
Gradient variance
5000
800
15000
10-5
9000
flog scale)
Reward
10-1
4000
10000
3000
5000
10
7000
10
SoftTreeMax
PPO Reward
Depth
Depth
Depth
Depth
SoftTreeMax Variance
PPO Variance
Kuasterr
NameTh
Phoenix
VideoPinball
Gradient variance
20000
75000
10
(log scale)
Reward
60000
15000
50000
600000
10-6i
400000
40000
10000
10-7 25000
10
200000
Depth
Depth
Depth
Depth
Figure 3. Reward and Gradient variance: GPU Maxe (single worker) PPO (256 GPU workers). The blue reward plots show the average of 50 evaluation episodes. The red variance plots show the average gradient variance of the corresponding training runs, averaged over five seeds. The dashed lines represent the same PPO. Note that the variance y-axis in log-scale.
gradient (Kakade, 2001). Multi-step returns for value esti- mation, such as n-step TD (Sutton et al., 1998) and GAE (Schulman et al., 2015), primarily affect learning targets rather than policy parameterization. Our approach differs by integrating planning directly into the policy structure, mak- ing it fundamentally different from both variance reduction techniques and multi-step return methods.
et al., 2021: Silva et al., 2019). It's important to distinguish between tree search and our approach of tree expansion. Traditional tree search methods like MCTS identify the best trajectory through selection process, while approach explores all possible trajectories up certain depth to compute improved policy. Recent work by Efroni et al. (2018) showed that multi-step greedy policies improve con- vergence rates, and Protopapas & Barakat (2024) combined policy mirror descent with lookahead planning, though us- ing different mechanisms than our SoftTreeMax policy.
Softmax Operator.
The softmax policy became a canonical part of PG the point where theoretical results of PG focus specifically it (Zhang al., 2021; Mei al., 2020b; Li et al., 2021; Ding et al., 2022). Even though we focus tree extension to the softmax policy, our methodology is general and can be easily applied to other discrete or contin- uous parameterized policies as in (Mei et al., 2020a; Miahi
Parallel Environments. In this work used accurate parallel models that are becoming more common with the increasing popularity of GPU-based simulation (Makoviy- chuk et al., 2021; Dalton et al., 2020; Freeman et al., 2021).
8

---

Policy Gradient with Tree Expansion
Alternatively, in relation to Theorem 4.8, one can rely on recent works that learn the underlying model (Ha & Schmid- huber, 2018; Schrittwieser et al., 2020) and use approxi- mation of the true dynamics, Risk Aversion. Previous work considered exponential utility functions for risk aversion (Chen et al., 2007; Garcia & Fernández, 2015; Fei et al., 2021). This utility function is the same as E-SoffTreMaax 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 Th are uniform. This is ex- pressed by the second eigenvalue of the transition matrix pT.. 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 & Brunskill, 2018; Tarbouriech & Lazaric, 2019).
based approaches like MuZero and policy gradient methods. Additionally, adapting the behavior policy H, time as learning progresses could lead more relevant exploration, though this requires careful management to maintain vari- ance reduction benefits. Finally, the E-SoftTieeMax variant offers potential applications in risk-sensitive RL by natu- rally emphasizing exceptional trajectories through its reward exponentiation structure.
A Acknowledgements
Gugan Thoppe's 's research is supported by the Walmart Cen- tre for Tech Excellence at IISc, the Indo-French Centre for the Promotion of Advanced Besearch-CEFFiPe (7102-1), DST SERB's Core Research Grant CRG/2021/0083300 and the Pratiksha Trust Young Investigator Award.
8. Diseussion and Future Work
Impact Statement
In this work, we introduced for the first time a differentiable parametric policy that combines TS with PG. We proved that Sof is essentially a variance reduction technique and explained how to choose the expansion policy mini- mize 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 & Szepesvári, 2006). We believe that this can be answered by analyzing the convergence rate of TreeMx,, relying on the bias and variance results we obtained here.
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential socictal consequences of our work, none which we feel must be specifically highlighted here.
Reproducibility and Limitations
The code for our implementation is available at https: //g1 t.ub.. com /omm//1la /VN / /SoftTreeMax We provide 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), 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.
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 represents deterministic learning, which is fastest).
There are additional limitations to consider. First, our anal- ysis focuses on MDPs with finite state and action spaces, and while we discuss theoretical extensions to continuous domains, practical implementations would require careful design choices to manage the sampling process efficiently. Second, while SofeTreeMax can be applied to any environ- ment, it provides most benefit in quasi-deterministic envi- ronments where accurate forward planning is possible. In highly stochastic environments, the exponential growth of the tree width with depth would present challenges even with pruning.
Our work can be extended to infinite action spaces, where the theoretical analysis would involve transition ker- nels rather than matrices while preserving the same non- expansive operator properties. For implementation, the tree of continuous actions can be expanded by maintaining a parametric distribution over actions depending on 6, similar to tree adaptation of MPPI (Williams et al., 2017) with a value function. This would significantly broaden the appli- cability of Max to important domains like robotics and continuous control.
Further important extensions include learning the forward model from interactions (Ha & Schmidhuber, 2018; Schrit- twieser et al., 2020) rather than using the true dynamics. Our analysis in Theorem 4.8 already provides theoretical guid- ance on the resulting gradient bias, but empirical validation with learned models would bridge the gap between model-
References
Agarwal, A., Kakade, S. M., Lee, J. D., and Mahajan, G. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. J. Mach. Learn. Res., 22 2021.
9

---

Policy Gradient with Tree Expansion
Bellemare, M. G., Naddaf, Y., Veness, J., and Bowling, M. The arcade learning environment: An evaluation plat- form for general agents. Journal of Artificial Intelligence Research, 47:253-279, 2013.
learning- Journal of Machine Learning Research, 5(9), 2004.
Ha, D. and Schmidhuber, World models. arXiv preprint arXiv: 10122, 2018.
Bertinetto, L., Valmadre, J., Henriques, F., Vedaldi, A., and Torr, P. H. Fully-convolutional siamese networks for object tracking. In European conference on computer vision, pp. 850-865. Springer, 2016.
Howard, R. A. and Matheson, E. Risk-sensitive markov decision processes. Management science, 18(7):356-369 1972.
Jin, C., Krishnamurthy, A., Simchowitz, M., and Yu, T. Reward- exploration for reinforcement learning. In International Conference on Machine Learning, pp. 4870- 4879. PMLR, 2020.
Bhatnagar, S., Sutton, R. S., Ghavamzadeh, M., and Lee, Natural actor-critic algorithms. Automatica, 45(11): 2471-2482, 2009.
Chatterjee. S. and Seneta, E. Towards consensus: Some convergence theorems on repeated averaging. Journal of Applied Probability, 14(1):89-99) 97, 1977.
Kakade, S. M. A natural policy gradient. Advances in neural information processing systems, 14, 2001.
Kim, S. W., Zhou, Y., Philion, J., Torralba, A., and Fi- dler, S. Learning simulate dynamic environments with gamegan. In Proceedings of the IEEF/CVF Conference Computer Vision and Pattern Recognition, pp. 1231- 1240, 2020.
Chen, X., Sim, M., Simchi-Levi, D., and Sun, P. Risk aversion in inventory management. Operations Research, 55(5):828-842, 2007.
Dalal, G., Hallak, A., Dalton, S., Mannor, S., Chechik, G., et al. Improve agents without retraining: Parallel tree search with off- offlpolliy correction. Advances in Neural Information Processing Systems, 34:5518- 8-5530, 2021.
Kocsis, L. and Szepesvári, C. Bandit based monte-carlo planning. In European conference on machine learning, pp. 282-293. Springer, 2006.
Dalton, S. et al. Accelcrating reinforcement learning through gpu atari emulation. Advances in Neural In- formation Processing Systems, 33:19773-19782, 2020.
Levin, D. A. and Peres, Y. Markov chains and mixing times, volume 107. American Mathematical Soc., 2017.
Li, G., Wei, Y., Chi, Y., Gu, Y., and Chen, Y. Softmax policy gradient methods can take exponential time to converge. In Conference on Learning Theory, pp. 3107- 3110. PMLR, 2021.
Ding, Y., Zhang, J., and Lavaei, J. On the global optimum convergence of momentum-based policy gradient. In International Conference on Artificial Intelligence and Statistics, pp. 1910-1934, PMLR, 2022.
Liu, Y. and Brunskill, E. When simple exploration is sample efficient: Identifying sufficient conditions for random exploration to yield pac rl algorithms. arXiv preprint arXiv: fath::880009944 2018.
Efroni, Y., Dalal, G., Scherrer, B., and Mannor, S. Beyond the one-step greedy approach in reinforcement learning. In International Conference on Machine Learning, pp- 1387-1396. PMLR, 2018.
Liu, Y., Zhang, K., Basar, T., and Yin, W. An improved analysis of (yariance-reducce)) policy gradient and natural policy gradient methods. Advances in Neural Information Prucessing Systems, 33:7624-7636, 2020.
Fei, Y., Yang. Z.., Chen, Y., and Wang, Z. Exponential bellman equation and improved regret bounds for risk- sensitive reinforcement learning. Advances in Neural Information Processing Systems, 34:20436-2044 2021.
Makoviychuk, V., Wawrzyniak, L., Guo, Y., Lu, M., Storey, K., Macklin, M., Hoeller, D., Rudin, N., Allshire, A., Handa, A., al. Isaac gym: High performance gpu-based physics simulation for robot learning. arXiv preprint arXiv:2108. 10470, 2021.
Freeman, D., Frey, E.,, Raichuk, A., Girgin, S., Mordatch, I,, and Bachem, O. Brax-a differentiable physics gine for large scale rigid body simulation. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021.
Mathkar, A. S. and Borkar, V. Nonlinear gossip. SIAM Journal on Control and Optimization, 54(3):1535-1557, 2016.
Garcia, J. and Fernández, F. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research,
Mei, J., Xiao, C., Dai, B., Li, L., Szepesvári, C., and Schu- urmans, D. Escaping the gravitational pull of softmax. Advances in Neural Information Processing Systems, 33: 21130-21140 2110, 2020a.
2015.
Greensmith, E., Bartlett, P. L., and Baxter, J. Variance reduc- tion techniques for gradient estimates reinforcement
10

---

Policy Gradient with Tree Expansion
Mei, J., Xiao, C., Szepesvari, C., and Schuurmans, D. On the global convergence rates of softmax policy gradient methods. In International Conference on Machine Learn- ing, Pp. 2820-6829) PMLR, 2020b.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and
Klimov, O. Proximal policy optimization algorithms. 2017.
arXiv preprint
Shen, Z., Ribeiro, A., Hassani, H., Qian, H., and Mi, C. Hes- sian aided policy gradient. In International conference on machine learning, pp. 5729-5736 PMLR, 2019.
Miahi, E., MacQueen, R., Ayoub, A., Masoumzadeh, A., and White, M. Resmax: An alternative soft-greedy oper- ator for reinforcement learning. 2021.
Silva, A., Killian, T., Rodriguez, I. J., Son, S.-H., and Gombolay, M. Optimization methods for interpretable differentiable decision trees in reinforcement learning. arXiv preprint arXiv: 1903:i::/000.000 993399338,,, 2019.
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A. A., Veness, J.. Bellemare, M. G., Graves, A., Riedmiller, M., Fidje- land, A. K., Ostrovski, G., et al. Human-level control through deep reinforcement learning. nature, 518(7540): 029-533,222 2015.
Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvame V., Lanetot, M., et Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
Noorani, and Baras, J. Risk-sensitive reinforce: A monte carlo policy gradient algorithm for exponential performance criteria. In 2021 60th IEEE Conference Decision and Control (CDC), pp. 1522-1527. IEEE, 2021.
Sutton, R. S., Barto, A. G., et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
Papini, M., Binaghi, D., Canonaco, G., Pirotta, M., and Restelli, M. Stochastic variance-reduced policy gradi- ent. In International conference on machine learning, pp- 4026-4035. PMLR, 2018.
Sutton, R. S., McAllester, D., Singh, S., and Mansour, Y.
Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Pelletier, M. On the almost sure asymptotic behaviour of stochastic algorithms. Stochastic processes and their applications, 78(2):217-244, 1998.
Szepesvári, C. Algorithms for reinforcement learning. Syn- thesis lectures artificial intelligence and machine learning, 4(1):1-103, 2010.
Pham, N., Nguyen, L., Phan, D., Nguyen, P. H., Dijk, M., and Tran-Dinh, Q. hybrid stochastic policy gradient algorithm for reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pp. 374-385. PMLR, 2020.
Tarbouriech, and Lazaric, A. Active exploration in markov decision processes. The 22nd International Confer- ence on Artificial Intelligence and Statistics, pp. 974-982. PMLR, 2019.
Protopapas, K. and Barakat, A. Policy mirror descent with lookahead. Advances in Neural Information Processing Systems, 37:26443 264811 2024.
Thomas, P. S. and Brunskill, E. Policy gradient meth- ods for reinforcement learning with function approxi- mation and action-dependent baselines. arXiv preprint arXiv: 17060663636, 2017.
Raffin, A., Hill, A., Ernestus, M., Gleave, A., Kanervisto, A., and Dormann, N. Stable baselines3, 2019.
Weaver, L. and Tao, N. The optimal reward baseline for gradient-based reinforcement learning. In Proceedings of the Seventeenth conference on Uncertainty in artificial intelligence, pp. 538-545, 2001.
Rosenberg, A., Hallak, A., Mannor, S., Chechik, G., and Dalal, G. Planning and learning with adaptive lookahead. arXiv preprint 220iv:2201.1 10:00...24.0 12403, 2022.
Williams, G., Wagener, N., Goldfain, B., Drews, P., Rehg, J. M., Boots, B., and Theodorou, E. A. Information theoretic mpe for model-based reinforcement learning. In 2017 IEEE International Conference on Robotics and Automation (ICRA), pp. 1714-1721. IEEE, 2017.
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, I.., Schmitt, S., Guez, A., Lockhart, E., Hassabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839): 604-609, 2020.
Wu, C., Rajeswaran, A., Duan, Y., Kumar, V., Bayen, A. M., Kakade, S., Mordatch, I., and Abbeel, P. Variance reduc- tion for policy gradient with action deeenddde factorized basclines. In International Conference on Learning Rep- resentations, 2018.
Schulman, J., Moritz, P., Levine, S., Jordan, M., and Abbeel, P. High-dimensional continuous control using generalized advantage estimation. arXiv preprint 176.:21506.02438, 2015.

---

Policy Gradient with Tree Expansion
Xu, P., Gao, F., and Gu, Q. An improved convergence analysis of stochastic variance-reduced policy gradient. In Uncertainty in Artificial Intelligence, pp. 541-551. PMLR. 2020.
Zhang, J., He, T., Sra, S., and Jadbabaie, A. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv preprint 0rXiv:1905.11880 2019.
Zhang. J., Ni, C., Szepesvari, C., Wang. M., al. On the convergence and sample efficiency variance-reduced policy gradient method. Advances in Neural Information Processing Systems, 34:2228-2242 2021.
12

---

Policy Gradient with Tree Expansion
Appendix
A. Proofs
A.1. Proof of Lemma 4.1 Bound on the policy gradient variance
For any parametric policy ne and function Q S x A - R
Var (Vo log (a(a)) a)) 2 max [Q(s, max IVo log To (+1)))
8,0
where V log (T(-)s) E RAxxim(0) is a matrix whose a-th row Vo =
Proof. The variance for parametric policy Te is given as follows:
Var (Vo log no (als)Q(a, s)) = -R goddng (-1e)
Vo log mo(als) Vo log To (als))((s (s)
E Banday Bo (-1))))1))
V log (o(a)s)((sss a)] Banday-amall Vo log (o(als)((s.aa) a)];
where (8, a) is the currently estimated Q-function and dm is the discounted state visitation frequency induced by the policy ne. Since the second term we subtract is always positive (it is of quadratic form v v) we can bound the variance by the first term:
Var (V log (o)loo/ s))
SE Endd a [v log (a(ss) Vo V (ogg
Σ da (8) Σ log TO (als)
Vo (allm
Σ max [1Q(s, a)) no(als)
Σ di (s)
Vo (als)
Vo log (a((s)
max
[Q(s, a) max V log (als))
Vo log
8,0
max
(Q(s, a)} max V (+1s)1
A.2. Proof of Lemma 4.2 Vector form of C-SoffTreeMax In vector form, (3) is given
3 Co.d + P. (pxs) (po))d-
exp
(-1s)
(8)
1 exp
Ps (Pao) P
where
d=1
C., = = RR, + Pa
Σn h-d (proyh-
Rev
(9)
Proof. Consider the vector laf- E IRIAA Its expectation satisfies
d=1 1
E L (d:
T
Σ
r-d rx + 0(ss)
d-1
R2 +
r
P. (P
Rr. + PP(
As required.
13

---

Policy Gradient with Tree Expansion
A.3. Proof of Lemma 4.3 Gradient of C-SoftTreeMax
The C-SoftTreeMax gradient of dimension A S is given by
Va log d.g - 3 [L - [n(F)))) I Ps (Pma)a-)
where for brevity, we drop the 8 index in the policy above, i.e., ~6,0 = (1)()
Proof. The (j, k) entry of Vo log og,0 satisifes
Vo log -5ppi
a0(s*)
En exp B Ca.a + Ps( (pmm)d-1
[P. (pos)d 0.k
= (Pre dd 1in
I exp
d-1 8) B Ca.d + Ps (P (P
= -|Ps(
C (ap(( [Pa (P Ja. 0.k 13.0 B Σ
- 13 [(F5,)) PP(PFU)dd)
= -P 1ز. [LA(T T Po (Pas)d-1]
j.k"
Moving back to matrix form, we obtain the stated result.
A.4. Proof of Theorem 4.4 Exponential variance decay of C-Sof The C-SoffTreeMax policy gradient is bounded by
3252 (1- (
Var (Vo log T. (ss,(()))((ssa) 2
Proof. We use Lemma 4.1 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 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-decompositionn on p
pr 1sp
Σ Nuae,
(10)
where y is the stationary distribution of Pra, and u; and vs are left and right eigenvectors correspondingly.
S
||3 (LA.- LAT Pa ( = B1l (LA.A - lAT P.
1su
+ +
x
dd
uuu
l|f
i=2
(Pa is stochastic)
= A1l (LA.A - LAT
Lau + Σ
x
Puuny
lle
1=2
(projection nullifies LAH²)
= B1| (LA.A - 1AT
Σ 1p- P.u., llf
(triangle inequality)
3
(IA.A - 11477) 1) (2--1 P.uu;)) |le
im2
(matrix norm sub-multiplicativiry
< 8/12
Σ
| - 1Ax Pallelal le
1=2
= 3/1 |1S 1) - T|
14

---

Policy Gradient with Tree Expansion
Now, we can bound the norm A 1 by direct calculation:
[AA LAT 1| Tr
[IA.A LAT (IA,A - LAT))
(11)
=Tr [A.A- - 114 T - + ④ N 1
(12)
- - 11 + AT
(13)
s 2A.
(14)
From the Cchwar-Schwartz inequality,
ΣΣ [Palaao]] Σ I|Palo,,112 S
Σ
A.
So,
Var (Vo log (c)((()))(s a)) S max
[Q(s, a))² max Vo log (-10(-))
s,a
1
|13 (TA, - LAT Ps (Pmo)d-111
7)2
1
82)2 ((d
(1 ya
which obtains the desired bound.
A.5. A lower bound on C-SoftTreeMax gradient (result not in the paper)
For completeness we also supply a lower bound on the Frobenius norm of the gradient. Note that this result does not translate the 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:
Ve log x (+)s))) 2 C (PTA(("))
(15)
Proof. We begin by moving to the induced l2 norm by norm-equivalence:
|| (LA.A LAx P. (P || 2 || (LA.A LAT P. 1|2-
Now, taking the vector be the eigenvector of the second eigenvalue of p
||3 (Ia,a 1Ax ( 2| 2 ||3 (Ia,a 1AT 1) Po ull2
= B1| (LA.A - 1AT P2ul12
(1aa(P))) (LA.A LAT Paul22.
Note that even though Psu can be O, that is not the common case since we freely change (and therefore the eigenvectors of pT)).
A.6. Proof of Lemma 4.5 Vector form of E-SoftTreeMax
For d 2 1, (4) is given by
(- = Es,d ex)( 1 Ead exp(
(16)
where
d-1
Edd = P.
II
D (exp]3><-4 4- R!) pm))
(17)
h=1
with R being the diy-eimmesll vector whose s-th coordinate r(s).
15

---

Policy Gradient with Tree Expansion
Proof. Recall that
(d,(d))) 7-
r(s) Σ xr((s) (0(()))
(18)
and, hence,
d-1
(r((s)) + +
exp]$6..0(d;) exp
By-d (r(s)
Σ t=l
(19)
Therefore,
+ d-1
Elexp Bls,a (d; 0)] = E
exp
By d
((s)
r5(())
E (exp 13 (0(sa))][s1,.... 86-11
(20)
1=1
II
exp
Bry
r(s) +
>rr(se)
pm (100-10811) exp(30)
(21)
t=1
E
exp
Byy By-
r(s) + Σ r
expl8y (eoj3y-r(ssu-1)l (( exp(50).
(22)
By repeatedly using iterative conditioning as above, the desired result follows. Note that exp(B7 exp(33-<rss)) does depend on the action and is therefore cancelled out with the denominator.
A.7. Proof of Lemma 4.6 Gradient of E-SeftTreeeMx
The E-SoftTreeMax gradient of dimension A S is given by
E D T2,0
E. Eexx(ex))(3
Vo log F = 3 ITa (A(
expe(3)((3
where for brevity, we drop the s index in the policy above, i.., =B = {)s)
Proof. The (j, k) entry of V log Th, satisfies
alog Plog(miio(a(0))
Vo logog F Telia
80(s*)
θ 00(s*)
(Ion((En,a]] exp(30) - -log(1/ Fe.d exp(30))
exp(30) B(Es,d)s, exp((((ss (E,,)),
31 A E.E.de exp((0(s)))
1 E. (30
B(E,,dek
exp(33(s* (Ea,)) exp(30) cf
B1 Ed.dek expl (30(s*))
Fag exp(
I7
Exaee exp(38(s**)
e E.dd exp((e)
1 Esdd exp(30)
Hence,
V log - 3 exp(,e))p(33)) (1, ex0( Ex,dck exp(30(s*))
From this, follows that The desired result is now easy to see.
E. ((ep(()((33))
V log - 3 ( - 1
(23)
1 exp(ee((33))
16

---

Policy Gradient with Tree Expansion
A.8. Proof of Theorem 4.7 Exponential variance decay of T-ee
There exists a E (0,1) such that, for any function Q Q: A IR,
Var (Vo log ( () E = If all rewards are equal (r const), then = [A2(Pma)]
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
Ve log rdlell
V log T0,02
==||==
with Va log T2 given in Lemma 4.6. Let E R be arbitrary vector such that ||^|| = 1. Then, = = Q24, where ci are scalar coefficients and zi are vectors spanning the S-dimensional space. In the full proof, we show our specific choice of 2 and prove they are linearly independent given that choice. We do note that
1s.
The first part of the proof relies on the fact that (Vo log F ( 1 = 0. This is easy to verify using Lemma 4.6 together with (6), and because [A [ is projection matrix whose null-space is spanned by 1s. Thus,
E Vo log 0, = Vo log 7
Σ
4z
1=2
In the second part of the proof, focus on Ex, from (6), which appears within Vo log ~ā,0- Notice that Ex,d consists of the product d-1 (D (exp(8)"~dR) pm)). Even though the elements in this product are not stochastic matrices, in the full proof we show how to normalize each of them to stochastic matrix Bg We thus obtain that
d-1
E. - P.D(M) II B
where Mi E R is some strictly positive vector. Then, we can apply result by Mathkar & Borkar (2016), which itself
builds (Chatterjee & Seneta, 1977). The result states that the product of stochastic matrices I B₁ of our particular form converges exponentially fast to a matrix of the form 1sp" s.t. ||p T -I=1 B. s Ca for some constant C. S.
Lastly, IsHT, gets canceled due to our choice of zi, ,. This observation along with the above fact that the remainder decays then shows that V log -d,0 E O(a"), which gives the desired result.
Full technical proof. Let d 2 2. Recall that
d-1
Exd = P. II (D (exp) (expp33*^-11)) prt).
(24)
and that R refers to the S-dimensional vector whose 8-th coordinate is r(s).
Define
p
if d 1, if = 1,..., d 2,
B, =
(25)
D (P D(M++1)
M1+1)
and the vector
Mj = eexp(3r RR)
if d 1, if 1,...,d- 2
(26)
exp((3r-d pmd Mi+1
where o denotes the dementtwise product. Then,
d-1
= P.D=P II
B
(27)
17

---

Policy Gradient with Tree Expansion
is easy to see that each B, row-stochastic matrix, i.e., all entries non-negative and B 1g.
Next, we prove that all non-zeros entries of Bi are bounded away from 0 by constant. This is necessary to apply the next result from (Chatterjee & Seneta, 1977). The j-th coordinate of M, satisfies
(M), expl87 Rz] E pmy 3a (M++1)a s ||exp(8>'' |oc Mi
(28)
Separately, observe that |MM-ll $ <eep(88> R)ll Plugging these relations in (26) gives
4-1
||exp(8>"<111"""" - s (exp|3>~<]]l0>> (29)
[M
II 1=1
expl8\/ h-4 12]1
II
h=1
Similarly, for every isis 1, we have that
|
II
-xppp3y- -R
(30)
h=1
The jk-th entry of B, D-1(P=MM PTo (TMM+++))) is
Pih [M++1]k |1 ElS P5P [M_ile
pR jk
(Bi)jk
jk
(31)
P
M4+lle
[xpp
Hence, for non-zero T the entries bounded away from by the same. We proceed with applying the following result.
Now, by (Chatterjee & Seneta, 1977, Theorem 5) (see also (14) in (Mathkar & Borkar, 2016)), lima- I B, exists and
of the form su T for some probability vector μ.. Furthermore, there is (0, 1) such that e(d) (11 B.)
1s u satisfies
[[<(d]] a
(32)
Pick linearly independent vectors W2, wg such that
# T w. for ......d
(33)
Since Ei a, is perpendicular to for any a2,... ag and because 4² exp(30) O, there exists no choice of a.2...., ag such that Σi= a,w; = exp( Hence, if we let 21 = 1g and 2; = D(exp(30)) - "w, for i 2,..., S, then follows that (21,..., 28) linearly independent. In particular, it implies that {21,..., 28) spans R.S.
18

---

Policy Gradient with Tree Expansion
Now consider an arbitrary unit norm vector = E Gaz E R s.t. ||=||2 1. Then,
s
Ve log m Volog T0,0
Σ C:4
(34)
D (1.00
E Ec.p((c))(5
8
=1 -1a (AT,0)
Σ
(35)
1 Edd ex0(()
Bad D 11.00
8
= 3 LA -LA 1
c:
(36)
A E exp((x((3 exp((3))
i=2
D (d.0
[1sH"
+d(d)
8
= B [TA -1a
C200g
(37)
exp(3()
2=2
s D T E(d)
= 3 [LA -1A( []
Σ
qwv
(38)
1 Ea,d exp(30)
1=2
D (1.00)
((d)D(exp(3e)) (= 1s),
= B [TA -1x [mg
(39)
1
where (34) follows from the fact that log =0 holds since zi = ((exp(30))) (38) because le is perpendicular wy for each i, while (39) follows by reusing 24 = ((exx(39))) -',,; ww relation along with the fact that 21 1s.
log _ 1 0, (35) follows from Lemma 4.6, (36)
From (39), it follows that
EE D d,0 1/ Fasd (xp
Vo log r.o0 & (l(E(d11
[TA -
[D(oxp((00)] || lsll
(40)
-1
D 12.0
< 388 + (1
(3 exp(8max (()) 1sll
(41)
Fs,d exp (xp(
D B
< 3Baa VA)
exp 100Bmaxe (m83mx ((s) cs
(42)
1 exp(())))
3.00 + VA)
D (Ea,d (xp( cxp(3 max ((s))| assl
(43)
( + v
min, EEa, (3)) exp(3 max 0 lsl
(44)
exp(3 max, 0(s)
$1 + VA)
exp(8 min, 0(s)) min, |M] || 1sll
(45)
cxp(3 max, (s))
s 3 (1 + VA)
exp ( min, 0(s) exp(3 min, r(s)) || ||= cls||
(46)
2a + VA) exp @ mmax (()) min min ((s) min (s))))| 1sll
(47)
Lastly, we prove that ||2-c 1s| bounded independently of First, denote by e = (ci,..., es) and (0, 02 cs)
19

---

Policy Gradient with Tree Expansion
Also, denote by Z the matrix with 2; as its i-th column. Now,
8
|| - c
Σ
(48)
1=2
Z
(49)
S / |||||
(50)
S 2 |||
(51)
|| 2 --
(52)
S ||
(53)
where the last relation is due being unit vector. All matrix norms here are l2-induced norms.
Next, denote by W the matrix with w; in its i-th column. Recall that in (33) we only defined w2..... ws. We now set wh ex)( Note that wi is linearly independent of {wz,- (wz,..... wss because of (33) together with the fact that μ T wi >> 0. We can now express the relation between Z and W by Z = D (ex) Substituting this in (53), we have
2 c1ss| [ID-(cpp((000///1/V000000)00))))
(54)
(55)
further holds that
[D(exp(30))I & maxexp (30(s)) < max(1, expl3 max ( (())))
(56)
where the last relation equals if ((s) 0 for all s. Similarly,
[exp(((( min, exp (30(s)) min expl min, (s))))}
(57)
Furthermore, by the properties of the lz-induced norm,
| 2 WS
(58)
max |w:]h Vs
(59)
isis
V max mex max [|w/||i)
(60)
25i5s
VS 1mxpf1 max (s)), max [[w;]l)))
(61)
25i5s
Lastly,
|W-1
(62)
G (W)
maxx (W)
(63)
(W)
Omin (W)
(a (W))--1
(64)
Is (W))
|[i]s-1
(65)
dW")
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];,1 that are independent of d, depending on their sign:
min(1, exp[3 min (s)))) [W]..1 < max exp[3 max (ss)))
(66)
Using this, together with (53), (55). (56), (57), and (61), we showed that || clss| is upper bounded by a constant independent d. This concludes the proof.
20

---

Policy Gradient with Tree Expansion
A.9. Bias Estimates Lemma A.2. For any matrix A and Å,
k (A-A)A -
A
A
A (Ā
hm1
Proof. The proof follows from first principles:
Σ A
(A-
A)A
A AA Σ AA
(67)
k
k Σ A
A A
(68)
h=1
h=1
k-1
A A + Σ
A Ah
Σ
Ak-
(69)
h=1
h=2
= A A
(70)
Henceforth, || will refer to -| ol,,, i.c. the induced infinity norm. Also, for brevity, we denote Th and #5.0 by Th and t, respectively. Similarly, use dr, and dx, to denote d,g and dæg As for the induced norm of the matrix P and its
perturbed counterpart P. which are of size Sx A S, we slightly abuse notation and denote P PI| ma [[l[),),)],, Pall), where P3 is defined in Section? 2.
Definition A.3. Let e be the maximal model mis-specification, i.e., max Pl1, |/r 1|r f|]] =. Lemma A.A. Recall the definitions of Rg, Pey Ru, and pas from Section 2, and respectively denote their perturbed counterparts by R2, Po, Ru, and sad Then, for defined in Definition A.3,
max - 1All..-Pllll,,,, 1| O(e).
(71)
Proof. The proof follows easily from the fact that the differences above are convex combinations of P P and
Lemma A.5. Let ne be as in (5), and let ñe also be defined as in (5), but with Rg Pa, pT replaced by their perturbed counterparts R,, P., pn throughout. Then,
#G, (3d
(72)
Proof. To prove the desired result, we work with (5) to bound the error between Rs, Ps, PKb, Rr, and their perturbed versions.
First, we apply Lemma A.2 together with Lemma A.4 obtain that [|(*))** (oa)**=(pppp (fx0)*M) O(ke). Next, denote by M the argument in the exponent in (5), i.e.
M == (C.. P. (P
Similarly, let M be the corresponding perturbed sum that relies on P and f. Combining the bounds from Lemma A.4, and using the triangle inequality, we have that |M M = (Bdy -e). The factor y appears because C..d includes the term -d as shown in Lemma 4.2.
Eq. (5) states that the C-SoftTreeMax policy in the environment is ne (XP exp(M)). Similarly define मं using M for the approximate model. Then,
Ao (To o exp(M M))1 exp (xp(M)/(1 exp(Mi))
21

---

Policy Gradient with Tree Expansion
where denotes element-wise multiplication.
Using the above relation, we have that || TTl IT exp M M-M) exp(M) 1|l. Using the relation |e- - 11 ((x) + 0, the desired result follows.
exp(M)
Theorem A.6. Let e be as in Definition A.3. Further let &6.0 being the corresponding approximate policy as given in Lemma 4.2. Then, the policy gradient bias is bounded by
8 80 (v V=o)
a 00
(v [ = (1-
(73)
We first provide proof outline for conciseness, and only after it the complete proof.
Proof outline. First, we prove that max{ - R.I1, IIp Pall. [R, - Rr, ll. Pm - 1 = ((e). This follows from the fact that the differences above are suitable convex combinations of either the rows of P - P or - f. We use the above observation along with the definitions of "2,0 and #2,0 given in (5) to show that ||=5,0 #5 ((3dd-4e) The proof for the latter builds upon two key facts: (a) [|( - ( S
Ek= || | P pm
Okke) for any
k 2 0, and (b) le" - 1| = (xx) as x - O. Next, we decompose the LHS of (7) to get
Σ II x.(+ II x
ΣΣ (1) x (Xi(s)- xi(3)) (.+1(8)....Axu))
where (s) das, (s) E R, X2(8) = (Vo log - 0,9 (-ss)) E RSXA, Xz(s) = D(F E RAXA, Xa(s) Q (s.)) E RAXA and (s)s)...... (s)) are similarly defined with =Ι,0 replaced by #G.0- Then, we show that, for i the desired result follows.
4, (i) (s) x.(()]| 46) and (ii) max [|l]]] is bounded by problem parameters. From this,
Proof. We have
(v
T VTT)
V=)
(74)
00
00
Egnddn (-1)) Vo Voloomaadss) Eonde 8) [ol
(75)
= (((4))))(((()))))aaoa00000000000 da, (o
(76)
Σ (1)
(77)
da (-1)) ((a
(78)
Σ II (x) II x1(00)
(79)
Σ (s) ... (s)) (x((s) - {((0)) Xi4.(s)....X111)
(80)
where Xy(s) dr. (s) E R, (s) = (Vo log T RSXA, X²(s) = ((To(-Is)) E RAXA, Xa(s) (s. RAXA, and
(s) are similarly defined with go replaced by At.
Therefore,
|
ag (N (vTrze) /| < (maxT(s) S. 012 (v v 00
(81)
where
r(s) = Σ (s) ... 1 (x1(s) -((()) Xi+1(s) ... Xx(I)]
(82)
22

---

Policy Gradient with Tree Expansion
Next, since aa d ne; and Fe all distributions, we have
max mooKXXX(elll.(((0)sssssso...). 1
(83)
Separately, using Lemma 4.3, have
||,||
VV lo - [([|aall + || P. (P
(84)
Since all rows of the above matrices have non-negative entries that add up to 1, we get
| I 2B.
(85)
In the rest of the proof, we bound each of || - Xill...., | - -aal-
Finally,
| 5 1-7
(86)
Similarly, the same bounds hold for X2, Xs and X1. From, we have
XX - ((1 Σ
79100 [|luu(ppo)) - u (PPR0)
(87)
S |]]|
7 tde
(88)
$ (1 de
71t
(89)
1=0
ydc
(90)
1 1-7
The last relation follows from the fact that (1 7) = Ei= 7",
which in turn implics
1
? Oy
ht
(91)
1-
0
From Lemma A.5, it follows that
||83 Xx11 = ((dde..
(92)
Next, recall that from Lemma 4.3 that
Xa(s,-)= = [LA IA(TT) Ps (Pru)d-1
Then,
[x2(s - (2( $|BB [TA (TT) 1 Palll (pms)u-1 _ (pm)) d-1
(93)
||3 [LA 1 7] I||.a - -sllll (isTa)
d-1
(94)
d-1
(TT (70)) TI (pms)
l1.
(95)
Following the same argument as in (85) and applying Lemma A.2, we have that (93) is (Bde). Similarly from the argument of (85), Eq. (94) is (Be)). Lastly, (95) is ((dd) due to Lemma A.5. Putting the above three terms together, we have that
(x,.) X2(s, ) ((3d).
(96)
23

---

Policy Gradient with Tree Expansion
Since the state-action value function satisfies the Bellman equation, we have
+=
(97)
and
Q = 7 7PQ00
(98)
Consequently,
- Q S -| + PQ PQ - <<+>|PP1IQ - -= + + P11110011
(99)
(100)
S s - 2 + 1-7
(101)
which finally shows that
|Xa - Xall = 11@00 - Q^01| s
(102)
(1 71-
B. Experiments
B.1. Implementation Details
The environment engine the highly efficient Atari-CuLE (Dalton et al., 2020), a CUDA- version of Atari that runs on GPU. Similarly, we use Atari-CuLE for the GPU-based breadth-first TS as done (Dalal et al., 2021).
We train MoaX for depths 1...8, 8, with single worker. We five seeds for each experiment.
For the implementation, extend Maale-Baseelines? (Raffin et al., 2019) with all parameters taken as default from the original PPO paper (Schulman et al., 2017). For depths d 2 3, we limited the tree maximum width of 1024 nodes and pruned con-promising trajectories in terms of estimated weights. Since the distributed PPO baseline advances significantly faster in terms of environment steps, for 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.
B.2. GPU-Based Tree Expansion Implementation
Figure 4 illustrates the GPU-based tree expansion mechanism used our implementation. We achieve efficient parallelization by duplicating and concatenating all states in the current level of the tree with each possible action, then advancing them simultaneously with a single forward pass through the simulator.
The implementation follows these steps:
1. Start with a root state so and expand across all actions A.
2. For each depth level from d:
Collect all states 8; from the previous level.
Duplicate each state for each action to create batch of state-action pairs.
Submit the entire batch to the GPU simulator in single forward pass.
Collect the resulting next states and corresponding rewards.
I < and pruning enabled, select the top-k most promising branches based on accumulated rewards.
3. For the leaf states (at depth d), compute the neural network outputs We(sa).
4. Combine the accumulated rewards along each trajectory with the corresponding leaf state values.
5. Compute the final policy logits according Equation (2) and apply the softmax operation.
This parallel implementation allows us to efficiently explore a larger number of trajectories compared to sequential tree expansion, making deeper tree depths practically feasible.
24

---

Policy Gradient with Tree Expansion
S.
state duplication
00
a
a2
S
a
S
a,
S
a2
batch
S S11
S |11
S2Mr
GPU simulation
state duplication
S
M+
a
S
a
S
a
S1 t+
a
S
a,
S
a
0
a1
22
0
a
a2
3
22
S
ap
S i+1
a,
S
i+1
a2
batch
142
S
S11u2
S
S.0
S
S2
GPU simulation
Logit inference
[000, 001, 002, 0 011, 012,020 020, 0 0
Figure 4. diagram of the expansion used by Max. In every step, the states in the current level of the tree duplicated and concatenated with each possible action. The resulting state-action pairs are then fed as 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 Wo and the logits are computed using the corresponding rewards along each trajectory.
Algorithm 1 C-SoftTreeMax
Input: GPU environment G, network 0, depth
Init tensors: state 5 [s], action Åg = [0, 1, 2, A 1). reward R = [0] for is to d do
5 S 5 x A, R R A (Replicate state and reward tensors A times)
F, F. = A(( Al) [Feed [S, to simulator and advance]
R + R + S S' (Accumulate discounted reward }
A + A x A (Replicate action tensor A times]
end for
Ix. + A. (R R 7 (5S))/ (Weighted average induced by m.}
Return maass)) exp 1 (d; 0)] [Return optimal action at the root)
B.3. Algorithms
This section provides the pseudocode for our SoftTreeMax implementation. Algorithm details the C-SoftTreeMax policy computation, which efficiently utilizes GPU parallelization to perform tree expansion. Algorithm shows how SofeTreeMax integrates with the PPO algorithm, distinguishing the usage of our new policy in red.
Note that in Algorithm 2, we use Generalized Advantage Estimation (GAE) with À = 0.95 for calculating advantage estimates, which is the standard configuration in the stable-baselines3 PPO implementation that we build upon.
B.4. Time-Based Training Curves
We provide the training curves in Figure 5. For brevity, we exclude a few of the depths from the plots. As seen, there a clear benefit for SoftTreeMax over distributed PPO with the standard softmax policy. In most games, PPO with the SoffTreeMax policy shows very high sample efficiency: it achieves higher episodic reward although it observes much less episodes, for the same running time.
25

---

Policy Gradient with Tree Expansion
Algorithm 2 SoftTreeMax-PPO
1: Initialize policy parameters dg Initialize value function parameters
3: for k = 0, .,,.....
......d. do
4: Collect set of trajectories DA: { by running policy ad,0, from Algorithm
5: Compute rewards-to-ge to-go RG
6: Compute advantage estimates A, using GAE with À 0.95
7: for each epoch do
8:
for each minibatch do
Compute policy ratio (0) = [s.
9:
ma,ay (8))
Compute clipped surrogate objective:
11:
I.CLIP (0) E, (m) Ac, clip(r, (0)) 6, + A))all
Update 0 with gradient step L CLLIP (0)
12:
13:
Compute value function loss: L (6) (Vo(8.) f2.)²
14:
Update o with gradient step on L vr (6)
15:
end for
16:
end for
end for
Asteroids
Breakout
Gopher
Krull
8000
6000
400
Reward
3000
6000
4000
200
2000
2000
4000
PPO
ToeeTre Depth
0
2000
SoftTreeMax Depth
100
100
100
100
SoftTreeMax Depth
Time [hours]
Time [hours]
Time [hours]
Time [hours]
SoftTreeMax Depth
SoftTreMMax Depth
KungFuMaster
ThisGame
Phoenix
VideoPinball
60000
30000
Reward
10000
200000
40000
20000
5000
10000
20000
100
100 Time [hours]
100 Time [hours]
100 Time [hours]
Time [hours]
Figure 5. Training curves: GPU Mrxx (single worker) vs PPO (256 GPU workers). The plots show average reward and standard deviation over seeds. The x-axis is the wall-clock time. The runs ended after one week with varying number of time-steps. The training curves correspond to the evaluation runs in Figure 3.
B.5. Step-Based Training Curves
In Figure 6 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 the reward with increasing tree depth is noticeable most games.
26

---

Policy Gradient with Tree Expansion
Asteroids
Breakout
Gopher
Krull
4000
8000
Reward
400
6000
3000
6000
4000
2000
2000
4000
PPO
SoltfreeMax Depth 2
10a
10"
10
10
10"
10a
10*
10
10a
10
10
10
Num of
Num of
Num of
Num of
SoftTreeMax Depth
SoftTreeMax Depth 5
online interactions
online interactions
online interactions
online interactions
SofthheMMxx Depth
SoftTreeMax Depth
KunsteMaster
NameTaie
Phoenix
VideoPinball
30000
300000
Reward
10000
20000
200000
40000
100000
5000
10000
10a
10
10a
10a
10a
10
10a
10a
10"
10f
10
10"
Num
Num
Num
Num of
online interactions
online interactions
online interactions
online interactions
Figure Training curves: GPU Mofx (single worker) vs PPO (256 GPU workers). The plats show average reward and standard deviation over 5 seeds. The x-axis is the number of online interactions with the environment. The runs ended after one week with varying number of time-steps. 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 obtain with larger depths, in which we apply the network states that are very different from the ones observed so far in the buffer (recall that in any deep RL algorithm, learned good enough value function yet, and there is no guiding dense reward along the trajectory, the policy becomes noisier, and can take more steps converge even more than those we run in week-long experiment.
train the model only on states in the buffer). If the model hasn't t
For a concrete example, let us compare Breakout to Gopher. Inspecting Fig. 6, we observe that Breakout quickly (and monotonically) gains from large depths since relies on the short term goal of simply keeping the paddle below the moving ball. In Gopher, however, for large depths (2=5), learning barely started even by the end of the training run. Presumably, this 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. 6). 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 the case.
C. Further discussion
C.1. The case of (2(PT
When p 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 + oo). 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 Na the optimal policy. However, a similar phenomenon was also discovered in the standard softmax with deterministic policies: 0(s, a) + 3o 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 Vo log no (a)s) ea T where e, E [0, 1 is the vector with 0 everywhere except for the a-th coordinate. I.e., it will be zero for a deterministic policy. Moax avoids these local optima by integrating the reward into the policy itself (but may get stuck in another, as discussed above).
27


Reviews:
[{"review_id": "duq0BcMi5n", "reviewer": "Reviewer_RjcM", "summary": "The paper introduces the SoftMaxTree algorithm, a softmax policy gradient algorithm extended with planning. The main idea of the extension is that estimation the gradient from longer paths reduces the variance of the gradient. Two variant are considered depending on the expectation being inside or outside of the exponent. The main theoretical results are bounds on the variance for the two variants considered. A bound is also provided for the bias introduced in the case of approximate forward model. Empirical evidence on several Atari benchmarks show that the proposed algorithms perform favorably compared to the standard PPO in terms of cumulative reward and variance.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and for recognizing the value of our theoretical and empirical contributions. We appreciate your feedback and address your questions below.\n\n**MCTS vs. our search approach**\n\nThank you for this important point of clarification. SoftTreeMax and MCTS represent fundamentally different approaches with distinct goals. While both involve planning, our work specifically focuses on how to directly integrate forward modeling into the policy gradient framework itself.\nThe key innovation of SoftTreeMax is that it incorporates planning directly into policy parameterization, creating a differentiable policy that naturally reduces gradient variance. Unlike MCTS, which operates as a separate planning algorithm that produces improved rollouts, our approach transforms the underlying architecture of policy gradient methods.\nThis enables us to maintain the theoretical elegance of policy gradient approaches while addressing their fundamental variance issues. Our analytical results show how this integration affects gradient properties in ways that would not be directly applicable to MCTS-based methods. We believe this approach opens new avenues for research at the intersection of planning and policy-based methods.\n\nThat said, we acknowledge the value of comparing against and potentially combining with MCTS-based approaches in future work.\n\n**Additional MCTS baselines**\n\nFollowing your comment, we added comparisons with a strong baseline to our experiments: EfficientZero [Ye et al., 2021]. We chose it because it is a highly sample-efficient version of MuZero that is also open source. MuZero is one of the best known RL algorithms to date. The results that were completed during the rebuttal are given here: https://ibb.co/zHG3KP8m, showing that SoftTreeMax surpasses EfficientZero in all the games we tested except one. We will add full experiments in the final version.\n\n\n**Variance scaling with eigenvalue assumptions**\n\nYou've identified an important theoretical-practical gap. While our theory depends on properties of transition matrices that may not always hold in practice, our empirical results consistently show variance reduction across various environments.\nIn cases where the eigenvalue assumptions are violated, we would expect diminished (but still positive) variance reduction benefits rather than variance increases. Figure 3 in our paper demonstrates this empirically - the variance consistently decreases with depth across different games, even though the theoretical assumptions likely vary between environments.\n\n\n**Bias-variance tradeoff with approximate models**\n\nThe bias-variance tradeoff we observed follows this pattern:\n\n- The closer $\\pi_b$ is to $\\pi_{d,t}$ at time t, the lower the bias\n\n- The closer $\\pi_b$ is to having similar rows in $P^{\\pi_b}$, the lower the variance\n\nFor practical solutions to this tradeoff, we recommend:\n\n1. Adaptive depth selection based on model certainty\n2. Progressive increase in depth as model accuracy improves during training\n3. Ensemble methods to better estimate model uncertainty\n\nOur Theorem 4.8 shows that the gradient bias diminishes with approximation error while retaining similar variance reduction properties. Even with approximation errors, the variance continues to decay exponentially, just at a rate dictated by $\\hat{P}^{\\pi_b}$ instead of $P^{\\pi_b}$.\n\n\nWe will expand this discussion in the revised paper to provide clearer guidance on managing this tradeoff in practical implementations.", "claims_and_evidence": "The claims are supported both theoretically and empirically.", "methods_and_evaluation_criteria": "The evaluation is sensible.", "theoretical_claims": "I checked the proof succinctly.", "experimental_designs_or_analyses": "The experimental setting is fine. Additional baselines that use planning in the training phase (e.g. using MCTS to generate traces) would have been useful.", "relation_to_broader_scientific_literature": "The bounds on the variance are valuable.", "essential_references_not_discussed": "The most relevant literature is mentioned, but a deeper comparision with policy gradient algorithms that are combined with MCTS would have been useful.", "other_strengths_and_weaknesses": "The search/planning part of the algorithm is awkward. It is attempting to do exhaustive search, but does some intensive pruning to keep the tree narrow (to achieve linear complexity scaling with depth). It is really unclear why not MCTS or at least Monte Carlo sampling. That would also have linear dependence, and both are extensively studied in RL. I assume the variance would have been more difficult to analyze theoretically in case of MCTS.", "other_comments_or_suggestions": "One of the advantages of PPO is that it is model free. Usually, one still needs a simulator for training, but there is less constrain on the simulator compared to the setting here that needs a forward model (e.g., we need to be able to reset the simulator to a certain state). If it is possible to use planning, using MCTS to generate traces improves the training (even with an approximate forward model as in MuZero). Therefore I would suggest to add baselines that use MCTS in the training phase (perhaps off-policy variants).", "questions_for_authors": "The variance scales favorably with depth, but a small eigenvalue is crucial in this dependence. The assumption on the transition matrix is common for the theoretical analyses, but often does not hold in practical applications. Do you expect an increase of variance with depth in such scenarios? A discussion on the topic could be useful.\n\nThe variance is reduced with search depth, however the bias increases with approximate forward model. Any solutions for the tradeoff?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GQW0BVx0t7", "reviewer": "Reviewer_587W", "summary": "This paper extends softmax policy gradient methods by integrating planning through tree expansion. The authors introduce two implementation variants—C-SoftTreeMax and E-SoftTreeMax—which differ in whether the expectation is computed inside or outside the exponent. They analyze the policy gradient variance for both approaches in the tabular case and demonstrate that the variance upper bound decreases with increasing tree depth at a rate determined by the second-largest eigenvalue of the behavior policy. Additionally, the paper examines the gradient bias introduced when employing an approximate model. The theoretical findings are further validated through experiments on several Atari games, where GPU based simulation was leveraged to mitigate the computational overhead introduced by tree expansion.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review. We appreciate the careful reading of our proofs and the constructive feedback. We address your concerns below.\n\n\n**Proof issues in Lemma A.5 and Theorem A.6**\n\n\nYou're correct that $\\gamma^{-d}$ from $C_{s,d}$ should appear in this bound. This oversight makes the proper bound $O(\\beta d \\gamma^{-d} \\epsilon)$. When the approximation error $\\epsilon$ is small enough and $\\beta$ properly controlled, this term still approaches 0 as required for the subsequent proof.\nWe will correct this in the revised version.\n\n\n**Variance bounds discrepancy**\n\n\nE-SoftTreeMax also depends on $S,A$, but Theorem 4.7 uses O-notation while C-SoftTreeMax shows explicit constants. Both establish the crucial exponential decay rate with tree depth.This exponential decay explains why deeper trees provide more stable policy gradient estimates and better overall performance, regardless of specific constants or dependencies. The key contribution is establishing this relationship between tree depth and variance reduction through the eigenstructure of the MDP's transition dynamics.\n\n\nWe'll highlight this insight more clearly.\n\n\n**Assumption that r(s,a) equals r(s)**\n\n\nThis simplification was made for theoretical analysis of E-SoftTreeMax. Our practical implementation handles general reward structures correctly in all experiments.\n\n\n**Definition of $\\theta$ and $\\Theta$**\n\n\nThank you for highlighting this inconsistency. In Section 2.1, we use $\\theta$ conventionally for policy parameters, but later use $\\Theta \\in \\mathbb{R}^S$ and $\\theta(s)$ as a scalar state function.\nWe'll replace the notation in Section 2.1 with a different symbol, reserving $\\theta$ and $\\Theta$ exclusively for our SoftTreeMax formulation to ensure consistency throughout the paper.\n\n\n**Verification for C-SoftTreeMax**\n\n\nWe analyzed both variants but included only E-SoftTreeMax results (Figure 1) to avoid redundancy. C-SoftTreeMax shows similar behavior and correspondence to theory. We focused on E-SoftTreeMax to verify that parameter $\\alpha$ relates to the transition matrix's second eigenvalue in the general case, which our experiments confirmed.\nWe'll add a brief discussion of C-SoftTreeMax's empirical variance characteristics for completeness.\n\n\n**Behavior policy choice**\n\n\nOur theory identifies properties of ideal behavior policies but doesn't provide concrete procedures for complex domains. In Atari experiments, we used uniform policies for simplicity. Preliminary experiments with other policies showed similar trends but weren't comprehensive enough to include.\n\n\nWe'll clarify our claims and add discussion about the practical impact of behavior policy choice.\n\n\n\n\n**Minor comments and title suggestion**\n\n\nWe appreciate your writing suggestions and will consider revising the title to better reflect our focus.\n\n\n**Question 1 regarding loose variance bounds**\n\n\nThis is indeed an artifact of our analysis approach. Our bound becomes loose in this regime because we use worst-case inequalities at multiple steps in the proof that don't capture the special structure that emerges in deterministic settings. The $\\beta^2$ term appears because we bound the gradient norm without accounting for how the policy structure changes as $\\beta$ increases. For practical parameter ranges used in training, our bound provides useful insights about variance reduction with tree depth. In the revised paper, we'll clarify this limitation and discuss how the bound could be refined to better handle the deterministic limit case.\n\n\nWe'll clarify this limitation and discuss potential refinements.\n\n\n**Question 2 on the potential issue with the proof**\n\n\nPlease see our response above regarding the missing $\\gamma^{-d}$ term.\n\n\n**Question 3 on variance bounds intuition**\n\n\nC-SoftTreeMax has a more decomposable gradient structure (Lemma 4.3):\n$$\\nabla_\\theta\\log \\pi_{d,\\theta}^C = \\beta\\left[I_{A} - 1_A (\\pi_{d,\\theta}^C)^\\top \\right]P_s \\left(P^{\\pi_b}\\right)^{d-1}.$$\nThis allows direct spectral decomposition. When the projection matrix interacts with the decomposed $(P^{\\pi_b})^{d-1}$, the stationary distribution term gets canceled, leaving terms scaled by $\\lambda_i^{d-1}$, with $|\\lambda_2|$ dominating.\n\n\nE-SoftTreeMax's analysis is more complex as its gradient involves matrix products and ratios with exponentiated rewards, requiring sophisticated techniques to establish the convergence rate $\\alpha$.\n\n\nThis explains why we provide an exact variance bound for C-SoftTreeMax, while for E-SoftTreeMax we give an asymptotic bound where $\\alpha = |\\lambda_2(P^{\\pi_b})|$ only under certain conditions.", "claims_and_evidence": "The main issue is in the bound for the policy gradient bias when using an approximate model. \n\nThe proof doesn’t entirely check out. In the proof for Lemma A.5, it is stated that $||\\hat{M}-M||=O(\\beta d \\epsilon)$. However, it is unclear why $\\gamma^{-d}$ from $C_{s,d}$ does not appear in the bound. Moreoever, if the gap between $\\hat{M}$ and $M$ is big, i.e., not approaching 0, the subsequent proof may not hold.\n\nAnother reservations that I have is regarding the difference between the variance upper bounds of C-SoftTreeMax and E-SoftTreeMax. \nSpecifically, the former one scales with $S^2$ and $A^2$ while the later one does not contain either $S$ and $A$. I’m not sure why there is such discrepancy. \n\nRelated to the bounds, while the variance of E-SoftTreeMax is verified in numerical experiment illustrated in Fig. 1,  there’s no verification provided for the result of C-SoftTreeMax.", "methods_and_evaluation_criteria": "The methods and evaluation mostly make sense to me. However, I have concerns about the assumption in E-SoftTreeMax that $r(s, a)$ equals $r(s)$, as this does not seem to hold in either Atari or MuJoCo—despite the paper's claim that this is typical for these environments. In Atari, for example, the reward is determined by transitions between consecutive states (i.e., $r(s, s^′)$). MuJoCo’s reward functions generally include a dedicated component determined by the action.", "theoretical_claims": "I carefully checked the proof for all results. \n\nThere seems to be an issue with the proof for Lemma A.5 and Theorem A.6, as I discussed in the “Claims And Evidence” section.\n\nIn addition, the mathematical analysis would benefit from improved clarity and proofreading.\n\nFor instance, the definition of $\\theta$ and $\\Theta$ are ambiguous. Sec 2.1 describes $\\theta$ as a mapping from S x A to R. However, the beginning of Sec 2 describes $\\Theta$ as a vector of $R^S$, a vector representation of $\\theta(s)$, which seems to imply that the later is a scalar. It is not clear how $\\theta(s)$ is defined, or what its dimensionality is.\n\nAlso please note the typos listed under “Other Comments Or Suggestions”.", "experimental_designs_or_analyses": "Yes, I checked the validity of both numerical and Atari experiments and found no major issues overall. However, there are two minor issues:\n\n1. variance bounds: the missing numerical result for C-SoftTreeMax which was brought up earlier. \n\n1. behavior policy: In Sec 8,  the paper claims to have “explained how to choose the expansion policy to minimize the gradient variance”. However, more precisely, I think what explained was instead on what property we want the induced transition matrix to have. It is not clearly explained the procedure of how to “choose” a behavior policy. In fact, in the deep RL experiment in Atari, only one type of policy (uniform) was used as the behavior policy. It is not clear what level of impact the behavior policy has on the learning performance.", "relation_to_broader_scientific_literature": "It extends the existing RL literature on two fronts:\n\n1, integrates planning with model-free policy gradient methods\n\n2, in the context of planning with a local simulator [Yin et. al 2022], most work were concerned with action value methods. This paper provides some empirical evidence supporting the benefits of combining policy gradient methods with planning.\n\nD. Yin, B. Hao, Y. Abbasi-Yadkori, N. Lazić, and C. Szepesvári. Efficient local planning with linear\nfunction approximation. The 33rd International Conference on Algorithmic Learning Theory, 2022.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "N.A.", "other_comments_or_suggestions": "I recommend changing the title to “Softmax” policy gradient with tree expansion, since the main results are specific to softmax parameterization.\n\nList of typos:\n\npage 1, the last sentence of the 2nd last paragraph, “we prove that the with …”\n\npage 2, Sec 2, the RHS of the definition of $\\mu_\\pi$ should be $\\mu_\\pi^\\top$ instead of $P^\\pi$\n\npage 2, Sec 2.1, in the variance expression of a random vector $X$,  $^\\top$ should be on the second parenthesis.\n\npage 24, Sec B.3, the symbol inside the parenthesis in “In Gopher, however, for large depths ()”", "questions_for_authors": "1. The bounds on the variance seems loose, in that, when the temperature \\beta is infinite t, the bound is infinite while the policy is deterministic, which should yield 0 variance in a deterministic MDP. Could you clarify if it is an artifact or if there’s another reason?\n2. Could you clarify regarding the potential issue with the proof for Lemma A.5 and Theorem A.6? A detailed discussion on the bound on $||\\hat{M}-M||$ would be particularly helpful.\n3. Could you provide intuitions regarding the differences between the variance upper bounds of C-SoftTreeMax and E-SoftTreeMax?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ZderQbOhrp", "reviewer": "Reviewer_q2bN", "summary": "The authors propose a generalization of softmax parametrization for policy gradient methods that utilizes the breadth-first search tree of the future states. This type of parametrization combines planning with policy gradient methods to reduce the latter's variance. The authors proved that given some assumption on the planning policy, the variance of the new gradient decreases exponentially with the depth of planning, allowing trade-off between computation efficiency and gradient variance. Finally, the authors propose the GPU implementation of their algorithm, which shows improvement against the PPO baseline regarding sample complexity.\n\n## After rebuttal\nI was happy to see the MCTS baseline. However, I still believe that the assumption on the presence of GPU simulators is very strong, so I decided to maintain my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and thoughtful review. We appreciate the careful reading of our work and the constructive feedback. Below, we address your concerns:\n\n**Additional MCTS baselines**\n\nFollowing your comment, we added comparisons with a strong baseline to our experiments: EfficientZero [Ye et al., 2021]. We chose it because it is a highly sample-efficient version of MuZero that is also open source. MuZero is one of the best known RL algorithms to date. The results that were completed during the rebuttal are given here: https://ibb.co/zHG3KP8m, showing that SoftTreeMax surpasses EfficientZero in all the games we tested except one. We will add full experiments in the final version.\n\n**Question 1: Does the number of online interactions include the samples generated during the tree search?**\n\nNo, the number of online interactions does not include the samples generated during tree search. We distinguish between two types of interactions:\n\n1. Online environment interactions: These are actual interactions with the environment during training (e.g., when collecting trajectories for PPO updates).\n\n2. Planning interactions (tree search): These are simulated forward passes using our GPU-based simulator to expand the tree.\n\nAs we explained in our experiments, the computational complexity of planning interactions is substantially lower than online interactions, especially with our GPU batching mechanism for tree expansion. This is why we track them separately, and our sample complexity plots show only the online interactions.\n\nFor a complete view of computational efficiency, we also provide wallclock time comparisons in Figure 4, which accounts for the total computation including tree expansion. As seen there, despite the additional computation for tree search, SoftTreeMax still demonstrates favorable performance per unit time for appropriate choices of depth.\n\n**Question 2: Why can increasing d to 1 or 2 harm performance in some experiments?**\n\nThis is an excellent observation. While theoretically more information should help, this illustrates the basic bias-variance tradeoff in the SoftTreeMax policy:\n\nAt small depths (d=1,2), we introduce a structural bias by incorporating tree search, but the variance reduction benefit hasn't fully kicked in yet. The theoretical variance reduction is exponential in d, so the effect becomes more pronounced at larger depths.\n\nWe will address this valuable insight in the revised paper.", "claims_and_evidence": "**Claim 1.** The gradient variance of SoftTreeMax decays exponentially with a rate depending on the spectral properties of the transitions matrix induced by a tree search policy;\n\nThis claim has been proven theoretically (Theorems 4.4 and 4.7) and practically given an experiment on a randomly generated MDP. \n\n**Claim 2.** In the case of an approximate forward model, the gradient bias is proportional to the approximation error.\n\nThis claim has been proven theoretically; however, there is no empirical evidence since all the experiments use the exact forward model. \n\n**Claim 3.** The algorithm is implementable on GPU and outperforms PPO in terms of both gradient variance reduction and achieving higher reward.\n\nThis claim was confirmed on 8 Atari games, implemented in a parallelizable GPU version of the environment.", "methods_and_evaluation_criteria": "The proposed method and the evaluation methodology make sense to me. However, given the focus on improving the algorithm's planning capabilities, the inclusion of the planning benchmark (e.g., the game of Go or any other strategic game) could be important in assessing the final improvement.", "theoretical_claims": "Theoretical statements make sense to me, although I have not carefully checked them in the appendix.", "experimental_designs_or_analyses": "- The experimental design lacks comparisons with other model-based approaches. \n\nThe proposed algorithm assumes a relatively strong assumption on the interaction model with MDP: the presence of a simulator or an approximate simulator to perform a BFS over the next states. Given this assumption, I believe the comparison with MCTS-based methods  (with a comparable tree search budget) should be present. \n\n- The lack of planning benchmarks.\n\nSince the method is assumed to improve planning capabilities, I expect to see a comparison on some planning benchmark, such as Go or any other strategic game (at least Connect-4).", "relation_to_broader_scientific_literature": "The improvement of the planning capabilities of policy-gradient methods is relevant to the current scientific literature.", "essential_references_not_discussed": "The paper lacks a discussion of another type of planning without deep search, such as k-step lookahead, during advantage estimation (e.g., Tree Backup (Precup et al, 2000), Retrace (Munos et al, 2016), application in Muesli (Hessel et al, 2021)). In particular, I found that the method resembles the idea of n-step Tree Backup (e.g., Sutton & Barto book) but from the perspective of policy parameterization, not Q-value estimation.\n\nPrecup, Doina, Richard S. Sutton, and Satinder Singh. \"Eligibility traces for off-policy policy evaluation.\" ICML. Vol. 2000. 2000.\n\nMunos, R., Stepleton, T., Harutyunyan, A., & Bellemare, M. (2016). Safe and efficient off-policy reinforcement learning. Advances in neural information processing systems, 29.\n\nHessel, M., Danihelka, I., Viola, F., Guez, A., Schmitt, S., Sifre, L., ... & Van Hasselt, H. (2021, July). Muesli: Combining improvements in policy optimization. In International conference on machine learning (pp. 4214-4226). PMLR.", "other_strengths_and_weaknesses": "Overall, I enjoy the idea of integrating planning directly into the policy's parametrization and the theoretical guarantees for variance reduction given the depth of the search. \n\nHowever, as a limitation acknowledged by the authors, I also found that the assumption of the presence of GPU-sumilator is very strong.", "other_comments_or_suggestions": "-", "questions_for_authors": "- Does the number of online interactions include the samples generated during the tree search? \n- For $d=0$, the method corresponds to a usual softmax parametrization (lines 147-157). Why, in your experiments, increasing $d$ to 1 or 2 can harm performance? In this case, it looks like the method gets more information on the environment and should help in training. I also found the case of $d = 1,2$ the most interesting from a practical perspective since it should not introduce a significant overhead on tree-search counterpart but decrease the variance (given all the experiments, the relative variance reduction between $d=0$ and $d=1$ is most significant).", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1aegWkloze", "reviewer": "Reviewer_9Djd", "summary": "The paper introduces SoftTreeMax, a new approach that combines policy gradient (PG) with tree search. The goal is to address the inherent high gradient variance in traditional PG algorithms. \n\nThe authors present theoretical analysis showing that the gradient variance of SoftTreeMax decays with the depth of the tree expansion, and that this decay rate is influenced by the second eigenvalue of the transition matrix induced by the tree expansion policy.\n\nOn the empirical side, they utilize a GPU-based simulator to efficiently manage the exponential computational cost of expanding the tree.\nEmpirical validation is done through experiments on the Atari benchmark suite, where the proposed approach demonstrated lower gradient variance and better performance compared to Proximal Policy Optimization (PPO).\n\n## after rebuttal\nThe rebuttal did not significantly change my view of the paper. I will maintain my score.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed and thoughtful review. \n\n**Dependency of bound on $S$ and lower bound**\n\nRL analysis on tabular MDPs often include $S$ terms, usually stemming from the triangle inequality applied on summation over states. These dependencies can sometimes be replaced by structural assumptions on the transition matrix (like branching). This is a common and important theoretical-practical gap. While the theory suggests meaningless bounds for large $S$ in practice we see exponential variance reduction with $d$ for the huge state space in Atari.\n\nOn a related note, a lower bound to the variance is given in the appendix A.5 (it does not include $S$).\n\n**Bound on E-SoftTreeMax does not contain $S$**\n\nThere is a dependence on $S,A$ also in the case of E-SoftTreeMax. However, notice that the bound in Thm 4.7 is presented in O-notation while the C-SoftTreeMax bound is expressed explicitly. The critical theoretical contribution in both cases is the identification of the exponential decay rate with respect to tree depth. \n\nRegarding the missing numerical variance measurements for C-SoftTreeMax: we conducted similar variance analyses for both SoftTreeMax variants, but included only the E-SoftTreeMax results in the paper to avoid redundancy. The numerical results for C-SoftTreeMax show similar behavior and correspondence to the theory. We will clearly state this in the revision.\n\nThe primary reason we included the detailed variance measurements for E-SoftTreeMax was to verify our theoretical hypothesis from Thm 4.7 that the parameter $\\alpha$ is related to the $\\lambda_2$. Our numerical experiments confirmed this in the general case by showing that variance decays at a rate determined by the mixing properties of the baseline policy.\n\nWe will add a brief discussion of C-SoftTreeMax's empirical variance characteristics in the revised manuscript to ensure completeness.\n\n**Comparison to MCTS**\n\nFollowing your comment, we added comparisons with a strong baseline to our experiments: EfficientZero [Ye et al., 2021]. We chose it because it is a highly sample-efficient version of MuZero that is also open source. MuZero is one of the best known RL algorithms to date. The results are given here: https://ibb.co/zHG3KP8m, showing that SoftTreeMax surpasses EfficientZero in all the games we tested except one. We will add full experiments in the final version.\n\n**Question 1: How can one handle a stochastic environment with an infinite number of states?**\n\nFor stochastic environments with infinite state spaces, we propose three complementary approaches:\n\n1. We can use Monte Carlo sampling to expand the tree. This approach maintains a parametric distribution over actions (potentially dependent on $\\theta$) and samples from it to build the tree. This method can be viewed as a tree adaptation of Model Predictive Path Integral control (MPPI) with a value function.\n\n2. Function approximation for leaf states: For infinite state spaces, we already use neural networks to approximate the value function at leaf states. This approach naturally extends to continuous state spaces.\n\n3. Theoretical extension: The key concepts in our analysis can be adapted to infinite state spaces by:\n\n    3a. Replacing transition matrices with kernels in the continuous case.\n\n    3b. Exploiting that $\\pi_b$'s transition kernel is a non-expansive operator.\n\n    3c. Leveraging that the eigenvalue 1 gets canceled for the policy gradient.\n\nThese properties can be shown to hold for decision models with infinite state and action spaces, though we leave the complete theoretical development for future work.\n\n**Question 2: Does it make sense to change the behavioral policy over time? Specifically, does it make sense to use the current policy as the tree-expansion policy?**\n\nYes, adapting the behavioral policy over time is a promising direction. As training progresses and the policy improves, using it as the tree expansion policy could lead to exploring more relevant parts of the state space. This creates a bootstrapping effect similar to what's done in MCTS algorithms like AlphaZero, where the current policy guides the search.\n\nStill, care must be taken as trained policies often become deterministic leading to unfavorable variance properties for semi-deterministic environments (when $P^{\\pi_b}$ approaches a permutation matrix, $\\lambda_2(P^{\\pi_b})$ could potentially approach 1, eliminating the variance reduction benefit). Also we can’t use the SoftTreeMax policy directly to open the tree as it requires expanding a tree while expanding a tree.\n\nA reasonable compromise might be to learn a proxy policy as a mixture of the current policy and a uniform policy, with the mixing weight potentially decreasing over time to balance exploration and exploitation while maintaining favorable variance properties.\n\nThank you for raising this excellent point, which provides fertile ground for future work. We will address this in the revised paper.", "claims_and_evidence": "Overall, the claims are clear. \n\nThe variance bound of C-SoftTreeMax in Theorem 4.4 depends on the number of states S, which is typically large (or even infinite). It is unclear whether this dependency is unavoidable in the worst case, as no lower bound is given, and a comparison between the bound and the empirical variance is not provided.\n\nFor E-SoftTreeMax, the bound in Theorem 4.7 does not depend on S, and an empirical comparison is included. Thus, this claim is more convincing.\n\nTheorem 4.8 provides a bound on the bias of the gradient estimator in cases where the dynamics model is inaccurate. This bound scales linearly with $S$ and $d$. Again, it is unclear whether this dependency is unavoidable, and no empirical comparison is provided. \nAssuming this dependency is accurate, the variance decreases with \n$d$ while the bias increases. This tradeoff on overall performance is not examined, as the experiments are conducted with a precise forward model (which does not exist in most practical applications).\n\nFinally, in the conducted experiment on the Atari environment, when a precise forward model does exists, the paper demonstrates that indeed the variance reduces with $d$ and the overall performance is better than PPO's.", "methods_and_evaluation_criteria": "The theoretical results provide bounds on the variance as well as on the bias of the estimator as a function of the tree depth \nd which definitely makes sense.\n\nThe comparison to PPO also makes sense, as this is one of the most popular PG algorithms. \n\nSince the method aims to combine PG with tree search, it would make sense to compare it to a benchmark that is based on MCTS (such as AlphaZero or something similar).", "theoretical_claims": "I read the proof outline in the main text and did not find anything that raised my suspicion. I have also read the proof of Lemma 4.1 in the appendix.", "experimental_designs_or_analyses": "The conducted experiments seem valid. As mentioned earlier, some additional experiments could be beneficial, such as a comparison to a tree search algorithm. Additionally, regarding the comment in footnote 1, I think it would be beneficial to compare the two suggested algorithms in a non-deterministic benchmark.", "relation_to_broader_scientific_literature": "The paper addresses the well-known issue of high gradient variance in policy gradient algorithms.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The paper presents a novel approach for reducing gradient variance.\n\nA major limitation is that the algorithm requires computing expectations over future states, which can be very challenging in stochastic environments with a large number of states.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "* How can one handle a stochastic environment with an infinite number of states?\n\n* Does it make sense to change the behavioral policy over time? Specifically, does it make sense to use the current policy (or a regularized version of it) as the tree-expansion policy?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper proposed a generalization of softmax, a tree based implementation for the softened max, to improve planning for RL.

The method reduces the variance of the policy gradient, and is well supported by experiments.

The tree depth is a hyper parameter that can influence the performance quite bit.

One reviewer pointed out that the main issue is the bounding of the variance given an approximate model. The authors need to improve the presentation of the result to be easier to follow.

Another reviewer suggested the exhaustive search is too time consuming, and can be improved by sampling based search.

These can be addressed to improve the quality of the paper.
Decision: Accept (poster)