Distributional Monte-Carlo Planning with Thompson Sampling in Stochastic Environments

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 14:27:29

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.

Main paper: 
 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. for value backup as in the Upper Confidence bounds applied to Trees method (UCT) (21) , suffer from underestimation of optimal values while maximum backup suffers from overestimation of optimal values (9) . The power mean estimator (12) offers a balanced solution by computing a mean between the average and maximum values. In our approach, we also use power mean for value operator as each V node stores the power mean of empirical means of succeeding Q-value nodes, eliminating the need for V to be modeled as a distribution.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.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. (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. 
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 t (s 0 ) an estimator of the optimal value in s 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:E [R(s 0 , t)] = E V ⋆ (s 0 ) -V t (s 0 ) ≤ r(t),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 allEQUATIONwhere 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.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 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) . 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.EQUATION 
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.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 deriveEQUATIONwith 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.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 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.return a SimulateV (s h , t) (Sec 3.2) a =SelectAction (s h ) SimulateQ (s h , a, t) T sh (t) = T sh (t) + 1 Q(s h , a) = i z i (s h , a)p i (s h , a) V (s h ) = a Ts h ,a(t) Ts 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 sh,a (t) = T sh,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 max (s h , a)] then Q max (s h , a) = max{Q t (s h , a), Q max (s h , a)} Q min (s h , a) = min{Q t (s h , a), Q min (s h , a)} △z = Qmax-Qmin 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 = arg max 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 Ts h ,a (t) Ts 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 sh,a (t) = T sh,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)• 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 anintermediate 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 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.V (s) = a Ts,a(n) Ts(n) Q p (s, a) 1 p , p ≥ 1,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.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 (α 0 , α 1 , . . . , α N ) by Dir(α 0 , α 1 , . . . , α N ), whose density function is given byΓ( N i=0 α i ) Π N i=0 Γ(α i ) Π N i=0 x α i -1 i for (x 0 , . . . , x N ) ∈ [0, 1] N +1 such that N i=0 x i = 1.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:a t = arg max a ϕ t (s, a)After taking action a t and get an intermediate rewardQ 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 CategoricalThompson 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 asa t = arg max a ϕ t (s, a) .After taking action a t and get an intermediate reward In the next section, we provide a theoretical analysis of the convergence of simple regret for CATS and PATS. 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,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 , 2R(t) N , • • • , R(t)] ⊤ L t a = 
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.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 limitV if the following two properties hold:(A) Concentration: For all n ≥ 1, for all ε > 0, ∃c > 0 thatP | V n -V | > ε ≤ cn -1 ε -1 . (B) Convergence: lim n→∞ E[ V n ] = V .In that case, we write plimn→∞ V n = V .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: 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.Assumption 1. Consider K arms that for a ∈ [K], let ( µ a,n ) n≥1We 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. 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.Before presenting the main results (Theorem 3 Theorem 4), we first show an important LemmaLemma 1. Let ( V m,n ) n≥1 , m ∈ [M ], be a sequence of estimator satisfying plim n→∞ V m,n = V m .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 estimatorQ n = 1 n n i=1 R i + γ M m=1 N n m n V m,N n m .Then plimn→∞ Q n = µ + M m=1 p m V m .The significance of Lemma 1 lies in demonstrating the concentration and convergence of an estimated 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 denotes an intermediate reward received by 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, plimn→∞ Q n (s h , a k ) = Q(s h , a k ).(ii) For any node s h at the depth h th in the tree, plimn→∞ V n (s h ) = V (s h ).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, plimn→∞ Q n (s h , a k ) = Q(s h , a k ).(ii) For any node s h at the depth h th in the tree, plimn→∞ V n (s h ) = V (s h ).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 ,E V ⋆ (s 0 ) -V n (s 0 ) ≤ O(n -1 ).Theorem 6. (Convergence of Simple Regret of PATS) We have at the root node s 0 ,E V ⋆ (s 0 ) -V n (s 0 ) ≤ O(n -1 ).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 Fixed-Depth-MCTS (33) . Furthermore, Our finding more broadly applies to the power mean estimator with p ∈ [1, +∞). 
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.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 ) rate of the fixed version of UCT (32) . Empirical findings conclusively demonstrate the effectiveness of our approach in stochastic environments.• Notations will be described in Section B.• 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.• Additional Experimental results are shown in Section H. We start with a result of the following lemma which plays an important role in the analysis of our MCTS algorithm.Lemma 1. For m ∈ [M ], let ( V m,n ) n≥1 be a sequence of estimator satisfying plim n→∞ V m,n = V m .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 estimatorQ n = 1 n n i=1 R i + γ M m=1 N n m n V m,N n m .Then there exists some constant c ′ (which depends on p i (i=1,2,...,M), γ, µ) such thatplim n→∞ Q n = µ + M m=1 p m V m . 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 (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 defineV = (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 iis the number of times that population i was observed. We haveQ n = R n + γ p n , V n . Therefore, P Q n -µ + γ ⟨p, V ⟩ ≥ ϵ ≤ P R n -µ ≥ 1 2 ϵ + P γ p n , V n -γ ⟨p, Y ⟩ ≥ 1 2 ϵ ≤ exp{-2n ϵ 2 4 } + P p n , V n -⟨p, Y ⟩ ≥ 1 2γ ϵ A .To upper bound A, let us consider p n , V -⟨p, V ⟩ = ( p n -p), V n + p, ( V -V ) . Then,A ≤ P ( p n -p), V n ≥ 1 4γ ϵ A1 + P p, ( V n -V ) ≥ 1 4γ ϵ A2 .By applying a Hölder inequality to p n -p and V , we obtain( p n -p), V n ≤∥ p n -p ∥ 1 ∥ V n ∥ ∞ =∥ p n -p ∥ 1 L,with L is the supremum of V . Then we can deriveA 1 = P ( p n -p), V n ≥ 1 4γ ϵ ≤ P ∥ p n -p ∥ 1 L ≥ 1 4γ ϵ = P ∥ p n -p ∥ 1 ≥ 1 4γLϵ .According to (39) , we have for any M ≥ 2 and δ ∈ [0, 1]P ∥ p n -p ∥ 1 ≥ 2M ln(2/δ) n ≤ δ. Define ϵ = 2M ln(2/δ) n, therefore δ = 2 exp{ -nϵ 2 2M }, we haveP ∥ p n -p ∥ 1 ≥ ϵ ≤ 2 exp{ -nϵ 2 2M }.Therefore,A 1 ≤ P ∥ p n -p ∥ 1 ≥ ϵ ≤ 2 exp{ -nϵ 2 32M γ 2 L 2 }.We also haveA 2 = P M m=1 p m ( V m,N n m -V m ) ≥ 1 4γ ϵ ≤ M m=1 E P 1 N n m N n m t=1 V m,t -V m ≥ 1 4γp m ϵ N n m ≤ M m=1 E c(N n m ) -1 ( ϵ 4γp m ) -1 .Let us define an event E = N n m ≥ npm. Therefore,A 2 ≤ M m=1 E c( np m 2 ) -1 ( ϵ 4γp m ) -1 + M m=1 E P(N n m < np m 2 ) = M m=1 (c2 1+2 γ 1 p -1+1 m )n -1 ϵ -1 + M m=1 E P(N n m -p m n ≤ - p m n 2 ) ≤ M m=1 (c2 3 γ)n -1 ϵ -1 + M m=1 exp -2n( p m n 2 ) 2We consider p m > 0 only since if p m = 0, p m ( V m,N n m -V m ) = 0, and has been eliminated.A ≤ A 1 + A 2 ≤ 2 exp{ -nϵ 2 32M γ 2 L 2 } + M m=1 (c2 3 γ)n -1 ϵ -1 + M m=1 exp -2n( p m n 2 ) 2 .That leads toP Q n -µ + γ ⟨p, V ⟩ ≥ ϵ ≤ exp{-2n ϵ 2 4 } + 2 exp{ -nϵ 2 32M γ 2 L 2 } + M m=1 (c2 3 γ)n -1 ϵ -1 + M m=1 exp -2n( p m n 2 ) 2 ≤ c ′ n -1 ϵ -1 ,with c ′ > 0 depends on c, M, p i . So thatP Q n -µ + γ ⟨p, V ⟩ ≥ ϵ ≤ c ′ n -1 ϵ -1 ,By following the same steps, we can deriveP Q n -µ + γ ⟨p, V ⟩ ≤ -ϵ ≤ c ′ n -1 ϵ -1 .Therefore, with n ≥ 1, ϵ > 0,P Q n -µ + γ ⟨p, V ⟩ ≥ ϵ ≤ c ′ n -1 ϵ -1 .Furthermore,Q n -µ + γ ⟨p, V ⟩ = ( R n -µ) + γ p n , V n -γ ⟨p, Y ⟩ = ( R n -µ) + γ ( p n -p), V n + p, ( V -V ) Therefore, ⇒ E[ Q n ] -µ + γ ⟨p, V ⟩ ≤ E[( R n -µ)] + γ |E[ p n -p]| V n + p E[ V -V ] ⇒ E[ Q n ] -µ + γ ⟨p, V ⟩ ≤ E[( R n -µ)] + γ L |E[ p n -p]| + p E[ V -V ] 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, lim n→∞ E[ Q n ] = µ + γ M m=1 p m V m . That mean plim n→∞ Q n = µ + γ M m=1 p m V m ,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(x + y) m ≤ x m + y m .Proof. With y = 0, or x = 0, the inequality (2) becomes correct. Let consider the case wherex > 0, y > 0, the inequality (2) can be written as( x y + 1) m ≤ x y m + 1.Let us define a functionf (t) = (t + 1) m -t m -1, (t > 0).We can see thatf ′ (t) = m(t + 1) m-1 -mt m-1 = m (t + 1) m-1 -t m-1 ≤ 0 with m ∈ [0, 1], t > 0, because g(x) = x m-1 is a decreasing function with m ∈ [0, 1], x > 0. Therefore, f (t) ≤ f (0) = 0 with t > 0.So that, (t + 1) m -t m -1 ≤ 0, (t > 0). 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 inequalityi (|x i + y i |) p 1 p ≤ i (|x i |) p 1 p + i (|y i |) p 1 p.Proof. This is a basic result.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:Pr(X > a) ≤ E[X] a .Proof. This is a well-known result.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.We introduce 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 asd(P, Q) :=∥ p -q ∥ ∞ = sup i∈[0,N ] |p i -q i |This represents the L ∞ distance between p and q in R N +1 .We also denotes KL(P ∥ Q) as the Kullback-Leibler divergence between P and Q, and denoteK inf (F a , µ ⋆ ) = inf G:E[G]>µ⋆ KL(F a ∥ G).In addition, we denoteK (N ) inf (F a , µ ⋆ ) = inf KL(F a ∥ G) the support of G ∈ 0, R N , 2R N , • • • , R , E[G] > µ ⋆ .We see that the definition of K inf (F a , µ ⋆ ) andK (N )inf (F a , µ ⋆ ) is only difference in the support set.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∼Fa [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 byp 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. 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 mostE[T a (n)] ≤ (1 + ϵ 0 ) log n K (N ) inf (F a , µ ⋆ ) -ϵ 1 + o(log n) + O(1), Proof. We have ϕ a,t = [0, R N , 2R N , • • • , R] ⊤ L a,t , with L a,t ∼ Dir(α 0 a (t), . . . , α N a (t)).To analyze the expectation associated with selecting a suboptimal arm a, we decompose it into two components:E n t=1 1(I(t) = a) = E n t=1 1(I(t) = a), ϕ a,t ≥ µ * -ϵ 1 , d( F I(t) , F I(t) ) ≤ ϵ 2 ) A1 + E n t=1 1(I(t) = a), ϕ a,t < µ * -ϵ 1 , d( F I(t) , F I(t) ) > ϵ 2 )We first find an upper bound for A 1 :A1 = n t=1 n m=1 1 I(t) = a, θ k (t) ≥ µ ⋆ -ϵ 1 ; ∥ α a (t) T k (t) + N + 1 -p a (t) ∥ ∞ ≤ ϵ 2 , T k (t) = mWe see that if the eventI(t) = a, θ k (t) ≥ µ ⋆ -ϵ 1 ; ∥ α a (t) T k (t) + N + 1 -p a (t) ∥ ∞ ≤ ϵ 2 , T k (t) = m 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 m ∈ [n] n t=1 1 I(t) = a, θ k (t) ≥ µ ⋆ -ϵ 1 ; ∥ α a (t) T k (t) + N + 1 -p a (t) ∥ ∞ ≤ ϵ 2 , T k (t) = m ≤ 1We can bound for anyEQUATIONBy applying results of Lemma 13 Appendix F (29), we have Pr θ k (t) ≥ µ ⋆ -ϵ 1 α a , T k (t) = m ≤ C(m + N + 1) N/) = K inf (P αa(t) , µ ⋆ -ϵ 1 ), therefore Pr θ k (t) ≥ µ ⋆ -ϵ 1 α a (t), T k (t) = m ≤ C(m + N + 1) N/2 exp{-(m + N + 1)K inf (P αa(t) , µ ⋆ -ϵ 1 )},whereC = exp{1/12} Γ(N +1) 1 √ 2π N . On the other hand, K inf (x, µ ⋆ -ϵ 1 ) is continuous in x ∈ [0, 1] N +1on the probability simplex with respect to the L ∞ distance from (( 19), Theorem 7) and Lemma 18 in Appendix H (29) . Therefore, for any ϵ 3 > 0, there exists ϵ 2 > 0 and constant C ′ > 0 such thatEQUATIONChoosingEQUATIONFurthermore, as from (( 19), Theorem 7), it is proven that µ → K inf (F, µ) is continuous for µ < 1,when we scale reward from [0,1] to [0, R] therefore µ from [0,1] to [0, R]. We have µ → K inf (F, µ)is continuous for µ < R. Therefore, ∀ϵ 4 > 0, ∃ϵ 1 > 0, such that|K inf (p a , µ * -ϵ 1 ) -K inf (p a , µ * )| ≤ ϵ 4 ⇒ K inf (p a , µ * -ϵ 1 ) -ϵ 3 ≥ K inf (p a , µ * ) -ϵ 3 -ϵ 4 Therefore, ∀ϵ 0 > 0 A1 ≤ (ϵ 0 + 1) log n K inf (p a , µ ⋆ ) -N -1 + C ′ 1Also According to Proposition 8 (29) , for any ϵ 0 > 0 we haveEQUATIONCombining inequality ( 5) and inequality (6) leads us toE[T a (n)] ≤ (1 + ϵ 0 ) log n K (N ) inf (F a , µ ⋆ ) + o(log n) + O(1).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 mostE[T a (n)] ≤ log n K inf (F a , µ ⋆ ) -ϵ 0 + o(log n) + O(1).Proof. In this Theorem, we use the Levy distance. Recall that the Levy distance between two cumulative distribution functions F and G on [0, 1] is defined asD L (F, G) = inf{ϵ > 0 : ∀x ∈ [0, 1], F (x -ϵ) -ϵ ≤ G(x) ≤ F (x + ϵ) + ϵ}.The proof follows the same steps as in Lemma 5. We also can deriveE n t=1 1(I(t) = a) = E n t=1 1(I(t) = a), ϕ a,t ≥ µ * -ϵ 1 , D L ( F I(t) , F I(t) ) ≤ ϵ 2 ) B1 + E n t=1 1(I(t) = a), ϕ a,t < µ * -ϵ 1 , D L ( F I(t) , F I(t) ) > ϵ 2 )We can use the same ways of derivations as in Lemma 5, equation ( 4) to have the same boundEQUATIONAccording to Lemma 15 in Appendix G.1 (29) on conditional probabilities, for any ν ∈ (0, 1) we havePr θ k (t) ≥ µ ⋆ -ϵ 1 D L F a (t), F a (t) ≤ ϵ 2 , T k (t) = m ≤ 1 ν exp -n K inf ( F a (t), µ ⋆ -ϵ 1 ) -ν µ ⋆ -ϵ 1 1 -(µ ⋆ -ϵ 1 )Because K inf (F, µ) is continuous in F with respect to the Levy distance from (19) , Theorem 7, for any ϵ 3 > 0 there exists ϵ 2 > 0 such thatD L ( F a (t), F a ) ≤ ϵ 2 ⇒ K inf ( F a (t), µ ⋆ -ϵ 1 ) -K inf (F a , µ ⋆ -ϵ 1 ) ≤ ϵ 3Therefore, ∀ν ∈ (0, 1) and for any ϵ 5 > 0, there exists ϵ 1 , ϵ 2 > 0 such thatPr θ k (t) ≥ µ ⋆ -ϵ 1 D L F a (t), F a (t) ≤ ϵ 2 , T k (t) = m ≤ 1 ν -m K inf (F a , µ ⋆ -ϵ 1 ) -ϵ 3 -ν µ ⋆ -ϵ 1 1 -(µ ⋆ -ϵ 1 ) (Theorem 6 (19) ) ≤ 1 ν -m K inf (F a , µ ⋆ ) ϵ 1 1 -µ ⋆ -ϵ 3 -ν µ ⋆ -ϵ 1 1 -(µ ⋆ -ϵ 1 )This implies that ∀ϵ 0 > 0, there exists ν ∈ (0, 1), ϵ 1 > 0 and ϵ 2 > 0 such thatPr θ k (t) ≥ µ ⋆ -ϵ 1 D L F a (t), F a (t) ≤ ϵ 2 , T k (t) = m ≤ 1 ν exp {-m(K inf (F a , µ ⋆ ) -ϵ 0 )}Therefore, according to inequality (7) and the fact thatPr D L F a (t), F a (t) ≤ ϵ 2 , T k (t) = m ≤ 1we haveB1 ≤ m 0 + n t=1 1 ν exp {-m(K inf (F a , µ ⋆ ) -ϵ 0 )} ≤ m 0 + 1 ν T exp {-m 0 (K inf (F a , µ ⋆ ) -ϵ 0 )}Because power mean is an increasing function of p, so thatK a=1 T a (n) n µ a,n ≤ K a=1 T a (n) n (µ a,n ) p 1/p. Furthermore, we observe that µ a,n ≤ µ a,Ta(n) + µ a,Ta(n) -µ a,n . So that, from equation ( 9) we haveµ ⋆,n -µ n (p) ≤ R K a=1,a̸ =a * T a (n) n + K a=1 T a (n) n (µ a,n ) p 1/p -µ n (p) ≤ R K a=1,a̸ =a * T a (n) n + K a=1 T a (n) µ a,Ta(n) + µ a,Ta(n) -µ a,n p 1 p - K a=1 T a (n) µ a,Ta(n) p 1 p n 1 p (Minkovski's inequality) ≤ R K a=1,a̸ =a * T a (n) n + K a=1 T a (n) µ a,Ta(n) -µ a,n p 1 p n 1 p (Properties of L p norm) ≤ R K a=1,a̸ =a * T a (n) n + K a=1 T a (n) µ a,Ta(n) -µ a,n n 1 p = R K a=1,a̸ =a * T a (n) n + K a=1 Ta(n) t R a,t -T a (n)µ a,n n 1 p Therefore |E[ µ n (p) -µ ⋆,n ]| ≤ R K a=1,a̸ =a * E[T a (n)] n + E K a=1 Ta(n) t R a,t -T a (n)µ a,n n 1 p = R K a=1,a̸ =a * E[T a (n)] nPlease note that because we study non-stationary bandits, E[n t R a,t ] = nµ a,n , therefore, E K a=1 Ta(n) t R a,t -T a (n)µ a,n n 1 p = 0According to Lemma 5, we have|E[ µ n (p) -µ ⋆,n ]| ≤ R K a=1,a̸ =a * E[T a (n)] n ≤ R n K a=1,a̸ =a * (1 + ϵ 0 ) log n K (N ) (F a , µ ⋆ ) + o(log n) + O(1) ,which concludes the proof. |E[ µ n (p)] -µ ⋆ | ≤ |δ ⋆,n | + R n K a=1,a̸ =a * log n K inf (F a , µ ⋆ ) -ϵ 0 + o(log n) + O(1)Proof. Similar to Lemma 7, we can derive|E[ µ n (p) -µ ⋆,n ]| ≤ |δ ⋆,n | + R K a=1,a̸ =a * E[T a (n)] n .And according to Lemma 6, we have|E[ µ n (p) -µ ⋆,n ]| ≤ R K a=1,a̸ =a * E[T a (n)] n ≤ R n K a=1,a̸ =a * log n K inf (F a , µ ⋆ ) -ϵ 0 + o(log n) + O(1) ,which concludes the proof. Then, for all p ∈ [1, ∞), the sequence of estimatorsµ n (p) = K a=1 T a (n) n µ p a,Ta(n) 1 p , where T a (n) = n-1 t=1 1(a t = a) is the number of selections of a prior to round n satisfies plim n→∞ µ n (p) = µ ⋆ .Proof. We first prove thatlim n→∞ E[ µ n (p)] = µ * .According to the result of Lemma 7, we have|E[ µ n (p)] -µ ⋆ | ≤ |δ ⋆,n | + R K a=1,a̸ =a * E[T a (n)] n ≤ |δ ⋆,n | + R n K a=1,a̸ =a * (1 + ϵ 0 ) log n K (N ) (F a , µ ⋆ ) + o(log n) + O(1)with δ ⋆,n = µ ⋆ -µ ⋆,n , and because lim n→∞ µ * ,n = µ ⋆ , we can concludes thatlim n→∞ E[ µ n (p)] = µ * .Second, we prove that∀n ≥ 1, ∀ε > 0, ∃c > 0 that P (| µ n (p) -µ ⋆ | > ε) ≤ cn -1 ε -1 .We observe that| µ n (p) -µ ⋆ | ≤ | µ n (p) -µ ⋆,n | + |µ ⋆ -µ ⋆,n | = | µ n (p) -µ ⋆,n | + |δ ⋆,n | =⇒P(| µ n (p) -µ ⋆ | ≥ ϵ) ≤ P(| µ n (p) -µ ⋆,n | ≥ ϵ/2) + P(|δ ⋆,n | ≥ ϵ/2). Because lim n→n |δ ⋆,n | = 0, therefore, ∃N 0 > 0 such that ∀n ≥ N 0 , we have |δ ⋆,n | < ϵ/2 that means ∀n > N 0 , P(|δ ⋆,n | ≥ ϵ/2) = 0.Next, according to Lemma 7,|E[ µ n (p)] -µ ⋆,n | ≤ R n K a=1,a̸ =a * (1 + ϵ 0 ) log n K (N ) (F a , µ ⋆ ) + o(log n) + O(1) = O(n -1 ),that leads toP(| µ n (p) -µ ⋆,n | ≥ ϵ/2) ≤ |E[ µ n (p)] -µ ⋆,n | ϵ/2 = O(n -1 ) ϵ/2 .Therefore, ∃c > 0 such thatP(| µ n (p) -µ ⋆,n | ≥ ϵ/2) ≤ cn -1 ϵ -1 , which means ∀n ≥ N 0 , ∀ε > 0, ∃c > 0 that P (| µ n (p) -µ ⋆ | > ε) ≤ cn -1 ε -1 . Now we see that | µ n (p) -µ ⋆ | ≤ R. With ϵ ≥ R, we have | µ n (p) -µ ⋆ | > ϵ ⇔ | µ n (p) -µ ⋆ | > R,therefore the inequality holds asP (| µ n (p) -µ ⋆ | > ε) = 0 ≤ cn -1 ε -1 . with 0 < ϵ < R, 1 ≤ n < N 0 ⇒ nϵ < RN 0 ⇒ n -1 ε -1 > 1/RN 0 . Therefore ∀C > 1/RN 0 ⇒ P (| µ n (p) -µ ⋆ | > ε) ≤ 1 < Cn -1 ε -1 , which means ∀n ≥ 1, ∀ε > 0, ∃C > 0 that P (| µ n (p) -µ ⋆ | > ε) ≤ Cn -1 ε -1 .That concludes the proof. Then, for all p ∈ [1, ∞), the sequence of estimatorsEQUATION, where T a (n) = According to the result of Lemma 8, we have|E[ µ n (p)] -µ ⋆ | ≤ |δ ⋆,n | + R K a=1,a̸ =a * E[T a (n)] n ≤ |δ ⋆,n | + R n K a=1,a̸ =a * log n K inf (F a , µ ⋆ ) -ϵ 0 + o(log n) + O(1)with δ ⋆,n = µ ⋆ -µ ⋆,n , and because lim n→∞ µ * ,n = µ ⋆ , we can concludes thatlim n→∞ E[ µ n (p)] = µ * .Second, we prove that∀n ≥ 1, ∀ε > 0, ∃c > 0 that P (| µ n (p) -µ ⋆ | > ε) ≤ cn -1 ε -1 .We observe that| µ n (p) -µ ⋆ | ≤ | µ n (p) -µ ⋆,n | + |µ ⋆ -µ ⋆,n | = | µ n (p) -µ ⋆,n | + |δ ⋆,n | =⇒P(| µ n (p) -µ ⋆ | ≥ ϵ) ≤ P(| µ n (p) -µ ⋆,n | ≥ ϵ/2) + P(|δ ⋆,n | ≥ ϵ/2). Because lim n→n |δ ⋆,n | = 0, therefore, ∃N 0 > 0 such that ∀n ≥ N 0 , we have |δ ⋆,n | < ϵ/2 that means ∀n > N 0 , P(|δ ⋆,n | ≥ ϵ/2) = 0.Next, according to Lemma 8,|E[ µ n (p)] -µ ⋆,n | ≤ R n K a=1,a̸ =a * log n K inf (F a , µ ⋆ ) -ϵ 0 + o(log n) + O(1) = O(n -1 ),that leads toP(| µ n (p) -µ ⋆,n | ≥ ϵ/2) ≤ |E[ µ n (p)] -µ ⋆,n | ϵ/2 = O(n -1 ) ϵ/2 .Therefore, ∃c > 0 such thatP(| µ n (p) -µ ⋆,n | ≥ ϵ/2) ≤ cn -1 ϵ -1 , which means ∀n ≥ N 0 , ∀ε > 0, ∃c > 0 that P (| µ n (p) -µ ⋆ | > ε) ≤ cn -1 ε -1 . Now we see that | µ n (p) -µ ⋆ | ≤ R. With ϵ ≥ R, we have | µ n (p) -µ ⋆ | > ϵ ⇔ | µ n (p) -µ ⋆ | > R,therefore the inequality holds asP (| µ n (p) -µ ⋆ | > ε) = 0 ≤ cn -1 ε -1 . with 0 < ϵ < R, 1 ≤ n < N 0 ⇒ nϵ < RN 0 ⇒ n -1 ε -1 > 1/RN 0 . Therefore ∀C > 1/RN 0 ⇒ P (| µ n (p) -µ ⋆ | > ε) ≤ 1 < Cn -1 ε -1 , which means ∀n ≥ 1, ∀ε > 0, ∃C > 0 that P (| µ n (p) -µ ⋆ | > ε) ≤ Cn -1 ε -1 .That concludes the proof.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 7for 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,plim n→∞ Q n (s h , a k ) = Q(s h , a k ).(ii) For any node s h at the depth h th in the tree,plim n→∞ V n (s h ) = V (s h ).Proof. We will prove this by induction on the depth D of the tree. If the tree only has depth (1) .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 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 )is the mean of the intermediate reward at state s 0 , after taking action a k . We recall the definition of Q(s 0 , a k ), with π 0 is the rollout policy to estimate the newly added node at the leaf,Q(s 0 , a k ) = R(s 0 , a k ) + γ s1∈As 0 P(s 1 |s 0 , a k ) V (s 1 )where V (s 1 ) is the value of the policy π 0 at state s 1 , A s0 is the set of feasible actions at state s 0 , |A s0 | = 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 ((1)), we have Q n (s 0 , a k ) = 1 n n t=1 r t (s 0 , a k ) + γ s1∼τ (s0,a k ) T s1 s0,a k (n) n V T s 1 s 0 ,a k (n) (s 1 ) (i) is a direct result ofplim n→∞ Q n (s 0 , a k ) = Q(s 0 , a k ), with a k ∈ A s0Because by definition:V (s 0 ) = max a k ∈As 0 Q(s 0 , a k ) V n (s 0 ) =   a∈As 0 T s0,a (n) n Q Ts 0 ,a (n) (s 0 , a) p   1 pfor some p ∈ [1, +∞)Then we haveplim n→∞ V n (s 0 ) = V (s 0 )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 H ) = V 0 (s H ) and, for all h ≤ H -1, Q(s h , a) = R(s h , a) + γ s h+1 ∈As P(s h+1 |s h , a) V (s h+1 ) V (s h ) = max a Q(s h , a)By the assumption of the induction the root node of a subtree with depth (D) at state s 1 we have plim n→∞ V n (s 1 ) = V (s 1 ) (i) Let'sQ n (s 0 , a k ) = Q(s 0 , a k ), with a k ∈ A s0(ii) follows the results of Theorem 1 as at the root node s 0 of depth D + 1, withV (s 0 ) = max a k ∈As 0 Q(s 0 , a k ) V n (s 0 ) = a∈As T s0,a (n) n Q Ts 0 ,a (n) (s 0 , a) p 1 p for some p ∈ [1, +∞) And because plim n→∞ Q n (s 0 , a k ) = Q(s 0 , a k ), with a k ∈ A s0Then, we haveplim n→∞ V n (s 0 ) = V (s 0 ).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,plim n→∞ Q n (s h , a k ) = Q(s h , a k ).(ii) For any node s h at the depth h th in the tree,plim n→∞ V n (s h ) = V (s h ).Proof. The proof follows the same steps as Theorem 3 by applying the results of Lemma 1 and Theorem 2.Theorem 5. (Convergence of Expected Payoff of CATS) We have at the root node s 0 ,E V n (s 0 ) -V ⋆ (s 0 ) ≤ O(n -1 ).Proof. We prove the result by induction and use the results of Theorem 3 to prove this Theorem. Let us assume that the depth of the tree is D = 1, as the results of Lemma 7, we haveE[ V n (s 0 )] -V ⋆ (s 0 ) ≤ |δ ⋆,n | + O( log n n ) = |δ ⋆,n | + O(n -1 ).And because the tree only have the depth D = 1, we have |δ ⋆,n | = 0, so that the result holds at the depth D = 1. Let us assume that we have the result of the tree at the depth D. Now when the depth of the tree is D + 1, at the root node s 0 , the conditions of Assumption 1 hold as the results of Theorem 3 then we haveE[ V n (s 0 )] -V ⋆ (s 0 ) (Lemma 7) ≤ |δ ⋆,n | + O( log n n ) = |δ ⋆,n | + O(n -1 ),where the bias|δ ⋆,n | = E[ Q n (s 0 , a ⋆ )] -Q ⋆ (s 0 , a ⋆ ) (contraction) ≤ γ ∥ E[ V (1) n ] -V ⋆ ∥ ∞ (by induction) ≤ γO(n -1 ).Therefore,E[ V n (s 0 )] -V ⋆ (s 0 ) ≤ O(n -1 ),that concludes the proof.Next, we present the results of Theorem 6. The proof follows the same steps as Theorem 5.Theorem 6. (Convergence of Expected Payoff of PATS) We have at the root node s 0 , • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).E V n (s 0 ) -V ⋆ (s 0 ) ≤ O(n -1 ).• 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.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). 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). 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).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.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. 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. 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.Since reward is scaled in [0, R], Qmin is not updated in our setup. 
 
 Appendix:  


