Title: Policy Gradient with Tree Expansion
Venue: NeurIPS 2024
Accepted: no
Keywords: Policy Gradient

Softmax policy

Tree expansion

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

Abstract

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

1 Introduction

Policy Gradient (PG) methods [Sutton et al., 1999] for Reinforcement Learning (RL) are often the 18 first choice for environments that allow numerous interactions at a fast pace [Schulman et al., 2017]. 19 Their success is attributed to several factors: they are easy to distribute to multiple workers, require 20 no assumptions on the underlying value function, and have both on-policy and off-policy variants. 21

Despite these positive features, PG algorithms are also notoriously unstable due to the high variance 22

of the gradients computed over entire trajectories [Liu et al., 2020, Xu et al., 2020]. As a result, PG 23 algorithms tend to be highly inefficient in terms of sample complexity. Several solutions have been 24 proposed to mitigate the high variance issue, including baseline subtraction [Greensmith et al., 2004, 25 Thomas and Brunskill, 2017, Wu et al., 2018], anchor-point averaging [Papini et al., 2018], and other 26 variance reduction techniques [Zhang et al., 2021, Shen et al., 2019, Pham et al., 2020]. 27

A second family of algorithms that achieved state-of-the-art results in several domains is based on 28

planning. Planning is exercised primarily in the context of value-based RL and is usually implemented 29 using a Tree Search (TS) [Silver et al., 2016, Schrittwieser et al., 2020]. In this work, we combine 30 PG with TS by introducing a parameterized differentiable policy that incorporates tree expansion. 31 Namely, our SoftTreeMax policy replaces the standard policy logits of a state and action, with the 32 expected value of trajectories that originate from these state and action. We consider two variants of 33 SoftTreeMax, one for cumulative reward and one for exponentiated reward. 34

Combining TS and PG should be done with care given the biggest downside of PG-its high gradient 35

variance. This raises questions that were ignored until this work: (i) How to design a PG method based 36

on tree-expansion that is stable and performs well in practice? and (ii) How does the tree-expansion 37

policy affect the PG variance? Here, we analyze SoftTreeMax, and provide a practical methodology 38 to choose the expansion policy to minimize the resulting variance. Our main result shows that a 39 desirable expansion policy is one, under which the induced transition probabilities are similar for 40 each starting state. More generally, we show that the gradient variance of SoftTreeMax decays at 41 a rate of | λ 2 | d , where d is the depth of the tree and λ 2 is the second eigenvalue of the transition 42 matrix induced by the tree expansion policy. This work is the first to prove such a relation between 43 PG variance and tree expansion policy. In addition, we prove that the with an approximate forward 44 model, the bias of the gradient is bounded proportionally to the approximation error of the model. 45

46

47

48

49

50

51

52

53

54

55

56

To verify our results, we implemented a practical version of SoftTreeMax that exhaustively searches the entire tree and applies a neural network on its leaves. We test our algorithm on a parallelized Atari GPU simulator [Dalton et al., 2020]. To enable a tractable deep search, up to depth eight, we also introduce a pruning technique that limits the width of the tree. We do so by sampling only the most promising nodes at each level. We integrate our SoftTreeMax GPU implementation into the popular PPO [Schulman et al., 2017] and compare it to the flat distributed variant of PPO. This allows us to demonstrate the potential benefit of utilizing learned models while isolating the fundamental properties of TS without added noise. In all tested Atari games, our results outperform the baseline and obtain up to 5x more reward. We further show in Section 6 that the associated gradient variance is smaller by three orders of magnitude in all games, demonstrating the relation between low gradient variance and high reward.

57

58

59

60

61

62

63

64

65

66

67

We summarize our key contributions. (i) We show how to combine two families of SoTA approaches: PG and TS by introducing SoftTreeMax: a novel parametric policy that generalizes softmax to planning. Specifically, we propose two variants based on cumulative and exponentiated rewards. (ii) We prove that the gradient variance of SoftTreeMax in its two variants decays exponentially with its tree depth. Our analysis sheds new light on the choice of tree expansion policy. It raises the question of optimality in terms of variance versus the traditional regret; e.g., in UCT [Kocsis and Szepesvári, 2006]. (iii) We prove that with an approximate forward model, the gradient bias is proportional to the approximation error , while retaining the variance decay. This quantifies the accuracy required from a learned forward model. (iv) We implement a differentiable deep version of SoftTreeMax that employs a parallelized GPU tree expansion. We demonstrate how its gradient variance is reduced by three orders of magnitude over PPO while obtaining up to 5x reward.

2 Preliminaries 68

Let ∆ U denote simplex over the set U. Throughout, we consider a discounted Markov Decision 69 Process (MDP) M = ( S , A , P, r, γ, ν ) , where S is a finite state space of size S , A is a finite action 70 space of size A , r : S × A → [0 , 1] is the reward function, P : S × A → ∆ S is the transition 71 function, γ ∈ (0 , 1) is the discount factor, and ν ∈ R S is the initial state distribution. We denote 72 the transition matrix starting from state s by P s ∈ [0 , 1] A × S , i.e., [ P s ] a,s ′ = P ( s ′ | a, s ) . Similarly, 73 let R s = r ( s, · ) ∈ R A denote the corresponding reward vector. Separately, let π : S → ∆ A be a 74 stationary policy. Let P π and R π be the induced transition matrix and reward function, respectively, 75 i.e., P π ( s ′ | s ) = ∑ a π ( a | s ) Pr( s ′ | s, a ) and R π ( s ) = ∑ a π ( a | s ) r ( s, a ) . Denote the stationary 76 distribution of P π by µ π ∈ R S s.t. µ ⊤ π P π = P π , and the discounted state visitation frequency 77 by d π so that d ⊤ π = (1 -γ ) ∑ ∞ t =0 γ t ν ⊤ ( P π ) t . Also, let V π ∈ R S be the value function of π 78 defined by V π ( s ) = E π [ ∑ ∞ t =0 γ t r ( s t , π ( s t )) | s 0 = s ] , and let Q π ∈ R S × A be the Q-function 79 such that Q π ( s, a ) = E π [ r ( s, a ) + γV π ( s ′ )] . Our goal is to find an optimal policy π ⋆ such that 80 V ⋆ ( s ) ≡ V π ⋆ ( s ) = max π V π ( s ) , ∀ s ∈ S . 81

82

83

84

85

86

For the analysis in Section 4, we introduce the following notation. Denote by Θ ∈ R S the vector representation of θ ( s ) ∀ s ∈ S . For a vector u , denote by exp( u ) the coordinate-wise exponent of u and by D ( u ) the diagonal square matrix with u in its diagonal. For a matrix A , denote its i -th eigenvalue by λ i ( A ) . Denote the k -dimensional identity matrix and all-ones vector by I k and 1 k , respectively. Also, denote the trace operator by Tr . Finally, we treat all vectors as column vectors.

2.1 Policy Gradient 87

PGschemes seek to maximize the cumulative reward as a function of the policy π θ ( a | s ) by performing 88 gradient steps on θ . The celebrated Policy Gradient Theorem [Sutton et al., 1999] states that 89

The variance of the gradient is thus 90

In the notation above, we denote the variance of a vector random variable X by

similarly as in [Greensmith et al., 2004]. From now on, we drop the subscript from Var in (1) for brevity. When the action space is discrete, a commonly used parameterized policy is softmax: π θ ( a | s ) ∝ exp( θ ( s, a )) , where θ : S × A → R is a state-action parameterization.

3 SoftTreeMax: Exponent of trajectories

We introduce a new family of policies called SoftTreeMax, which are a model-based generalization of the popular softmax. We propose two variants: Cumulative (C-SoftTreeMax) and Exponentiated (E-SoftTreeMax). In both variants, we replace the generic softmax logits θ ( s, a ) with the score of a trajectory of horizon d starting from ( s, a ) , generated by applying a behavior policy π b . In C-SoftTreeMax, we exponentiate the expectation of the logits. In E-SoftTreeMax, we first exponentiate the logits and then only compute their expectation.

Logits . We define the SoftTreeMax logit ℓ s,a ( d ; θ ) to be the random variable depicting the score of a trajectory of horizon d starting from ( s, a ) and following the policy π b :

In the above expression, note that s 0 = s, a 0 = a, a t ∼ π b ( ·| s t ) ∀ t ≥ 1 , and r t ≡ r ( s t , a t ) . For brevity of the analysis, we let the parametric score θ in (2) be state-based, similarly to a value function. Instead, one could use a state-action input analogous to a Q-function. Thus, SoftTreeMax can be integrated into the two types of implementation of RL algorithms in standard packages. Lastly, the preceding γ -d scales the θ parametrization to correspond to its softmax counerpart.

C-SoftTreeMax . Given an inverse temperature parameter β , we let C-SoftTreeMax be

C-SoftTreeMax gives higher weight to actions that result in higher expected returns. While standard softmax relies entirely on parametrization θ, C-SoftTreeMax also interpolates a Monte-Carlo portion of the reward.

E-SoftTreeMax . The second operator we propose is E-SoftTreeMax:

here, the expectation is taken outside the exponent. This objective corresponds to the exponentiated reward objective which is often used for risk-sensitive RL [Howard and Matheson, 1972, Fei et al., 2021, Noorani and Baras, 2021]. The common risk-sensitive objective is of the form log E [exp( δR )] , where δ is the risk parameter and R is the cumulative reward. Similarly to that literature, the exponent in (4) emphasizes the most promising trajectories.

SoftTreeMax properties . SoftTreeMax is a natural model-based generalization of softmax. For d = 0 , both variants above coincide since (2) becomes deterministic. In that case, for a state-action parametrization, they reduce to standard softmax. When β → 0 , both variants again coincide and sample actions uniformly (exploration). When β → ∞ , the policies become deterministic and

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

greedily optimize for the best trajectory (exploitation). For C-SoftTreeMax, the best trajectory is defined in expectation, while for E-SoftTreeMax it is defined in terms of the best sample path.

SoftTreeMax convergence. Under regularity conditions, for any parametric policy, PG converges to local optima [Bhatnagar et al., 2009], and thus also SoftTreeMax. For softmax PG, asymptotic [Agarwal et al., 2021] and rate results [Mei et al., 2020b] were recently obtained, by showing that the gradient is strictly positive everywhere [Mei et al., 2020b, Lemmas 8-9]. We conjecture that SoftTreeMax satisfies the same property, being a generalization of softmax, but formally proving it is subject to future work.

