Title: Distributional Monte-Carlo Planning with Thompson Sampling in Stochastic Environments
Venue: NeurIPS 2024
Accepted: no
Keywords: Monte-Carlo Tree Search

Planning

Reinforcement Learning
Abstract: We focus on a class of reinforcement learning algorithms, Monte-Carlo Tree Search (MCTS), in stochastic settings. While recent advancements combining MCTS with deep learning have excelled in deterministic environments, they face challenges in highly stochastic settings, leading to suboptimal action choices and decreased performance. Distributional Reinforcement Learning (RL) addresses these challenges by extending the traditional Bellman equation to consider value distributions instead of a single mean value, showing promising results in Deep Q Learning. In this paper, we bring the concept of Distributional RL to MCTS, focusing on modeling value functions as categorical and particle distributions. Consequently, we propose two novel algorithms: Categorical Thompson Sampling for MCTS (CATS), which uses categorical distributions for Q values, and Particle Thompson Sampling for MCTS (PATS), which models Q values with particle-based distributions. Both algorithms employ Thompson Sampling to handle action selection randomness. Our contributions are threefold: We introduce a distributional framework for Monte-Carlo Planning to model uncertainty in return estimation. We prove the effectiveness of our algorithms by achieving a non-asymptotic problem-dependent upper bound on simple regret of order $O(n^{-1})$, where $n$ is the number of trajectories. We provide empirical evidence demonstrating the efficacy of our approach compared to baselines in both stochastic and deterministic environments.

Abstract

We focus on a class of reinforcement learning algorithms, Monte-Carlo Tree Search (MCTS), in stochastic settings. While recent advancements combining MCTS with deep learning have excelled in deterministic environments, they face challenges in highly stochastic settings, leading to suboptimal action choices and decreased performance. Distributional Reinforcement Learning (RL) addresses these challenges by extending the traditional Bellman equation to consider value distributions instead of a single mean value, showing promising results in Deep Q Learning. In this paper, we bring the concept of Distributional RL to MCTS, focusing on modeling value functions as categorical and particle distributions. Consequently, we propose two novel algorithms: Categorical Thompson Sampling for MCTS (CATS), which uses categorical distributions for Q values, and Particle Thompson Sampling for MCTS (PATS), which models Q values with particle-based distributions. Both algorithms employ Thompson Sampling to handle action selection randomness. Our contributions are threefold: We introduce a distributional framework for Monte-Carlo Planning to model uncertainty in return estimation. We prove the effectiveness of our algorithms by achieving a non-asymptotic problemdependent upper bound on simple regret of order O ( n -1 ) , where n is the number of trajectories. We provide empirical evidence demonstrating the efficacy of our approach compared to baselines in both stochastic and deterministic environments.

1 Introduction

Online planning in Markov decision processes (MDPs) involves making real-time decisions based on the current state of the environment. It requires balancing exploration and exploitation while handling uncertainty and partial observability. Monte Carlo Tree Search (MCTS) is a highly effective online planning method for tackling complex MDPs. MCTS has shown impressive performance in various tasks, including traditional board games like Chess and Go, video games, and real-world challenges. Notable successes include advancements in Chess (35) and Go (34; 36; 30), video game strategy (28), robot assembly (16), robot path planning (15; 13), and autonomous driving (24).

Despite these achievements, current MCTS methods are primarily effective in deterministic environments, often overlooking the significant impact of randomness in real-world scenarios. In highly stochastic and partially observable environments, conventional MCTS approaches face substantial challenges due to widespread randomness and limited observability. This leads to compromised value estimates, suboptimal decisions, and diminished overall performance. Therefore, there is a clear need for improved methods capable of navigating the complexities of randomness and partial observability in value estimation.

We now review related works to understand the advancements and limitations in these areas.

Related work In MCTS, value estimation methods and action selection rules are critical factors for 36 algorithm performance. Traditional value estimation methods, such as using empirical average mean 37

for value backup as in the Upper Confidence bounds applied to Trees method (UCT) (21), suffer from 38 underestimation of optimal values while maximum backup suffers from overestimation of optimal 39 values (9). The power mean estimator (12) offers a balanced solution by computing a mean between 40 the average and maximum values. In our approach, we also use power mean for value operator as 41 each V node stores the power mean of empirical means of succeeding Q-value nodes, eliminating the 42 need for V to be modeled as a distribution. 43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

For action selection in MCTS, strategies from Multi-Armed Bandits (MAB) are commonly employed. For instance, UCT extends the UCB1 strategy from bandits to the tree by computing confidence intervals at each step. However, original UCT's performance is hindered by the incorrect choice of logarithmic bonus constant (32). Shah et al. (32) propose an adapted version of UCT incorporating a polynomial bonus term instead of the "logarithmic" bonus term in UCT and show the non-asymtotic convergence of rate O ( n -1 / 2 ) , with n is the number of rollout trajectories. On the other hand, our method improves over this rate with theoretical guarantee of O ( n -1 ) . Although Thompson sampling has been less explored in MCTS, some approaches like those by Bai et al. (1) and Bai et al. (2) incorporate it for exploration. However, these methods lack convergence rate analysis. Furthermore, in the article Bai et al. (1), authors model value functions as a mixture of Normal distributions, which may lack the generality of complex real-world scenarios. Our approach adopts Thompson sampling for action selection but introduces a novelty by modeling the uncertainty of action value estimates over the tree as arbitrary categorical and particle-based distributions. This modification enhances our ability to handle more generality in highly stochastic environments effectively.

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

Entropy regularization techniques in RL modify value and action selection functions to balance exploration and exploitation, leading to improved value estimation (25; 17; 31; 18). Several works have applied these techniques in MCTS. Maximum Entropy Tree Search (MENTS) (40) emphasizes exploration by integrating MCTS with maximum entropy policy optimization. MENTS aims to maximize cumulative rewards and policy entropy concurrently, regulated by a temperature parameter. Dam et al. (14) extend MENTS by incorporating Relative and Tsallis entropy, leading to the RENTS and TENTS algorithms. However, the effectiveness of MENTS/RENTS/TENTS hinges on the temperature parameter, which may impede convergence. Furthermore, the value estimation converges exponentially to the regularized value not the optimal one. In contrast, Painter et al. (27) utilize a similar action selection approach but employ a maximum backup operator for value estimation. Although their method exhibits exponential decay of simple regret, it heavily relies on the sensitivity of the temperature parameter for Boltzmann Exploration, limiting its practicality.

Distributional Reinforcement Learning (RL) (6; 11; 22) addresses the randomness of the value estimation by introducing a distributional perspective to the traditional Bellman equation. This approach views the value function as a distribution rather than a single mean, providing a comprehensive understanding of uncertainties in rewards and the stochasticity from environments. Through discretization (26), parameterization (6), and quantization (10), it allows for efficient and effective approximation of value distributions, leading to improved performance in various RL tasks. However, these results are only for learning not for planning .

Outline and contribution In this work, we integrate the distributional approach from reinforcement learning (RL) into the planning framework to tackle the challenges of planning in stochastic environments. We focus on modeling value functions as categorical and particle distributions. Consequently, we propose two novel algorithms: Categorical Thompson Sampling for MCTS (CATS) and Particle Thompson Sampling for MCTS (PATS). CATS represents each Q value function as a categorical distribution and uses Thompson Sampling for action selection to manage uncertainty. PATS models each Q value function with a particle-based distribution, using a nuanced Thompson Sampling approach to handle action selection randomness.

Our contributions are threefold:

(i) In section 3, we introduce a distributional framework for planning to model uncertainty in return estimation, enhancing the robustness of value estimation in stochastic environments.

(ii) In section 4 Theorem 5 and Theorem 6, we prove the effectiveness of our algorithms by achieving a non-asymptotic problem-dependent upper bound on simple regret of O ( n -1 ) , which significantly improves upon the current state-of-the-art theoretical analysis of regret, previously established at O ( n -1 / 2 ) by Shah et al. (33).

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

(iii) In section 5, we provide comprehensive empirical evidence demonstrating the efficacy of our approach compared to baselines, showcasing competitive performance in stochastic settings and the Atari benchmark.

In the next section, we describe the problem setting addressed in this paper.

2 Setting

In our study, We address the dynamics of an agent navigating an infinite-horizon discounted Markov decision process (MDP), defined formally as M = ⟨S , A , R , P , γ ⟩ . Here, S represents the state space, A denotes the set of actions, and R quantifies the Reward function of the MDP ( R : S × A×S → R ). Transition dynamics are governed by P ( S × A → S ), with γ ∈ (0 , 1] as the discount factor. The agent interacts with the environment via a policy π ∈ Π : S → A , guiding action selection based on observed states. This yields an action-value function Q π , indicating the expected cumulative discounted reward from a state-action pair under π . The agent seeks the optimal policy maximizing the action-value function, adhering to the Bellman equation (7), given by Q ( s, a ) ≜ ∫ S P ( s ′ | s, a )[ R ( s, a, s ′ ) + γ max a ′ Q ( s ′ , a ′ )] ds for all states s and actions a . Upon acquiring the optimal action-value function, we derive the optimal value function V ( s ) ≜ max a ∈A Q ( s, a ) for all states s in S .

Monte-Carlo tree search (MCTS) (20; 8) is a planning approach for complex Markov decision processes (MDPs). It employs an iterative approach:

Selection : It begins by selecting an action using a specified strategy, followed by executing this action through Monte Carlo simulation.

Expansion : Subsequently, it assesses the resulting state, either by recursively evaluating if it already exists in the search tree or by inserting it into the tree.

Simulation : Or employing a rollout policy via simulations. This iterative process continues until certain termination criteria are met, allowing traversal through the search tree.

Backpropagation : Finally, the outcomes of the simulations are propagated backward through the chosen nodes to update their statistical metrics.

Simple Regret An MCTS algorithm dynamically gathers trajectories within an MDP starting from an initial state s 0 . After processing t trajectories, it provides two outputs:

̂ a t , a guess for the best action to take at state s 0

V ( s ) an estimator of the optimal value in s ,

̂ t 0 0

where s 0 is the state at the root node. The algorithm's performance can be assessed by its convergence rate r ( t ) of the simple regret, formulated as:

Here, R ( s 0 , t ) = V ⋆ ( s 0 ) -̂ V t ( s 0 ) is the simple regret of the algorithm at the root node with V ⋆ ( s 0 ) representing the optimal value at state s 0 .

In this article, we analyze an MCTS algorithm employing a maximal planning horizon H and a playout policy π 0 with value V 0 . We define ˜ V ( s H ) = V 0 ( s H ) recursively as follows: for all h ≤ H -1 ,

where r ( s h , a ) defined formally as the mean intermediate reward at state s h after taking action a . The primary objective of an MCTS algorithm is to estimate a tied rate r ( t ) by constructing estimates of ˜ Q ( s h , a ) and ˜ V ( s h ) to ultimately estimate ˜ Q ( s 0 , a ) and consequently Q ⋆ ( s 0 , a ) . In practical implementations of the MCTS algorithm, the maximal depth H can sometimes be set to + ∞ . However, for theoretical analysis, the maximal depth H is crucial as we will analyze the algorithm that always collects trajectories of length H.

Distributional Reinforcement Learning The mathematical framework used in reinforcement learning is based on the Bellman equation (37), which aims to find an agent to maximize the expected utility Q value. However, the single expected value function cannot encapsulate the stochasticity in the reward function and the dynamic of the environments. Recently, in the article (5), authors shed light on the distributional perspective of the Bellman equation by modeling each Q value function as a distribution instead of a single expected value. The main objective is to study the random return Q at the state s , action a , and is defined recursively as

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

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

where X ( s, a ) is the reward distribution at the state s , action a , Q ( s, a ) is the Q value distribution at state s , action a , and Q ( s ′ , a ′ ) is the Q value distribution at state s ′ , action a ′ . s ′ distributed according to P ( ·| s, a ) , a ′ distributed according to a policy π ( ·| s ′ ) . A D = B denotes that two random variables A and B have equal probability laws.

This distributional approach offers a deeper understanding of uncertainty and variability, especially in complex, stochastic systems where traditional expected value representations may fail to capture the true dynamics of the problem. which has been successfully used in Deep Q Learning (5). Categorical Value Distribution Based on the distributional Bellman equation, In the article (5), authors approximate the Q value distribution Q ( s, a ) as a discrete categorical distribution parametrized by N ∈ N , which denotes the number of atoms (N+1) at fixed-sized locations. This method effectively divides the Q value function into a set of equally spaced atoms z i ( s, a ) = Q min + i △ z : 0 ≤ i ≤ N , where Q min and Q max are respectively the minimum and maximum values at state s , action a . The size of each atom is set as △ z := Q max -Q min N .

This discrete distribution approach is highly expressive and computationally efficient, making it ideal for practical applications. For instance, in the article (5), authors successfully used this representation in Deep Q Learning (C51), showing promising results in several Atari games. In the next section, we demonstrate how to apply this idea to MCTS.

3 Distributional Thompson Sampling in Tree Search

In this section, we introduce two novel distributional approaches for MCTS based on Thompson sampling. The first method represents each Q-value node as a categorical distribution, while the second uses particle-based distributions for greater flexibility. Both methods integrate Thompson sampling for improved exploration and performance.

3.1 Distributional Monte-Carlo Tree Search

We leverage the success of distributional reinforcement learning (4; 3; 6) and apply this concept to MCTS. In MCTS, there are two types of nodes: V-nodes and Q-value nodes. Instead of treating each V value and Q value as a single expected value, we model these functions as distributions. Based on equation (2), we can derive

with s ′ ∼ P ( ·| s, a ) , where ¯ π ( . | s ′ ) is formally defined as the tree policy at state s ′ . We can model any Q distribution with equal law distributed as the sum of the distributions of the next reward and the Q distributions of the next states actions. We further model each V distribution, having equal probability law to the expectation of the chosen policy of the next Q-value distributions (3).

Our method follows the same four basic steps of MCTS but is different in Value Backup and Action selection steps. We introduce two distinct methodologies: categorical-based and particle-based. In the categorical based approach, we parameterize each V value and Q value function in the tree as a categorical distribution. In contrast, in the particle-based approach, we model each value distribution as a set of sampling particles, representing the values observed during the tree planning. We provide a detailed explanation for the value backup and action selection of each method in the next section.

3.2 Value Backup

In this work, we employ two approaches to represent the Q value distribution.

Categorical distribution : we represent each node in the tree as a categorical distribution. In each Q-value node, we: (1) store the empirical mean value of that Q-value node (same as in UCT), and (2) maintain a categorical distribution of the Q value function. To define a categorical distribution Q function, we require three essential pieces of information:

The number of atoms ( N+1 ): We choose a consistent number of atoms ( N+1 ) that remains the same for all Q distributions along the tree.

Minimum and maximum values ( min and max ): Each node in the tree may have different ranges for its minimum ( Q min ) 1 and maximum ( Q max ) values, depending on its state/action in the environment. When a new Q-value node is added to the tree, we initially set Q min to 0 (assuming we have scaled the reward range to [0, R]) and initialize Q max to a small

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

Algorithm 1

CATS

SelectAction

(

s

h

)

(Sec 3.2)

for

a

∈

[

A

]

do

L

(

s

h

, a

)

∼

Dir

(

α

0

(

s

h

, a

)

, . . . , α

N

(

s

h

, a

))

ϕ

(

s

h

, a

) = [

z

0

(

s

h

, a

)

, . . . , z

N

(

s

h

, a

)]

⊤

L

(

s

h

, a

)

a

a

{

ϕ

(

s

h

, a

)

}

return

a

SimulateV

(

s

h

, t

)

(Sec 3.2)

a

=

SelectAction

(

s

h

)

SimulateQ

(

s

h

, a, t

)

T

s

h

(

t

) =

T

s

h

(

t

) + 1

̂

Q

(

s

h

, a

) =

∑

i

z

i

(

s

h

, a

)

p

i

(

s

h

, a

)

̂

V

(

s

h

) =

( ∑

a

T

s h ,a

(

t

)

T

s h

(

t

)

̂

Q

p

(

s

h

, a

)

)

1

p

SimulateQ

(

s

h

, a, t

)

(Sec 3.2)

s

h

+1

∼

P

(

·|

s

h

, a

)

,

r

t

(

s

h

, a

)

∼ R

(

s

h

, a, s

h

+1

)

if

Node

s

h

+1

not expanded

then

Rollout

(

s

h

+1

)

else

SimulateV

(

s

h

+1

, t

)

T

s

h

,a

(

t

) =

T

s

h

,a

(

t

) + 1

Q

t

(

s

h

, a

) =

r

t

(

s

h

, a

) +

γ

̂

V

(

s

h

+1

)

if

Q

t

(

s

h

, a

)

̸∈

[

Q

min

(

s

h

, a

)

, Q

(

s

h

, a

)]

then

Q