Decision: I really appreciate the authors for conducting the rebuttal and the following discussions. There are still a few weaknesses after the rebuttal. For instance, reviewer 1ykv is unsure why PATS would not face approximation errors. Most of reviewers have concern about the experiments considered in the paper and some real-world stochastic experiments would amplify the contribution of the current study. 

I personally spotted an issue which might be critical for the analysis. The paper claims it achieved the better non-asymptotic rate of $O(1/n)$, this is counter-intuitive since the standard minimax lower bound in statistical learning requires $\Omega(1\sqrt{n})$ complexity. According to the proof, Line 611-612 is the key for achieving $O(1/n)$, but I am not sure why the equation under line 612 is zero. This seems to be a critical point that should be mentioned clearly. In general, I don't think it is correct. Because of that, I recommend the paper be rejected. The paper should absorb the suggestions and make better presentations in the next version.

Reviews: review: paper_summary: paper_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})$.
strength_weakness: strength_weakness: 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: questions: 1- In the numerical experiments shown in Figure 3, Why does PATS perform much better than CATS?
limitations: limitations: 1- The added high computational complexity from maintaining a distribution for each node in the MCTS.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: Despite 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:

1. Empirical Evaluation
Experiments are limited to a toy domain and results on the Atari benchmark reported in the appendix.
The toy domain is a tabular environment that is being generated randomly and contains stochasticity in both the final reward and transitions.
For an illustrating example this makes it hard to judge the combined effect on the overall return distribution to be approximated. 

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