SoftTreeMax gradient. The two variants of SoftTreeMax involve an expectation taken over S d many trajectories from the root state s and weighted according to their probability. Thus, during the PG training process, the gradient ∇ θ log π θ is calculated using a weighted sum of gradients over all reachable states starting from s . Our method exploits the exponential number of trajectories to reduce the variance while improving performance. Indeed, in the next section we prove that the gradient variance of SoftTreeMax decays exponentially fast as a function of the behavior policy π b and trajectory length d . In the experiments in Section 6, we also show how the practical version of SoftTreeMax achieves a significant reduction in the noise of the PG process and leads to faster convergence and higher reward.

4 Theoretical Analysis

In this section, we first bound the variance of PG when using the SoftTreeMax policy. Later, we discuss how the gradient bias resulting due to approximate forward models diminishes as a function of the approximation error, while retaining the same variance decay.

We show that the variance decreases exponentially with the tree depth, and the rate is determined by the second eigenvalue of the transition kernel induced by π b . Specifically, we bound the same expression for variance as appears in [Greensmith et al., 2004, Sec. 3.5] and [Wu et al., 2018, Sec. A, Eq. (21)]. Other types of analysis could instead have focused on the estimation aspect in the context of sampling [Zhang et al., 2021, Shen et al., 2019, Pham et al., 2020]. Indeed, in our implementation in Section 5, we manage to avoid sampling and directly compute the expectations in Eqs. (3) and (4). As we show later, we do so by leveraging efficient parallel simulation on the GPU in feasible run-time. In our application, due to the nature of the finite action space and quasi-deterministic Atari dynamics [Bellemare et al., 2013], our expectation estimator is noiseless. We encourage future work to account for the finite-sample variance component. We defer all the proofs to Appendix A.

We begin with a general variance bound that holds for any parametric policy.

Lemma 4.1 (Bound on the policy gradient variance) . Let ∇ θ log π θ ( ·| s ) ∈ R A × dim( θ ) be a matrix whose a -th row is ∇ log π ( a | s ) ⊤ . For any parametric policy π and function Q π θ : S × A → R ,

Hence, to bound (1), it is sufficient to bound the Frobenius norm ∥∇ θ log π θ ( ·| s ) ∥ F for any s .

Note that SoftTreeMax does not reduce the gradient uniformly, which would have been equivalent to a trivial change in the learning rate. While the gradient norm shrinks, the gradient itself scales differently along the different coordinates. This scaling occurs along different eigenvectors, as a function of problem parameters ( P , θ ) and our choice of behavior policy ( π b ), as can be seen in the proof of the upcoming Theorem 4.4. This allows SoftTreeMax to learn a good 'shrinkage' that, while reducing the overall gradient, still updates the policy quickly enough. This reduction in norm and variance resembles the idea of gradient clipping Zhang et al. [2019], where the gradient is scaled to reduce its variance, thus increasing stability and improving overall performance.

A common assumption in the RL literature [Szepesvári, 2010] that we adopt for the remainder of the section is that the transition matrix P π b , induced by the behavior policy π b , is irreducible and aperiodic. Consequently, its second highest eigenvalue satisfies | λ 2 ( P π b ) | < 1 .

From now on, we divide the variance results for the two variants of SoftTreeMax into two subsections. For C-SoftTreeMax, the analysis is simpler and we provide an exact bound. The case of E-SoftTreeMax is more involved and we provide for it a more general result. In both cases, we show that the variance decays exponentially with the planning horizon.

4.1 Variance of C-SoftTreeMax 173

We express C-SoftTreeMax in vector form as follows. 174

Lemma 4.2 (Vector form of C-SoftTreeMax) . For d ≥ 1 , (3) is given by 175

The vector C s,d ∈ R A represents the cumulative discounted reward in expectation along the trajectory of horizon d. This trajectory starts at state s, involves an initial reward dictated by R s and an initial transition as per P s . Thereafter, it involves rewards and transitions specified by R π b and P π b , respectively. Once the trajectory reaches depth d, the score function θ ( s d ) is applied,.

Lemma 4.3 (Gradient of C-SoftTreeMax) . The C-SoftTreeMax gradient is given by

in R A × S , where for brevity, we drop the s index in the policy above, i.e., π C d,θ ≡ π C d,θ ( ·| s ) .

We are now ready to present our first main result:

Theorem 4.4 (Variance decay of C-SoftTreeMax) . For every Q : S × A → R , the C-SoftTreeMax policy gradient variance is bounded by

We provide the full proof in Appendix A.4, and briefly outline its essence here.

Proof outline. Lemma 4.1 allows us to bound the variance using a direct bound on the gradient norm. The gradient is given in Lemma 4.3 as a product of three matrices, which we now study from right to left. The matrix P π b is a row-stochastic matrix. Because the associated Markov chain is irreducible and aperiodic, it has a unique stationary distribution. This implies that P π b has one and only one eigenvalue equal to 1; all others have magnitude strictly less than 1 . Let us suppose that all these other eigenvalues have multiplicity 1 (the general case with repeated eigenvalues can be handled via Jordan decompositions as in [Pelletier, 1998, Lemma1]). Then, P π b has the spectral decomposition P π b = 1 S µ ⊤ π b + ∑ S i =2 λ i v i u ⊤ i , where λ i is the i -th eigenvalue of P π b (ordered in descending order according to their magnitude) and u i and v i are the corresponding left and right eigenvectors, respectively, and therefore ( P π b ) d -1 = 1 S µ ⊤ π b + ∑ S i =2 λ d -1 i v i u ⊤ i .

The second matrix in the gradient relation in Lemma 4.3, P s , is a rectangular transition matrix that translates the vector of all ones from dimension S to A : P s 1 S = 1 A . Lastly, the first matrix [ I A -1 A ( π C d,θ ) ⊤ ] is a projection whose null-space includes the vector 1 A , i.e., [ I A -1 A ( π C d,θ ) ⊤ ] 1 A = 0 . Combining the three properties above when multiplying the three matrices of the gradient, it is easy to see that the first term in the expression for ( P π b ) d -1 gets canceled, and we are left with bounded summands scaled by λ i ( P π b ) d -1 . Recalling that | λ i ( P π b ) | < 1 and that | λ 2 | ≥ | λ 3 | ≥ . . . for i = 2 , . . . , S, we obtain the desired result.

Theorem 4.4 guarantees that the variance of the gradient decays exponentially with d. It also provides a novel insight for choosing the behavior policy π b as the policy that minimizes the absolute second eigenvalue of the P π b . Indeed, the second eigenvalue of a Markov chain relates to its connectivity and its rate of convergence to the stationary distribution [Levin and Peres, 2017].

Optimal variance decay . For the strongest reduction in variance, the behavior policy π b should be chosen to achieve an induced Markov chain whose transitions are state-independent. In that case, P π b where 176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

is a rank one matrix of the form 1 S µ ⊤ π b , and λ 2 ( P π b ) = 0 . Then, Var( ∇ θ log π θ ( a | s ) Q ( s, a )) = 0 . 210 Naturally, this can only be done for pathological MDPs; see Appendix C.1 for a more detailed 211 discussion. Nevertheless, as we show in Section 5, we choose our tree expansion policy to reduce the 212 variance as best as possible. 213

214

215

216

217

Worst-case variance decay . In contrast, and somewhat surprisingly, when π b is chosen so that the dynamics is deterministic, there is no guarantee that it will decay exponentially fast. For example, if P π b is a permutation matrix, then λ 2 ( P π b ) = 1 , and advancing the tree amounts to only updating the gradient of one state for every action, as in the basic softmax.

218

219

220

221

222

223

224

225

226

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

247

4.2 Variance of E-SoftTreeMax

The proof of the variance bound for E-SoftTreeMax is similar to that of C-SoftTreeMax, but more involved. It also requires the assumption that the reward depends only on the state, i.e. r ( s, a ) ≡ r ( s ) . This is indeed the case in most standard RL environments such as Atari and Mujoco.

Lemma 4.5 (Vector form of E-SoftTreeMax) . For d ≥ 1 , (4) is given by

The vector R above is the S -dimensional vector whose s -th coordinate is r ( s ) .

The matrix E s,d ∈ R A × S has a similar role to C s,d from (5), but it represents the exponentiated cumulative discounted reward. Accordingly, it is a product of d matrices as opposed to a sum. It captures the expected reward sequence starting from s and then iteratively following P π b . After d steps, we apply the score function on the last state as in (6).

Lemma 4.6 (Gradient of E-SoftTreeMax) . The E-SoftTreeMax gradient is given by

where for brevity, we drop the s index in the policy above, i.e., π E d,θ ≡ π E d,θ ( ·| s ) .

This gradient structure is harder to handle than that of C-SoftTreeMax in Lemma 4.3, but here we also can bound the decay of the variance nonetheless.

Theorem 4.7 (Variance decay of E-SoftTreeMax) . There exists α ∈ (0 , 1) such that,

for every Q. Further, if P π b is reversible or if the reward is constant, then α = | λ 2 ( P π b ) | .

Theory versus Practice. We demonstrate the above result in simulation. We draw a random finite MDP, parameter vector Θ ∈ R S + , and behavior policy π b . We then empirically compute the PG variance of E-SoftTreeMax as given in (1) and compare it to | λ 2 ( P π b ) | d . We repeat this experiment three times for different P π b : (i) close to uniform, (ii) drawn randomly, and (iii) close to a permutation matrix. As seen in Figure 1, the empirical variance and our bound match almost identically. This also suggests that α = | λ 2 ( P π b ) | in the general case and not only when P π b is reversible or when the reward is constant.

4.3 Bias with an Approximate Forward Model

The definition of the two SoftTreeMax variants involves the knowledge of the underlying environment, in particular the value of P and r. However, in practice, we often can only learn approximations of the dynamics from interactions, e.g., using NNs [Ha and Schmidhuber, 2018, Schrittwieser et al., 2020]. Let ˆ P and ˆ r denote the approximate kernel and reward functions, respectively. In this section, we study the consequences of the approximation error on the C-SoftTreeMax gradient.

where

2

4

6

8

10

Depth d

10

34

10

28

10

22

10

16

10

10

10

4

10 2

SoftTreeMax

Gradient variance

Permutation: Empirical variance

Permutation: Variance bound

Random: Empirical variance

Random: Variance bound

Uniform: Empirical variance

Uniform: Variance bound

𝑊(𝑆

!"#

$,$

)

logits

for

logits

for

logits

for

𝑎

!"$

($)

Policy

network

𝑎

!"$

