Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Authors: Tuan Quang Dam
Pascal Stenger
Lukas Schneider
Joni Pajarinen
Carlo D'Eramo
Odalric-Ambrym Maillard
Venue: ICML 2025 spotlightposter (2025)
Keywords: Monte-Carlo Tree Search
Planning under Uncertainty
Primary area: reinforcement_learning->planning
TL;DR: We develop a new planning algorithm that handles uncertainty by tracking probability distributions instead of single values, achieving 20-80% better performance in unpredictable environments.

Abstract: This paper introduces a novel backup strategy for Monte-Carlo Tree Search (MCTS) tailored for highly stochastic and partially observable Markov decision processes. We adopt a probabilistic approach, modeling both value and action-value nodes as Gaussian distributions, to introduce a novel backup operator that computes value nodes as the Wasserstein barycenter of their action-value children nodes; thus, propagating the uncertainty of the estimate across the tree to the root node. We study our novel backup operator when using a novel combination of $L^1$-Wasserstein barycenter with $\alpha$-divergence, by drawing a crucial connection to the generalized mean backup operator. We complement our probabilistic backup operator with two sampling strategies, based on optimistic selection and Thompson sampling, obtaining our Wasserstein MCTS algorithm. We provide theoretical guarantees of asymptotic convergence of $\mathcal{O}(n^{-1/2})$, with $n$ as the number of visited trajectories, to the optimal policy and an empirical evaluation on several stochastic and partially observable environments, where our approach outperforms well-known related baselines.

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Tuan Dam
Pascal Stenger
Lukas Schneider Odalric-Ambrym Maillard
Joni Pajarinen Carlo D'Eramo" 2.56
Abstract
Reinforcement Learning (RL) problems, particularly after its widespread success in deterministic games like Go and Chess (Silver et al., 2016a 2017b). However, moving be- yond these deterministic settings toward highly stochastic or partially observable Markov Decision Processes (MDP- s/POMDPs) reveals major difficulties. In these cases, two key obstacles arise: Uncertainty in Value Estimates: In problems with substantial randomness or limited observ- ability, naive value backups may lead to erroneous or un- stable estimates, which propagate through the tree and de- grade overall performance. applooration-Exploitttti ancing: Traditional UCT-based exploration bonuses (Koc-) sis et al.) 2006) can falter under high variance transitions, often causing either over- or under-exploration. Recent works (Tesauro al.) 2012 Bai et al. 2013 2014) have suggested Bayesian distributional methods for MCTS to better quantify uncertainty. Meanwhile, Metelli et al. (2019) leveraged - Wasserstnnn barycenters to propagate distributional information in temporal-diffrrent ing. Yet, several open questions remain on how to unify distribution- backups and flexible exploration strate- gies within a single MCTS framework that provably han- dles high stochasticity and partial observability. Our Approach. In this paper, we propose new MCTS al- gorithm, Wasserstein MCTS, that models each node's 's value as a Gaussian distribution and propagates both mean and variance estimates throughout the tree. Crucially, we intro- duce novel backup operator that computes value nodes as I .--Sasserstein barycenters of their action-value children, using an a-divergence the distance measure. This yiclds:
This paper introduces a novel backup strat- egy for Monte-Carlo Tree Search (MCTS) tai- lored for highly stochastic and partially observ- able Markov decision processes. We adopt a probabilistic approach, modeling both value and action-value nodes as Gaussian distributions, to introduce novel backup operator that com- putes value nodes as the Wasserstein barycenter of their action-vale value children nodes; thus, prop- agating the uncertainty of the estimate across the tree to the root node. We study our novel backup operator when using novel combination of Wasssesstt barycenter with a-divergence, by drawing a crucial connection to the gener- alized mean backup operator. We complement our probabilistic backup operator with two pling strategies, based on optimistic selection and Thompson sampling, obtaining our Wasserstein MCTS algorithm. We provide theoretical guar- antees of asymptotic convergence of (mn-1/2), with n as the number visited trajectories, to the optimal policy and an empirical evaluation on several stochastic and partially observable envi- ronments, where our approach outperforms well- known related baselines.
Bal-
learn-
1. Introduction
Monte-Carke Tree Search (MCTS) has become a crucial al- gorithmic paradigm for tackling challenging planning and
Distributional Value Backups: By tracking distributions (rather than point estimates), our method captures the in- herent uncertainty of each node's value, especially valu- able in stochastic or partially observable domains.
'Hanoi University of Science and Technology, Hanoi, Viet- nam *Department of Computer Science, Technical Univer- sity of Darmstadt, Germany 'ETHZ. ETH Zurich, Switzer- land 'Department of Electrical Engineering and Automa- tion, Aalto University, Finland 'Center for Artificial Intelli- gence and Data Science. University of Würzzurg, Germany Hesssia Center for Artificial Intelligence (Hlessian.ai,, Ger- many 'niiy Lille, Inria, CNRS, Centrale Lille. UMR 9189- CRIStAL, F-59000 Lille, France. Correspondence to: Tuan Dam <tuandq
Generalized Mean Operator: The a-divergence ties natu- rally to the power-mean backup (Dam al.) [2019] 2024a), letting interpolate between average-like and max-like updates to mitigate the overestimation often seen in RL (Hasselt 2010).
We complement these distributional backups with two ex- ploration mechanisms an optimistic UCT bonus, and a Thompson sampling approach that selects actions by sam-
Proceedings of the 1nnd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
pling from the node's Gaussian posterior.
P SSxx A - S is the transition kernel, and ? E [0. 1) is the discount factor. A policy л IT 8 A defines the action selection probabilities based on states. The action-value function QF is given by Q" (s. - E [= r*;++++1 8; = 8, - -.a==a..] representing the expected cumulative discounted reward for executing action a in state 8 and fol- lowing policy т. The objective is to find the op- timal policy that maximizes Q satisfying the Bellman equation (Bellman) 1954): (*(s,a) [g P(s"ls, [R(s, a, (,, + Q"(s',''))) (d'',a')ds'' Vs
Our Key Contributions. 1. Uncertainty Propagation via L'-Wasserstein Barycenters. We provide a principled way to back up distributions in an MCTS, unifying L'- Wasserstein geometry and a-divergences to handle high variance and partial observability, 2. Connection to Gener- alized Mean Backup. Our backup operator yields a power- mean update for node values, enabling controllable con- tinuum between overly optimistic (max-like) and risk- averse (average-like) estimates. 3. Polynomial Convergence Analysis. We prove that Wasserstein MCTS with Thomp- son sampling converges to the optimal policy at a rate Onn-1/3) 1/2, matching known lower bounds. This is in con- trast to prior distributional MCTS methods that lacked ex- plicit convergence guarantees. Extensive Empirical Val- idation. On a suite of highly stochastic MDPs (e.g. River- Swim, Taxi) and partially observable tasks (Pocman, Rock- sample), our approach outperforms established baselines, including UCT, Power-UCT, and Bayesian MCTS vari- ants.
(averag--ikk)eestiatte.33
S,a E A. From the optimal action-value func- tion, we derive the optimal value function as (*()) - maxae Q"(ss Va E.
3.2. Monte-Carlo Tree Search
Monte-Carlo Tree Search (MCTS) combines Monte-Carlo sampling, tree search, and exploration strategies from multi-armed bandits (Auer et al.] (002b) to solve MDPs. It builds search tree where states are nodes and actions are edges. MCTS involves four key steps: Selection: Nav- igate from the root to leaf node using tree-policy, Expansion: Expand the reached node based on the tree policy. Simulation: Perform a rollout (Monte-Care Carlo simu- lation) from the child node to estimate its value, or use a pretrained neural network (Silver et al. 2016a) for this es- timation. Backup: Update the action-values Q(-) along the visited trajectory using the collected rewards.
Overall, Wasserstein MCTS offers a flexible and theo- retically grounded framework for handling uncertainty within MCTS. By combining Gaussian node models, L,-- Wasserstein barycenters, and a-divergences, it effectively balances exploration and exploitation in domains where noise or partial observability make traditional MCTS meth- ods brittle.
2. Related Work
4. Formalization
Metelli al. (2019) use I.2- -asssrsesntnin barycenters to propagate uncertainty in cemporal-different learning. In
Problem Setup Monte Carlo Tree Search (MCTS) is an algorithm for exploring and evaluating trajectories in an MDP. Starting from an initial state 8g, MCTS incremen- tally builds planning tree by simulating trajectories. Each trajectory either reaches a leaf node or terminates when predetermined maximum depth H is reached. At the end of each trajectory, playout policy (which may be determin- istic or stochastic) is executed from the final node reached, allowing the algorithm to evaluate the associated state. Af- ter running for t trajectories. the MCTS algorithm provides the following outputs:
MCTS, Bayesian methods handle uncertainty by treating values as Gaussian distributions (Tesauro et al., 2012) or Dirichlet-NormaGaamma posteriors (Bai et al.) 2013
2014). Unlike these, we propagate uncertainty throughout the tree via L Wassersttein barycenters and a-divergences. linking to generalized-meet -mani backups (Dam et al.] 2019) and maintaining both mean and variance estimates. This distributional perspective is effective in highly stochas- tic or partially observable tasks. In multi-armed bandits, optimism (Auer et al., 2002a) and Thompson sampling [Thompson] 1933) are standard; we combine these with our uncertainty propagation scheme to guide action selection in MCTS.
ae: estimate of the optimal action to take in state Sp, Ve(so): estimate of the optimal value function at so-
Evaluating MCTS Performance The performance of the MCTS algorithm is assessed based its convergence rate, r(t), which quantifies how quickly the algorithm ap- proaches the optimal policy. Specifically, the following bounds hold:
3. Background
3.1. Markov Decision Process
We consider an agent in an infinite-horizon discounted Markov decision process (MDP) M = AS R, P,7), where S is the state space, A is the finite action space, R : 8 SxxXx + R is the reward function,
E (sV) (so. U.)] < r()),
| [V*(s)) - V v.())]] <
2

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
where (s0) and (*0,,, a) are the true optimal value and action-value functions at state so, respectively.
X' whose q-th moment is finite. For two distributions M, v E (2)),, the I T.--asserstein distance is
Recursive Value Estimation To analyze the MCTS algo- rithm, we consider planning horizon H and a playout pol- icy To with an associated value function Vo. For each node 9), at depth (i.e., the state reached after / steps from sg), we recursively define the value function V V(Sn) as follows. At the leaf nodes (h- = H). the value function is simply the playout policy's value:
W)(u,v) = -(inf inf (xx.)) (PET())))
where I'u,, v) is the set of joint couplings whose marginals match and v. Given n distributions {u: and weights {w:} summing to I, the L Wassersttnn barycenter is
F = arg min Σ w. (o
V(SH)= Vo(HH).
For all other depths h < H - 1, we compute the action- value function Q88h; a) and value function V(*)) ) as:
Our work focuses on =
5.2. a-divergence and the L Wasserstein Barycenter
(5) r(sh: +
Σ P(sh+1 Sh. () aht1EA,
In many distribution-base backup schemes, the Wasser- stein distance is natural choice to quantify how "far apart" two distributions are. A commonly used approach (Metelli et al. 2019) employ the L2=Wasserstein metric. In con- trast, we consider L -assersteii formulation coupled with an a-divergence for two main reasons:
((%)) = max Q(s,,a),
where r(s), a) is the mean immediate reward obtained by taking action a in state 8n: P(sh+1 8%; a) is the probabil- ity of transitioning to state Sh+1 from Sh given action a and y is the discount factor.
Robusiness & Aggregation Control. An L -b-se metric can be more robust to outliers and large deviations than T.2. Furthermore, combining it with the a-divergence allows a continuous interpolation between averaging and max-like backups (through the a parameter).
Bounding the Error The recursive structure of the value estimates gives rise to a bound on the error between the true optimal action-value function (So. a) and the MCTS es- timate (ss. a). Specifically, we have:
( a) @(su. a) < 7 # #IV* - Volloo:
Connection to Power-Mean Updates. Modeling nodes as Gaussians (or particle distributions) and relying L Wasserstein with an a-divergence yields closed-form up- dates that coincide with the power-mean operator. This uni- fies average and maximum backups in a single formula and lets us propagate both means and variances (uncertainty) through the tree.
where the supremum norm V Volloo can be restricted to states reachable within H steps from 8g-
Goal of MCTS The ultimate aim of the MCTS algorithm is to minimize the convergence rate r(f) by constructing accurate estimates of Q(5a, a) and (so), which in turn approach the true optimal functions Q* (so, a) and (*oo,,
f-divergences and the a-divergence. An f-divergence (Csisáár 1964 between two points X and Y over a Mani- fold M defined as
and then identify the best action at the root node:
= arg max Q"(ss,a))
(XII) ES ( fa fa(x)=
f((()) xa -1-1-a(a-
a(a-1)
5. Wasserstein Barycenter With a-Divergence
where varying a controls how aggressively conserva- tively we measure the "distance" between X and Y. Constructing the L -assersteein Barycenter. In our ap- proach, the L Wassssssnnnn distance between and vs is defined via
We introduce the key notions behind our distribution-basee backups: the Wasserstein barycenter and the a-divergence. Unlike prior works that use 1.2-based Wasserstein dis- tances (Metelli et al., 2019), we adopt an La-Wasserstein distance combined with a-divergences. This combination yields more robust value backups in stochastic and partially observable settings.
W1 (M( = inf PET(MM)
Exx,r--(PPsxxx,1)] (1)
The L Wasserstein barycenter then solves
5.1. Wasserstein Baryeenter
n
Let (X', d) be a Polish (complete, separable metric) space. For q 2 1, define (t) as the set of probability measures
-argini inf
wc WW(())))
3

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
i... we seek the single distribution i that jointly minimizes its I ---aasserstein distance (defined via the a-divergence) to all the Vi.
divergence.
Direct Theoretical Connections. Under mild assumptions, 1 I--Wasserstein geometry paired with a-divergences ad- mits closed-form near-closed- las (Dam et al., 2019). This not only streamlines theoret- ical analysis but also simplifies implementation by allow- ing straightforward computation of mean and variance up- dates at each node.
Why L' instead of L Using the L' distance in equation naturally leads to a backup rule resembling the power mean operator Proposition [ This power-mean update is more robust to high-variance samples and connects smoothly to both the average backup (when a + or p 1) and the max backup (as + oo p + oc). Hence, L'- Wasserstein with a-divergences offers principled way to blend distributions in highly stochastic environments while controlling the balance between underestimation and over- estimation in the final backup.
|
power-mean formu-
In practice, these properties make a-divergences well- suited for uncertainty propagation within MCTS: despite not being a metrie, their adaptability and connection to gen- eralized means allow them to effectively handle complex, high-variance environments.
Why Use an a-Divergence Instead of L Although a- divergences are not strict metries (they can be asymmet- ric and need not satisfy the triangle inequality), their use within an 7, Ta-sassersteinn framework provides distinct ben- efits for MCTS under stochastic or partially observable conditions:
5.3. V-posterior
It is natural to define value node the V-posterior computed with Wasseestnnn barycenters of the chil- dren nodes Q-posteriors, following a procedure inspired by Metelli et al. 2019 (Metelli et al. 2019) and tailored to MCTS.
Greater Flexibility via Generalized Means. When com- bined with the L'-wasserstein distance, an a-divergence naturally yields power-mean style backup operator (Dam et al.) 2019). By adjusting the parameter a, one smoothly interpolates between average-like and max-like backups, allowing precise control over how conservative or aggressive the updates should be. This stands in con- trast to 1.2-based distances, which only yield fixed (e.g. purely quadratic) aggregation behavior.
Definition (V-posterior). Given a policy v and a state E S, we define the V-posterior V(s) induced by Q- posteriors Q(s, a) with a A as the L -assereeenn burycenter of the Q(s, a):
V(s) E arginf En-a(ss) [W1(), R(s.a)))
In this work, we model each node in the tree as a Gaussian distribution. We define 1 - a and derive the following. Proposition 1. Consider the V-posterior value function ((s) as a Gaussian: N ((m (f (s)). Define each ( a) as the action-value function child node of V(s). Each ((s, a) is assumed as a Gaussian distributions ((s, a) Nmm(s, o(s as the Wasserstein barycenter of the action-value function Q(s, a), given the policy n, we have
Robustness to Stochastic Variations. Because a- divergences can emphasize or de-emphasize portions of the distribution differently depending on a, they help mitigate overestimation or underestimation in highly stochastic settings. Empirical studies in distributional RL (Metelli et al. (2019) suggest that more adaptive divergence measures can significantly improve stability and performance when the underlying dynamics involve heavy noise.
If the value function
V(s) is defined
m(s)) = (Ba~* [mo,, a)
No Need Symmetry in Backups. MCTS requires cost functional to aggregate posterior distributions across chil- dren nodes, rather than a strict metric. Hence, the lack of symmetry or the triangle inequality does not under- a-
8(s) = (Eawn [6(s, ))
Proposition 1 [shows the closed form solutions of the mean and standard deviation of the Gaussian value function (>)) considering it as the L!--Waseerstein barycenter Q- posteriors. In detail, the mean of (ss) are the power mean of all mean values of all the Q(s, a) function, considering the finite set of actions. When p = 1, we derive the ex- pected form solutions.
mine its validity here. An aedivergence-includee divergences is sufficient to drive consistent updates of belief distributions in the tree.
Unified Framework for Various Divergences. The divergence family subsumes and generalizes many stan- dard divergences (c.g. KL, reverse KL). This single- parameter approach enables users to easily switch or fine- tune the update behavior for different problem character- istics, rather than designing separate algorithms for each
We point out that our approach is not restricted to the Gaus- sian distribution model. We get the following result by con- sidering each tree node as a particle model.
4

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Proposition 2. Consider the V-posterior value function V(s) as an equally weighted Particle model: sı(8)= E [1, M). M is an integer and M 2 1. Assume each aation--vluee function Q(s, a) has M particles zz a)
where n(ss,) a) is the visit count of action state 8, and N(s) = E. n(s, a). For p = 1, this reduces to the standard average, whereas > induces more "max-like" backup Dam et al. 2019).
[1, M. If the value function V(s) is defined as the Wasser- stein barycenter of the action-value function ((s,a) given the policy n, each particle o(( E [1 can be esti- mated as
Under the Bellman-style backup for each Q- node,
Q-nodes.
(m(s,a) = (]](s,a)]++B[Va.. Qua
(s, a) = y Vsta(s"),
[T(s) = (aaw (x(s,))))""
we replace expectations by empirical sums and visitation counts:
Proposition shows that each particle of the V-posterior value function V(s) can be derived as the power mean of the respective particles of all the ((s, a) function. p 1, we again get the closed-form solutions as the expectation of the respective particles of all the Q(s, a) functions. The results in Proposition [ and Proposition 2 can be consid- ered as the generalized result of Proposition A.3 in Metelli ett al.] (2019). In the next section, we present Wasser- stein Monte-Carlo tree search (W-MCTS) ) algorithm, as- suming each tree node is a Gaussian distribution.
t E +x
m(s,a)
B,als,,,n(s,,0
Th N(6
a(x,a)
Here, the sums range over transitions and children states s', weighted by their visit counts N(s'). As n(s, a) grows large, both the variance and mean estimators stabilize, eventually converging to deterministic values.
6.2. Action Selection
6. Wasserstein Monte-Carlo Tree Search
Monte Carlo Tree Search can adopt variety of exploration strategies based on the original UCT framework (Kocsis) et al. 2006). In practice, multiple refinements exist, such as the variants used in AlphaGo (Silver et al. 2016b).
We introduce our Wasserstein Monte-Carlo Tree Search (W-MCTS), where V-posteriors are modeled as Wasserstein barycenters of action-value distributions. With Gaussian distributions at each node, we define backup operators for mean and variance. Additionally, we propose two action se- lection strategies: optimistic selection and Thompson sam- pling.
AlphaZero (Silver et al., 2017c a). MuZero (Schrittwieser et al. 2020). Stochastic MuZero (Antonoglou al. 2021). )), and Boochastic-Power-Uue different choices of the exploration constant or bonus lead to different performance characteristics, we retain the stan- dard, state-of-the-ar designs described below. In our theo- retical analysis, however, we focus specifically on Thomp- son sampling, since the UCT-like optimistic selection can be viewed as a special case of the well-studied Pow
(Dam et 20044). Although
6.1. Backup Operator
We model each V-node and Q-node as a Gaussian with mean and standard deviation:
algorithm (Dam al. 2019 2024b).
(s), (s) and (s, (s).
We denote (s, Ns)) as the empirical mean estimate of the V-node at state after N(s) total visits, and Qi (s, a, n(s,, a)) as the empirical mean estimate of the Q- node at (s, a) after nt(s, a) visits. Likewise, V (Na(s,NN))) and Quta(s, a, n(ss, a)) are their corresponding empirical standard deviation estimates.
Optimistic Selection. A classic UCT-style selection picks actions using upper confidence bounds on Q-values.
argmnax mmss,a)) as) + C
log (e) m(3,a;)
where m(s, aq) is the empirical mean, n(ss, 0;) is the visit count of action o, and N(s) is the total visit count state 8. Replacing the Va(a.a.) term by the empirical standard deviation a as, a;) yields an optimistic variant of Wasser- stein MCTS (N
V-nodes. From Proposition I the mean and the standard deviation of a V-node is a power-mean aggregation of its Q-children:
a(s.a) [() [Qm (s, (l,,, n(s,)))))) 1/p {
V ((ss) +
Σ
argmax [m(ssaa) + Ca(s,as) Vlog N(*)
V "snd(sN(s))
+
Σ
m(s.a N(s)
Qsta(s, a, n(8, ((s,))))
The factor a(,, a:) ~ 1/ 1/vn(s,a;) follows from a CLT- based argument.
5

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Thompson Sampling. In contrast, Thompson sampling stochastically samples action from the Qpposterior:
Proof Sketch
MCTS a Hierarchical Bandit Structure. The Monte Carlo Tree Search (MCTS) algorithm can be viewed as hierarchy of multi-armed bandits (MABs), where each node in the search tree represents an independent handit problem. In this framework, the reward for each node, or current bandit, is influenced by the performance of the bandit algorithms applied to its child nodes. Since the W-MCTS policy adapts dynamically to balance exploitation and exploration, the rewards at each node are inherently non-stationary. The proof of Theorem | unfolds through three essential steps:
argmna {0, ~ (M(((,))))) o² a.
We refer to this Thompson variant as Wasserstein MCTS- TS (W-MCTS-IS). In Section 7 we analyze its conver- gence properties under non-stationary multi-armed bandits and then leverage these results to establish convergence in the planning tree.
7. Theoretical Analysis
7.1. Analysis Setup
1. Analyzing Non-stationary Bandits The initial step focuses the analysis of non-stationary multi-armed bandit, which reflects the behavior of MABs at each MCTS node. We establish that if the rewards of these non- stationary bandits meet specific concentration properties, the regret induced by the W-MCTS algorithm will exhibit corresponding concentration guarantees. This outcome formally stated in Theorem
We define the setting for our theoretical analysis using a class of non-stationary Multi-Armed Bandit (MAB) prob- lems at each state 8 in the MCTS tree. Consider K arms (actions), each with mean reward μıx. for & E [K]. At time step 4, pulling arm k yields a random reward Xg.t> bounded within [0, R]. The average reward for arm after n trials is:
2 XR.ts with
X
Pk.n = E[X km
2. Induction Argument Next, we utilize an inductive ar- gument to transfer the convergence and concentration prop- erties from the lower tree levels to the root node. As the rewards from one level inform those of the next, the find- ings from Step can be recursively applied. We begin depth H - and move upward, demonstrating inductively that the bandit rewards at each level H of the MCTS satisfy the criteria required by Theorem 2 This process propagates the desired properties up to the root node, completing the induction.
72
Let represent quantities related to the optimal arm, and denote Ti& (n) as the number of times arm le has been played by step n. We assume the following concentration condi- tion holds:
Assumption 1. We assume that the reward sequence, (XRa= :l > 1}, is a non-stationary process satisfying the assumption: for all
> 3c > 0 that
Pr (| k. - pxl > e) E cn te (2)
3. Error Analysis from the Oracle The final step exam- ines the error introduced by the leaf node estimator, rep- resented by the value function oracle Vo. With this oracle, the depth-/f MCTS can be interpreted as performing I steps of value iteration, starting from Vo at the leaf nodes (as mentioned in (Dam al., 2024b)). Importantly, the or- acle's error decreases geometrically contraction mapping property of value iteration, leading to diminishing еттоr as we ascend from the leaf nodes to the root. The complete proof for Proposition 3 can be found in the supplemental material. Finally, we get the main result.
7.2. Main Results
We show the polynomial convergence of the expected esti- mated mean value function at the root node in Theorem
7.2.1. CONVERGENCE OF W-MCTS
rate of y due the
We start with an important result as shown below Proposition 3. Applying W-MCTS to an MCTS tree of depth (H), at any depth h of the tree, we have
(i) At any depth h, 3 constant Co > that for any < O, 2 1, we can derive
Theorem 1. We have at the root node 80.
L M (s0. m)] [ss)) < OO(n 1/2).
Pr V Vm n) (ish, at) 2E) & 1on-
Our proposed method, W-MCTS, achieves polynomial convergence rate of On221/), matching the results of Dam et al.] (2024b). In contrast, Xiao et al. (2019) intro- duced MENTS, followed by RENTS and TENTS from Dam al. (002)). which leverage exponential convergence to a
(ii) At any depth h. 3 constant Co > that for any < 8 < O, > 1, we can derive
Pr (mh. ax)) 20)= 1og-2
6

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
regularized value function through maximum entropy reg- ularization. However, these methods face bias due errors in the regularized value function, potentially leading to in- correct action selection. Conversely, Painter et al. (2024) employ a similar action selection strategy with a maximum backup operator for value estimation, resulting in exponen- tial reductions in simple regret. However, their method's effectiveness heavily relies on the temperature parameter in Boltzmann exploration, limiting its practical use.
agent aims to reach a goal in the bottom-right corner. Due to frequent slips, each move has high uncertainty. Fig- ure I shows that W-MCTS- -S (Thompson sampling) out- performs DNG, UCT, Power-UCT, and W-MCTS (opti- mistic selection), with W-MCTS at = 1 performing com- parably to W W-MCTS-TTT
NChain. An agent can move forward or backward along chain of length 5. Actions may reverse with 20% proba- bility, making consistent forward progress difficult. In Fig- ure | both W-MCTS-TT and W-MCTS-B8 OS exceed UCT and Power-UCT in convergence speed and final returns. RiverSwim. Similar to NChain but more complex tran- sitions: sometimes the agent remains in the same state or only partially moves. This rewards long-term planning to reach high-value states. As in Figure I W-MCTS-OO converges fastest and attains the best performance, while Power-UCT eventually reaches similar returns more slowly.
7.2.2. WASSERSTEIN NON- STATIONARY MULTI-ARMED BANDIT
A crucial part of the proof for Theorem is to derive the following result for the W-MCTS in bandit setting. Under the Assumption T. we consider applying Thompson Sam- pling strategy as the action selection method for the non- stationary multi-armed bandit (MAB) problems describes above. At each time step n. an action is selected
SixArms. A 7-state chain with 6 possible arms (actions) leading to different rewards that scale inversely with their success probabilitics. This environment demands high ex- ploration. Figure shows that W-MCTS is the only method consistently securing strong returns.
a = argmax {0,- ~ NN( k, Ve/Ta(n)))- (1..
(3)
Let's define (p) -Ex [a(m) n (n)
1/2
Taxi. A 7 x 6 grid where the agent must pick up three passengers, then reach goal region. Slips occur 10% of the time, adding further uncertainty. Only W-MCTSSTS manages to collect all passengers reliably, outperforming Power-UCT and W WCMCTS with optimistic selection.
the power mean value backup at the root node, Ta (n) =
E (a, a) is the number of selections of prior to round n. We show theoretical results of our method as follows. Under the Assumption 1 we establish the con- centration properties of the power mean backup operator (p) towards the mean value of the optimal arm μux = maxa } a E [K], as shown in Theorem 2
8.2. Partially Observable, Highly Stochastic Tasks
Theorem 2. Consider a non- non-stationary bandit problem de- scribed as in 7.1 with action selection as Equation Then,
We also test W-MCTS against POMCP (UCT), D2NG, and DESPOT in classic POMDP benchmarks: rocksample, poc- man, Tag, and LaserTag. Code for POMCP (UCT) (Silver
& Venees 2010b), D2NG (Bai al.. 2014), and DESPOT (Somani et al. 2013) is used released by the original au- thors.
(r( - ma| > E) Cn --2
Theorem 2 states the concentration properties of the power mean estimation by W-MCTS for non-stationary continuous-arman bandit problem, and play an important role for the induction proof of Proposition 3 leading to the main result presented at Theorem
Rocksample. A robot on an n grid can sample or ignore k rocks, then exit. We test three variants: (11,11), (15,15), and (15,35). Figure that W-MCTS- TS consistently outperforms both UCT and D2NG.
Poeman. partially observed maze (Silver & Veness 2010a) where the agent must collect pellets while avoiding ghosts. Table I
indicates that R7-MCTS-PR with p = 100 outperforms UCT and D2NG across most rollout- budget set- tings, and W-MCTS-OS OS also matches or surpasses these baselines in some configurations.
8. Experiments
8.1. Fully Observable, Highly Stochastic Tasks
We compare W-MCTS to UCT Kocsis et al., 2006).
Comparison with DESPOT. We additionally compare W-MCTS to DESPOT across Tag. LaserTag. rocksample (15 x 15), and Pocman. Table 2 shows that W-MCTS-OS and W-MCTS-TS achieve higher returns than AB-DESPOT and AR-DESPOT in rocksample. Similarly, W-MCTS-TS surpasses DESPOT in Pocman, Tag, and LaserTag, while W-MCTS-OS outperforms AB-DESPOT in Pocman. Role
Power-UCT (Dam et al., 2019), and DNG (Bai et al. 2013) in five benchmark environments: FrozenLake, NChain, RiverSwim, SixArms, and Taxi. These tasks all feature significant stochasticity or long-horizon exploration chal- lenges.
FrozenLake. A 4 x grid with slippery transitions, im- plemented in OpenAl Gym (Brockman et al. 2016). The
7

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
FrozenLake
NChain
RiverSwim
Discoureed Reeium
Disourtad Recum
0.2
06
14t
SixArms
Taxi
Discrurted Returm
Disccurted
6.00
Simullations
Simulations
DNG
Powe-UUT
UCT
W-MCTS-D5
W-MCTS-TS
.... W-METS--S,
p-1
Figure 1: Performance W-MCTS vs. DNG, Power-OCT, and UCT on five MDPs. Each curve shows the mean discounted return (averaged over 50 runs), with shaded regions indicating standard error
rocksample 11x11 (16 actions)
rocksample 15x15 (20 actions)
rocksample 15x35 (4D actions)
Discounted Return
Discounted Return
Discounted
Simulations
Simulations
Simulations
W-MCTS-TS
Figure 2: Performance of W-MCTS D2NG in rocksample, averaged over 1000 runs (except UCT, 100 runs). Shaded areas denote standard error.
of o- a-Divergence. We explored several values of a to vary how aggressively our backups shift between average-like and max-like behavior. When approaches or oc, the update becomes nearly a pure average (p 1) or nearly max backup, respectively. In practice, we found that moderate values often provide suitable balance be- tween these extremes, and we report results with the best- performing choices. Although a more extensive sensitivity analysis could be conducted, the core takeaway that com- bining power-mean backups with variance propagation sig- nificantly enhances performance highly stochastic tasks.
Table 1: Discounted total reward in pocman. Mean stan- dard error are computed from 1000 random seeds.
HOTE 06
ME3220000 0.6
02.2 = 0.29
54.8
1.8
N-M078-T8
67.38 0.63
75.64 0.51
77.68
0.777
77.70
1.22
000
23.4
1.39
28.8 1.00
21.00
3.30
25.5 3.8
Table 2: Average total discounted reward. The results for POMCP, and DESPOT are taken from (Somani et al. 2013).
Tog
Lawer Tay
AST15 15)
Porman
R-801S-C8
0.53
19,76
R-MOTS-13 73
-5.90 0.66
-8.76 0.5
20.20
-2 0.22
315.46 2.16
8.3. Key Performance Factors
FONCP
0..11
0.06 -19.58
12.23
0.32
291.16: 1.06
DESPOT
6.57 0.26
11.133
18.18
290.34
4.12
The superior performance our method stems from two complementary components that address fundamental lim-
AR-DESPOT
6.28 - 0.24
9.34 4 0.28
18.57
0.30
307.06
8

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
itations in existing MCTS approaches for stochastic and partially observable environments:
Munos 2010) for even broader applicability.
Explicit Variance Propagation. Unlike previous methods that only propagate point estimates or use fixed variance models, our approach dynamically updates both means and variances at each node through the L*-Wasserstein barycenter formulation. This capability is particularly cru- cial in highly stochastic and partially observable environ- ments where uncertainty quantification directly impacts de- cision quality. Our experimental results demonstrate con- sistent improvements over Bayesian MCTS methods: we achieve up to 80% improvement over DNG in Frozen- Lake, and significant gains over POMCP across all POMDP environments, with particularly notable improvements of 55.31% in LaserTag and 65.90% in 15cksample(15,15) Additionally, observe improvements of up to 21.38% over AB-DESPOT in LaserTag. highlighting the effective- ness of our distributional approach.
Impact Statement
Our proposed Wasserstein MCTS algorithm offers princi- pled way to tackle complex, stochastic tasks in both fully and partially observable domains. Potential applications in- clude robotics, autonomous systems, and large-scale re- source management, all of which require adaptive plan- ning strategies to handle real-world variability. While we do not anticipate immediate negative societal implications, responsible deployment remains essential. As with any AI- driven technology, understanding ethical, economic, and security ramifications- as autonomy in safety-critical gystems-shoulee guide practical use.
Acknowledgments
This research is funded by Hanoi University of Science and Technology (HUST) under Project No. T2024-TD-024, the French Ministry of Higher Education and Research, the Hautsde-France region, Inria, the MEL, the French National Research Ageney under PEPR IA FOUNDRY project (ANR-23-PEIA-00003)
Our approach's ability to interpolate between average-like and max-like backups through the a-divergence parameter al- lows adaptive behavior across varying levels of stochastic- ity. In highly stochastic environments such FrozenLake and NChain, we found that moderate a values (leading to more average-like updates with p closer to 1) performed optimally by preventing overestimation bias. Conversely, in environments with more deterministic regions of the state space, larger a values (yielding more max-like behavior) proved beneficial for faster convergence to optimal poli- cies. This flexibility, combined with Thompson sam- pling strategy, enables our algorithm to automatically adapt its iaploration-exploittiei variance observed at each node.
Flexibility in Balancing
References
Antonoglou, I., Schrittwieser, J., Ozair, S., Hubert, T. K., and Silver, D. Planning in stochastic environments with learned model. In International Conference on Learning Representations, 2021.
Auer, P., Cesa- Desa-Bianchh, N. and Fischer, P. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47(2):235-2566 2002a.
balance based on the empirical
The synergy between these two components--pricciplee uncertainty propagation and adaptive backup operators- explains why W-MCTS consistently outperforms both clas- sical MCTS variants and existing Bayesian approaches across our diverse set of benchmark environments.
Auer, P., Cesa-Bianchi, N., and Fischer, P. Finite-time analysis of the multiarmed bandit problem. Mach. Learn., 47(2-3):235-256, may 2002b. ISSN 0885-61 25. doi: 10. 1023/A:10136897043322 URL http:///doo.or/1000022/// 101 1013689704352
9. Conclusion
Bai, A., Wu, F., and Chen, X. Bayesian mixture modelling and inference based thompson sampling in monte-carlo tree search. Advances in neural information processing systems, 26, 2013.
We proposed Wasserstein MCTS, an algorithm that mod- els node values as Gaussian distributions and employs unify
L Ja-sassersteiin barycenters with a-divergences average- and max-like backups. Coupled with Thomp- son sampling or optimistic selection, our method achieves strong empirical performance while offering O(n convergence guarantees. Experiments in both stochastic MDPs and POMDPs show significant improvements over classic baselines and Bayesian MCTS variants. Future work includes extending these Wasserstein-based ideas to open-loop planning (Leurent& & Maillard 2020; Bubeck &
Bai, A., Wu, F., Zhang, Z., and Chen, X. Thompson sam- pling based monte-carlo planning in pomdps. the Interna- tional Conference on Automated Planning and Schedul- ing, 24(1), 2014.
Bellman, R. The theory of dynamic programming. Tech- nical report, Rand corp santa monica ca, 1954.

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J., and Zaremba, W. Openai gym. arXiv preprint arXiv:16000 01 01540, 2016.
Metelli, A. M., Likmeta, A., and Restelli, M. Propagat- ing uncertainty in reinforcement learning via wasserstein barycenters. Advances in Neural Information Processing Systems, 32, 2019.
Bubeck, and Munos, R. Open loop optimistic planning. In COLT 2010-The 23rd Conference on Learning Theory, 2010.
Painter, M., Baioumy, M., Hawes, N., and Lacerda, B. Monte carlo tree search with boltzmann exploration. Ad- vaRces in Neural Information Processing Systems, 36, 2024.
Cichocki, A. and Amari, S.-i. Families of alpha- beta- and gamma- divergences: Flexible and robust measures similarities. Entropy,
2010. ISSN 1099-4300. doi: 10.3390/e 12611533 URL https://www. mdpi.com//0099-45000/02000
Perlman, M. D. Jensen's inequality for convex vector- valued function an infinite-dimensionall space. Jour- nal of Multivariate Analysis, 4(1):52-65, 1974. ISSN 0047-259X. doi: http:://oo.ooo///0000000000/2030
Csiszár, I. Eine informationsthoorrtisshh ungleichung und ihre anwendung auf beweis der ergodizitaet von markoff- schen ketten. Magyer Tud. Akad. Mat. Kutato Int. Koezl., 8:85-108, 1964.
90005-0. URL httpss//www...dddneeedder......s article/ #tttec//pi/004755xx77900000
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E., Has- sabis, D., Graepel, T., et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588 (7839):604-009, 2020.
Dam, T., Klink, P., D'Eramo, C., Peters, J., and Pajari- nen, J. Generalized mean estimation in monte-carlo tree search. arXiv preprint 19e:::1111.0,,4 2019.
Dam, T. D'Eramo, C., Peters, J., and Pajarinen, J. A unified perspective value backup and exploration in monte-carlo tree search. Journal of Artificial Intelligence Research, 11:511-577,8 2024a.
Silver, D. and Veness, J. Monte-carlo planning in large pomdps. In Advances in neural information processing systems, 2010a.
Dam, T., Maillard, O.-A., and Kaufmann, Power mean estimation in stochastic monte-carlo tree_search. The 40th Conference on Uncertainty in Artificial Intelligence (UAI), 2024b.
Silver, D. and Veness, J. Monte-carlo planning in large pomdps. In Lafferty, J., Williams, C., Shawe-Taylor, J., Zemel, R., and Culotta, A. (eds.), Advances in Neural Information Processing Sys- tems, volume 23. Curran Associates, Inc., 2010b. URL http:://ppoeeeddnggsssuusssspppp edibe aac12244bb0040004440000000000000
Dam, T. Q., D' Eramo, C., Peters,, J., and Pajarinen, J. Con- vex regularization in monte-carlo tree search. In Interna- tional Conference on Machine Learning, pp. 2365-2375. PMLR, 2021.
Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., van den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershellamm, V., Lanetot, M., Dieleman, S., Grewe, D., Nham, J., Kalchbrenner, N., Sutskever, I., Lillicrap, T., Leach, M., Kavukcuoglu, K., Graepel, T., and Hass- abis, D. Mastering the game of Go with deep neural net- works and tree search. Nature, 529(7587):444-88, Jan- uary 2016a. doi: 10. 1031//0ttrreeeeeeeeeeeeeeeeeee 16961.
Gerchinovitz, S., Ménard, P., and Stoltz, G. Fano's in- equality for random variables. Statist. Sci, 2020.
Hasselt, H. V. Double q-learning. In Advances in Neural Information Processing Systems, 2010.
Jin, T., Xu, P., Xiao, X., and Anandkumar, A. Finite-time regret of thompson sampling algorithms for exponential family multi-armed bandits. Advances in Neural Infor- mation Processing Systems, 35:38475-386 -34477, 2022.
Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484, 2016b.
Kocsis, L., Szepesvári, C., and Willemson, J. Improved monte-carlo search. Univ. Tartu, Estonia, Tech. Rep, I, 2006.
Leurent, F. and Maillard, O.-A. Practical open-loop op- timistic planning. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECML PKDD 2019, Würzburg, Germany, September 16-20, 2019, Proceedings, Part III, pp. 69-85, Springer, 2020.
Silver, D., Hubert, T., Schrittwieser, J., Antonoglou, I., Lai, M., Guez, A., Lanctot, M., Sifre, L., Kumaran, D., Graepel, T., et al. Mastering chess and shogi by self-play with a general reinforcement learning algorithm. arXiv preprint arXiv: 2017a.
10

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang, A., Guez, A., Hubert, T., Baker, I.,, Lai, M., Bolton, A., Chen, Y. Lillierap, T., Hui, F,, Sifree L., van den Driessche, G., Graepel, T., and Hassabis, D. Mas- tering the game of go without human knowledge. Na- ture, 550:354-, October 2017b. URL http.//dx.doo..or/ 10.1038/nature222700
Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M., Bolton, A., et al. Mastering the game of go without human knowledge. Nature, 550(7676):354-35,, 2017c
Soch, J. The book of statistical proofs. URL
http::/ssutpnoofoooo.giucce io, 2020.
qf.html
Somani, A., Ye, N., Hsu, D., and Lee, W. S. Despot: Online pomdp planning with regularization. Advances in neural information processing systems, 26, 2013.
Tesauro, G., Rajan, V., and Segal, R. Bayesian in- ference in monte-carlo tree search. arXiv preprint arXiv:/203.3511 2012.
Thompson, W. R. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-22 1933.
Xiao, C., Huang, R., Mci, J., Schuurmans, D., and Müller, M. Maximum entropy monte-carlo planning. Ad- vances in Neural Information Processing Systems, pp. 9516-9524, 2019.

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Outline
Notations will be described in Section A.
Hyperparameters are provided in Section B.
Derivation of Wasserstein barycenter with Gaussian and particle filter distributions will be described in Section C.
Supporting Lemmas will be provided in Section D.
Full proof for the convergence of Wasserstein Non-stationary multi-armed bandit will be provided in Section E.
Full proof for the convergence of Wasserstein Monte-Carlo tree search will be provided in Section F.
A. Notations
Table 3: List of all notations of Wasserstein barycenter with Gaussian and particle filter distributions.
Notation
Type
Description
N(m, 82)
R
Gaussian distribution with mean m, standard deviation 6
(X
complete separable metric (Polish) space
Wy(uv)
(P-Wasserstein distance between 14,
Wi(u,v)
L' Wassersstii distance between μ4,
(1)
quantile function of distribution p(z)
r(u,v)
X' x y
set of measures on X' x y with marginals μt, V
d(X,Y)
distance between X and
Dy, (XI(Y)
a-divergence distance between X and Y the inverse of the function
erf (t)
v So exef (oof--2]ec dd
2
B. Experimental setup and Parameters selection
All the experiments were done on an Intel(R) Core(TM) 19-14900K 3. 3.20 GHz 24 cores/CPU.
To compare the performance of W-MCTS to other state-of- ofrthe-rrt w planning algorithms, we run several experiments on stan- dard MDP as well as POMDP environments. For comparison, we consider UCT (Koesis et al., 2006), Power-UCT (Dam) al. 2019), DNG (Bai al. 2013) and D2NG (Bai et al. 2014). The hyperparameters are tuned using grid-search. Except for the case of Poeman environment, we scale the rewards into the range [0. We use the discount factor 7 0.95. For DNG, D2NG, set hyperparameters recommended in the paper and from the author's source code (Bai et al. 2013 2014). We set exploration constant for UCT, Power-UCT to 22. We set initial standard deviation value to std = 30. In all Rocksample and Pocman environments, we set the heuristic for rollouts as treeknowledde = O, rorllouukknowweeeee = 1. For all environments, we increase the value of p and choose the best power mean p value for Power-Uee UCT, and W-MCTS. Details can be found in Table 6 For POMDP environments such as Rocksample, Pocman we get the source code released
from the author of DNG (aai et 2013) and D2NG (Bai al. 2014
12

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Table 4: List of all notations of Wasserstein Non-stationary multi-armed bandits.
Notation
Type
Description
K
N
number of arms/actions
Pk
R
mean value of arm At
R
optimal mean value
A
R
A p p.
A
= [
X
R
average reward of the optimal arm after s visitations
F2
CDF of Gaussian with mean X
(n)
N
number of visitations of arm A at timesteps power mean backup operator with power p
(p(
R
R
average rewards of arm k after (n) visits
C. Derivation of Wasserstein barycenter with Gaussian and particle filter distributions
We revisit the definition of Wasserstein distance: The Wa-Sassennniin distance (with q > 0) between two distributions μt, with the cost function d(æ, y) x y + R is defined as
1/9
W Wg(,.v) = inf
E (E) (
[d(x,
(4)
here [(u, u) the set of measures on X' with marginals 14, (t) as the quantile function of distribution
Define F p(x))
p(x) (t) (t) inf x R.t Fy(2)).
(5)
With d(x,d( X Y| as the Euclidean distance, we can derive
(w, = (1) With d(x,)) D,,(( as the a-divergence distance (defined in section 4.1), we can derive
(6)
W2(, = ( D1, (0)))
C.I. L'- Wassserssnninn barycenter with a-divergence distance
We have
W((pp)= inf E [d(X,Y)] inf pETUNe)XXYY (ET(.))X,
DD (X, Y)).
13

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Table 5: List of all notations of Wasserstein Monte-Carlo Tree Search.
Notation
Type
Description
KL
KL divergence
(sh)
R
optimal mean of V value at root state Sh, at depth (h)
Qm (m)
R
mean of Q value function at state action de, at depth (h)
V ())
empirical estimated mean of value at state 8 after n visitations at depth (h)
Qm (Sh, nx n)
empirical estimated mean of Q value at root at state 8h, action dg after n visitations at depth (h)
(8h)
optimal mean of V value at depth (h) state 8h
(n)))aak)
R
mean of Q value function at depth (h) state sy, action ag
V (m(sn,n)
R
empirical estimated mean of V value depth (h) state s after visitations
Qm (Sh, n)
R
empirical estimated mean of Q value depth (h) at state sh, action a after visitations number of plays of action at state at timestep
Ten. Tan (n)
N
Ts' (n)
N
number of plays of action dg at state s to state at timestep n
5,ax
Table 6: List of all hyperparameters.
Environments
p Value Search
Best Value
FrozenLake
1,2, 4, 10, 100
W-MCTS-03 (p=100 P00
NChain
p= 1, 4,8, 15, 100
M-MCTS-OS N-CSSSSS (p= 100)
RiverSwim
p= 1,2,4, 4,8, 15, 100
W-MCTS 100
R-MCTS-OS MCTS TS P
SixArms
1,2,4,8, 15, 100
100
Taxi
p 1,2,4,8,
100
W-MOTS-OS -P htwp://000111111__0ooooooooC1110
Rocksamplet Ix11)
10, 50, 80, 100, 150
W-MCTS-OS 007 (p=100)
Rocksample( 15x15)
10, 50, 80, 100, 150
W-MCTS-OS (p= 100) MC
(p-10
15:33))
p 10, 80,
(p= 150) N-MCTS-OS
(p=10)))
Poeman
1, 2,8, 10, 100
W-MCTS-05 (p= Mw0.MCMS---- MM (p=10)
We find the lower bound of W (p,,v) with a-divergence as a measure cost function. Let denote N(m, 52) as a Gaussian distribution with mean m and standard deviation With = A(m,, o7). N (m2, 52) We first want to show that by applying Data Processing Inequalities (Lemma 2.1 (Gerchinovitz et al. 2020
14

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
with [(x= = mi, and g(X) X m2, we derive
W)(m, inf E pErlam))x 1D (X,
inf PET p
E DD (X m2, Y mm))
-W,(V(,, (f) m1, 62)),
and Wi( inf DET(mN)x,E<P
E DDs (X, inf
E peT(p.nxxX
DD m2, Y m2))
inf pEI(HN))X 1D (m2 X, m2 {) with the transform function f(X) -X)
E
WW( - m1, NO, 62)).
(10)
Now according to the L." -Wassersteii distance with a-diveegencc distance is defined as
(i(m, (L (0,1
(11)
We show that the quantile function Gaussian distribution (ochh 2020 F (2,
(1) 28errf (21 1) + 14,
(12)
where erf Therefore, the L' Wasserstein distance with a o-divergence distance as the cost function between two Gaussian distributions = [W (f), v = ((m,, 83) can be measured as
erf=(()) (l) is the inverse of the function 3 S0 expl-x²)ddr
(i(,, v) (L. D1220eer (21 - 1) + mi mwlll/22eeee
222eeff (21 - 11) 1
Applying the convexity properties a-divergence (Cichocki & Amari 2010), and from I0 we have
Wi(H,D) > 1 3/2 (L. Dra Dia (V26eet 1(t 1) 1(t - 1) + m2 mı dt
+ ! Dja (V28 (21
(21 -1)]dt)
(f) Dj, (v2 (2t
m2
m1
25 1 1(t-1) 1))+ m2 m1
+
m2 mi
m2
m
WW(N(
2
2
Applying Data Processing Inequalities (Lemma 2.1 Gerchinovitz et al. 2020)), with I(X) X ma-my
we can derive
Wi (i(,, v) > 67) 83)) = ( Dj. (225 25 (21 1)1l 2222222z2 (211-122)))))) 1))dt
Let us consider the sequences tu < t & ... tN there exists & E [t,, t++1 that
i=N
Wi(u.v)
Σ
(t++ 1, (v2djerf ' (25, 1)llv28gerf '(2;: 1))
i=N
Σ
A,D 2228 - 1)ll '(5 1)),
=
with A, = (t++1 - t.). Since D,, (cPlIcQ) = D5. (PIIQ) where e is a constant. We can derive
im N
Wi(Av)) >
Σ (6110))) = D,, (81
(13)
We start with the first Proposition about the closed solutions of mean and variance of Gaussian value function V(s) as V-posterior T,1 Wasserstnn barycenter of all action value function distributions Q(s, a).
15

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Proposition 1. Consider the V-posterior value function ((s) as a Gaussian: ((m(ss) d² (s)). Let's define each Q a) as
the function child node of V()). Each ((s, a) assumed as Gaussian distributions ((s, a) [[m(s, a), 8(s, a a)2) If the value function (s) is defined the Wasserstein barycenter of the funetion (s a) given the policy n, we will have:
( (Ba- [m a)?))
(14)
(() = (Ba (6(8 ()))
(15)
with -1 - O.
Proof. By the definition of the V-posterior value function, we have:
(m(ss, 8())) arg min (V.
(16)
We first compute the standard deviation 8(s).
From 13). and (16), want find (s) that is the minimizer of
(()) arg min [B=]) 8(s)
derive 8(5) is the solution of
VE Dr (8(s)))0(s,a)) 0.
(17)
Vo(s)
Since
(z(())
a(ra-1-1 -1
x a-1
(18)
V
a(a - Q-1
With Dj, (xllv) = Ey (.) we can have
VD fa
(x (v)
(g)a-1-1
(19)
x
a-1
y
We can derive
B(a) ja-1 1
3(s) ja-1 ((ssaa)
Ean
E
((..a)
Exnt
( (8(3,0))-11
(20)
(a (a-1) 1)
Now define p that leads
8(s) = (Ea= (88, a) P))
(21)
To compute pis). Let's revisit here again the definition of Wasserstein distance between two Gaussian distributions μ(m;, 67), v(m2,68) 62).
Wi v) = in((E/Ds
(22)
According to Jensen's 's inequalityPPerlman 1974) we can derive
E[D;, (u][b)) > D, [[]]]]][]] = D (m l|m)).
(23)
Therefore, according to the definition of Wasserstein barycenter, the mean of a Gaussian V-posterior value function V(s) can be derived as
(T(s) arg min Eg E [D,, (m(s)]|m(s,o)) m(s)
(24)
Following the same steps as to compute 5(s), we can get
m((s) = (Eaw [m((s, aa a)
(25)
with = a that concludes the proof.
16

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Next, we consider each node equally weighted Particle model and derive the following proposition.
Proposition 2. Let's assume the V-posterior value function V(s) as a equally weighted Particle model: xF(88) i [L, J. M is an integer and M > 1. Let's assume each Q function ((s, a) has M particles æ4(s, a), E [1, M. If the value function V(s) is defined as the Wasserstein barycenter of the function (ss, a) given the policy nt each particle (8)()) E [1, M)) can be estimated as
(7(s) (Bawn (1:
(26)
with p
Proof. We compute the quantile function of y and
(1)
Σ 3q 11 (1)
Σ yb (1).
(27)
Therefore from
we can get
w)(k, - (f( Dja D2,00,"000))000000
(28)
M
(
(f Dr. (FT (F (F
(29)
M
(L, Dra (zillvs)dt
Σ
(30)
(Le dt
Dsa (a:
(31)
M
w. W, (zillys).
(32)
We can see that for each particle (Zr()), [1, MD, we can derive
(7) = arg min Ea (1)())(())))
(33)
=((s)
(3) = (Euwa as)
(34)
with p = 1 - a.
D. Supporting Lemmas
We will make use of the following basic results.
Lemma 1. (Minkowski's inequality) Given p 1, (x,, vi) E R, = 1,2, n, then we have the following inequality
1/
1. + | [((e)))))
E(k, un)P)
(35)
Proof. This is basic result.
Lemma 2. (Markov's 's inequality) If X is nonnegative random variable and > then the probability that X is least is at most the expectation of X divided by a:
Pr(X > a) < E
(36)
17

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
E. Convergence of Wasserstein Non-stationary multi-armed bandits
We note that an MCTS tree, each node is considered non-stationary multi-armed bandit where average mean drifts due the given action selection strategy. Therefore, we first study the convergence of Wasserstein -on-stationary multi- armed bandits where the action selection Thompson sampling, with the power mean backup operator the root node. Detailed descriptions of the Wasserstein Non-stationary multi-armed bandits settings can be found in the main article in the Theoretical Analysis section.
We briefly summarize the theoretical results below. Lemma is about the upper bound on the expectation of the number of suboptimal arms playing, following the corresponding Theorem 4.2 in (Jin et al. 2022). Lemma i about the bias the expected value of the power mean backup operator, which follows the result as Theorem
Bepchastic-Power-UUss al. 2024b). Theorem deals with the polynomial concentration of the power mean backup operator around the optimal mean the root node of the non-stationary Wasserstein problem for multi-armed bandits. This theorem plays important role in deriving the polynomial convergence of the choice of the optimal action the root node in the Wasserstein MCTS tree, described in the next section.
(Dam)
Now, we will find an upper bound for the expectation of numbers of pulling a suboptimal arm. Let us define the event E = (t) & -==== for all E [K].e > 0, 0 (t) is sampled from (X (/Tx(n)) at timestep t. Let us consider the decomposition
[[Tx(n)) =1 =1+
Σ
[ ok; (1))
Σ
I(A, 0k, (1){
(37)
(=K+1
K+1
=1+B E
1(4, on Eg (t))
(t)
+E
Σ
1f4
ax. E (t,
(38)
(=K+1
I=K+1
8
Here E is the complement of an event E, > 8 V/n is an arbitrary constant.
Bounding Term Let's define
KL(H"-E- -E-z, < Alog( th )/5}
a₂ = sup xE(0,n -E)
(39)
Lemma 3. (Lemma A. (Jin al.
2022 ) Let M [16V log(ne? Loggne2/VV/e" and ar be the same defined in 39 then
M
Σ t=
1(A = ab. En (t)}
F7 (n- - =) - 11), e - -
1
ae. ac. 11}] + (
(40)
where F2 is the CDF of Gaussian with mean X", 5. X is the average reward of the optimal arm after s visitations.
Lemma 4. (Lemma A.2 (Jin el al. 2022 ) Let M- = [16V log(ne" /V)/=2V Then
(Vlog(ne//V)
Σ E x (1
1[x; (μ¹ a. 113 θ F2 -c)
(41)
62
Bounding Term B:
Lemma 5, (Lemma C.I (Jin et al. 2022
Let N = min
M 25 2}. lng(ne
For any &k, >> O that satisfies e Ph < Ai.
then
2V
V
Nlog(ne2 /V) KL(U + Pke Pk. n² =))
Σ
=
+ +
+
(42)
2
22
t= K+1
18

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
From Assumption 1, derive the upper bound for the expectation of the number of plays of suboptimal arm.
Lemma 6. Consider Thompson Sampling strategy (using power mean estimator) applied where the pay-off sequence satisfies Assumption 1. Fix e 2 O. Let Tu(n) denote the number of plays of arm k. Then if lo is the index of a suboptimal arm, then each sub-optimal arm k is played in expectation at most
stan-ttat problem
[TTx())] & 0
0 1+
(log(nA22))
(43)
4
Proof. The proof Lemma 6 closely follows Theorem 4.2((Jin et al. 2022 by observing results from Lemma 3
From equation 38 putting all Lemma
we have
V Vlog(ne2/V) (An - ph))
Vlog(ne"1) (V))))e2/)))))))))))))))) M
( e e
+
+
(44)
Pk
2
g²
Set E = Pk- Ax/4, we derive
[[]((N)] & <
V Vlog(nAi/V))
(45)
A
Lemma 7. Consider a non-stationary problem where the pay-off sequence satisfies Assumption We consider bandit algorithm that selects each arm as
a argmax ~ NN a..i (1...)
k. V/Tx(n))).
Let us define the power mean estimator (p) as xn(p) = (EK= Ta(n) X a.Ta(()) )² and dan Pa,n For any p E 1, O, we have
E[X pl| < 10 +
R
Σ a=1,aça,
K
@
V "llo((n4 Vlog(nAg/V)
(46)
A
n
Proof. We observe that
xm(P) -Hal Xn(P) pa,al + 1m pr,ll Xx Mx,nl + +10mm
(47)
Furthermore,
Xa.T,(a) & Ha, + X Haml
(48)
Since Mr,n maXxce[K] {M we have
K Σ a=1
Ta(n) (Xanan)))
K Ta(n) (Ham,) )2
-x) -een - -x,()) Σ Talmiaan &
(49)
a=l
K Ta(n) (X (N..T
DK (n( (Ma, )p)
(50)
n
Applying Minkowski's inequality from Lemma 1 and the result of equation
48
we have
ΣK Ta(n) (Ma + X XaTaon - Ma,n 1)
EK Ta(n) (Han)))
xx( Mn.n x
(51)
n
ΣK Ta(n) ((x Ma,n 1)²)
(52)
nr
19

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
On the other hand,
Mr, (p(p)= nm n
nfpnn - (E Ta (n) + Ea= Ta(n)Man mX
(53)
n
n
Ex= Lapa (n) luen Ma,al + Ea=1 Ta (n)) nX
(54)
n
K
(n(n)
K
Ta(n)
R
Σ
+
Ma,n - xx
(55)
o=1
Because power mean an increasing function of p, so that
K
Ta(n)
Ta(m) (Ha,,)) Σ
1/8
Ma,n
n
Furthermore, we observe that
Ha,n X (a)To()) + (a,TT(n) Mam
So that, from equation 55 we have
1/8
Ha,n x
&
Σ
K
Ta(n)
+
K
Ta(n)
(Ham))) - -x((p)
(56)
a 1 Layo.
a=1
K
Ta(n)
<
R
(57)
a=
ΣK Ta(n) (X (a.To()) + (a.T.()) - )))=
(IK Ta(n) (X (aTT.()))))
(58)
n2
ΣK
(a(n) ( Man
(Minkovski's incquality)
Ta(n)
R
Σ
(59)
e=1,aşa.
(Properties norm)
K
Ta(n)
( Ta(n) (|Xa,I,(m) Maml))
R
Σ
+
(60)
To((m)
Ta(n)
Xo,t Ta (n))a
R
(61)
Therefore
[[>* ΣT Ta(n))
K
E[Ta(n)]
E
Xa,t
Ta (n))aam
x MM,n]] < R
(62)
n
11
K
([Ta(n))
R
(63)
Please note that because study non-stationary bandits, EE 7 Xa.1] -MHon therefore,
[|Em] Σ (a(n) Xa,t - (a( )] 0
E
n
20

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
According to Lemma have
K
(lT,(n)) < 16
K Σ
0(1
V
Vlog(nAg/V)
E x (p) - pn & 16,ml1 +
R a=1aa
(64)
n
71
4
which concludes the proof.
Theorem 2. For E [K], let (X (Xa,))23 be sequence of estimator satisfying Assumption I and let μ4 max
Assume that all the estimators are bounded in [0, R). We consider bandit algorithm that selects each arm as
argmax (...(1
{0e ~ NN( k, V/Te(n)))-
Then, for all p E [1, o.), the sequence of estimators
K Σ x
(n(n)
(x)
(n))(
n
where Ta(n) = En-1 (ae a) the number of selections of prior to round n satisfies
Pr([Xn(p) nal > 8) < On 12-
Proof. We first prove that lim E = According to the result of Lemma 7
we have
E[Ta(n)
B nal m6 R
Σ
(65)
a=1,aşa
R
K
(1 (1+ Eo) logn o(logn) + +011
(66)
KIN) v.)
a=
with 6.,n = 11. - Me,,, and because
lim usnn = Ma, we can concludes that n-000
lim [[X,(p)) = M4.
n-000
Second, we prove that We observe that
Vn > V. 0, 3, 0 that P ([X - nal > e) & -t--2
(p) /4 < Xx, M4,1 + 1 Mr,al = (p) Me,m| + 18
(67)
X u.l > $) < X - p.m > s/2) + P(66.,n > 8/2).
(68)
Because lim 5. = 0, therefore, 3No > such that Vn > No. we have 0o.nl < 8/2 that means n1-p
An No, 1( > =
Next, according to Lemma 6
B.X,, (p)] - Mnnn <
R
Σ
K
Eo) logn (IN)(Fa,H..)
+ o(logn) + +0(1] +ollgn)+ =0((=0)) O(n-1),
(69)
a=La
21

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
that leads
X (p))
O(n -((='))
P( (p) - Ma.10 E/12) <
(70)
E/2
E/2
Therefore, 3c 0, 44 < e < such that
(() - uamal 3 E/2) & on -
(71)
which means
Vn > No. < 0 that (I) 14l 8) & cn t
Now we see that (p) n.| < R. With E max 1, R), we have (p((p) - M. SE xx(P) - Hal > R, therefore the inequality holds as
(() 14| > 2) 0 cn
with 0 < max R). 12 No max (1, R).oo -- 1/ max R).o Therefore
AC 1/ max 1, No P P (p) p41 =) < 1 -n- < On Cn-
which means
Vn > 1, 0A < < 1,30 > 0 that I (I) - n.1 =) & Cn Cn 1
That concludes the proof.
F. Convergence of Wasserstein Monte-Carlo tree search
We start with Lemma 8 which shows the concentration of empirical Q value at any internal node in the tree. This plays an important role in the analysis of our MCTS algorithm.
From the results of Lemma and Theorem 2 we derive Propostion 3 which shows the concentration of any internal V-node and Q-node in the tree. Finally, we get the expected simple bias with convergence rate of (nn11/2) in Theorem
Let us start with Lemma
Lemma 8. (Lemma 11 el al. 2024bi) For m E [M], let (mm, n))n21 be sequence of estimator satisfying
Pr (V n) ( (m)) 8) < Cn 1--
Assume that there exists a constant L > O such that supremum V (m, n) hns Let Ra an iid sequence with mean I and St be an iid sequence from a distribution p (pı...., .PM) supported on (1....; M}. Introducing the random variables No. #lli < nn:5. smil }, we define the sequence of estimator
1 R
M
Nu ((m, N
((n)
n
n
m=
Then there exists some constant d' (which depends on p; (i=1,2....,M) 7, μ)) such that
Pr (mm) - Σ Pm (m) de < Cn Cn-
Based the results of the described nonstationary multi-armed bandit problem, derive theoretical results for W-MCTS.
We derive Proposition which shows the polynomial concentration of the estimated mean of the Q-value function at the root node. In Proposition 3, we also show that the estimated mean of the V-value function at the root node converges polynomially to the optimal mean. Based on Proposition 3 we derive the result in Theorem [ which shows the bias the expected payoff of the power mean backup at the root node.
22

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
At any node of state depth h in the tree, the mean of the value function, and the mean value of the optimal value function are defined
((sn, a) - R(s, a) + (V(++)) ((s)) argmax a)),
(72)
(73)
with h - 1,,,, 1, (sx) is the value return from rollouts at state R(sh: a) is the mean reward received at state 8/. after taking action a. Let us denote og- as the optimal action at the root node.
Proposition 3. When we apply the W-MCTS algorithm to an MCTS tree of depth (H), any depth h of the tree, we have
(i) At any depth h constant Co > thr any 0 < < 1, 2 1, we can derive
Pr m/ n) - V (sh, ax) > < Con-le
(74)
(ii) At any depth h, constant Co that for any < 1, 2 1, we can derive
Pr Pr( ( as, n) (Sn, an) _ < Con
(75)
Proof. We will prove this by induction the depth D of the tree.
Base case (depth // = I):
At depth I, the tree consists of only the root node. The state at the root is denoted by so- At time step t, suppose the agent takes action dg at so. resulting in an intermediate reward re (S0. ax), and transitions to the next state 81-
We assume that the reward R(50, ax) represents the mean reward received at state so after taking action dg-
We recall the definition of (8, ag), defined as
( = R(s,a) + (s))
(76)
where (s1) is the value of the rollout policy at state Aso is the set of feasible actions at state |Asul = M, P(s1|s0, ax) is the probability transition of taking action ag at state sy to state 81. We have
n
T (n) V m (81, Tei (n))
2m (30, ak, n)
re ax) +7
Σ
71
n
(30-())
Equation 774 a direct result of Lemma 8 where X, represents the intermediate reward re(so, ax) at time step The (x)) where P(-|so, ax) is the transition probability dynamic for
probability distribution p (p..P2.....PP. taking action aş in state so-
For each m [M], the sequence (V mm,aa) at time step & corresponds to the deterministic initial value function V. (31), where:
> 2) < Cn- Cn-
Pr
with given action dj.
1,.,, M, and 8m T(-150, Here, T(- 0x) denotes the transition kernel from state Sm,
Equation 75 is the direct results from Theorem 2 In detail, we have from equation 74
Pr (0 (50, ax, n) - (s0, ax)
> >e) < C' 11 with as E Aso
23

---

Monte-Carlo Tree Search with Uncertainty Propagation via Optimal Transport
Because by definition:
(so) max ((50: ax) og ErAop
(77)
$
Tsp.a(n) (2m (50, a, (x)))
Vn (s). =
with E +00)
(78)
n
acd.o
Then we have
(m -((0)) > 2)) Cn Cn-'ee
Pr
that concludes for Equation 75
Let us assume that for tree of depth H 1, the theorem holds for its children.
Now, consider tree with depth H. When an action is taken the root node, where the state is sp, the tree transitions into subtree of depth H. By the induction hypothesis, the results hold for any internal node of the tree after taking the first action.
We have 81 T(S,, ax), where T(s,, au) denotes the transition dynamics. By definition, the value function at the leaf nodes is V(SH) ), and for all H 1, the following holds:
(sh,aa) a) = R(ss, a) +
Epstt1 (sn+)))
P(snt1
Sh+1EA.
(sN) max (sh,
where R(sk,e) a) represents the immediate reward at state 8h after taking action a, 7 is the discount factor, and P(81+1 Sh: a) is the probability of transitioning to state sh-1 from sh by taking action a.
By the assumption of the induction the root node of subtree with depth (H 1) state have
Pr (V.
((s)) > 2) (
CCn-ee
75 Let's apply Lemma g with (X,) is the intermediate reward (re(80, ax)). p (p1.P2;...Px) P(-|so, ax). For [M], each (V m, time step is the empirical Value function (s().). We will have
Pr (n (50, ag, n) ((50, an) > 2)) < Cn -t- with a Aso
74
follows the results of Theorem as at the root node so of depth H, with
((s))= max 2(0) (su: ax))
(79)
Ex
(n)
(Om (So, a, Tao. (m)))p)
V m (So, n)
for some p +00)
(80)
n
AaEAD
And because
Pr (00 ax, n) - Q(00, as) > >)) < Snn-ts- with as E Aso
Then, we have
Pr (s. V(so) >) & Cn -2--
that concludes for 74
The results of Proposition 3 hold for any node in the tree with the tree of depth (H). By induction, we can conclude the proof.
24

---

Monte-Carke Tree Search with Uncertainty Propagation via Optimal Transport
Theorem 1. We have the rool node Su,
V (50. n)] (s0) < O(n
Proof. Using the convexity of ((x) |a]| and applying Jensen's inequality we have
BV m (80: n)] V(so)] < El[V m (s0, n)] [=
V (s - V(su) > s) ds P
+ox
1ds +
nn Con ds
(g-211 to0
E n
+Con
-2+1 n
Co
+ +1)n* -
2-1
Then,
V.V (80. -V(so) < 00nn 1/3.
That concludes the proof.
25


Reviews:
[{"review_id": "3HT2VihjZu", "reviewer": "Reviewer_EvfR", "summary": "This paper introduces Wasserstein Monte-Carlo Tree Search (W-MCTS), a novel approach that represents value nodes as Gaussian distributions (mean and variance), allowing explicit uncertainty propagation throughout the search tree. The method employs a backup operator based on the Wasserstein barycenter and α-divergence to aggregate value estimates from child nodes, enabling a more flexible backup strategy that interpolates between averaging and maximization. W-MCTS incorporates both an optimistic UCT-like action selection strategy and a Thompson Sampling-based approach, achieving an asymptotic convergence rate of $O(n^{−1/2})$ under the latter. Extensive experiments on highly stochastic and partially observable domains demonstrate superior performance over SOTA methods.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the Reviewer for their thoughtful review and insightful questions that help us improve our paper. Below are our responses to your specific concerns:\n\n## Key Performance Factors\n\nThank you for this insightful question. We would like to point out that the superior performance of our method stems from two complementary components:\n\n - Explicit variance propagation: Unlike previous methods that only propagate point estimates or use fixed variance models, our approach dynamically updates both means and variances at each node. This explains the superior performance gain of our method over Bayesian MCTS methods in highly stochastic and partially observable environments. For example, our algorithms demonstrated consistent improvements over POMCP across all environments, with particularly notable gains of 55.31% in LaserTag, 65.90% in RS(15,15), and with improvements of up to 21.38% over AB-DESPOT in LaserTag. In FrozenLake, we obtain an 80% improvement over DNG.\n\n - Flexibility in balancing exploration-exploitation: Our approach's ability to interpolate between average-like and max-like backups (through parameter α) allows it to adapt to varying levels of stochasticity. In highly stochastic environments, we found moderate α values (leading to more average-like updates) performed best, while in more deterministic regions of the state space, larger α values (more max-like) were optimal.\n\nWe will add these insights to Section 8.4 in the revised manuscript\n\n## Scope of Uncertainty Propagation\n\nThank you for highlighting this ambiguity. In our method, uncertainty propagation is indeed more extensive than just parent-child relationships, though it follows the tree structure.\n\nOur phrase “uncertainty is propagated throughout the tree” means that each node not only stores a (mean, variance) pair (or a Gaussian distribution) but also uses these distributions in the backup operator at its parent. That is, any child’s updated distribution is reflected at the parent level.\n\nTo clarify:\n\nW-MCTS propagates uncertainty bi-directionally in the tree:\n\n - \"Bottom-up\" through Q-nodes to V-nodes during backups. \"Top-down\" through action selection, where uncertainty influences exploration\n\n - Sibling nodes indirectly influence each other's visitation rates through their parent's uncertainty estimation, creating a form of lateral uncertainty influence. We do not directly share distributions among sibling or cousin nodes. However, each node’s distribution indirectly influences its parent’s distribution—and transitively influences siblings as the parent’s updated distribution affects how siblings are compared and selected in subsequent iterations.\n\nUnlike previous methods where uncertainty is often reset or approximated at each level, our approach maintains consistent distributional representations throughout the entire search process.\n\nThus, “throughout” means that over multiple rollouts, the uncertainty from leaf nodes consistently flows upward, eventually impacting the root node’s estimates, while updated root estimates modulate exploration down into deeper levels.\n\nWe will add a new paragraph in Section 5 to make this distinction clearer, showing how uncertainty flows throughout multiple levels of the tree simultaneously.\n\n## Experimental Reporting\n\nYou are correct that the x-axis in Figure 2 represents the number of simulations, and we apologize for not stating this explicitly. We will update the figure caption to clearly indicate this.\n\nRegarding fixed-time performance, we agree that a time-based comparison could be informative in practice—particularly for large or real-time systems. We have run standard, iteration-based comparisons, as is common in research prototypes. For the final version, we can add a note clarifying the computational budget used, and that a direct time-limited comparison in the revised version.\n\n## Other Corrections\n\nWe will fix the missing equation reference on Line 155 and explicitly define all terms in Proposition 1 to improve clarity.", "claims_and_evidence": "The empirical results support the claim that W-MCTS outperforms existing methods in highly stochastic and partially observable domains. However, as acknowledged in the paper, previous methods such as Tesauro (2012) and Bai (2013) also model uncertainty in value estimates. It is unclear what specifically enables W-MCTS to achieve superior performance. Is it the improved approximation of uncertainty, a more effective backup operation, or another factor? A clearer discussion could help isolate the key contributing components.\n\nThe paper states that uncertainty is propagated \"throughout\" the tree. Does this mean it extends beyond the parent-child relationships to include sibling and cousin nodes, or is it strictly along the path from leaf to root? Clarifying this distinction would enhance the reader’s understanding of the method's impact on tree-wide exploration and value estimation.", "methods_and_evaluation_criteria": "The selection of benchmark problems is comprehensive and relevant. However, the paper does not explicitly mention the number of simulations (or iterations) for each method in the results section. The x-axis in Figure 2 appears to represent the number of simulations, but this should be clearly stated in the figure description. Additionally, a comparison based on a fixed wall-clock timeout (e.g., 1 second per decision) or by reporting the time taken per search trial would provide a clearer picture of the expected performance.", "theoretical_claims": "The theoretical analysis provides sound justification for the claimed convergence properties. Proposition 1 establishes the mean and variance update rules under the Wasserstein barycenter framework. However, a minor presentation issue remains: the term $\\delta$ is not explicitly defined in Proposition 1. While this does not impact correctness, explicitly stating its definition would improve clarity.", "experimental_designs_or_analyses": "N.A.", "relation_to_broader_scientific_literature": "The paper's contributions are significant for the planning and search community, particularly in AI applications requiring robust decision-making under uncertainty.", "essential_references_not_discussed": "N.A.", "other_strengths_and_weaknesses": "N.A.", "other_comments_or_suggestions": "- Line 155 (right side) - Missing equation reference.", "questions_for_authors": "- Key Performance Factors: What is the primary driver of W-MCTS’s superior performance compared to previous uncertainty-aware MCTS methods? Is it the backup operation, improved uncertainty estimation, or another factor?\n\n- Scope of Uncertainty Propagation: When the paper states that uncertainty is propagated \"throughout\" the tree, does this extend beyond the standard parent-child relationships?\n\n- Experimental Reporting: Does Figure 2’s x-axis represent the number of simulations? If so, could this be explicitly mentioned? Additionally, would a fixed-time performance comparison provide further insights?\n\nOverall, I am leaning toward acceptance, pending clarification of the above concerns.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bBN65kWj2O", "reviewer": "Reviewer_mdxn", "summary": "This paper takes a distributional approach to Monte Carlo tree search. The\nauthors propose a framework for planning in environments with uncertainty and/or\npartial observability. The proposed framework models state and state-action\nvalues as distributions. They also introduce a backup operator that propagates\nuncertainty through nodes in tree by estimating backup values as Wasserstein\nbarycenters. Additionally, they propose two sampling techniques for use with the\ntree policy taking the uncertainty into account.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the Reviewer for your positive review and for recommending acceptance of our paper. We appreciate their recognition of our work's novelty and clarity.\n\n## On Our Claims and Evidence\n\nWe thank the Reviewer for acknowledging that our experimental results support our claims about robustness to stochastic variations. Regarding the comment on \"no need for symmetry in backups,\" we agree that this is more accurately described as an enabling factor rather than an advantage in itself. We will revise this framing in our final version to more precisely articulate the relationship between the α-divergence's properties and our method's effectiveness.\n\n## On Theoretical Claims\n\nWe appreciate the Reviewer's assessment that our mathematical analysis of the convergence rate appears sound, even while acknowledging the specialized nature of this theoretical work.\n\n## On LaTeX Compilation Issue\n\nWe will fix the LaTeX compilation issue on Pg. 3, Col. 2, Line 156 as you noted. We thank the Reviewer for bringing this to our attention.", "claims_and_evidence": "The authors make several claims:\n1. Their approach is robust to stochastic variations.\n- The experimental results mostly back up this claim.\n2. No need for symmetry in backups.\n- I do not think this is an advantage more than it is a reason why they are\n  able to make their approach work.\n3. Their approach has convergence rate of $\\mathcal{O}(n^{-1/2})$\n- See further.", "methods_and_evaluation_criteria": "The methodology is founded on ideas that have worked prior. Update values based\non Wasserstein barycenters are a reasonable method to model uncertainty over the\nvalues of action outcomes. The set of evaluation domains are appropriate for\ndemonstrating the performance of their proposed method.", "theoretical_claims": "The paper has the theoretical claim that the proposed search method has a finite\ntime convergence rate of $\\mathcal{O}(n^{-1/2})$. While I do not have\ntheoretical background to verify the correctness of the mathematics, the math\nlooks sound.", "experimental_designs_or_analyses": "The authors perform experiments on two types of domains:\n    - Long-horizon, highly stochastic domains\n    - Partially observable, stochastic domains\n\nThe results support the performance claims made by the authors. One exception\nbeing the results for river swim does not line up with the analysis provided.\nIn this setting, DNG clearly converges faster than the other algorithms.", "relation_to_broader_scientific_literature": "I believe this work would be of interest to the planning community on the whole.", "essential_references_not_discussed": "I do not think so.", "other_strengths_and_weaknesses": "**Novelty**\n\nThe work is a novel technique for computing update values.\n\n**Clarity**\n\nThe paper was well written and moderately easy to understand.", "other_comments_or_suggestions": "- Pg. 3, Col. 2, Line 156: There seems to be a LaTeX compilation issue.", "questions_for_authors": "No questions.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R0LW6BX6A8", "reviewer": "Reviewer_Nijy", "summary": "This paper introduces Wasserstein Monte-Carlo Tree Search (W-MCTS), a new MCTS variant designed for highly stochastic and partially observable environments. The key innovation lies in propagating uncertainty through the search tree using L1-Wasserstein barycenters combined with alpha-divergences, enabling robust distributional backups (prior work used L2-Wasserstein barycenters in the context of temporal-difference learning). The main idea is to aggregate distributions from child nodes using L1-Wasserstein barycenters, paired with alpha-divergences to interpolate between average-like and max-like backups. This balances exploration-exploitation and mitigates overestimation.The work bridges distributional RL and MCTS, offering a theoretically grounded framework for decision-making under uncertainty. The approaches can have applications in robotics, autonomous systems, and other domains requiring adaptive planning in noisy environments.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the Reviewer for the thorough review of our paper on Wasserstein Monte-Carlo Tree Search (W-MCTS).\n\n## On Key Innovations\n\nWe are pleased that the Reviewer recognizes our main contribution in propagating uncertainty through the search tree using L1-Wasserstein barycenters combined with alpha-divergences. As correctly highlighted, this approach enables robust distributional backups that effectively balance exploration-exploitation while mitigating the overestimation problem that often occurs in reinforcement learning.\n\n## On Theoretical Analysis\n\nWe thank the Reviewer for acknowledging our theoretical analysis regarding the convergence rate of O(n^(-1/2)). This theoretical foundation provides important guarantees about the behavior of our algorithm in the asymptotic case.\n\n## On Empirical Results\n\nWe appreciate the recognition of our comprehensive empirical evaluation across both stochastic MDPs and partially observable tasks. These experiments have been designed to demonstrate the practical advantages of our approach in environments where traditional MCTS methods struggle with high variance or limited observability.\n\n## On Broader Impact\n\nThe Reviewer's comment about the potential applications in robotics, autonomous systems, and other domains requiring adaptive planning aligns with our insight for this work. We believe that the ability to handle uncertainty in a principled way is crucial for deploying decision-making systems in real-world scenarios.", "claims_and_evidence": "The papers emphasized the advantages of using L1-Wasserstein and alpha-divergences vs. L2-Wasserstein in terms of robustness and connection with power-mean updates. Theoretical analysis is done in terms of convergence bound for the value with a polynomial\nconvergence rate of O(1/n^2).", "methods_and_evaluation_criteria": "NA", "theoretical_claims": "Theoretical analysis is done in terms of convergence bound for the value with a polynomial\nconvergence rate of O(1/n^2).", "experimental_designs_or_analyses": "Empirical comparisons are done to show the proposed methods outperform baselines (UCT, Power-UCT, Bayesian MCTS) in stochastic MDPs (e.g., RiverSwim, Taxi) and partially observable tasks (e.g., Pocman, Rocksample).", "relation_to_broader_scientific_literature": "NA", "essential_references_not_discussed": "NA", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "n2rJgQENY1", "reviewer": "Reviewer_6RvP", "summary": "The paper introduces Wasserstein-MCTS, an uncertainty-aware version of MCTS based on Wasserstein barycenter updates for the optimal value estimates and their variance. Due to the power-mean like updates the version is particularly suitable for highly stochastic environments. A theoretical analysis is included, that shows that the method has a polynomial convergence rate. \nSuperior performance compared to standard and Bayesian MCTS variants is demonstrated on common planning benchmarks.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank the Reviewer for their thoughtful and constructive feedback and comments, and appreciate the reviewer's acknowledgment of the strengths of our paper. We have provided detailed responses to address each of their concerns.\n\n## On reference suggestions. \n\nWe thank the Reviewer for these excellent suggestions. We will expand our related work section to include both the works of Hennig et al. 2010, Grosse et al. 2021 and add a brief discussion of their relationship to our approach, noting that while they share the concept of Gaussian distributions over optimal values, our method differs in its focus on stochastic/partially observable MDPs and its use of L1-Wasserstein barycenters with α-divergences.\n\nAdditionally, we will add a paragraph discussing the broader context of uncertainty quantification in reinforcement learning, referencing the suggested papers from O'Donoghue 2023, O'Donoghue et al. 2020 and Tarbouriech et al. 2023.\n\n## Last paragraph on page 3: There is a reference missing (\"??\").\n\nWe thank the Reviewer for catching this error. We will fix the missing reference, which should have pointed to Dam et al. (2019) when discussing the connection to power-mean updates.", "claims_and_evidence": "* The paper claims an asymptotic convergence of the proposed method, which is supported by the theoretical analysis in Section 7, in particular Theorem 1. \n* The paper claims a connection to POWER-UCT, which is given in Section 5.2.\n* The paper claims that their method can handle partial observability and high variance, which is supported by the Experiments in Section 8.", "methods_and_evaluation_criteria": "yes, see below.", "theoretical_claims": "Unfortunately, I could not find the time to check the proofs in the supplementary material.", "experimental_designs_or_analyses": "Yes, I reviewed both experiments in Section 8:\n\n* It makes sense to me to evaluate the method on the benchmarks used as they are a fairly standard choice. Since a claim in the paper is that the method is particularly suitable for stochastic environments, it also makes sense to evaluate them in highly stochastic environments. The method is compared to UCT, which is probably the standard MCTS baseline, as well as a Bayesian MCTS variant which similar in spirit to the proposed method due to also tracking uncertainty. The ablation with POWER-UCT is also interesting: POWER-UCT does not seem to be much better than UCT, providing evidence that the additional propagation of the uncertainty does indeed benefit the search.  \n\nThe experimental evaluation could be strengthened further by evaluating against newer MCTS variants, e.g. one of those with exponential convergence rate, that are cited in Section 5. \n\n* I am not very familiar with the partially observable setting, therefore I do not want to comment on the validity of this experiment.", "relation_to_broader_scientific_literature": "**Relation to other methods:**\n* The idea is related to Bayesian tree search algorithms in the sense that distributions instead of only point estimates are propagated through the tree. The acquisition functions (Thompson sampling and UCB) are standard choices in Bandits/Bayesian Optimization/ Bayesian Tree Search.\n* The method is also related to power mean back-ups in MCTS (Dam et al. 2019), but the proposed method additionally propagates uncertainty estimates in form of standard deviations through the tree. \n* The method also shares similarities with Wasserstein Q Learning (Metelli et al. 2019), but relies on the MCTS framework instead of Q-Learning and replaces the $L^2$ Wasserstein distance with the $L^1$ Wasserstein distance in order to be more robust in stochastic environments\n\n**Relation to other theoretical findings:**\n* The method shares a polynomial convergence rate with POWER-UCT (Dam et al. 2019)\n* There are MCTS variants with exponential convergence rate in the Maximum Entropy or Boltzmann framework (Xiao et al. (2019), Dam\net al. (2021), Painter et al. (2024)), but they suffer from biases (e.g. due to their additional regularization term).", "essential_references_not_discussed": "There is more work regarding Bayesian MCTS and uncertainty propagation (in form of Gaussian distributions over the optimal values) than cited in the paper:\n\n* _Coherent inference on optimal play in game trees_, Hennig et al. 2010\n* _Probabilistic DAG search_, Grosse et al. 2021\n\nThese methods are not designed for stochastic or partially observable MDPs, but since Tesauro et al. 2012 is also cited, I would suggested to add them to the list of cited methods in the introduction. \n\nI am not sure if one wants to keep the related work section restricted to MCTS, but there is more work on uncertainty quantification in context of the general RL setting, e.g. [1-3] . But maybe this is a judgement call.\n\n[1]  _Efficient Exploration via Epistemic-Risk-Seeking Policy Optimization_, O’Donoghue 2023.\n\n[2]  _Making sense of reinforcement learning and probabilistic inference_, O’Donoghue et al. 2020\n\n[3]  _Probabilistic Inference in Reinforcement Learning Done Right_, Tarbouriech et al. 2023", "other_strengths_and_weaknesses": "I appreciated the proof sketch for Theorem 1, as well as the intuition for the choice to replace the $L^2$ Wasserstein distance with $L^1$ Wasserstein distance + $\\alpha$-divergence.", "other_comments_or_suggestions": "Last paragraph on page 3: There is a reference missing (\"??\").", "questions_for_authors": "-", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper discusses a novel Monte-Carlo tree search method that uses $L^1$-Wasserstein barycenters to incorporate the uncertainty of state and state-action values of MDPs. The reviewers came to the consensus that the paper has good contribution to tackle highly stochastic and partially observable MDPs with a solid theory and convincing empirical results. While some concerns are raised such as the details of the experiments and literature review, they are mostly solved in the rebuttal. Thus I recommend acceptance and expect that the authors polish the paper incorporating the discussions with the reviewers in the final version.
Decision: Accept (spotlight poster)