Also CATS appears to be doing consistently worse than some of the other methods despite having the same theoretical properties as PATS.

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

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

2. Content division
The 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.

3. Overall presentation
There 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.
questions: questions: -
limitations: limitations: The 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.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: 1. **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. 

2. **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.
questions: questions: 1. **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?

2. **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?

3. **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?

4. **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?

5. **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.
limitations: limitations: ### Computational Demands
The 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.

### Fixed Precision
The 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.

### Number of Atoms
The 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.

### Addressing Limitations
While 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.

### Societal Impact
The 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.

### Suggestions for Improvement
1. **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.
   
2. **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.

Overall, 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.
paper_summary: paper_summary: The 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.
strength_weakness: strength_weakness: - Lack of referencing of existing work and novelty relative to existing work
    - 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.
- 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
- Key results in appendix and lack of empirical results in the main paper
- CATS never outperforms fixed depth MCTS on the synthetic tree task
- Unable to find code despite checklist saying it is provided

- Small issues
    - Figure 2 algorithms alignment off
    - Indication of Atari results in section 5 which are not there
    - Adding bold to best performing method in the Atari table would be useful for readability
questions: questions: - How does your approach differ to existing methods (Hayes et al)?
- Why do the baselines investigated change between the synthetic experiments and the atari experiments?
- Why do you think CATS is consistently outperformed by other baselines on the synthetic tree task?
limitations: limitations: - 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
    - “faces challenges in managing computational demand” : this does not say anything meaningful
    - “Our approach’s performance is slightly influenced by hyperparameters”, this can be said for essentially any method