(

s

h

, a

{

Q

t

(

s

h

, a

)

, Q

(

s

h

, a

)

}

Q

min

(

s

h

, a

) = min

{

Q

t

(

s

h

, a

)

, Q

min

(

s

h

, a

)

}

△

z

=

Q

-

Q

min

N

z

i

(

s

h

, a

) =

Q

min

+

i

△

z

: 0

i

N

Update

p

(

s

h

, a

) = [

p

0

(

s

h

, a

)

, . . . , p

N

(

s

h

, a

)]

Algorithm 2

PATS

SelectAction

(

s

h

)

(Sec 3.2)

for

a

∈

[

A

]

do

L

(

s

h

, a

)

∼

Dir

(

α

(

s

h

, a

))

ϕ

(

s

h

, a

) =

S

(

s

h

, a

)

⊤

L

(

s

h

, a

)

a

a

{

ϕ

(

s

h

, a

)

}

return

a

SimulateV

(

s

h

, t

)

(Sec 3.2)

a

=

SelectAction

(

s

h

)

SimulateQ

(

s

h

, a, t

)

T

s

(

t

) =

T

s

(

t

) + 1

̂

Q

(

s

h

, a

) =

∑

α

t

(

s

h

, a

)

Q

t

(

s

h

, a

)

̂

V

(

s

h

) =

( ∑

a

T

s h ,a

(

t

)

T

s h

(

t

)

̂

Q

p

(

s, a

)

)

1

p

SimulateQ

(

s

h

, a, t

)

(Sec 3.2)

s

h

+1

∼

P

(

·|

s

h

, a

)

,

r

t

(

s

h

, a

)

∼ R

(

s

h

, a, s

h

+1

)

if

Node

s

h

+1

not expanded

then

Rollout

(

s

h

+1

)

else

SimulateV

(

s

h

+1

, t

)

T

s

h

,a

(

t

) =

T

s

h

,a

(

t

) + 1

Q

t

(

s

h

, a

) =

r

t

(

s

h

, a

) +

γ

̂

V

(

s

h

+1

)

if

Q

t

(

s, a

)

∈ {S

(

s

h

, a

)

}

then

α

t

(

s

h

, a

) += 1

//

α

t

(

s

h

, a

) :

weight of

Q

t

(

s

h

, a

)

else

S

(

s

h

, a

) := (

S

(

s

h

, a

)

, Q

t

(

s

h

, a

))

α

(

s

h

, a

) := (

α

(

s

h

, a

)

,

1)

number, e.g., Q max = 0 . 001 . Since the min and max values are unknown, we start with a small range, that will get updated accordingly to the scale of the observed values.

Probabilistic parameterization: The probability of each atom ( p i ( s, a ) ) is determined based on the visitation count ratio. In detail, each atom stores statistical information about the visitation count, and the probability of that atom will be calculated as the visitation count divide with the total visitation count of that Q-value node. When we backpropagate the r t ( s, a ) + γ ̂ V t ( s ′ ) value to a specific node, we identify the atom whose value range includes the r t ( s, a ) + γ ̂ V t ( s ′ ) value. At this point, we increase its visitation count.

Additionally, as we backpropagate Monte-Carlo Q values over time, we empirically adjust the Q min and Q max values to account for the dynamic range of Q values observed in the tree. This dynamic scaling ensures that the atom locations are effectively rescaled to adapt to the changing conditions. This representation method allows us to encapsulate the knowledge gained through exploration in the form of categorical distributions, which helps in making informed decisions during the tree search.

Paricle based distribution : We represent each Q value distribution as a collection of sampling particles, which encapsulate the observed values during tree planning. Initially, we maintain an empty set of particles for the Q value distribution, denoted as S ( s, a ) . At time step t , upon receiving an intermediate reward Q t ( s, a ) = r t ( s, a ) + γ ̂ V t ( s ′ ) , with s ′ ∼ P ( ·| s, a ) , we add Q t ( s, a ) to the set S ( s, a ) if the particle does not already exist within it. If the particle Q t ( s, a ) already exists in S ( s, a ) , we increase the visitation count ratio associated with that particle.

Value function: The Q-value node is crucial in the tree because its representation influences action selection, as detailed in the next section. We now discuss modeling each V-value node. The V-value distribution is based on the expected outcomes of the chosen policy and the subsequent Q-distributions. Thus, the mean of the V-function corresponds to the tree policy's expectation of the means of all

214

215

216

217

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

248

249

250

251

252

253

254

255

256

257

258

259

260

succeeding Q-value nodes. The common approach is to use empirical average mean for the value backup, as in UCT (21). However, this approach underestimates the optimal value, while using the maximum value overestimates it (9). The power mean estimator (12) provides a balanced solution, falling between the average and maximum values. In our methods, each V node stores the power mean of the empirical means of all succeeding Q-value nodes, eliminating the need to model V as a distribution.

where T s ( n ) , T s,a ( n ) are the number of visitations at s and s, a at timestep n respectively. Next, we show how to select actions in the tree based on the categorical distribution of Q-value nodes.

3.3 Action Selection

Thompson sampling has shown promising results in real bandit scenarios due to the randomness of action selection. Taking advantage of the established categorical based distribution and particle based distribution, we use the Thompson sampling method for action selection. We maintain a Dirichlet distribution of parameter of the Q value distribution. We denote the Dirichlet distribution of parameters

Categorical distribution : The probability mass function of the discrete categorical distribution at each Q-value node at state s , action a : p ( s, a ) = [ p 0 ( s, a ) , p 1 ( s, a ) , . . . , p N ( s, a )] , where p i ( s, a ) represents the probability of selecting the i -th atom z i ( s, a ) , N +1 is the number of atoms. We maintain a Dirichlet distribution Dir ( α 0 ( s, a ) , α 1 ( s, a ) , . . . , α N ( s, a )) as the prior for the Q-value node at state s , action a . At each time step t we sample L t ( s, a ) ∼ Dir ( α 0 ( s, a ) , α 1 ( s, a ) , . . . , α N ( s, a )) and compute ϕ t ( s, a ) = [ z 0 ( s, a ) , z 1 ( s, a ) , . . . , z N ( s, a )] ⊤ L t ( s, a ) . Then, the action a t is selected as follows:

After taking action a t and get an intermediate reward Q t ( s, a t ) = r t ( s, a t ) + γ ̂ V t ( s ′ ) . The posterior is also a Dirichlet: Dir ( α 0 ( s, a ) , . . . , α t ( s, a ) + 1 , . . . , α N ( s, a )) with the intermediate reward at time step t : Q t ( s, a t ) is in the range of the atom z t ( s, a ) . We denote this mechanism as Categorical Thompson sampling for Tree Search (CATS) method.

Paricle based distribution : In the particle-based approach, the prior Dirichlet distribution of the Q-value node at state s , action a is Dir ( α ( s, a )) , with α ( s, a ) is initiated as [1] . Considering each Q value distribution at state s , action a has a set of particle { Q t ( s, a ) } with the corresponding weighted α ( s, a ) = { α t ( s, a ) } At each time step t we also sample L t ( s, a ) ∼ Dir ( α ( s, a )) and compute ϕ t ( s, a ) = [1 , Q 0 ( s, a ) , Q 1 ( s, a ) , . . . , Q N ( s, a )] ⊤ L t ( s, a ) . Then the action a t is chosen as

After taking action a t and get an intermediate reward Q t ( s, a t ) = r t ( s, a t ) + γ ̂ V t ( s ′ ) . We update α t ( s, a ) = α t ( s, a ) + 1 if Q t ( s, a t ) is in the set { Q t ( s, a ) } . If not, we add Q t ( s, a t ) to the set { Q t ( s, a ) } and add 1 to the set { α t ( s, a ) } = { α t ( s, a ) , 1 } .

We call this method as Paricle Thompson sampling for Tree Search (PATS) method. Detailed pseudocode and a comparison of CATS and PATS can be seen in Fig 1. The two methods are identical in all procedures except for the Q value function backup ( SimulateQ ) and the action selection function ( SelectAction ).

Remark 1. CATS and PATS both use similar action selection strategies within a bandit setting, specifically referring to Multinomial Thompson Sampling and Non-Parametric Thompson Sampling, respectively (29). While CATS action selection heavily depends strictly on Thompson Sampling by maintaining parameters of posterior Q-value distribution, PATS is not based on the posterior sampling in the strict sense. At each step, it computes an average of the observed rewards with random weight and is a Non-Parametric approach. Furthermore, CATS maintains a fixed set of atoms, whereas in PATS, the number of particles increases depending on the observed Q values.

In the next section, we provide a theoretical analysis of the convergence of simple regret for CATS and PATS.

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

Algorithm 3

CATS in Non-stationary bandits

Require:

K arms; n: number of plays;

N

+1

support size of categorical distributions

Init

(

α

0

a

, . . . , α

N

a

) = (1

, . . . ,

1)

for each

a

∈

[

K

]

Main

()

for

t = 0,1,2,. . . , n

do

for

a

∈

[

A

]

do

L

a,t

∼

Dir

(

α

0

a

, . . . , α

N

a

)

ϕ

a,t

= [0

,

R

(

t

)

N

,

2

R

(

t

)

N

,

· · ·

, R

(

t

)]

⊤

L

t

a

a

{

ϕ

a,t

}

Pull arm

a

and observe reward

R

a,t

=

mR

(

t

)

N

where

m

∈ {

0

,

1

, . . . N

}

Update

α

m

a

=

α

m

a

+1

Algorithm 4

PATS in Non-stationary bandits

Require:

K arms; n: number of plays;

Init

α

a

= (1)

;

S

a

= (1)

for each

a

∈

[

K

]

Main

()

for

t = 0,1,2,. . . , n

do

for

a

∈

[

A

]

do

L

a,t

∼

Dir

(

α

a

)

ϕ

a,t

=

S

⊤

a

L

a,t

a

a

{

ϕ

a,t

}

Pull arm

a

and observe reward

R

a,t

if

R

a,t

∈ {S

a

}

then

α

t

a

+= 1

//

α

t

a

:

weight of

R

a,t

else

S

a

:= (

S

a

, R

a,t

)

α

a

:= (

α

a

,

1)

4 Theoretical analysis

Planning in MCTS involves making a sequence of decisions along the tree, where each internal node functions as a non-stationary bandit, with the empirical mean drifting due to the action selection strategy. Therefore, we first study the non-stationary multi-armed bandit settings using the action selections of CATS and PATS, examining the concentration properties of the power mean backup for each arm relative to the optimal arm. We then apply these results to MCTS.

4.1 Non-stationary multi-armed bandit

We consider a class of non-stationary multi-armed bandit (MAB) problems with K ≥ 1 arms. Let R a,t denote the random reward obtained by playing arm a ∈ [ K ] at the time step t bounded in [0 , R ] . We consider ̂ µ a,n = 1 n ∑ n t =1 R a,t as the average rewards collected at arm a after n plays. We first define:

Definition 1. A sequence of estimators ( ̂ V n ) n ≥ 1 is concentrated and convergent towards some limit V if the following two properties hold:

(A)

Concentration: For all

(B)

Convergence:

lim

[

n

V

E

→∞

n

In that case, we write plim

n

→∞

̂

We assume that the reward sequence { R a,t } , t ≥ 1 is a non-stationary process satisfying the convergence and concentration properties from Definition 1, by making the following assumption:

Assumption 1. Consider K arms that for a ∈ [ K ] , let ( ̂ µ a,n ) n ≥ 1 be a sequence of estimator satisfying

The action selection of CATS and PATS follows closely as in Section 3.3 and pseudocode are shown

1

in Fig. 2. Let us define ̂ µ n ( p ) = ( ∑ K a =1 T a ( n ) n ̂ µ p a,T a ( n ) ) p as the power mean value backup operator after n rounds. Here 1 ≤ p < ∞ is a constant. We denote T a ( n ) is the number of visitations of the arm a .

We define µ ⋆ = max a ∈ [ K ] { µ a } and assume that µ ⋆ is unique. Then, we establish the concentration and convergence properties of the power mean backup operator ̂ µ n ( p ) towards the optimal value µ ⋆ , as shown in Theorem 1 and Theorem 2, respectively for CATS and PATS.

Theorem 1. For a ∈ [ K ] , let ( ̂ µ a,n ) n ≥ 1 be a sequence of estimator satisfying plim n →∞ ̂ µ a,n = µ a and let

µ ⋆ = max a { µ a } . Assume that all the estimators are bounded in [0 , R ] . We consider a bandit algorithm that selects each arm according to CATS once in each round n ≥ K . Then, plim ̂ µ n ( p ) = µ ⋆ .

Theorem 2.

µ

µ

̂

a,n

)

n

1

For

µ

a

∈

[

K

]

, let

(

. Assume that all the estimators are bounded in

[0

⋆

a

{

a

}

that selects each arm according to PATS once in each round

n

, R

be a sequence of estimator satisfying plim

µ

a,n

=

µ

a

and let

n

→∞

]

̂

. We consider a bandit algorithm

K

. Then, plim

n

µ

(

p

) =

µ

n

→∞

̂

⋆

.

n

→∞

V

n

=

̂

V

n

.

1

] =

, for all

V

.

ε >

0

,

∃

c >

0

that

(

|

V

-

V

|

)

> ε

cn

-

1

ε

-

1

.

P

n

̂

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

318

319

320

321

322

323

324

325

326

327

328

329

330

331

332

333

334

Detailed proofs of the two Theorems can be found in the appendix. Based upon these results we analyse the concentration properties for any internal node and convergence of the simple regret in the MCTS in the next section.

4.2 Monte-Carlo Tree Search

Before presenting the main results (Theorem 3 Theorem 4), we first show an important Lemma

Assume that there exists a constant L > 0 such that L = supremum { ̂ V m,n } n ≥ 1 . Let R i be an iid sequence with mean µ and S i be an iid sequence from a distribution p = ( p 1 , . . . , p M ) supported on { 1 , . . . , M } . Introducing the random variables N n m = # |{ i ≤ n : S i = s m }| , we define the sequence of estimator

The significance of Lemma 1 lies in demonstrating the concentration and convergence of an estimated denotes an intermediate reward received by

Q value, conditioned on the concentration and convergence of a child V-value node. Here, ̂ V · ,n represents the value estimation at time step n , and R i taking a specific action at a particular state.

Next, we first start with Theorem 3 to show the convergence and concentration of any V-Node and Q-node in the tree for CATS.

Theorem 3. When we apply the CATS algorithm, we have

(i) For any node s h at the depth h th in the tree, plim ̂ Q n ( s h , a k ) = ˜ Q ( s h , a k ) .

n

→∞

(ii) For any node s h at the depth h th in the tree, plim ̂ V n ( s h ) = ˜ V ( s h ) .

n

→∞

We can derive a similar result for PATS as shown in Theorem 4.

Theorem 4. When we apply the PATS algorithm, we have

(i) For any node s h at the depth h th in the tree, plim ̂ Q n ( s h , a k ) = ˜ Q ( s h , a k ) .

(ii) For any node

s

h

at the depth

h

th

n

→∞

in the tree, plim

n

→∞

̂

˜

The results of Theorems 4 and 4 demonstrate that, at any node in the tree, both the V-value and Q-value nodes are convergent and concentrated. These results are applicable to any power mean backup operator of V-value nodes with p ∈ [1 , + ∞ ) . Finally, we show important results in Theorem 5, and Theorem 6, since they show the convergence of simple regret of CATS and PATS, respectively.

Theorem 5. (Convergence of Simple Regret of CATS) We have at the root node s 0 ,

Theorem 6. (Convergence of Simple Regret of PATS) We have at the root node s 0 ,

Remark 2. These results demonstrate that both CATS and PATS share the same convergence rate for value estimation at the root node of O ( n -1 ) , which improves over the rate O ( n -1 / 2 ) of FixedDepth-MCTS (33). Furthermore, Our finding more broadly applies to the power mean estimator with p ∈ [1 , + ∞ ) .

5 Experiments

We compare our methods with UCT (21), Fixed-Depth-MCTS (33), MENTS (40), RENTS, TENTS (14), BTS (27) and DNG (1) in a stochastic setting ( SyntheticTree ) to highlight the benefits of CATS and PATS in stochastic environments. Additionally, we test on 17 Atari games, comparing our algorithms with DQN (base network without planning) and other non-distributional planning methods (Power-UCT (12), MENTS (40), TENTS (14)) to demonstrate CATS and PATS' competitiveness and put results in Appendix. In all settings, we use 100 atoms for CATS, and set the discount factor γ to 0.99 for Atari, and γ to 1 for SyntheticTree .

V

n

(

s

h

) =

V

(

s

h

)

.

335

336

337

338

339

340

341

342

343

344

345

346

347

348

349

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

SyntheticTree : We evaluate CATS and PATS using the synthetic tree toy problem (14). This problem involves a tree with depth d and branching factor k . Each tree edge has a random value between 0 and 1. Returns at the leaf nodes are simulated using Gaussian distributions with means equal to the sum of edge values from the root to the leaf, and a standard deviation of 0 . 5 . Means are normalized between 0 and 1. An agent traverses the tree from the root, aiming to find the leaf node with the highest mean value. Internal nodes give zero reward, while leaf nodes provide a reward sampled from their Gaussian distribution. We introduce stochasticity into the environment by altering the transition probabilities: there is a 50% chance of moving to the intended node and a 50% chance of moving to a different node with equal probability. We conduct 25 experiments on five trees with five runs each, covering all combinations of branching factors k = { 2 , 4 , 6 , 8 , 10 , 12 , 14 , 16 , 100 , 200 } and depths d = { 1 , 2 , 3 , 4 } . We compute the value estimation error at the root node. Fig. 3 shows the convergence of the value estimations of CATS and PATS at the root node in the Synthetic Tree environment which shows they archives faster convergence compared to other methods.

0    500   1000

# Simulations

0.0

0.5

Value Estimation Error

k=16  d=1

0    500   1000

# Simulations

0

1

0    500   1000

# Simulations

0

1

k=14  d=3

0    500   1000

# Simulations

0

1

k=16  d=3

0    500   1000

# Simulations

0

2

k=16  d=4

0    500   1000

# Simulations

0.0

0.5

k=200  d=2

UCT

Power-UCT

DNG

Fixed-Depth-MCTS

MENTS

RENTS

TENTS

BTS

CATS

PATS

6 Conclusion

To conclude, our work introduces Categorical Thompson Sampling for MCTS (CATS) and Particle Thompson Sampling for MCTS (PATS), distributional planning approaches specifically designed to tackle complexities arising from stochasticity. CATS uses a categorical distribution, while PATS uses a particle-based distribution to represent and model the uncertainty inherent in return outcomes. We also propose exploration strategies based on Thompson Sampling that leverage this distributional modeling. Our methods come with a rigorous theoretical convergence guarantee, achieving a simple regret polynomial decay of the order O ( n -1 ) , which improves over the O ( n -1 / 2 ) rate of the fixed version of UCT (32). Empirical findings conclusively demonstrate the effectiveness of our approach in stochastic environments.

References

A. Bai, F. Wu, and X. Chen. Bayesian mixture modelling and inference based thompson sampling in monte-carlo tree search. Advances in neural information processing systems , 26, 2013.

A. Bai, F. Wu, Z. Zhang, and X. Chen. Thompson sampling based monte-carlo planning in pomdps. the International Conference on Automated Planning and Scheduling , 24(1), 2014.

M. Bellemare, S. Srinivasan, G. Ostrovski, T. Schaul, D. Saxton, and R. Munos. Unifying count-based exploration and intrinsic motivation. In Advances in neural information processing systems , pages 1471-1479, 2016.

M. 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.

M. G. Bellemare, W. Dabney, and R. Munos. A distributional perspective on reinforcement learning. In International Conference on Machine Learning , 2016.

M. G. Bellemare, W. Dabney, and R. Munos. A distributional perspective on reinforcement 373 learning. In Proceedings of the 34th International Conference on Machine Learning-Volume 70 , 374 pages 449-458. JMLR. org, 2017. 375

R. Bellman. The theory of dynamic programming. Technical report, Rand corp santa monica 376 ca, 1954. 377

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

R. Coulom. Efficient selectivity and backup operators in monte-carlo tree search. In Interna381 tional conference on computers and games . Springer, 2006. 382

W. Dabney, G. Ostrovski, D. Silver, and R. Munos. Implicit quantile networks for distributional 383 reinforcement learning. In International conference on machine learning , pages 1096-1105. 384 PMLR, 2018. 385

W. Dabney, M. Rowland, M. G. Bellemare, and R. Munos. Distributional reinforcement learning 386 with quantile regression. In Thirty-Second AAAI Conference on Artificial Intelligence , 2018. 387

T. Dam, P. Klink, C. D'Eramo, J. Peters, and J. Pajarinen. Generalized mean estimation in 388 monte-carlo tree search. arXiv preprint arXiv:1911.00384 , 2019. 389

T. Dam, G. Chalvatzaki, J. Peters, and J. Pajarinen. Monte-carlo robot path planning. IEEE 390 Robotics and Automation Letters , 7(4):11213-11220, 2022. 391

T. Q. Dam, C. D'Eramo, J. Peters, and J. Pajarinen. Convex regularization in monte-carlo tree 392 search. In International Conference on Machine Learning , pages 2365-2375. PMLR, 2021. 393

S. Eiffert, H. Kong, N. Pirmarzdashti, and S. Sukkarieh. Path planning in dynamic environments 394 using generative rnns and monte carlo tree search. In 2020 IEEE International Conference on 395 Robotics and Automation (ICRA) , pages 10263-10269. IEEE, 2020. 396

N. Funk, G. Chalvatzaki, B. Belousov, and J. Peters. Learn2assemble with structured repre397 sentations and search for robotic architectural construction. In Conference on Robot Learning , 398 pages 1401-1411. PMLR, 2022. 399

400

401

M. Geist, B. Scherrer, and O. Pietquin. A theory of regularized markov decision processes. In International Conference on Machine Learning , pages 2160-2169, 2019.

402

403

404

405

406

407

408

409

410

T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning , pages 1861-1870, 2018.

J. Honda and A. Takemura. An asymptotically optimal bandit algorithm for bounded support models. In COLT , pages 67-79. Citeseer, 2010.

L. Kocsis and C. Szepesvári. Bandit based monte-carlo planning. In Proceedings of the 17th European Conference on Machine Learning , ECML'06, page 282-293, Berlin, Heidelberg, 2006. Springer-Verlag. ISBN 354045375X. doi: 10.1007/11871842_29. URL https://doi. org/10.1007/11871842_29 .

L. Kocsis, C. Szepesvári, and J. Willemson. Improved monte-carlo search. Univ. Tartu, Estonia, 411 Tech. Rep , 1, 2006. 412

B. Mavrin, H. Yao, L. Kong, K. Wu, and Y. Yu. Distributional reinforcement learning for 413 efficient exploration. In International conference on machine learning , pages 4424-4434. 414 PMLR, 2019. 415

V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, 416

Riedmiller, A. K. Fidjeland, G. Ostrovski, et al. Human-level control through deep rein417

forcement learning. Nature , 518(7540):529-533, 2015. 418

S. Mo, X. Pei, and C. Wu. Safe reinforcement learning for autonomous vehicle using monte 419 carlo tree search. IEEE Transactions on Intelligent Transportation Systems , 23(7):6766-6773, 420 2021. 421

422

423

G. Neu, A. Jonsson, and V. Gómez. A unified view of entropy-regularized markov decision processes. arXiv preprint arXiv:1705.07798 , 2017.

424

425

426

427

428

429

430

431

432

433

434

435

436

437

438

439

440

441

442

443

444

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

461

462

T. Nguyen-Tang, S. Gupta, and S. Venkatesh. Distributional reinforcement learning via moment matching. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 9144-9152, 2021.

M. Painter, M. Baioumy, N. Hawes, and B. Lacerda. Monte carlo tree search with boltzmann exploration. Advances in Neural Information Processing Systems , 36, 2024.

D. Perez, S. Samothrakis, and S. Lucas. Knowledge-based fast evolutionary mcts for general video game playing. In 2014 IEEE Conference on Computational Intelligence and Games , pages 1-8. IEEE, 2014.

C. Riou and J. Honda. Bandit algorithms based on thompson sampling for bounded reward distributions. In Algorithmic Learning Theory , pages 777-826. PMLR, 2020.

J. 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.

J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz. Trust region policy optimization. In International Conference on Machine Learning , pages 1889-1897, 2015.

D. Shah, Q. Xie, and Z. Xu. Non-asymptotic analysis of monte carlo tree search. In Abstracts of the 2020 SIGMETRICS/Performance Joint International Conference on Measurement and Modeling of Computer Systems , pages 31-32, 2020.

D. Shah, Q. Xie, and Z. Xu. Nonasymptotic analysis of monte carlo tree search. Operation Research , 70(6):3234-3260, 2022.

D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. van den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot, S. Dieleman, D. Grewe, J. Nham, N. Kalchbrenner, I. Sutskever, T. Lillicrap, M. Leach, K. Kavukcuoglu, T. Graepel, and D. Hassabis. Mastering the game of Go with deep neural networks and tree search. Nature , 529(7587): 484-489, Jan. 2016. doi: 10.1038/nature16961.

D. Silver, T. Hubert, J. Schrittwieser, I. Antonoglou, M. Lai, A. Guez, M. Lanctot, L. Sifre, D. Kumaran, T. Graepel, et al. Mastering chess and shogi by self-play with a general reinforcement learning algorithm. arXiv preprint arXiv:1712.01815 , 2017.

D. Silver, J. Schrittwieser, K. Simonyan, I. Antonoglou, A. Huang, A. Guez, T. Hubert, L. Baker, M. Lai, A. Bolton, Y. Chen, T. Lillicrap, F. Hui, L. Sifre, G. van den Driessche, T. Graepel, and D. Hassabis. Mastering the game of go without human knowledge. Nature , 550:354-, Oct. 2017. URL http://dx.doi.org/10.1038/nature24270 .

R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction . MIT press, 2018.

H. Van Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double q-learning. In Thirtieth AAAI conference on artificial intelligence , 2016.

T. Weissman, E. Ordentlich, G. Seroussi, S. Verdu, and M. J. Weinberger. Inequalities for the l1 deviation of the empirical distribution. Hewlett-Packard Labs, Tech. Rep , 2003.

C. Xiao, R. Huang, J. Mei, D. Schuurmans, and M. Müller. Maximum entropy monte-carlo planning. In Advances in Neural Information Processing Systems , pages 9516-9524, 2019.

A Outline 463

464

Notations will be described in Section B.

465

466

467

468

469

470

471

Supporting Lemmas are presented in Section C.

The Convergence of CATS and PATS in Non-stationary multi-armed bandits is shown in Section D.

Section E presents the concentration and convergence guarantee of CATS and PATS in MCTS.

Section F discusses about Limitations and possible improvements.

Experimental setup is provided in Section G.

472

Additional Experimental results are shown in Section H.

B Notations 473

C Supporting Lemmas 474

We start with a result of the following lemma which plays an important role in the analysis of our 475 MCTS algorithm. 476

477

Lemma 1. For m ∈ [ M ] , let ( ̂ V m,n ) n ≥ 1 be a sequence of estimator satisfying plim n →∞ ̂ V m,n = V m .

478

479

480

481

Assume that there exists a constant L > 0 such that L = supremum { ̂ V m,n } n ≥ 1 . Let R i be an iid sequence with mean µ and S i be an iid sequence from a distribution p = ( p 1 , . . . , p M ) supported on { 1 , . . . , M } . Introducing the random variables N n m = # |{ i ≤ n : S i = s m }| , we define the sequence of estimator

Then there exists some constant c ′ (which depends on p i (i=1,2,...,M), γ , µ ) such that 482

Proof. Let p = ( p 1 , p 2 , ...p M ) , p ∈ △ M where △ M = { x ∈ R M : ∑ M i =1 R i = 1 , R i ≥ 0 } is the 483 ( M -1) -dimensional simplex. Let us study a random vector ̂ p n = ( N n 1 n , N n 2 n , ..., N n M n ) . Let us define 484

V = ( V 1 , V 2 , ...V M ) . Let ̂ R n = 1 n ∑ n i =1 R i , ̂ V n = ( ̂ V 1 ,N n 1 , ̂ V 2 ,N n 2 , ..., ̂ V M,N n M ) , ∑ M i =1 N n i = n , N n i 485 is the number of times that population i was observed. We have ̂ Q n = ̂ R n + γ 〈 ̂ p n , ̂ V n 〉 . Therefore, 486

To upper bound A, let us consider 〈 ̂ p n , ̂ V 〉 -⟨ p, V ⟩ = 〈 ( ̂ p n -p ) , ̂ V n 〉 + 〈 p, ( ̂ V -V ) 〉 . Then, 487

By applying a Hölder inequality to ̂ p n -p and ̂ V , we obtain 488

with L is the supremum of ̂ V . Then we can derive 489

According to (39), we have for any M ≥ 2 and δ ∈ [0 , 1] 490

Define ϵ = √ 2 M ln(2 /δ ) n , therefore δ = 2exp { -nϵ 2 2 M } , we have 491

Let us define an event E = { N n m ≥ np m 2 } . Therefore, 494

We consider p m > 0 only since if p m = 0 , p m ( ̂ V m,N n m -V m ) = 0 , and has been eliminated. 495 Therefore, 496

Therefore, 492

We also have 493

That leads to 497

with c ′ > 0 depends on c, M, p i . So that 498

By following the same steps, we can derive 499

Therefore, with n ≥ 1 , ϵ > 0 , 500

Furthermore, 501

Therefore, 502

Also because lim n →∞ E [ ̂ V m,n ] = V m , lim n →∞ ̂ N n m n = p m , and E [( ̂ R n -µ )] = 0 so that, 503

That mean 504

505

506

507

508

which concludes the proof.

Results from Lemma 1 is important as it shows the concentration for the Q value estimation given the concentration of V value of the children nodes.

Lemma 2. Let consider non-negative variables x, y ∈ R + , and a constant m that 0 ≤ m ≤ 1 . Then

Proof. With y = 0 , or x = 0 , the inequality (2) becomes correct. Let consider the case where 509 x > 0 , y > 0 , the inequality (2) can be written as 510

Let us define a function 511

We can see that 512

513

514

515

516

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

532

533

534

535

536

537

538

539

because g ( x ) = x m -1 is a decreasing function with m ∈ [0 , 1] , x > 0 . Therefore,

So that, with t = x y ≥ 0 , we can derive the inequality (2).

We use Minkowski's inequality as shown below

Lemma 3. (Minkowski's inequality) Given p ≥ 1 , { x i , y i } ∈ R , i = 1 , 2 , ..., n , then we have the following inequality

Proof.

Lemma 4. (Markov's inequality) If X is a nonnegative random variable and a > 0 , then the probability that X is at least a is at most the expectation of X divided by a:

Proof. This is a well-known result.

D Convergence of CATS and PATS in Non-stationary multi-armed bandits

We note that in an MCTS tree, each node is considered a non-stationary multi-armed bandit where the average mean drifts due to the given action selection strategy. Therefore, we first study the convergence of CATS and PATS in non-stationary multi-armed bandits where the action selection is Thompson sampling, with the power mean backup operator at the root node. Detailed descriptions of the CATS and PATS in Non-stationary multi-armed bandits settings can be found in the main article in the Theoretical Analysis section.

We first establish the convergence and concentration properties for the power mean backup operator in non-stationary bandits, detailed in Theorem 1 for CATS and Theorem 2 for PATS.

To achieve these results, we demonstrate that the expected payoff of the power mean backup operator decays polynomially at a rate of O ( log n n ) . This is supported by Lemma 7 for CATS and Lemma 8 for PATS. Critical to this analysis are Lemma 5 and Lemma 6, which establish an upper bound of log( n ) for the expected number of suboptimal arm pulls.

Weintroduce some important definitions. F n a represents the empirical cumulative distribution function of arm a after n visitations, and F a represents the cumulative distribution function of arm a . We employ the following distance measure: If P and Q are two distributions characterized by parameters p = ( p 0 , p 1 , · · · , p N ) and q = ( q 0 , q 1 , · · · , q N ) respectively, then the distance is defined as

This represents the L ∞ distance between p and q in ❘ N +1 . We also denotes 540 KL ( P ∥ Q ) as the Kullback-Leibler divergence between P and Q , and denote 541 K inf ( F a , µ ⋆ ) = inf G : E [ G ] >µ ⋆ KL ( F a ∥ G ) . In addition, we denote K ( N ) inf ( F a , µ ⋆ ) = 542 inf { KL ( F a ∥ G ) ∣ ∣ ∣ ∣ the support of G ∈ { 0 , R N , 2 R N , · · · , R } , E [ G ] > µ ⋆ } . 543

We see that the definition of K inf ( F a , µ ⋆ ) and K ( N ) inf ( F a , µ ⋆ ) is only difference in the support set. 544

545

546

547

548

549

550

551

552

553

554

555

556

557

We denote the true parameter of arm a by p a = ( p 0 a , p 1 a , . . . , p N a ) with p i a = Pr X ∼ F a [ X = i N ] . We denote the parameter of the posterior distribution of arm a as α a = ( α 0 a , α 1 a , . . . , α N a ) . Since each arm a is non-stationary, we also denote the parameter of arm a after n visitations by p a ( n ) = ( p 0 a ( n ) , p 1 a ( n ) , . . . , p N a ( n )) with p i a ( n ) = Pr X ∼ F n a [ X = i N ] . The parameter of the posterior distribution of arm a denoted as α a ( n ) = ( α 0 a ( n ) , α 1 a ( n ) , . . . , α N a ( n )) We first show the results of an important Lemma 5. The proof follows closely to the Proof of Proposition 7 (29). The only difference is that in our settings, we study non-stationary bandits.

Lemma 5. Consider Categorical Thompson Sampling(CATS) strategy applied to a non-stationary problem where the pay-off sequence satisfies Assumption 1. Let T a ( n ) denote the number of plays of arm a up to timestep n .

If a is the index of a suboptimal arm, Then for any ϵ 0 , ϵ 1 ≥ 0 , each sub-optimal arm a is played in expectation at most

To analyze the expectation associated with selecting a suboptimal arm a , we decompose it into two 558 components: 559

We first find an upper bound for A 1 : 560

561

We see that if the event

occurs at step t for a certain m ∈ [1 , n ] , then T k ( t ′ ) > T k ( t ) = m for any t ′ > t . Therefore, for any 562 m ∈ [ n ] 563

We can bound for any m 0 ∈ [ n ] 564

By applying results of Lemma 13 Appendix F (29), we have 565

where P ∗ µ ⋆ -ϵ 1 = arg min x : u ⊤ x ≥ µ ⋆ -ϵ 1 KL ( P α a ∥ x ) and P α a ( t ) = 1 n + N +1 α a ( t ) . And by definition 566 KL ( P α a ( t ) ∥ P ∗ µ ⋆ -ϵ 1 ) = K inf ( P α a ( t ) , µ ⋆ -ϵ 1 ) , therefore 567

where C = exp { 1 / 12 } Γ( N +1) ( 1 √ 2 π ) N . On the other hand, K inf ( x, µ ⋆ -ϵ 1 ) is continuous in x ∈ [0 , 1] N +1 568 on the probability simplex with respect to the L ∞ distance from ((19), Theorem 7) and Lemma 18 in 569 Appendix H (29). Therefore, for any ϵ 3 > 0 , there exists ϵ 2 > 0 and constant C ′ > 0 such that 570

And because Pr ( ∥ α a ( t ) T k ( t )+ N +1 -p a ( t ) ∥ ∞ ≤ ϵ 2 , T k ( t ) = m ) ≤ 1 . Therefore, 571

Choosing m 0 = log n K inf ( p a ,µ ⋆ -ϵ 1 ) -ϵ 3 -N -1 , we have 572

Furthermore, as from ((19), Theorem 7), it is proven that µ →K inf ( F, µ ) is continuous for µ < 1 , 573 when we scale reward from [0,1] to [0 , R ] therefore µ from [0,1] to [0 , R ] . We have µ →K inf ( F, µ ) 574 is continuous for µ < R . Therefore, ∀ ϵ 4 > 0 , ∃ ϵ 1 > 0 , such that 575

Also According to Proposition 8 (29), for any ϵ 0 > 0 we have 577

Combining inequality (5) and inequality (6) leads us to 578

Therefore which concludes the proof.

Lemma 6. Consider Particle Thompson Sampling(PATS) strategy applied to a non-stationary problem where the pay-off sequence satisfies Assumption 1. Then for any ϵ 0 ≥ 0 . Let T a ( n ) denote the number of plays of arm a up to timestep n . Then if a is the index of a suboptimal arm, then each sub-optimal arm a is played in expectation at most

Therefore, ∀ ϵ 0 > 0 576

579

580

581

582

583

Proof. In this Theorem, we use the Levy distance. Recall that the Levy distance between two 584 cumulative distribution functions F and G on [0 , 1] is defined as 585

The proof follows the same steps as in Lemma 5. We also can derive 586

We can use the same ways of derivations as in Lemma 5, equation (4) to have the same bound 587

According to Lemma 15 in Appendix G.1 (29) on conditional probabilities, for any ν ∈ (0 , 1) we 588 have 589

Because K inf ( F, µ ) is continuous in F with respect to the Levy distance from (19), Theorem 7, for 590 any ϵ 3 > 0 there exists ϵ 2 > 0 such that 591

Therefore, ∀ ν ∈ (0 , 1) and for any ϵ 5 > 0 , there exists ϵ 1 , ϵ 2 > 0 such that 592

This implies that ∀ ϵ 0 > 0 , there exists ν ∈ (0 , 1) , ϵ 1 > 0 and ϵ 2 > 0 such that 593

Therefore, according to inequality (7) and the fact that 594

we have 595

Choose m 0 = log n K inf ( F a ,µ ⋆ ) -ϵ 0 we have 596

Also According to Proposition 10 (29), for any ϵ 0 > 0 we have 597

which concludes the proof.

That leads us to 598

599

600

601

602

̂

µ

n

(

p

)

as

̂

µ

n

(

p

) =

K

a

=1

a

(

n

n

p

a,T

̂

µ

a

(

n

)

, and

̸

δ

⋆,n

=

µ

⋆

-

µ

⋆,n

For any

p

1

, ϵ

0

>

Proof. We observe that 603

604

Furthermore,

Since µ ⋆,n = max a ∈ [ K ] { µ a,n } , we have 605

Applying Minkowski's inequality from Lemma 3, and the result of (8), we have 606

On the other hand, 607

̸

Lemma 7. Consider Categorical Thompson Sampling(CATS) strategy applied to a non-stationary problem where the pay-off sequence satisfies Assumption 1. Let us define the power mean estimator

( ∑

)

1

p

̸

T

)

0

, we have

Because power mean is an increasing function of p , so that 608

Furthermore, we observe that 609

̸

So that, from equation (9) we have 610

̸

̸

Therefore 611

̸

Please note that because we study non-stationary bandits, E [ ∑ n t R a,t ] = nµ a,n , therefore, 612

According to Lemma 5, we have 613

̸

which concludes the proof. 614

615

616

Lemma 8. Consider Particle Thompson Sampling(PATS) strategy applied to a non-stationary problem where the pay-off sequence satisfies Assumption 1. Let us define the power mean estimator

̸

̸

̸

̸

̸

Proof. Similar to Lemma 7, we can derive 618

And according to Lemma 6, we have 619

̸

620

621

622

623

which concludes the proof.

Theorem 1. For a ∈ [ K ] , let ( ̂ µ a,n ) n ≥ 1 be a sequence of estimator satisfying plim n →∞ ̂ µ a,n = µ a and let µ ⋆ = max a { µ a } . Assume that all the estimators are bounded in [0 , R ] . We consider a bandit algorithm that selects each arm according to CATS once in each round n ≥ K .

Then, for all p ∈ [1 , ∞ ) , the sequence of estimators 624

where T a ( n ) = ∑ n -1 t =1 1 ( a t = a ) is the number of selections of a prior to round n satisfies 625

Proof. We first prove that lim n →∞ E [ ̂ µ n ( p )] = µ ∗ . According to the result of Lemma 7, we have 626

̸

̸

with δ ⋆,n = µ ⋆ -µ ⋆,n , and because lim n →∞ µ ∗ ,n = µ ⋆ , we can concludes that 627

Second, we prove that 628

We observe that 629

630

Next, according to Lemma 7, 631

that leads to 632

Because lim n → n | δ ⋆,n | = 0 , therefore, ∃ N 0 > 0 such that ∀ n ≥ N 0 , we have | δ ⋆,n | < ϵ/ 2 that means

̸

̸

̸

Therefore, ∃ c > 0 such that 633

which means 634

Now we see that | ̂ µ n ( p ) -µ ⋆ | ≤ R . With ϵ ≥ R , we have | ̂ µ n ( p ) -µ ⋆ | > ϵ ⇔| ̂ µ n ( p ) -µ ⋆ | > R , 635 therefore the inequality holds as 636

with 0 < ϵ < R, 1 ≤ n < N 0 ⇒ nϵ < RN 0 ⇒ n -1 ε -1 > 1 /RN 0 . Therefore 637

638

639

640

641

642

which means

That concludes the proof.

Theorem 2. For a ∈ [ K ] , let ( ̂ µ a,n ) n ≥ 1 be a sequence of estimator satisfying plim n →∞ ̂ µ a,n = µ a and let µ ⋆ = max a { µ a } . Assume that all the estimators are bounded in [0 , R ] . We consider a bandit algorithm that selects each arm according to PATS once in each round n ≥ K .

Then, for all p ∈ [1 , ∞ ) , the sequence of estimators 643

where T a ( n ) = ∑ n -1 t =1 1 ( a t = a ) is the number of selections of a prior to round n satisfies 644

Proof. The proof follows the same steps as Theorem 1. We first prove that lim n →∞ E [ ̂ µ n ( p )] = µ ∗ . 645 According to the result of Lemma 8, we have 646

̸

̸

with δ ⋆,n = µ ⋆ -µ ⋆,n , and because lim n →∞ µ ∗ ,n = µ ⋆ , we can concludes that 647

Second, we prove that 648

We observe that 649

Because lim n → n | δ ⋆,n | = 0 , therefore, ∃ N 0 > 0 such that ∀ n ≥ N 0 , we have | δ ⋆,n | < ϵ/ 2 that means 650

Next, according to Lemma 8, 651

652

653

654

655

656

657

658

659

660

661

662

663

664

665

666

667

668

669

670

671

672

673

674

that leads to

Therefore, ∃ c > 0 such that

which means

Now we see that | ̂ µ n ( p ) -µ ⋆ | ≤ R . With ϵ ≥ R , we have | ̂ µ n ( p ) -µ ⋆ | > ϵ ⇔| ̂ µ n ( p ) -µ ⋆ | > R , therefore the inequality holds as

with 0 < ϵ < R, 1 ≤ n < N 0 ⇒ nϵ < RN 0 ⇒ n -1 ε -1 > 1 /RN 0 . Therefore

which means

That concludes the proof.

E Convergence of CATS and PATS in Monte-Carlo Tree Search

Based upon the results of CATS and PATS using power mean as the value backup operator on the described non-stationary multi-armed bandit problem, we derive theoretical results for CATS in an MCTS tree.

We derive Theorem 3 for CATS and Theorem 4 for PATS, which show concentration and convergence for any internal node in the tree. These proofs utilize induction, leveraging the results of Lemma 7 for CATS and Lemma 8 for PATS, and Lemma 5 for CATS and Lemma 6 for PATS. Additionally, we use Lemma 1, which demonstrates the concentration and convergence of an estimated Q-value based on the child V-value node, applying it recursively throughout the tree.

Our main results, Theorem 5 for CATS and Theorem 5 for PATS, show that the simple regret converges non-asymptotically at a rate of O ( n -1 ) .

Theorem 3. When we apply the CATS algorithm, we have

(i) For any node s h at the depth h th in the tree,

(ii) For any node s h at the depth h th in the tree,

Proof. We will prove this by induction on the depth D of the tree. If the tree only has depth (1) . 675 The state at the root node is s 0 , let us assume that at time step t , after taking action a k , the MCTS tree 676 gets an intermediate reward r t ( s 0 , a k ) and traverses to the next state s 1 . Let us assume that R ( s 0 , a k ) 677 is the mean of the intermediate reward at state s 0 , after taking action a k . We recall the definition of 678 ˜ Q ( s 0 , a k ) , with π 0 is the rollout policy to estimate the newly added node at the leaf, 679

̸

where ˜ V ( s 1 ) is the value of the policy π 0 at state s 1 , A s 0 is the set of feasible actions at state s 0 , 680 |A s 0 | = M , P ( s 1 | s 0 , a k ) is the probability transition of taking action a k at state s 0 to state s 1 . From 681 ((1)), we have 682

( i ) is a direct result of Lemma 1 with X t is the intermediate reward r t ( s 0 , a k ) at time t , p = 683 ( p 1 , p 2 , ...p M ) ∼ P ( ·| s 0 , a k ) , where P ( ·| s 0 , a k ) is the probability transition dynamic of taking action 684 a k at state s 0 . For m ∈ [ M ] , each ( ̂ V m,t ) t ≥ 1 at time step t is the deterministic initial Value function 685 ˜ V ( s 1 ) . We have 686

( ii ) Direct results from Theorem 1. In detail, we have from ( i ) , 687

Because by definition: 688

Then we have 689

690

691

692

693

694

695

that concludes for ( ii )

Let us assume that with the tree of depth D , the theorem holds for all its children.

Now let's consider the tree with depth ( D +1) . When we take one action at the root node at the state s 0 , it comes to a subtree with depth ( D ) . According to the induction assumption, the results hold for any internal node in the tree after we take the first action. We have s 1 ∼ τ ( s 0 , a k ) . By the definition, V ( s ) = V ( s ) and, for all h ≤ H -1 ,

˜ H 0 H

By the assumption of the induction the root node of a subtree with depth ( D ) at state s 1 we have 696

( i ) Let's apply Lemma 1 with { X t } is the intermediate reward { r t ( s 0 , a k ) } , p = ( p 1 , p 2 , ...p M ) ∼ 697 P ( ·| s 0 , a k ) . For m ∈ [ M ] , each ( ̂ V m,t ) t ≥ 1 at time step t is the empirical Value function ̂ V t ( s 1 ) . We 698 will have 699

( ii ) follows the results of Theorem 1 as at the root node s 0 of depth D +1 , with 700

And because 701

Then, we have 702

703

704

705

706

707

708

709 710

that concludes for ( ii )

The results of Theorem 3 hold for any node in the tree with the tree of depth ( D +1) . By induction, we can conclude the proof.

Similarly we can derive the following Theorem

Theorem 4. When we apply the PATS algorithm, we have

(i) For any node s h at the depth h th in the tree,

(ii) For any node s h at the depth h th in the tree,

Proof. The proof follows the same steps as Theorem 3 by applying the results of Lemma 1 and 711 Theorem 2. 712

Theorem 5. (Convergence of Expected Payoff of CATS) We have at the root node s 0 , 713

Proof. We prove the result by induction and use the results of Theorem 3 to prove this Theorem. Let 714 us assume that the depth of the tree is D = 1 , as the results of Lemma 7, we have 715

And because the tree only have the depth D = 1 , we have | δ ⋆,n | = 0 , so that the result holds at 716 the depth D = 1 . Let us assume that we have the result of the tree at the depth D . Now when the 717 depth of the tree is D +1 , at the root node s 0 , the conditions of Assumption 1 hold as the results of 718 Theorem 3 then we have 719

where the bias 720

Therefore, 721

that concludes the proof. 722

Next, we present the results of Theorem 6. The proof follows the same steps as Theorem 5. 723

Theorem 6. (Convergence of Expected Payoff of PATS) We have at the root node s 0 , 724

725

726

727

728

729

730

731

732

733

734

735

736

737

738

739

740

741

742

743

744

745

746

747

F Limitations

Computational Demands : The CATS distributional Monte Carlo Tree Search (MCTS) faces challenges in managing computational demands while maintaining and updating probability distributions, leading to a slightly increased complexity.

Fixed precision : The PATS set of particles can increase in size if the observed value are different. We prevent this in the implementation by fixing the float precision.

Number of atoms : Our approach's performance is slightly influenced by hyperparameters, with the number of atoms being a critical factor. Suboptimal choices may affect performance.

G Experimental setup

All the experiments were done on 8 Intel Xeon Gold 6130 (Skylake), x86_64, 2.10GHz, 2 CPUs/node, 16 cores/CPU. Whenever feasible, we opted for open-source implementations of algorithms and environments.

Parameters selection We search the number of atoms from {10,20,...,100} and choose the results with best performances. We set the discount factor γ = . 99 for MDPs, and γ = . 95 for POMDPs. For UCT, we use the exploration constant C = √ 2 × ( R max -R min ) .

Atari hyperparameters We run CATS in Atari with 10 random seeds, where each seed with 512 samples and collect the average score. We found that only 512 simulations were necessary due to the utilization of a pretrained neural network. We run CATS with 100 atoms. The temperature parameter τ of MENTS and TENTS is tuned from {0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0}. The selected parameter τ are shown in Table 4. The exploration constant ϵ for MENTS and TENTS are set to 0 . 01 . For Power-UCT, we select the power mean p = 2 .

Atari

Atari environments (4) provide diverse video game-inspired scenarios commonly used in reinforce748 ment learning research. These environments offer challenges based on classic Atari 2600 games 749 (23; 38; 6). To explore enhanced exploration in deep reinforcement learning, we employ a Deep 750 Q-Network pre-trained following the experimental setup outlined in (23). This pre-trained network 751 initializes action-values for each node, combined with a Monte-Carlo Tree Search method similar to 752 the AlphaGo one. Here, P prior represents the Boltzmann distribution derived from the action-values 753 Q ( s, . ) computed by the network. The results in Table 3 show that CATS and PATS outperform UCT, 754 DQN, Power-UCT, TENTS and MENTS in most of the games. For example, CATS is significant 755 better than other methods in Breakout, Enduro , while PATS is significant better than other methods 756 in MsPacman, Solaris . Our intention in this experiment is not to assert exceptional superiority, but 757 rather to emphasize that CATS and PATS actually work in complicated Atari benchmark. 758

759

760

761

762

763

764

765

766

767

768

769

770

771

772

773

774

775

776

777

778

779

780

781

782

783

784

785

786

787

788

789

790

791

792

793

NeurIPS Paper Checklist

The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and precede the (optional) supplemental material. The checklist does NOT count towards the page limit.

Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:

You should answer [Yes] , [No] , or [NA] .

[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.

Please provide a short (1-2 sentence) justification right after your answer (even for NA).

The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.

The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.

IMPORTANT, please:

Delete this instruction block, but keep the section heading 'NeurIPS paper checklist" ,

Keep the checklist subsection headings, questions/answers and guidelines below.

Do not modify the questions and only use the provided macros for your answers .

(i) Claims

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

Answer: [Yes] ,

Justification: We discuss the problem of planning in stochastic environments and we present a method to tackle problem with clear contributions.

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.

(ii) Limitations

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

Answer: [Yes]

Justification: We discuss the limitation in Section 6

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.

(iii) 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: We provide the main theorems in the main paper and proofs 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.

(iv) 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: Code and reproducibility steps are provided in 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.

(v) 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: Full code is available in 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.

(vi) 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: The experimental setting is detailed in the appendix.

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.

(vii) 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: We provide error bars for the plots. For Atari, we report the standard deviation. 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.

(viii) 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: We provide the details about the computer resources used (CPU and number of cores).

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).

(ix) 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: The research conducted in the paper conforms the Code of Ethics.

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).

(x) Broader Impacts

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

Answer: [NA]

Justification: The research conducted in the paper has no societal 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.

1000

1001

1002

1003

1004

1005

1006

1007

1008

1009

1010

1011

1012

1013

1014

1015

1016

1017

1018

1019

1020

1021

1022

1023

1024

1025

1026

1027

1028

1029

1030

1031

1032

1033

1034

1035

1036

1037

1038

1039

1040

1041

1042

1043

1044

1045

1046

1047

1048

1049

1050

1051

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).

(xi) 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 research proposed in this 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.

(xii) 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: [NA]

Justification: We do not use existing assets.

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.

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

(xiii) New Assets

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

Answer: [Yes]

Justification: The provided code is well documented.

1052

1053

1054

1055

1056

1057

1058

1059

1060

1061

1062

1063

1064

1065

1066

1067

1068

1069

1070

1071

1072

1073

1074

1075

1076

1077

1078

1079

1080

1081

1082

1083

1084

1085

1086

1087

1088

1089

1090

1091

1092

1093

1094

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.

(xiv) 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.

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.

(xv) 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:\nThe authors explore the use of distributional reinforcement learning within Monte Carlo Tree search. They propose two algorithms CATS and PATS a categorical distribution and particle distribution based approach respectfully. They perform a theoretical analysis of the methods and show analysis of regret. They then evaluate on a synthetic planning tasks and evaluate it in combination with a pre-trained network on the atari benchmark.\n\nSTRENGTHS:\n- Thorough theoretical analysis\n- The authors address distributional RL applied to planning which is a clear important direction of research\n\nWEAKNESSES:\n- Lack of referencing of existing work and novelty relative to existing work\n    - Hayes, C.F., Reymond, M., Roijers, D.M., Howley, E. and Mannion, P., 2023. Monte Carlo tree search algorithms for risk-aware and multi-objective reinforcement learning. *Autonomous Agents and Multi-Agent Systems*, *37*(2), p.26.\n- Unjustifiable statement: “For example, CATS is significantly better than other methods in Breakout, Enduro” There is no significance testing performed so this statement cannot be made and in fact the Confidence intervals overlap\n- Key results in appendix and lack of empirical results in the main paper\n- CATS never outperforms fixed depth MCTS on the synthetic tree task\n- Unable to find code despite checklist saying it is provided\n\n- Small issues\n    - Figure 2 algorithms alignment off\n    - Indication of Atari results in section 5 which are not there\n    - Adding bold to best performing method in the Atari table would be useful for readability\n\nQUESTIONS:\n- How does your approach differ to existing methods (Hayes et al)?\n- Why do the baselines investigated change between the synthetic experiments and the atari experiments?\n- Why do you think CATS is consistently outperformed by other baselines on the synthetic tree task?\n\nLIMITATIONS:\n- Limitations are not included in the main body of the paper which they should be especially considering there is space. The limitations are also not thoroughly discussed for example\n    - “faces challenges in managing computational demand” : this does not say anything meaningful\n    - “Our approach’s performance is slightly influenced by hyperparameters”, this can be said for essentially any method\n- It seems that the distributional approach has an additional memory cost which if correct should be added to the limitations\n- Given CATS and PATS do not massively outperform all baselines on the synthetic task I think the limitations should be where this is addressed and perhaps some insight given into why this is and why performance on Atari is also not particular strong relative to methods such as MENTS.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper propose two algorithms, Categorical Thompson Sampling for MCTS (CATS) and Particle Thompson Sampling for MCTS (PATS). These algorithms extend Distributional Reinforcement Learning (RL) to Monte-Carlo Tree Search (MCTS) by modeling value functions as categorical and particle distributions, respectively to improve the performance of MCTS in highly stochastic settings.\n\nSTRENGTHS:\n- **Originality:** The integration of Distributional RL into MCTS using categorical and particle distributions is innovative and addresses a critical need in handling stochastic environments (Sections 3.1-3.3).\n- **Quality:** The theoretical analysis is rigorous, with well-defined proofs and clear methodology (Sections 4.1 and 4.2).\n\nWEAKNESSES:\n1. **Empirical Validation**: While the paper presents a comprehensive set of experiments demonstrating the efficacy of the proposed methods (CATS and PATS) in synthetic scenarios, there is an evident lack of diversity in the benchmarks used. \n\n2. **Algorithm Complexity and Overhead**: Both CATS and PATS introduce additional complexity by incorporating distributional approaches and Thompson Sampling into MCTS. The paper does not sufficiently address the computational overhead or the scalability of these methods when applied to environments with larger state or action spaces. This could be crucial for understanding the practical deployment of these algorithms in real-world applications.\n\nQUESTIONS:\n1. **Generalization to Complex Environments**: As we know, UCT is used as default choice in AlphaGo series. How do CATS and PATS perform in highly dynamic or unstructured environments beyond SyntheticTree? Can the authors provide insights or empirical evidence on the performance of these methods in such scenarios?\n\n2. **Theoretical vs. Practical Performance**: The theoretical improvements in regret bounds are notable. Can the authors discuss how these improvements manifest in practical scenarios? Are there specific environments or settings where the reduced regret significantly enhances decision-making?\n\n3. **Computational Overhead**: What is the computational overhead introduced by the distributional and sampling methods in CATS and PATS compared to traditional MCTS? How do these methods scale with the size of the state and action spaces?\n\n4. **Hyperparameter Sensitivity Analysis**: Given the potential impact of hyperparameters like the number of particles or categories on the algorithms' performance, can the authors provide a detailed sensitivity analysis? How robust are CATS and PATS to variations in these parameters?\n\n5. **Related Work**: Could you provide any related work in bandit literature that may use Thompson sampling for categorical posterior distribution. Discuss the technical novelty if any such literature exists.\n\nLIMITATIONS:\n### Computational Demands\nThe authors recognize that the Categorical Thompson Sampling (CATS) distributional Monte Carlo Tree Search (MCTS) involves increased complexity due to the management and updating of probability distributions. This acknowledgment is crucial as it highlights a potential scalability issue, especially in environments where computational resources are limited or real-time responses are required.\n\n### Fixed Precision\nThe approach used in the Particle Thompson Sampling (PATS) to manage the growth in the number of particles by fixing the float precision is a practical solution to prevent computational overload. However, this method may introduce limitations in the precision and adaptiveness of the model, potentially affecting the accuracy of value estimations in environments with high variability.\n\n### Number of Atoms\nThe performance sensitivity to the number of atoms indicates a hyperparameter dependency, which could impact the effectiveness and robustness of the model. The authors mention that suboptimal choices in this hyperparameter may affect performance, suggesting a need for careful tuning and validation to optimize the model's accuracy and efficiency.\n\n### Addressing Limitations\nWhile the authors have outlined these limitations, the discussion could be expanded to include more detailed strategies for mitigating these issues, particularly the computational demands and fixed precision aspects. For instance, strategies to optimize computational efficiency or adaptive techniques to dynamically adjust precision based on the context could further strengthen the approach.\n\n### Societal Impact\nThe paper does not explicitly address the potential negative societal impacts of the research. In the realm of reinforcement learning and AI planning, concerns such as the deployment in sensitive or critical environments, where errors may have significant consequences, should be considered. Discussions around ethical implications, misuse, and long-term effects would be beneficial.\n\n### Suggestions for Improvement\n1. **Enhanced Computational Strategies**: The authors could explore methods to reduce computational overhead, such as parallel processing or optimizing algorithmic efficiency, to make the model more practical for real-time applications.\n   \n2. **Dynamic Precision Adjustment**: Introducing mechanisms to adjust the precision of particle distributions dynamically based on the observed variability in the environment could help maintain balance between computational efficiency and model accuracy.\n\nOverall, the authors should be commended for their upfront discussion of the limitations, but there is room for deeper analysis and additional strategies to address these limitations comprehensively.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces distributional return estimates to MCTS-based planning. For this the authors borrow from work on distributional Q-Learning and show how to adapt the MCTS value back-up and action selection steps to compute and utilise these distributions. They formulate two approaches based on different distribution representations (quantile and particle based) for which they provide some theoretical convergence analysis as well as first experimental results.\n\nSTRENGTHS:\nThe paper combines two well-established ideas in MCTS and distributional value approximation resulting in a new algorithm with better theoretical guarantees. The overall approach and implementation of this combination makes sense and should at least in theory overcome limitations in stochastic environments.\n\nThough I was unable to verify all proofs in detail, the theoretical analysis seems to make sense and establish the advantages of the proposed methods.\n\nWEAKNESSES:\nDespite the soundness of the overall proposed method, I found the paper very hard to follow and felt details were missing due to an overall lack of focus. Contributing to this were the following issues:\n\n1. Empirical Evaluation\nExperiments are limited to a toy domain and results on the Atari benchmark reported in the appendix.\nThe toy domain is a tabular environment that is being generated randomly and contains stochasticity in both the final reward and transitions.\nFor an illustrating example this makes it hard to judge the combined effect on the overall return distribution to be approximated. \n\nHow the combinations of branching factor and depths that were plotted were chosen is unclear to me. Beyond this I am not sure how meaningful these plots are. In the right most plot it appears as if the PATS approximates the root value almost correctly in under 100 simulations - at which point it could not even have tried all k = 200 actions available to it.\n\nAlso CATS appears to be doing consistently worse than some of the other methods despite having the same theoretical properties as PATS.\n\nFor the Atari baseline the authors make use of Q-networks and point to a related paper. However, the exact implementation details and hyperparameters are not discussed making it hard to reproduce this work based on the paper alone.\n\nWhile stochasiticity and the exploration challenges this causes form one of the main motivations for this paper, no further ablations how the proposed methods improve here are presented.\n\n2. Content division\nThe author devote a significant amount of space to the summarisation of MCTS and distributional RL. While the theoretical analysis is arguably the strongest part of the presented work only the main theorems are found in the main body of the paper with very little contextualization.\n\n3. Overall presentation\nThere are several presentation issues in overall formatting, grammar and spelling. The former includes, but is not limited to overlapping lines, inconsistent / in-text section headers and wrong section references.\n\nQUESTIONS:\n-\n\nLIMITATIONS:\nThe discussion of limitations is restricted to a short paragraph in the appendix listing generic points such as increased computational demand and sensitivity to hyperparameter choices. However, no further investigation or explanation as to their severity is provided.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-12"}

comments: SUMMARY:
This paper introduces Categorical Thompson Sampling for MCTS (CATS) and Particle Thompson Sampling for MCTS (PATS) algorithms, which incorporate distributional reinforcement learning into Monte Carlo Tree Search (MCTS) to handle value estimation in stochastic settings. By modeling value functions as categorical and particle-based distributions and applying Thompson Sampling for action selection, the proposed algorithms aim to improve the robustness and accuracy of value estimates. The paper proves the theoretical effectiveness of these methods by achieving a non-asymptotic problem-dependent upper bound on simple regret of $O(n^{−1})$.

STRENGTHS:
The idea is interesting and original and the non-asymptotic problem-dependent upper bound on simple regret of $O(n^{−1})$ significantly advances the state-of-the-art from the previous $O(n^{−1/2})$.

WEAKNESSES:
1- While using distributional RL in MCTS to do Thompson sampling is interesting, it introduces much computation complexity hindering the applicability of the proposed algorithms.

2- The numerical experiments for the stochastic environments that are the main motivation of this work are done on a toy problem.


Minor comments

1- The presentation of the paper can be improved, specifically the parentheses () citation style can be confused with equations reference. 

2- Line 42, the authors mention V node for the first time without properly defining what is a V node.

QUESTIONS:
1- In the numerical experiments shown in Figure 3, Why does PATS perform much better than CATS?

LIMITATIONS:
1- The added high computational complexity from maintaining a distribution for each node in the MCTS.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-12