(()

𝑎

!"$

($)

𝑎

!"$

(()

𝑎

!"$

($)

𝑎

!"$

(()

𝑎

!")

($)

𝑎

!")

(#)

𝑎

!")

(()

𝑎

!"#

(%)

𝑎

!"#

(')

𝑎

!"#

(()

𝑆

!")

𝑆

!"$

($)

𝑆

!"$

(#)

𝑆

!"$

(()

𝑆

!"#

($,$)

𝑆

!"#

(#,$)

𝑆

!"#

($,()

𝑆

!"#

(#,()

𝑆

!"#

((,$)

𝑆

!"#

((,()

𝑊(𝑆

!"#

$,(

)

𝑊(𝑆

!"#

#,$

)

𝑊(𝑆

!"#

#,(

)

𝑊(𝑆

!"#

(,$

)

𝑊(𝑆

!"#

(,(

)

248

249

250

251

252

253

Let ˆ π C d,θ be the C-SoftTreeMax policy defined given the approximate forward model introduced above. That is, let ˆ π C d,θ be defined exactly as in (5), but using ˆ R s , ˆ P s , ˆ R π b and ˆ P π b , instead of their unperturbed counterparts from Section 2. Then, the variance of the corresponding gradient again decays exponentially with a decay rate of λ 2 ( ˆ P π b ) . However, a gradient bias is introduced. In the following, we bound this bias in terms of the approximation error and other problem parameters. The proof is provided in Appendix A.9.

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

Theorem 4.8. Let ϵ be the maximal model mis-specification, i.e., let max {∥ P -ˆ P ∥ , ∥ r -ˆ r ∥} = ϵ. Then the policy gradient bias due to ˆ π C d,θ satisfies

To the best of our knowledge, Theorem 4.8 is the first result that bounds the bias of the gradient of a parametric policy due to an approximate model. It states that if the learned model is accurate enough, we expect similar convergence properties for C-SoftTreeMax as we would have obtained with the true dynamics. It also suggests that higher temperature (lower β ) reduces the bias. In this case, the logits get less weight, with the extreme of β = 0 corresponding to a uniform policy that has no bias. Lastly, the error scales linearly with d : the policy suffers from cumulative error as it relies on further-looking states in the approximate model.

5 SoftTreeMax: Deep Parallel Implementation

Following impressive successes of deep RL [Mnih et al., 2015, Silver et al., 2016], using deep NNs in RL is standard practice. Depending on the RL algorithm, a loss function is defined and gradients on the network weights can be calculated. In PG methods, the scoring function used in the softmax is commonly replaced by a neural network W θ : π θ ( a | s ) ∝ exp( W θ ( s, a )) . Similarly, we implement SoftTreeMax by replacing θ ( s ) in (2) with a neural network W θ ( s ) . Although both variants of

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

314

315

316

317

SoftTreeMax from Section 3 involve computing an expectation, this can be hard in general. One approach to handle it is with sampling, though these introduce estimation variance into the process. We leave the question of sample-based theory and algorithmic implementations for future work.

Instead, in finite action space environments such as Atari, we compute the exact expectation in SoftTreeMax with an exhaustive TS of depth d . Despite the exponential computational cost of spanning the entire tree, recent advancements in parallel GPU-based simulation allow efficient expansion of all nodes at the same depth simultaneously [Dalal et al., 2021, Rosenberg et al., 2022]. This is possible when a simulator is implemented on GPU [Dalton et al., 2020, Makoviychuk et al., 2021, Freeman et al., 2021], or when a forward model is learned [Kim et al., 2020, Ha and Schmidhuber, 2018]. To reduce the complexity to be linear in depth, we apply tree pruning to a limited width in all levels. We do so by sub-sampling only the most promising branches at each level. Limiting the width drastically improves runtime, and enables respecting GPU memory limits, with only a small sacrifice in performance.

To summarize, in the practical SoftTreeMax algorithm we perform an exhaustive tree expansion with pruning to obtain trajectories up to depth d. We expand the tree with equal weight to all actions, which corresponds to a uniform tree expansion policy π b . We apply a neural network on the leaf states, and accumulate the result with the rewards along each trajectory to obtain the logits in (2). Finally, we aggregate the results using C-SoftTreeMax. We leave experiments E-SoftTreeMax for future work on risk-averse RL. During training, the gradient propagates to the NN weights of W θ . When the gradient ∇ θ log π d,θ is calculated at each time step, it updates W θ for all leaf states, similarly to Siamese networks [Bertinetto et al., 2016]. An illustration of the policy is given in Figure 2.

6 Experiments

We conduct our experiments on multiple games from the Atari simulation suite [Bellemare et al., 2013]. As a baseline, we train a PPO [Schulman et al., 2017] agent with 256 GPU workers in parallel [Dalton et al., 2020]. For the tree expansion, we employ a GPU breadth-first as in [Dalal et al., 2021]. We then train C-SoftTreeMax 1 for depths d = 1 . . . 8 , with a single worker. For depths d ≥ 3 , we limited the tree to a maximum width of 1024 nodes and pruned trajectories with low estimated weights. Since the distributed PPO baseline advances significantly faster in terms of environment steps, for a fair comparison, we ran all experiments for one week on the same machine. For more details see Appendix B.

In Figure 3, we plot the reward and variance of SoftTreeMax for each game, as a function of depth. The dashed lines are the results for PPO. Each value is taken after convergence, i.e., the average over the last 20% of the run. The numbers represent the average over five seeds per game. The plot conveys three intriguing conclusions. First, in all games, SoftTreeMax achieves significantly higher reward than PPO. Its gradient variance is also orders of magnitude lower than that of PPO. Second, the reward and variance are negatively correlated and mirror each other in almost all games. This phenomenon demonstrates the necessity of reducing the variance of PG for improving performance. Lastly, each game has a different sweet spot in terms of optimal tree depth. Recall that we limit the run-time in all experiments to one week The deeper the tree, the slower each step and the run consists of less steps. This explains the non-monotone behavior as a function of depth. For a more thorough discussion on the sweet spot of different games, see Appendix B.3.

7 Related Work

Softmax Operator. The softmax policy became a canonical part of PG to the point where theoretical results of PG focus specifically on it [Zhang et al., 2021, Mei et al., 2020b, Li et al., 2021, Ding et al., 2022]. Even though we focus on a tree extension to the softmax policy, our methodology is general and can be easily applied to other discrete or continuous parameterized policies as in [Mei et al., 2020a, Miahi et al., 2021, Silva et al., 2019]. Tree Search. One famous TS algorithm is Monte-Carlo TS (MCTS; [Browne et al., 2012]) used in AlphaGo [Silver et al., 2016] and MuZero [Schrittwieser et al., 2020]. Other algorithms such as Value Iteration, Policy Iteration and DQN were also shown to

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

Asteroids

Breakout

Gopher

Krull

Gradientvariance

5000

10-5

800

15000

10-5

9000

(log scale)

Reward

4000

10-7

600

10-5

10-6

10000

10-6

8000

10-7

3000

10-7

5000

10-7

7000

10-9

SoftTreeMaxReward

2

4

6

8

400

2

4

6

8

2

4

6

8

N

4

6

8

PPO Reward

Depth

Depth

Depth

Depth

SoftTreeMaxVariance

PPOVariance

KungFuMaster

10-5

20000

NameThisGame

10-5

75000

Phoenix

10-5

800000

VideoPinball

Gradientvariance

(log scale)

Reward

60000

15000

10-6

50000

600000

10-6

40000

10-7

10000

10-725000

10-7

400000

200000

10-8

2

4

6

8

2

4

6

8

2

4

6

8

2

4

6

8

Depth

Depth

Depth

Depth

give an improved performance with a tree search extensions [Efroni et al., 2019, Dalal et al., 2021]. Parallel Environments. In this work we used accurate parallel models that are becoming more common with the increasing popularity of GPU-based simulation [Makoviychuk et al., 2021, Dalton et al., 2020, Freeman et al., 2021]. Alternatively, in relation to Theorem 4.8, one can rely on recent works that learn the underlying model [Ha and Schmidhuber, 2018, Schrittwieser et al., 2020] and use an approximation of the true dynamics. Risk Aversion. Previous work considered exponential utility functions for risk aversion [Chen et al., 2007, Garcıa and Fernández, 2015, Fei et al., 2021]. This utility function is the same as E-SoftTreeMax formulation from (4), but we have it directly in the policy instead of the objective. Reward-free RL. We showed that the gradient variance is minimized when the transitions induced by the behavior policy π b are uniform. This is expressed by the second eigenvalue of the transition matrix P π b . This notion of uniform exploration is common to the reward-free RL setup [Jin et al., 2020]. Several such works also considered the second eigenvalue in their analysis [Liu and Brunskill, 2018, Tarbouriech and Lazaric, 2019].

8 Discussion

In this work, we introduced for the first time a differentiable parametric policy that combines TS with PG. We proved that SoftTreeMax is essentially a variance reduction technique and explained how to choose the expansion policy to minimize the gradient variance. It is an open question whether optimal variance reduction corresponds to the appealing regret properties the were put forward by UCT [Kocsis and Szepesvári, 2006]. We believe that this can be answered by analyzing the convergence rate of SoftTreeMax, relying on the bias and variance results we obtained here.

As the learning process continues, the norm of the gradient and the variance both become smaller. On the face of it, one can ask if the gradient becomes small as fast as the variance or even faster can there be any meaningful learning? As we showed in the experiments, learning happens because the variance reduces fast enough (a variance of 0 represents deterministic learning, which is fastest).

Finally, our work can be extended to infinite action spaces. The analysis can be extended to infinitedimension kernels that retain the same key properties used in our proofs. In the implementation, the tree of continuous actions can be expanded by maintaining a parametric distribution over actions that depend on θ . This approach can be seen as a tree adaptation of MPPI [Williams et al., 2017].

9 Reproducibility and Limitations

In this submission, we include the code as part of the supplementary material. We also include a docker file for setting up the environment and a README file with instructions on how to run both training and evaluation. The environment engine is an extension of Atari-CuLE [Dalton et al., 2020], a CUDA-based Atari emulator that runs on GPU. Our usage of a GPU environment is both a novelty and a current limitation of our work.

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

386

387

388

389

390

391

392

393

394

395

396

397

398

References

Agarwal, S. M. Kakade, J. D. Lee, and G. Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. J. Mach. Learn. Res. , 22(98):1-76, 2021.

G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. Journal of Artificial Intelligence Research , 47:253-279, 2013.

Bertinetto, J. Valmadre, J. F. Henriques, A. Vedaldi, and P. H. Torr. Fully-convolutional siamese networks for object tracking. In European conference on computer vision , pages 850-865. Springer, 2016.

Bhatnagar, R. S. Sutton, M. Ghavamzadeh, and M. Lee. Natural actor-critic algorithms. Automatica , 45(11):2471-2482, 2009.

B. Browne, E. Powley, D. Whitehouse, S. M. Lucas, P. I. Cowling, P. Rohlfshagen, S. Tavener, D. Perez, S. Samothrakis, and S. Colton. A survey of monte carlo tree search methods. IEEE Transactions on Computational Intelligence and AI in games , 4(1):1-43, 2012.

Chatterjee and E. Seneta. Towards consensus: Some convergence theorems on repeated averaging. Journal of Applied Probability , 14(1):89-97, 1977.

Chen, M. Sim, D. Simchi-Levi, and P. Sun. Risk aversion in inventory management. Operations Research , 55(5):828-842, 2007.

Dalal, A. Hallak, S. Dalton, S. Mannor, G. Chechik, et al. Improve agents without retraining: Parallel tree search with off-policy correction. Advances in Neural Information Processing Systems , 34:5518-5530, 2021.

Dalton et al. Accelerating reinforcement learning through gpu atari emulation. Advances in Neural Information Processing Systems , 33:19773-19782, 2020.

Ding, J. Zhang, and J. Lavaei. On the global optimum convergence of momentum-based policy gradient. In International Conference on Artificial Intelligence and Statistics , pages 1910-1934. PMLR, 2022.

Efroni, G. Dalal, B. Scherrer, and S. Mannor. How to combine tree-search methods in reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 33, pages 3494-3501, 2019.

Fei, Z. Yang, Y. Chen, and Z. Wang. Exponential bellman equation and improved regret bounds for risk-sensitive reinforcement learning. Advances in Neural Information Processing Systems , 34: 20436-20446, 2021.

D. Freeman, E. Frey, A. Raichuk, S. Girgin, I. Mordatch, and O. Bachem. Brax-a differentiable physics engine for large scale rigid body simulation. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1) , 2021.

Garcıa and F. Fernández. A comprehensive survey on safe reinforcement learning. Journal of Machine Learning Research , 16(1):1437-1480, 2015.

Greensmith, P. L. Bartlett, and J. Baxter. Variance reduction techniques for gradient estimates in reinforcement learning. Journal of Machine Learning Research , 5(9), 2004.

Ha and J. Schmidhuber. World models. arXiv preprint arXiv:1803.10122 , 2018.

A. Howard and J. E. Matheson. Risk-sensitive markov decision processes. Management science , 18(7):356-369, 1972.

Jin, A. Krishnamurthy, M. Simchowitz, and T. Yu. Reward-free exploration for reinforcement learning. In International Conference on Machine Learning , pages 4870-4879. PMLR, 2020.

W. Kim, Y. Zhou, J. Philion, A. Torralba, and S. Fidler. Learning to simulate dynamic environments with gamegan. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 1231-1240, 2020.

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

Kocsis and C. Szepesvári. Bandit based monte-carlo planning. In European conference on machine learning , pages 282-293. Springer, 2006.

A. Levin and Y. Peres. Markov chains and mixing times , volume 107. American Mathematical Soc., 2017.

Li, Y. Wei, Y. Chi, Y. Gu, and Y. Chen. Softmax policy gradient methods can take exponential time to converge. In Conference on Learning Theory , pages 3107-3110. PMLR, 2021.

Liu and E. Brunskill. When simple exploration is sample efficient: Identifying sufficient conditions for random exploration to yield pac rl algorithms. arXiv preprint arXiv:1805.09045 , 2018.

Liu, K. Zhang, T. Basar, and W. Yin. An improved analysis of (variance-reduced) policy gradient and natural policy gradient methods. Advances in Neural Information Processing Systems , 33: 7624-7636, 2020.

Makoviychuk, L. Wawrzyniak, Y. Guo, M. Lu, K. Storey, M. Macklin, D. Hoeller, N. Rudin, A. Allshire, A. Handa, et al. Isaac gym: High performance gpu-based physics simulation for robot learning. arXiv preprint arXiv:2108.10470 , 2021.

S. Mathkar and V. S. Borkar. Nonlinear gossip. SIAM Journal on Control and Optimization , 54 (3):1535-1557, 2016.

Mei, C. Xiao, B. Dai, L. Li, C. Szepesvári, and D. Schuurmans. Escaping the gravitational pull of softmax. Advances in Neural Information Processing Systems , 33:21130-21140, 2020a.

Mei, C. Xiao, C. Szepesvari, and D. Schuurmans. On the global convergence rates of softmax policy gradient methods. In International Conference on Machine Learning , pages 6820-6829. PMLR, 2020b.

Miahi, R. MacQueen, A. Ayoub, A. Masoumzadeh, and M. White. Resmax: An alternative soft-greedy operator for reinforcement learning. 2021.

Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 518(7540):529-533, 2015.

Noorani and J. S. Baras. Risk-sensitive reinforce: A monte carlo policy gradient algorithm for exponential performance criteria. In 2021 60th IEEE Conference on Decision and Control (CDC) , pages 1522-1527. IEEE, 2021.

Papini, D. Binaghi, G. Canonaco, M. Pirotta, and M. Restelli. Stochastic variance-reduced policy gradient. In International conference on machine learning , pages 4026-4035. PMLR, 2018.

Pelletier. On the almost sure asymptotic behaviour of stochastic algorithms. Stochastic processes and their applications , 78(2):217-244, 1998.

Pham, L. Nguyen, D. Phan, P. H. Nguyen, M. Dijk, and Q. Tran-Dinh. A hybrid stochastic policy gradient algorithm for reinforcement learning. In International Conference on Artificial Intelligence and Statistics , pages 374-385. PMLR, 2020.

Raffin, A. Hill, M. Ernestus, A. Gleave, A. Kanervisto, and N. Dormann. Stable baselines3, 2019.

Rosenberg, A. Hallak, S. Mannor, G. Chechik, and G. Dalal. Planning and learning with adaptive lookahead. arXiv preprint arXiv:2201.12403 , 2022.

Schrittwieser, I. Antonoglou, T. Hubert, K. Simonyan, L. Sifre, S. Schmitt, A. Guez, E. Lockhart, D. Hassabis, T. Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature , 588(7839):604-609, 2020.

Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017.

Shen, A. Ribeiro, H. Hassani, H. Qian, and C. Mi. Hessian aided policy gradient. In International conference on machine learning , pages 5729-5738. PMLR, 2019.

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

463

464

465

466

467

468

469

470

471

472

Silva, T. Killian, I. D. J. Rodriguez, S.-H. Son, and M. Gombolay. Optimization methods for interpretable differentiable decision trees in reinforcement learning. arXiv preprint arXiv:1903.09338 , 2019.

Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature , 529(7587):484-489, 2016.

S. Sutton, D. McAllester, S. Singh, and Y. Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems , 12, 1999.

Szepesvári. Algorithms for reinforcement learning. Synthesis lectures on artificial intelligence and machine learning , 4(1):1-103, 2010.

Tarbouriech and A. Lazaric. Active exploration in markov decision processes. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 974-982. PMLR, 2019.

S. Thomas and E. Brunskill. Policy gradient methods for reinforcement learning with function approximation and action-dependent baselines. arXiv preprint arXiv:1706.06643 , 2017.

Williams, N. Wagener, B. Goldfain, P. Drews, J. M. Rehg, B. Boots, and E. A. Theodorou. Information theoretic mpc for model-based reinforcement learning. In 2017 IEEE International Conference on Robotics and Automation (ICRA) , pages 1714-1721. IEEE, 2017.

Wu, A. Rajeswaran, Y. Duan, V. Kumar, A. M. Bayen, S. Kakade, I. Mordatch, and P. Abbeel. Variance reduction for policy gradient with action-dependent factorized baselines. In International Conference on Learning Representations , 2018.

Xu, F. Gao, and Q. Gu. An improved convergence analysis of stochastic variance-reduced policy gradient. In Uncertainty in Artificial Intelligence , pages 541-551. PMLR, 2020.

Zhang, T. He, S. Sra, and A. Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. arXiv preprint arXiv:1905.11881 , 2019.

Zhang, C. Ni, C. Szepesvari, M. Wang, et al. On the convergence and sample efficiency of variance-reduced policy gradient method. Advances in Neural Information Processing Systems , 34: 2228-2240, 2021.

473

Appendix

A Proofs 474

475

A.1 Proof of Lemma 4.1 - Bound on the policy gradient variance

For any parametric policy π θ and function Q : S × A → R , 476

477

where ∇ θ log π θ ( ·| s ) ∈ R A × dim( θ ) is a matrix whose a -th row is ∇ θ log π θ ( a | s ) ⊤ .

Proof. The variance for a parametric policy π θ is given as follows: 478

where Q ( s, a ) is the currently estimated Q-function and d π θ is the discounted state visitation frequency 479 induced by the policy π θ . Since the second term we subtract is always positive (it is of quadratic form 480 v ⊤ v ) we can bound the variance by the first term: 481

482

A.2 Proof of Lemma 4.2 - Vector form of C-SoftTreeMax 483

In vector form, (3) is given by 484

where 485

Proof. Consider the vector ℓ s, · ∈ R |A| . Its expectation satisfies 486

As required. 487

A.3 Proof of Lemma 4.3 - Gradient of C-SoftTreeMax 488

489

The C-SoftTreeMax gradient of dimension A × S is given by

where for brevity, we drop the s index in the policy above, i.e., π C d,θ ≡ π C d,θ ( ·| s ) . 490

Proof. The ( j, k ) -th entry of ∇ θ log π C d,θ satisifes 491

Moving back to matrix form, we obtain the stated result. 492

A.4 Proof of Theorem 4.4 - Exponential variance decay of C-SoftTreeMax 493

The C-SoftTreeMax policy gradient is bounded by 494

Proof. We use Lemma 4.1 directly. First of all, it is know that when the reward is bounded in [0 , 1] , 495 the maximal value of the Q-function is 1 1 -γ as the sum as infinite discounted rewards. Next, we 496 bound the Frobenius norm of the term achieved in Lemma 4.3, by applying the eigen-decomposition 497 on P π b : 498

where µ is the stationary distribution of P π b , and u i and v i are left and right eigenvectors correspond499 ingly. 500

Now, we can bound the norm ∥ I A,A -1 A π ⊤ ∥ F by direct calculation: 501

From the Cauchy-Schwartz inequality, 502

So, 503

which obtains the desired bound. 504

505

506

507

508

509

510

511

512

513

514

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 to the a lower bound on the variance since we have no lower bound equivalence of Lemma 4.1.

Lemma A.1. The Frobenius norm on the gradient of the policy is lower-bounded by:

Proof. We begin by moving to the induced l 2 norm by norm-equivalence:

Now, taking the vector u to be the eigenvector of the second eigenvalue of P π b :

Note that even though P s u can be 0 , that is not the common case since we can freely change π b (and therefore the eigenvectors of P π b ).

A.6 Proof of Lemma 4.5 - Vector form of E-SoftTreeMax

515

For d ≥ 1 , (4) is given by where 516

517

with R being the | S | -dimensional vector whose s -th coordinate is r ( s ) .

Proof. Recall that 518

and, hence, 519

Therefore, 520

By repeatedly using iterative conditioning as above, the desired result follows. Note that 521 exp( βγ -d r ( s )) does not depend on the action and is therefore cancelled out with the denomi522 nator. 523

A.7 Proof of Lemma 4.6 - Gradient of E-SoftTreeMax 524

The E-SoftTreeMax gradient of dimension A × S is given by 525

where for brevity, we drop the s index in the policy above, i.e., π E d,θ ≡ π E d,θ ( ·| s ) . 526

Proof. The ( j, k ) -th entry of ∇ θ log π E d,θ satisfies 527

Hence, 528

From this, it follows that 529

The desired result is now easy to see. 530

A.8 Proof of Theorem 4.7 - Exponential variance decay of E-SoftTreeMax 531

532

533

534

535

There exists α ∈ (0 , 1) such that, for any function Q : S × A → R ,

If all rewards are equal ( r ≡ const), then α = | λ 2 ( P π b ) | .

Proof outline. Recall that thanks to Lemma 4.1, we can bound the PG variance using a direct bound on the gradient norm. The definition of the induced norm is

with ∇ θ log π E d,θ given in Lemma 4.6. Let z ∈ R S be an arbitrary vector such that ∥ z ∥ = 1 . Then, 536 z = ∑ S i =1 c i z i , where c i are scalar coefficients and z i are vectors spanning the S -dimensional space. 537 In the full proof, we show our specific choice of z i and prove they are linearly independent given that 538 choice. We do note that z 1 = 1 S . 539

The first part of the proof relies on the fact that ( ∇ θ log π E d,θ ) z 1 = 0 . This is easy to verify using Lemma 4.6 together with (6), and because [ I A -1 A ( π E d,θ ) ⊤ ] is a projection matrix whose null-space is spanned by 1 S . Thus,

In the second part of the proof, we focus on E s,d from (6), which appears within ∇ θ log π E d,θ . Notice 540 that E s,d consists of the product ∏ d -1 h =1 ( D ( exp( βγ h -d R ) P π b ) . Even though the elements in this 541 product are not stochastic matrices, in the full proof we show how to normalize each of them to a 542 stochastic matrix B h . We thus obtain that 543

where M 1 ∈ R S is some strictly positive vector. Then, we can apply a result by Mathkar and Borkar [2016], which itself builds on [Chatterjee and Seneta, 1977]. The result states that the product of stochastic matrices ∏ d -1 h =1 B h of our particular form converges exponentially fast to a matrix of the form 1 S µ ⊤ s.t. ∥ 1 S µ ⊤ -∏ d -1 h =1 B h ∥ ≤ Cα d for some constant C.

Lastly, 1 S µ ⊤ π b gets canceled due to our choice of z i , i = 2 , . . . , S. This observation along with the above fact that the remainder decays then shows that ∇ θ log π E d,θ ∑ S i =2 z i = O ( α d ) , which gives the desired result.

Full technical proof. Let d ≥ 2 . Recall that

and that R refers to the S -dimensional vector whose s -th coordinate is r ( s ) . Define

and the vector

where ◦ denotes the element-wise product. Then, 554

It is easy to see that each B i is a row-stochastic matrix, i.e., all entries are non-negative and 555 B i 1 S = 1 S . 556

Next, we prove that all non-zeros entries of B i are bounded away from 0 by a constant. This is 557 necessary to apply the next result from Chatterjee and Seneta [1977]. The j -th coordinate of M i 558 satisfies 559

Separately, observe that ∥ M d -1 ∥ ∞ ≤ ∥ exp( βγ -1 R ) ∥ ∞ . Plugging these relations in (26) gives 560

Similarly, for every 1 ≤ i ≤ d -1 , we have that 561

The jk -th entry of B i = D -1 ( P π b M i +1 ) P π b D ( M i +1 ) is 562

544

545

546

547

548

549

550

551

552

553

Hence, for non-zero P π b jk , the entries are bounded away from zero by the same. We can now proceed 563 with applying the following result. 564

Now, by [Chatterjee and Seneta, 1977, Theorem 5] (see also (14) in [Mathkar and Borkar, 2016]), 565 lim d →∞ ∏ d -1 i =1 B i exists and is of the form 1 S µ ⊤ for some probability vector µ. Furthermore, there 566 is some α ∈ (0 , 1) such that ε ( d ) := ( ∏ d -1 i =1 B i ) -1 S µ ⊤ satisfies 567

Pick linearly independent vectors w 2 , . . . , w S such that 568

Since ∑ S i =2 α i w i is perpendicular to µ for any α 2 , . . . α S and because µ ⊤ exp( β Θ) > 0 , there 569 exists no choice of α 2 , . . . , α S such that ∑ S i =2 α i w i = exp( β Θ) . Hence, if we let z 1 = 1 S and 570 z i = D (exp( β Θ)) -1 w i for i = 2 , . . . , S, then it follows that { z 1 , . . . , z S } is linearly independent. 571 In particular, it implies that { z 1 , . . . , z S } spans R S . 572

Now consider an arbitrary unit norm vector z := ∑ S i =1 c i z i ∈ R S s.t. ∥ z ∥ 2 = 1 . Then, 573

where (34) follows from the fact that ∇ θ log π E d,θ z 1 = ∇ θ log π E d,θ 1 S = 0 , (35) follows from 574 Lemma 4.6, (36) holds since z i = D (exp( β Θ)) -1 w i , (38) because µ is perpendicular w i for each i, 575 while (39) follows by reusing z i = D (exp( β Θ)) -1 w i relation along with the fact that z 1 = 1 S . 576

From (39), it follows that 577

Lastly, we prove that ∥ z -c 1 1 S ∥ is bounded independently of d. First, denote by c = ( c 1 , . . . , c S ) ⊤ 578 and ˜ c = (0 , c 2 , . . . , c S ) ⊤ . Also, denote by Z the matrix with z i as its i -th column. Now, 579

where the last relation is due to z being a unit vector. All matrix norms here are l 2 -induced norms. 580

Next, denote by W the matrix with w i in its i -th column. Recall that in (33) we only defined 581 w 2 , . . . , w S . We now set w 1 = exp( β Θ) . Note that w 1 is linearly independent of { w 2 , . . . , w S } 582 because of (33) together with the fact that µ ⊤ w 1 > 0 . We can now express the relation between Z 583 and W by Z = D -1 (exp( β Θ)) W. Substituting this in (53), we have 584

It further holds that 585

where the last relation equals 1 if θ ( s ) < 0 for all s. Similarly, 586

Furthermore, by the properties of the l 2 -induced norm, 587

The determinant of W is a sum of products involving its entries. To upper bound (65) independently 589 of d, we lower bound its denominator by upper and lower bounds on the entries [ W ] i, 1 that are 590 independent of d, depending on their sign: 591

Using this, together with (53), (55), (56), (57), and (61), we showed that ∥ z -c 1 1 S ∥ is upper bounded 592 by a constant independent of d. This concludes the proof. 593

A.9 Bias Estimates 594

Lemma A.2. For any matrix A and ˆ A,

Proof. The proof follows from first principles: 595

Lastly, 588

596

Henceforth, ∥ · ∥ will refer to ∥ · ∥ ∞ , i.e. the induced infinity norm. Also, for brevity, we denote π C d,θ 597 and ˆ π C d,θ by π θ and ˆ π θ , respectively. Similarly, we use d π θ and d ˆ π θ to denote d π C d,θ and d ˆ π C d,θ . As for 598 the induced norm of the matrix P and its perturbed counterpart ˆ P, which are of size S × A × S, 599 we slightly abuse notation and denote ∥ P -ˆ P ∥ = max s {∥ P s -ˆ P s ∥} , where P s is as defined in 600 Section 2. 601

Definition A.3. Let ϵ be the maximal model mis-specification, i.e., max {∥ P -ˆ P ∥ , ∥ r -ˆ r ∥} = ϵ. 602

603

604

605

606

607

608

609

610

Lemma A.4. Recall the definitions of R s , P s , R π b and P π b from Section 2, and respectively denote their perturbed counterparts by ˆ R s , ˆ P s , ˆ R π b and ˆ P π b . Then, for ϵ defined in Definition A.3,

Proof. The proof follows easily from the fact that the differences above are convex combinations of P -ˆ P and r -ˆ r.

Lemma A.5. Let π θ be as in (5) , and let ˆ π θ also be defined as in (5) , but with R s , P s , P π b replaced by their perturbed counterparts ˆ R s , ˆ P s , ˆ P π b throughout. Then,

Proof. To prove the desired result, we work with (5) to bound the error between R s , P s , P π b , R π b and their perturbed versions.

First, we apply Lemma A.2 together with Lemma A.4 to obtain that ∥ ( P π b ) k -( ˆ P π b ) k ∥ = O ( kϵ ) . Next, denote by M the argument in the exponent in (5), i.e.

Similarly, let ˆ M be the corresponding perturbed sum that relies on ˆ P and ˆ r . Combining the bounds 611 from Lemma A.4, and using the triangle inequality, we have that ∥ ˆ M -M ∥ = O ( βdϵ ) . 612

Eq. (5) states that the C-SoftTreeMax policy in the true environment is π θ = exp( M ) / (1 ⊤ exp( M )) . Similarly define ˆ π θ using ˆ M for the approximate model. Then,

where ◦ denotes element-wise multiplication. Using the above relation, we have that ∥ ˆ π θ -π θ ∥ = ∥ π θ ∥∥ exp( M -ˆ M )1 ⊤ exp( M ) 1 ⊤ exp( ˆ M ) -1 ∥ . Using the relation | e x -1 | = O ( x ) as x → 0 , the desired result follows.

Theorem A.6. Let ϵ be as in Definition A.3. Further let ˆ π C d,θ being the corresponding approximate policy as given in Lemma 4.2. Then, the policy gradient bias is bounded by

613

614

615

616

617

618

We first provide a proof outline for conciseness, and only after it the complete proof. 619

Proof outline. First, we prove that max {∥ R s -ˆ R s ∥ , ∥ P s -ˆ P s ∥ , ∥ R π b -ˆ R π b ∥ , ∥ P π b -ˆ P π b ∥} = O ( ϵ ) . This follows from the fact that the differences above are suitable convex combinations of either the rows of P -ˆ P or r -ˆ r. We use the above observation along with the definitions of π C d,θ and ˆ π C d,θ given in (5) to show that ∥ π C d,θ -ˆ π C d,θ ∥ = O ( βdϵ ) . The proof for the latter builds upon two key facts: (a) ∥ ( P π b ) k -( ˆ P π b ) k ∥ ≤ ∑ k h =1 ∥ ˆ P π b ∥ h -1 ∥ ˆ P π b -P π b ∥∥ p π b ∥ k -h = O ( kϵ ) for any k ≥ 0 , and (b) | e x -1 | = O ( x ) as x → 0 . Next, we decompose the LHS of (7) to get

where X 1 ( s ) = d π C d,θ ( s ) ∈ R , X 2 ( s ) = ( ∇ θ log π C d,θ ( ·| s )) ⊤ ∈ R S × A , X 3 ( s ) = D ( π C d,θ ( ·| s )) ∈ 620 R A × A , X 4 ( s ) = Q π C d,θ ( s, · ) ∈ R A × A , and ˆ X 1 ( s ) , . . . , ˆ X 4 ( s ) are similarly defined with π C d,θ re621 placed by ˆ π C d,θ . Then, we show that, for i = 1 , . . . , 4 , (i) ∥ X i ( s ) -ˆ X i ( s ) ∥ = O ( ϵ ) and (ii) 622 max {∥ X i ∥ , ∥ ˆ X i ∥} is bounded by problem parameters. From this, the desired result follows. 623

Proof. We have 624

where X 1 ( s ) = d π θ ( s ) ∈ R , X 2 ( s ) = ( ∇ θ log π θ ( ·| s )) ⊤ ∈ R S × A , X 3 ( s ) = D ( π θ ( ·| s )) ∈ R A × A , 625 X 4 ( s ) = Q π θ ( s, · ) ∈ R A × A , and ˆ X 1 ( s ) , . . . , ˆ X 4 ( s ) are similarly defined with π θ replaced by ˆ π θ . 626

Therefore, 627

where 628

Next, since d π θ , d ˆ π θ , π θ , and ˆ π θ are all distributions, we have 629

Separately, using Lemma 4.3, we have

630

Since all rows of the above matrices have non-negative entries that add up to 1 , we get 631

In the rest of the proof, we bound each of ∥ X 1 -ˆ X 1 ∥ , . . . , ∥ X 4 -ˆ X 4 ∥ . 632 Finally, 633

Similarly, the same bounds hold for ˆ X 1 , ˆ X 2 , ˆ X 3 and ˆ X 4 . 634

From, we have 635

The last relation follows from the fact that (1 -γ ) -1 = ∑ ∞ t =0 γ t , which in turn implies 636

Next, recall that from Lemma 4.3 that 638

Following the same argument as in (85) and applying Lemma A.2, we have that (93) is O ( βdϵ ) . 640 Similarly, from the argument of (85), Eq. (94) is O ( βϵ ) . Lastly, (95) is O ( βdϵ ) due to Lemma A.5. 641

642

Putting the above three terms together, we have that

Since the state-action value function satisfies the Bellman equation, we have 643

B Experiments

B.1 Implementation Details

The environment engine is the highly efficient Atari-CuLE [Dalton et al., 2020], a CUDA-based version of Atari that runs on GPU. Similarly, we use Atari-CuLE for the GPU-based breadth-first TS as done in Dalal et al. [2021]: In every tree expansion, the state S t is duplicated and concatenated with all possible actions. The resulting tensor is fed into the GPU forward model to generate the tensor of next states ( S 0 t +1 , . . . , S A -1 t +1 ) . The next-state tensor is then duplicated and concatenated again with all possible actions, fed into the forward model, etc. This procedure is repeated until the final depth is reached, for which W θ ( s ) is applied per state.

We train SoftTreeMax for depths d = 1 . . . 8 , with a single worker. We use five seeds for each experiment.

From Lemma A.5, it follows that 637

Then, 639

and 644

Consequently, 645

which finally shows that 646

647

648

649

650

651

652

653

654

655

656

657

658

For the implementation, we extend Stable-Baselines3 [Raffin et al., 2019] with all parameters taken 659 as default from the original PPO paper [Schulman et al., 2017]. For depths d ≥ 3 , we limited the 660 tree to a maximum width of 1024 nodes and pruned non-promising trajectories in terms of estimated 661 weights. Since the distributed PPO baseline advances significantly faster in terms of environment 662 steps, for a fair comparison, we ran all experiments for one week on the same machine and use the 663 wall-clock time as the x-axis. We use Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz equipped with 664 one NVIDIA Tesla V100 32GB. 665

B.2 Time-Based Training Curves 666

We provide the training curves in Figure 4. For brevity, we exclude a few of the depths from the plots. 667 As seen, there is a clear benefit for SoftTreeMax over distributed PPO with the standard softmax 668 policy. In most games, PPO with the SoftTreeMax policy shows very high sample efficiency: it 669 achieves higher episodic reward although it observes much less episodes, for the same running time. 670

Asteroids

Breakout

Gopher

Krull

8000

400

6000

Rewar

3000

4000

6000

200

2000

2000

4000

PPO

SoftTreeMax Depth 2

0

2000

SoftTreeMaxDepth3

0

100

0

100

0

100

0

100

SoftTreeMax Depth5

Time [hours]

Time [hours]

Time [hours]

Time [hours]

SoftTreeMax Depth 6

KungFuMaster

NameThisGame

Phoenix

VideoPinball

60000

30000

200000

Reward

10000

40000

20000

0

5000

10000

20000

0

100

0

100

0

100

0

100

Time [hours]

Time [hours]

Time [hours]

Time [hours]

B.3 Step-Based Training Curves 671

In Figure 5 we also provide the same convergence plots where the x-axis is now the number of online 672 interactions with the environment, thus excluding the tree expansion complexity. As seen, due to the 673 complexity of the tree expansion, less steps are conducted during training (limited to one week) as 674 the depth increases. In this plot, the monotone improvement of the reward with increasing tree depth 675 is noticeable in most games. 676

We note that not for all games we see monotonicity. Our explanation for this phenomenon relates to 677 how immediate reward contributes to performance compared to the value. Different games benefit 678 differently from long-term as opposed to short-term planning. Games that require longer-term 679 planning need a better value estimate. A good value estimate takes longer to obtain with larger depths, 680 in which we apply the network to states that are very different from the ones observed so far in the 681 buffer (recall that as in any deep RL algorithm, we train the model only on states in the buffer). If 682 the model hasn't learned a good enough value function yet, and there is no guiding dense reward 683 along the trajectory, the policy becomes noisier, and can take more steps to converge - even more 684 than those we run in our week-long experiment. 685

686

687

688

689

690

691

692

693

694

For a concrete example, let us compare Breakout to Gopher. Inspecting Fig. 5, we observe that Breakout quickly (and monotonically) gains from large depths since it relies on the short term goal of simply keeping the paddle below the moving ball. In Gopher, however, for large depths (>=5), learning barely started even by the end of the training run. Presumably, this is because the task in Gopher involves multiple considerations and steps: the agent needs to move to the right spot and then hit the mallet the right amount of times, while balancing different locations. This task requires long-term planning and thus depends more strongly on the accuracy of the value function estimate. In that case, for depth 5 or more, we would require more train steps for the value to 'kick in' and become beneficial beyond the gain from the reward in the tree.

695

696

697

698

699

700

701

The figures above convey two key observations that occur for at least some non-zero depth: (1) The final performance with the tree is better than PPO (Fig. 3); and (2) the intermediate step-based results with the tree are better than PPO (Fig. 5). This leads to our main takeaway from this work - there is no reason to believe that the vanilla policy gradient algorithm should be better than a multi-step variant. Indeed, we show that this is not the case.

C Further discussion

C.1 The case of λ 2 ( P π b ) = 0

When P π b is rank one, it is not only its variance that becomes 0 , but also the norm of the gradient 702 itself (similarly to the case of d →∞ ). Note that such a situation will happen rarely, in degenerate 703 MDPs. This is a local minimum for SoftTreeMax and it would cause the PG iteration to get stuck, 704 and to the optimum in the (desired but impractical) case where π b is the optimal policy. However, 705 a similar phenomenon was also discovered in the standard softmax with deterministic policies: 706

θ ( s, a ) → ∞ for one a per s . PG with softmax would suffer very slow convergence near these 707

local equilibria, as observed in Mei et al. [2020a]. To see this, note that the softmax gradient is 708

∇ θ log π θ ( a | s ) = e a -π θ ( ·| s ) , where e a ∈ [0 , 1] A is the vector with 0 everywhere except for the 709

a -th coordinate. I.e., it will be zero for a deterministic policy. SoftTreeMax avoids these local optima 710

by integrating the reward into the policy itself (but may get stuck in another, as discussed above). 711

712

713

714

715

716

717

718

719

720

721

722

723

724

725

726

727

728

729

730

731

732

733

734

735

736

737

738

739

740

741

742

743

744

745

746

747

748

749

750

751

752

753

754

755

756

757

758

759

760

761

762

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We included a relevant section at the end of the paper.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: All proofs can be found in the appendix.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: Yes. We also attached the repository, together with a docker environment, as supplementary material.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We attached the repository, together with a docker environment, as supplementary material.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: All relevant information can be found in the paper and the appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: The paper has no scietal impact.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: The paper poses no such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: [NA]

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

970

971

972

973

974

975

976

977

978

979

980

981

982

983

984

985

986

987

988

989

990

991

992

993

994

995

996

997

998

999

1000

1001

1002

1003

1004

1005

1006

1007

1008

1009

1010

1011

1012

1013

1014

1015

1016

1017

1018

1019

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDisclaimer: I do not have the mathematical background required to check all proofs. It is also my first time reviewing for NeurIPS. However, I have contributions in classical (deep) RL and finishing my PhD.\n\nIn this paper, authors use a novel policy parametrization for RL. Namely, the SoftTreeMax policy that replaces the traditional logit values by small horizon trajectories rewards values. The authors claim and prove that SoftTreeMax policy gradient has less variance than traditional policy gradient. They do experiments with PPO on Atari. \n\nI summarize my review as follows. The idea and work in this paper are novel, strong and well-motivated. However, it is in my opinion poorly presented, and poorly situated compared to existing related work. \n\nI very slightly lean towards accepting this paper as is but give feedbacks for the authors to increase their score in the following sections.\n\nSTRENGTHS:\nOriginality: it is the first time that I see someone use look-ahead information in the policy parametrization. I have already seen it done in the critic part of Policy Gradient (PG) methods, e.g. for advantages estimation in PPO using n-step returns with n>1 and generalized advantage estimations.\n\nContribution: if the proofs of variance reduction as a function of the depth of the tree expansion in the logit values and as a function of the approximation error of the forward models are correct, I believe this work opens a whole new avenue for future work (which, by the way, is missing from the paper).\n\nWEAKNESSES:\nIn the following comments, I assume that the techincal results are correct.\n \nFirst weakness: clarity.\n\nThe proposed algorithm (PPO with SoftTreeMaxPolicy) makes heavy use of three major branches of reinforcement learning namely policy gradient methods (core of the paper), model-based RL (use of a forward model to compute small trajectories), efficient implementations of tree search (control the exponential cost of look-ahead search). Those components are all central to your work but the paper lacks a summary of how those components work together.  I recommend to summarize those in a schematic (see Questions section).  \n\nSecond weakness: baselines and related work.\n\nThe related work mentionned in this work does not clearly help to understand the significance of your work. There are well-studied tools reducing the variance in the gradient estimates (n-steps returns and gae-lambda), and some work combining TS and PG. It would have been nice to see comparisons with those work. I recommend the authors to rewrite and extend their related work section (see Questions section). \n\nREFERENCES.\n\nTS + PG references:\n[An Actor-Critic Algorithm Using a Binary Tree Action Selector, Action Guidance with MCTS for Deep Reinforcement Learning, Policy Gradient Algorithms with Monte-Carlo Tree Search for Non-Markov Decision Processes, Policy Gradient Search: Online Planning\nand Expert Iteration without Search Trees]\n\nBetter sample efficiency and variance reduction:\n[Schulman et al., 2015b; Mnih et al., 2016; Munos\net al., 2016; Schulman et al., 2017; Hessel et al., 2018;\nSchrittwieser et al., 2020; Wurman et al., 2022; Chebotar\net al., 2023; Schwarzer et al., 2023, Brett Daley, Martha White, 2024]\n\nQUESTIONS:\nQuestions: \n\n- How do you learn the forward model? If it is not learned, how do you plan to use your algorithm for environments for which you have to learn the model?\n\n- Is there an Atari game named \"NameThisGame' or is this a typo in your plots ?\n\nSuggestions:\n\nI provide some suggestions to increase the score of the paper. I believe increasing my score by 1 for every three implemented suggestions seems fair.\n\n- Could you please briefly conclude on if it is worth spending all this compute on tree expansion with respect to how much variance is reduced ? \n\n- Could you motivate looking at the variance only at the end of training of PPO (line 301)? \n\n- Could you please add a schematic of the forward pass and gradient accumulation in the fashion of fig 1 from [Fully-Convolutional Siamese Networks for Object Tracking] ? Basically combining your figure 2 with the text in lines 282-285 and lines 649-658 (appendix B1) in a summary figure OR in an algorithm description with pseudo code ?\n\n- Could you please move to appendix results on one of C-SoftTreeMax or E-SoftTreeMax and move in the main paper the walltime plots and sample complexity plots (B2, B3)?\n\n- Could please add a dedicated future work section with dedicated paragraphs on efficient tree expansion implems (better pruning, parallelism, gpu, ...) and on forward model learning ?\n\n- Could you please have a dedicated related work section (you could again move some theory to the appendix) with dedicated paragraphs  on other variance reduction method in PG RL (you could have a look at the introduction of Averaging n-step Returns Reduces Variance in Reinforcement Learning 2024 ), on combining TS and PG and an other paragraph on GPU-based environment that have differentiable models ? \n\nAdvice:\n\nIf I may, having also worked with expensive tree structures, I recommend you to limit your experiments to e.g. depths 2, 4, 6,8 rather than 2,3,4,5,6,7,8. \n\nThank you\n\nLIMITATIONS:\nNo limitations.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-06-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces a model-based online planning method called SoftTreeMax. The method acts as an extension of the softmax, by replacing the logit in softmax with a n-step return. Based on SoftTreeMax, the paper proposes two learning algorithms, C-SoftTreeMax and E-SoftTreeMax. The work includes a mathematical analysis of both algorithms to prove the policy gradient variance is bounded and exponentially decays when the planning horizon becomes longer. Then, the paper provides an empirical test on C-SoftTreeMax to compare the learned policy and the variance with PPO.\n\nSTRENGTHS:\n* The method is theoretically sound. The paper provides a detailed theoretical analysis to show the advantage of the decayed policy gradient variance of the proposed algorithm. \n\n* The paper reports the implementation and computation in detail and ensures reproducibility.\n\n* The paper empirically checks C-SoftTreeMax. The empirical results compare between the performance and the change of variance as the agent learns. Curves clearly suggest the inverse relationship between the performance and the variance, thus supporting the importance of having a method with guaranteed decaying variance.\n\nWEAKNESSES:\n* The main concern I have is the difference between the proposed method and the n-step return. \n\n    In Section 3, the paper defines the SoftTreeMax logit as \n\n    $$l_{s,a}(d; \\theta) = \\gamma^{-d} \\left[ \\sum^{d-1}_{t=0} \\gamma^t r_t + \\gamma^d \\theta(s_d) \\right] .$$\n\n    In the Actor-Critic method, it is normal to use $Q(s,a)$ for the scoring function $\\theta$, and learn a policy $\\pi(a|s) \\propto exp(Q(s,a))$, for example, SAC, a commonly known online learning algorithm, with state-of-the-art performance. In this case, $l_{s,a}(d; \\theta)$ can be rewritten to \n\n    $$ \\gamma^{-d} \\left[ \\sum^{d-1}_{t=0} \\gamma^{t-1} r_t + \\gamma^d E_a Q(s_d, a_d) \\right] .$$\n \n    This equation is exactly the same as the equation in n-step TD, except the normalizer $\\gamma^{-d}. However, the normalizer can be written as part of the temperature in softmax. \n\n    In C-SoftTreeMax, the learned policy is written as $\\pi^C_{d,\\theta} (a|s) \\propto \\exp [\\beta E^{\\pi_b} l_{s,a}(d; \\theta)] $ in formula (3). \n\n    Replacing $l_{s,a}$ with the equation in (2), we have $\\pi^C_{d,\\theta} (a|s) \\propto \\exp [\\beta E^{\\pi_b} [\\gamma^{-d}  \\sum^{d-1}_{t=0} \\gamma^t r_t + \\gamma^d \\theta(s_d) ]] $. \n\n     As the expectation is taken on $\\pi_b$, the normalizer $\\gamma^{-d}$ does not depend on $\\pi_b$, thus could be moved out of the expectation. Then (3) is changed to $$ \\pi^C_{d,\\theta} (a|s) \\propto \\exp [(\\beta\\gamma^{-d}) \\mathbb{E}^{\\pi_b} [ \\sum^{d-1}_{t=0} \\gamma^t r_t + \\gamma^d \\theta(s_d)) ]]  .$$ \n\n     Now the temperature becomes $\\beta\\gamma^{-d}$, which is a tunable parameter, and the expectation term becomes n-step TD when using $Q(s,a)$ for $\\theta$.\n     Similarly, when using the state value $V(s) = \\theta$, the normalizer can be written as part of the temperature, and the expectation term becomes n-step TD.\n     In this case, using the proposed method seems to be no different from simply changing the value update from TD(0) to n-step TD. \n\n* The empirical test seems to be incomplete. The paper proposes two algorithms, C-SoftTreeMax and E-SoftTreeMax. But only C-SoftTreeMax was empirically tested. The paper explains the empirical test for E-SoftTreeMax was left for future work on risk-averse RL. But it does not make sense to me that the empirical test is completely omitted while the method is listed as a new algorithm in the paper. \n\n* There could be a stronger baseline. Given the similarity between SAC policy ($\\pi(a|s) \\propto exp(Q(s,a))$) and the policy learned by C-SoftTreeMax ($\\pi(a|s) \\propto exp(W_\\theta(s,a))$), it may worth testing SAC, both 1-step TD and n-step TD in critic update.\n\n* The empirical test could include more seeds. Currently there are only five seeds. The variance calculated based on 5 seeds could be highly inaccurate.\n\n* The empirical test is limited to discrete action space. It is acceptable and understandable to leave continuous control experiments to future work, but still, adding continuous control results could make the paper more convincing.\n\n* The new method introduces a new parameter, the planning length $d$. According to the empirical results (Figures 3, 4, and 5), different $d$ gave very different performances. The best setting is different across tasks. In 3 out of 8 tasks, (NameThisGame, Phoenix, and VideoPinball), longer planning hurts the performance. It is acceptable to say the performance is sensitive to the value of $d$ and leave this parameter as a tunable one, but it could be better to give an indication of how to pick this value when the reader would like to apply this method to a new task.\n\n\n**Typo:**\n* Line 307, “one week The” -> “one week. The”\n\nQUESTIONS:\n* In Section 5, the paper says $\\theta$ is replaced by a neural network $W_\\theta$. It would be good if the authors could further explain what the target is when updating $W_\\theta$, i.e. is $W_\\theta$ a value function estimation or not?\n\n\n* It would be better to explain the empirical setting in more detail. In Section 6, the paper explains the tree is pruned according to an estimated weight. It would be more clear if the paper could include a description of how the weights are updated.\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new policy parameterization called SoftTreeMax, which can reduce the variance of the stochastic policy gradient. The authors consider finite state and action spaces throughout the paper. They start by taking a softmax tabular policy and replacing the logit $\\theta(s,a)$ with a score from a trajectory of horizon $d$ sampled from a behavior policy. They then address the normalization accordingly. Because of the randomness of the trajectory of size $d$, the authors take an expectation with respect to the randomness of the trajectory of size $d$ to well-define the policy parameterization. By taking the expectation before or after the exponent operator, the authors propose two policy parameterizations: C-SoftTreeMax and E-SoftTreeMax. It turns out that under either C-SoftTreeMax or E-SoftTreeMax parameterization, the variance of the stochastic gradient of the value function decays exponentially with the planning horizon $d$. The theoretical findings are well supported by simulations. Finally, the authors propose a parallel GPU-based simulator for the practical implementation of SoftTreeMax. By applying the C-SoftTreeMax parameterization to PPO and adapting $\\theta(s)$ with a neural network, as well as using the parallel GPU-based implementation, leads to better performance compared to distributed PPO.\n\nSTRENGTHS:\n- The concept of integrating tree search (planning) directly into policy parameterization is both intriguing and innovative. It demonstrates two key advantages: 1) With the SoftTreeMax parameterization, the variance of the stochastic gradient of the value function decreases exponentially with the planning horizon. 2) It can be easily combined with various policy gradient algorithms, such as REINFORCE and PPO, to enhance their performance.\n- Although computing the parameterization is computationally intensive, the authors offer a practical solution through parallel GPU-based simulation.\n- SoftTreeMax shows significant potential for extension to infinite or large spaces, making it a promising method for addressing complex problems.\n\nWEAKNESSES:\n- Not really a weakness. Since the SoftTreeMax can be beneficial to different policy gradient methods, I expect to see its improvement applied on some (original) policy gradient methods for the experiments, such as REINFORCE and actor-critic, to demonstrate its generality, which is another strength of the paper.\n- Since SoftTreeMax is inspired from tree search, the authors can highlight the difference between the \"tree search\" literature and the \"tree expansion\" used in this paper in order to avoid confusion. The tree search method uses a max to find the best policy, while here the authors use a fixed behavior policy.\n- For the tree search literature, the authors can cite Efroni et al. [2018] and Protopapas et Barakat [2024]. In particular, Protopapas et Barakat [2024] combine a variant of PG, the policy mirror descent, with tree search.\n- For the theoretical PG literature in Line 312, the authors can cite Yuan et al. [2022], where they provide a general analysis of PG, including softmax tabular parameterization as a special case.\n\nEfroni, Y., Dalal, G., Scherrer, B., and Mannor, S. Beyond the one-step greedy approach in reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1387–1396. PMLR, 10–15 Jul 2018.\n\nKimon Protopapas, Anas Barakat (2024). Policy Mirror Descent with Lookahead.\n\nRui Yuan, Robert M. Gower, and Alessandro Lazaric. A general sample complexity analysis of vanilla policy gradient. In Proceedings of The 25th International Conference on Artificial Intelligence and Statistics, volume 151 of Proceedings of Machine Learning Research, pages 3332–3380. PMLR, 28–30 Mar 2022.\n\nQUESTIONS:\nNo specific questions.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a model based policy gradient (PG) algorithm that tries to decrease the variance of the gradient updates in PG methods and thus (hopefully) improving their sample complexity. The proposed approach works by modifying the softmax policy to work not just not with simple score functions (mapping states to real numbers), but rather with the truncated expected discounted return from a state. The resulting policy is called \"SoftTreeMax\". As such, this paper seems a little bit similar to decision time planning (or MCTS type) algorithms.\n\nThe paper proposes two variants of SoftTreeMax and analyzes their variance properties; in particular, the results show that as the size of truncation increases, the variance decreases exponentially. The paper also has a theorem which shows that SoftTreeMax can estimate the gradients well even with an approximate model. There are also some experiments on certain Atari games which show the efficacy of SoftTreeMax against the softmax policy using the popular PPO algorithm.\n\nSTRENGTHS:\nOriginality: The proposed algorithm seems novel to me in that I have not seen tree search used with policy gradient. As far as I remember, most tree search methods employ value function based RL algorithms (most popular being AlphaZero). In that way, this paper is quite novel. Further, the modification of logits to instead include the expected return seems novel too. (Although, it might be related to the \"soft-Q learning\" papers a little bit; something which could be mentioned in the paper. For instance, see https://arxiv.org/pdf/1704.06440, Eq. 2.)\n\nQuality and Clarity: The paper is well written and has clear theorem statements that are easy to interpret. The setting considered is the popular discounted finite MDP, making the results very easy to follow. The experiments are also conducted using popular benchmarks and thus help the reader understand and compare the proposed method. Another nice thing is that the experimental results clearly align with the theoretical statements (such as the variance reduction with depth) and the benefits in terms of achieving better sample complexity are quite clear. Overall, the paper seems like a logical next step: using policy gradient methods with model based and tree search based methods. (In particular, for continuous action spaces, this would be highly useful and relevant.)\n\nSignificance: The paper combines, a very popular class of RL methods - policy gradients, with tree search like planning algorithms. For many difficult problems, particularly those with very large search spaces, Tree search has been shown to be very effective. On the other hand, policy gradient methods are broadly the most popular way of dealing with continuous action RL problems. Therefore, an algorithm combining these two solution methods can have immense applications, especially as AI methods continue to be adopted for solving more and more real-world decision making tasks (robotics being a good representative application area).\n\nWEAKNESSES:\n- The major weaknesses of this paper are not enough discussion on the exact mechanisms for implementing SoftTreeMax. So SoftTreeMax is akin to a policy parameterization (such as softmax). As such, when implementing an RL agent, we need a large number of other details (some details being how is the exploration policy chosen, what optimizer is used to update the gradients, how are the function approximators initialized, is there any processing of the observation features, etc.). These implementation details are important since they can affect the performance of the RL algorithms more than the core idea itself (for instance, recall RainbowDQN https://arxiv.org/abs/1710.02298 and PPO https://arxiv.org/abs/2005.12729).\n\n- Can you guess how SoftTreeMax run on a continuous action space MDP, such as those from Mujoco?\n\n- Does SoftTreeMax solve the problems of Softmax policies? I know it helps with variance, but what about softmax's dependence on the initialization? Or the possibly slow convergence of sotmax PG in MDPs due to an exponential dependence on the state space size (https://arxiv.org/abs/2102.11270)? Do you have any intuition? From the nice discussion in Appendix C.1, it seems that the answer is no. In particular, choosing the behavior policy while generating the tree could be quite important. Maybe some comments on how this policy is chosen in the implementation could be helpful (I saw the high-level comments about controlling for the eigenvalues of transition matrix; but those seem high-level ideas).\n\n- What does SoftTreeMax really do? Is it just an MCTS like decision time planning algorithm that is helped a tiny bit by the theta values in Eq. 2? For instance, what if I do not update the theta values at all (i.e. set learning rate to zero). In that case, would the performance decrease drastically? Related to that, do you have any comments on how it relates with MCTS? I remember Bertsekas (http://web.mit.edu/dimitrib/www/LessonsfromAlphazero.pdf) said that as the depth increases, the utility / importance of the learned values (in this case theta values) in decision making decreases. Would that apply here as well?\n\n- Once the policy is learned, how do you sample from it? I guess, you still need a model to sample from this policy. This seems like a clear disadvantage (although, this is true for all tree search based / decision time planning methods). Maybe mention this as a limitation.\n\nline 53: \"In all tested Atari games, our results outperform the baseline and obtain up to 5x more reward.\" --> this is not true? For instance, results on VideoPinball or Breakout show that PPO is pretty much as good. Also, the paper doesn't specify how the hyperparameters were chosen. So maybe PPO with a better choice of hyper-params might outperform SoftTreeMax?\n\nQUESTIONS:\nThe method seems to be highly computationally expensive. Thus, a major question is whether the convergence is sped up by the variance reduction? The paper claim that SoftTreeMax algorithm would result in a faster convergence than traditional policy gradient methods.\n\nNow, it looks like Figure 4 (page 24) is talking about exactly that. (Consider moving it into the main paper? This seems like a very important experiment.)\nAlso, in Figure 4 caption \"The x-axis is the wall-clock time. The runs ended after one week with varying number of time-steps.\", what does 'varying number of steps' mean? Does this mean that the algorithms run for the same amount of time but different number of (non-simulated) environment interactions?\n\n\nMinor comments (feel free to ignore)\n\nline 22: \"PG algorithms are also notoriously unstable due to the high variance of the gradients computed over entire trajectories\" --> I don't think actor-critic type methods do that.\n\nline 34: \"for cumulative reward and one for exponentiated reward.\" --> reward is not exponentiated, right? The SoftTreeMax logits (Eq. 2) are exponentiated.\n\nline 55: \"demonstrating the relation between low gradient variance and high reward\" --> seems like a strong statement? We don't know if low gradient is necessarily good, right? As Appendix C.1 says, this could also be sometimes harmful.\n\nline 69-81: Nice preliminaries section! Also consider specifying the agent-MDP interaction: $s_0 \\sim \\nu, a_t \\sim( \\cdot | s_t), s_{t+1} \\sim p( \\cdot | s_t, a_t), r_{t+1} = r(s_t, r_t)$, etc..\n\nline 76: define $P^{\\pi'}(s' | s)$ as $P^{\\pi'}(s, s')$ instead?\n\nline 82: consider writing bold $\\Theta \\in \\mathcal{R}^S$ as bold $\\theta \\in \\mathcal{R}^S$ --- capital letters are often reserved for matrices.\n\nline 84: consider replacing $D(u)$ with $\\textrm{diag}(u)$ throughout the paper -- midway, I got confused by $D$.\n\nline 127: \"by showing that the gradient is strictly positive everywhere [Mei et al., 2020b, Lemmas 8-9].\" --> Gradient being positive is unclear, since it is a vector. Also, I'm not sure if Lemmas 8 and 9 actually show this.\n\nline 165: \"where the gradient is scaled to reduce its variance,\" --> do you mean that gradient is clipped to reduce variance?\n\nline 181: \"the C-SoftTreeMax gradient is given\" --> try using 'Jacobian' instead of 'gradient'\n\nline 214-217: would the variance zero in this case?\n\nFigure 2: Consider writing \"I.e., instead of the\" as --> \"That is, instead of the\"\n\nFigure 254: what norm is being used?\n\nline 254, Theorem 4.8: It is quite interesting that bound is unaffected by the the behavior policy and the tranistion dynamics $\\mathcal{P}^\\pi$. Any comments?\n\nline 273: \"SoftTreeMax with an exhaustive TS of depth d\" --> Is the search really going through all possible state-actions? If not, consider dropping the work exhaustive.\n\nline 282: maybe figure 2 should be moved somewhare closer.\n\nlines 311-326: maybe put a colon, instead of a period, after the bolded letters. So do \"true dynamics. Aversion Risk. Previous wor\"\n\nLIMITATIONS:\nThe major limitation seems to be that the paper doesn't provide\n\n1.  a written algorithm - while I understand that this is almost straightforward, just putting a clear algorithm (specifying order of planning/simulation steps, how to take real world actions with SoftTreeMax policy, computing gradients, and updating the policy) would be highly useful.\n\n2. The paper doesn't run experiments on continuous action space environments, which is arguably the more relevant domain for PG methods. (For discrete environments like Atari, we already have a plethora of value function + tree search based methods - like AlphaZero or MuZero.) Further, there are no hyper-parameter settings specified; it would be difficult to reproduce these results without looking at the codebase.\n\n\nFurther, a limitations section is missing\n- if the PG is biased, would that affect convergence? In particular, note that Theorem 4.8 shows that the bias would scale with the state space size S (which is generally huge). Maybe setting d = \\log(S) / \\log(1 / gamma) helps with this and reduces this error to \\log(S). But the computational complexity is probably exponential in d, so each gradient update becomes O(S), which is also bad..\n\nMaybe a discussion between these factors would be helpful\n---- how does the increase in bias affect the convergence speed (and the quality of the final policy obtained). What bias does the pruning approach introduce?\n---- how does the decrease in the variance of gradient update affect the convergence speed (and the quality of the final policy obtained). Does the additional variance (in computing Lemma 4.6) introduced by pruning remove counteract the variance reduction (as given in Theorem 4.4)?\n---- how does changing the bias / variance affect d and thus the computational complexity (edited)", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-08-01"}