- It seems that the distributional approach has an additional memory cost which if correct should be added to the limitations
- 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.

Comments: comment: Thank you for your response. In your own words, one of the main motivations behind your work is "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." However, your stochastic environments are toy problems. 
On the other hand, it seems that the Atari benchmark you study is deterministic but I wanted to double-check that. The Atari benchmark from OpenAI framework Gym can be made stochastic if sticky actions are allowed "Instead of always simulating the action passed to the environment, there is a small probability that the previously executed action is used instead" and on top of that there is the option of stochastic frame skipping -- "In each environment step, the action is repeated for a random number of frames". Although you do not specify these parameters of the Atari environments in your Experimental setup (Appendix G), I have checked the code you provided and noticed that you are using the "NoFrameskip-v4" versions of the games which suppress frame skipping and "v4" has repeat_action_probability= 0. This means that the Atari benchmark is deterministic indeed.

Based on the above, I recommend the authors do further testing on larger stochastic environments to properly validate their algorithms' performance and claims. I am keeping my score.
I thank the authors for taking the time to answer the questions raised by myself and the other reviewers.

With regards to the comparison between PATS and CATS, I am unsure why PATS would not face approximation errors. Surely, PATS has to manage a similar trade-off between computational effort and precision as CATS governed by the number of particles. Additionally, there is the precision cut-off which is required to collect more than a single rollout per particle which represents a discretization (and thereby approximation).

Overall, I still believe that the paper requires major revision and would benefit from another round of reviews.
Consequently, I will maintain my original score.
Thank you for your detailed rebuttal. I appreciate the efforts to address the concerns raised in my review. While your responses provide some clarification, I believe there are still some important points to consider:

1. **Empirical Validation**: Your commitment to include more diverse benchmarks is welcome. However, the current lack of diverse empirical validation remains a significant limitation of the paper.

2. **Algorithm Complexity and Overhead**: The explanation about CATS' increased computation with more atoms is helpful. I look forward to seeing a more detailed analysis in the camera-ready version, particularly regarding PATS and categorical reward distributions.

3. **Generalization and Practical Performance**: Your plan to include results from more complex environments and provide specific examples of practical improvements is promising. These additions will be crucial for demonstrating the real-world applicability of your methods.

4. **Computational Overhead and Hyperparameter Sensitivity**: The commitment to include detailed comparisons and a thorough sensitivity analysis is appreciated. These will be important for understanding the practical implications of implementing your algorithms.

I encourage you to make these improvements and include more discussions on the related works on the regret analysis for distributional rl, as they could significantly strengthen your contribution to the field.
Regarding "An explanation of why different baselines were used in different experiments will be included". As part of the rebuttal I would expect you to provide this explanation and therefore the response is disappointing, since as a reviewer I would like to know this information as I am concerned as to the reason behind this. I also do not agree with the dismissal of Hayes et al. (2023) as very different purely since they consider the multi-objective domain. I have no further questions and will raise my score to 3 to reflect some alleviated concerns.