Fast Bellman Updates for Wasserstein Distributionally Robust MDPs
Authors: Zhuodong Yu
Ling Dai
Shaohang Xu
Siyang Gao
Chin Pang Ho
Venue: NeurIPS 2023 poster (2023)
Keywords: Markov decision processes
distributionally robust optimization
TL;DR: We propose fast algorithms to solve distributionally robust Markov decision processes with Wasserstein ambiguity sets.

Abstract: Markov decision processes (MDPs) often suffer from the sensitivity issue under model ambiguity. In recent years, robust MDPs have emerged as an effective framework to overcome this challenge. Distributionally robust MDPs extend the robust MDP framework by incorporating distributional information of the uncertain model parameters to alleviate the conservative nature of robust MDPs. This paper proposes a computationally efficient solution framework for solving distributionally robust MDPs with Wasserstein ambiguity sets. By exploiting the specific problem structure, the proposed framework decomposes the optimization problems associated with distributionally robust Bellman updates into smaller subproblems, which can be solved efficiently. The overall complexity of the proposed algorithm is quasi-linear in both the numbers of states and actions when the distance metric of the Wasserstein distance is chosen to be $L_1$, $L_2$, or $L_{\infty}$ norm, and so the computational cost of distributional robustness is substantially reduced. Our numerical experiments demonstrate that the proposed algorithms outperform other state-of-the-art solution methods.

effective Fast Bellman Updates for Wasserstein Distributionally Robust MDPs
Zhuodong Yu'
Shaohang Xu' Siyang Gao? Chin Pang Ho' Schooo of Data Science, City University of Hong Kong, Hong Kong Department of Systems Engineering, City University of Hong Kong, Hong Kong (zhuodonyu2-c, lingdai5-c, shaohanxu2-c) htt.:cyymmymtyy....mmmmmmmmmm... {siyangao, clint. how.uiuy..eeeeeeeeoeooieee edu. hk
Ling Dai'
cityu.edu. hk
Abstract
Markov decision processes (MDPs) often suffer from the sensitivity issue under model ambiguity. In recent years, robust MDPs have emerged framework to overcome this challenge. Distributionally robust MDPs extend the robust MDP framework by incorporating distributional information of the uncertain model parameters to alleviate the conservative nature of robust MDPs. This paper proposes computationally efficient solution framework for solving distributionally robust MDPs with Wasserstein ambiguity sets. By exploiting the specific problem structure, the proposed framework decomposes the optimization problems associated with distributionally robust Bellman updates into smaller subproblems, which can be solved efficiently. The overall complexity of the proposed algorithm is quasi-linear in both the numbers of states and actions when the distance metric of the Wasserstein distance is chosen to be I,, T., or To norm, and so the computational cost of distributional robustness is substantially reduced. Our numerical experiments demonstrate that the proposed algorithms outperform other state-of-the-art solution methods.
Introduction
Markov Decision Processes (MDPs) provide flexible and powerful modeling framework for sequen- tial decision-making problems under uncertainty (Puterman, 2014; Sutton and Barto, 2018). However, the standard MDP model assumes that the exact knowledge of the transition kernel is available, which is not the case for most real-world applications. While these model parameters can be estimated from data, is well-known that the optimal policy of MDP is sensitive to estimation errors because the dynamic nature of the problem. In particular, small errors in estimating transition kernels could lead to catastrophic failures in practice (lyengar, 2005; Nilim and El Ghaoui, 2005; Le Tallee, 2007).
To overcome the aforementioned challenge in model ambiguity, robust MDPs assume that the transition kernels belong to ambiguity sets (Iyengar, 2005; Nilim and El Ghaoui, 2005; Xu and Mannor, 2006). Through optimizing the worst-case performance, robust MDPs compute robust policies to avoid disappointing performance due the sensitivity issue caused by inaccurate transition kernels. However, robust MDPs are often too conservative because they always adopt the worst possible transition kernels without any consideration of their likelihood. To mitigate this undesirable side effect of robust MDPs, distributionally robust MDPs are proposed to incorporate distributional information of model ambiguity and maximize the expected reward under the most adversarial probability distribution (Xu and Mannor, 2010; Yu and Xu, 2015).
In distributionally robust MDP, we assume that all the plausible distributions of the transition kernel belong to a prescribed ambiguity set, which could be formalized in different types, such as moment ambiguity sets (Delage and Ye, 2010) and Wasserstein ambiguity sets (Mohajerin Esfahani and
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Kuhn, 2018). In this paper, we focus on the Wasserstein ambiguity sets, which have been a popular choice for distributionally robust data-driven optimization in recent years because of their outstanding empirical performance as well as nice theoretical properties, such as consistency in optimality and finite-sample bounds (Mohajerin Esfahani and Kuhn, 2018; Gao and Kleywegt, 2022). By using Wasserstein distance to formulate the ambiguity set, Yang (2017) shows that there exists an optimal policy that stationary and Markovian for the corresponding Wasserstein distributionally robust MDP.
While (distributionally) robust MDPs can be solved by extending standard solution methods in classical MDPs to their (distributionally) robust counterparts, these solution methods become much more computationally demanding. For example, each Bellman update for (distributionally) robust MDPs can be formulated as a convex optimization problem. Without making use of any specific problem structure, one would need to use generic convex optimization solvers to compute these Bellman updates, which have to be evaluated numerous times for computing the (distributionally) robust value function. This computational challenge restricted the application of (distributionally) robust MDPs to small medium size of problems. In recent years, many efficient algorithms are proposed for solving robust MDPs to address this issue (lyengar, 2005; Nilim and El Ghaoui, 2005; Ho et al., 2018; Behzadian et al., 2021; Grand-Clément and Kroer, 2021b; Ho et al., 2022).
On the other hand, however, distributionally robust MDPs have received limited attention on their computational efficiency. To the best of our knowledge, (Grand-Clémen Clément and Kroer, 2021a) is the only pioneer work that proposes a first-order method to efficiently solve Wasserstein distributionally robust MDPs. However, being restricted by the nature of first-order methods, the proposed algorithm in Grrend-Clément and Kroer, 2021a) a) struggles when high accuracy is needed or when the discount factor is close to one, as it would take unsatisfactorily many numbers of iterations to compute the distributionally robust value function.
In this paper, we take an alternative approach to develop fast algorithms for solving Wasserstein distributionally robust MDPs. In particular, reformulate the optimization problem of the distribu- tionally robust Bellman update to a structural form, and by exploiting the specific problem structure, we propose decomposition schemes in which the problem of interest can be decomposed into smaller subproblems, which can be solved by our customized algorithms in quasi-linear time. As we will show later, the time complexities the proposed algorithms are linear in the numbers of actions and kernels and quasi-linear in the number of states, which are significantly better than the state-of-the-art solution methods (Xu and Mannor, 2010; GraandCCéémen and Kroer, 2021a).
The rest of this paper is organized as follows. Section 2 introduces the related work, and Section 3 describes the setting of Wasserstein distributionally robust MDPs. In Section we illustrate our decomposition scheme for computing the distributionally robust Bellman updates when the Wasserstein distance in the ambiguity set is chosen to be Ly norm. In Section 5, customized fast algorithms and their time complexities are particularly provided for the three most well-known cases: L1, L2, and Loe. The numerical experiments in Section 6 verify that our methods outperform other state-of-the-an algorithms.
Notations We use boldface lowercase and uppercase letters to denote vectors and matrices, re- spectively. We denote by [N] the set of natural numbers from 1 to N: that is, [N] = (1. N.. The vector e is denoted as the vector of all ones while the dimension depends on the context. The vector e, represents the vector of all zeros except the j-th component is one. Probability simplex denoted as AN {x RV T x x 0). The smallest component of a vector x is denoted by min(x). Given set X, we call the set of all the Borel probability measures P(X'). The 1/4
general Lg norm for vector æ E R. is denoted by ||>||, (EM los1") for 1 S q < co, and the Loo norm is || maX |a,]] The Dirac distribution at E g is denoted by 8z. For optimization problems with multiple constraints, we indicate the decision variables on the last line of the constraints.
2 Related Work
Research on Markov decision processes under model ambiguity can be traced back to the seventies (Satia and Lave Jr, 1973). In recent years, much progress has been made in the development of

---

robust MDPs. The most fundamental solution scheme to solve robust MDPs is robust value iteration, which is the robust counterpart of the standard value iteration and is first developed in (Givan et al., 2000; Iyengar, 2005; Nilim and El Ghaoui, 2005) for robust MDPs with (s, a) d)-rectangular ambiguity sets. Wiesemann et al. (2013) introduce the s-rectangular ambiguity sets; the authors reformulate the robust Bellman updates as convex optimization problems and solve them by using off-the-shelf solvers that have polynomial-time complexity. Researchers have been investigating the theoretical properties of robust MDPs, such as the error bound of applying robust MDPs with state aggregation (Petrik and Subramanian, 2014), the relationship between robustness and regularization (Derman et al., 2021), and the geometry of value function in robust MDPs (Wang et al., 2022). Other than the 8- and (s, -)-ambiguity sets, different types of ambiguity sets have been also proposed to mitigate the side effect of conservatism, such as k-rectangularity (Mannor et al., 2016) and r- r-rectnngularit (Goyal and Brénn-Clément, 2023). Different related models are also proposed to incorporate robustness in different settings, such as robust baseline regret (Ghavamzadeh et al., 2016), imitation learning (Brown et al., 2020), and soft-robust model (Lobo et al., 2020). In terms of algorithmic development, Ho al. (2018) propose fast algorithms for solving s-rectangular robust MDPs with weighted L norm ambiguity sets and propose to compute robust Bellman updates using the combination of bisection method and homotopy method, and Behzadian al. (2021) further extend the solution scheme for the unweighted La norm cases (Delgado et al., 2016). A first order method is introduced and Kroer, 2021b) to approximate robust Bellman updates for robust MDPs with ellipsoidal and Bullback-Leible (KL) s-rectangular ambiguity sets. Recently, Ho et al. (2022) propose the ø-divergence ambiguity set, which generalizes multiple popular choices of ambiguity sets, and they propose an unified solution scheme to compute robust Bellman updates efficiently. On the other hand, robust policy iteration alternates between robust policy evaluation and policy improvement (lyengar, 2005) and it has better empirical performance compared to robust value iteration. Kaufman and Schaefer (2013) propose a modified policy iteration for (s, reetrnctgn robust MDPs. Ho et al. (2021) introduce partial policy iteration for and (s, d)-rectangular robust MDPs with L, norm ambiguity sets. Kumar et al. (2022) study the connection between robust MDPs and regularized MDPs. and propose robust policy iteration for s-rectangular robust MDPs. Other than robust policy iteration, Li et al. (2022) propose a policy gradient method for solving rectangular robust MDPs. In order to solve large-scale problems, Tamar al. (2014) apply value function approximation and propose robust approximate dynamic programming for robust MDPs.
in
While this paper is focused on the model-based setting, recently there is an active line of research focusing on robust reinforcement learning (RL) algorithms (Roy et al., 2017: Badrinath and Kalathil, 2021; Wang and Zou, 2021, 2022; Panaganti and Kalathil, 2022). Liu et al. (2022) propose a robust Q-learning algorithm for (8, a)-rectangular robust MDPs. A function approximation approach is proposed in (Panaganti et al., 2022) for (s, a)-rentang robust MDPs with offline datasets. The combination of robust RL and deep RL also studied in (Pinto et al., 2017: Mankowitz et al., 2019; Zhang et al., 2020) although they do not provide theoretical guarantees of the learned policies. We would like to emphasize that both model- based and model-free methods have their own merits: while model-free methods are powerful for applications where data could be generated without much risk and cost, model-based methods are suitable for applications with very limited data or settings where the (time and financial) cost of obtaining new data is high.
For distributionally robust MDPs with the nested-set set structured parameter ambiguity sets, Xu and Mannor (2010) propose a Bellman type backward induction to compute the optimal policy, and Yu and Xu (2015) further extend to a more general state-wise ambiguity set following the unified framework proposed in (Wiesemann et al., 2014). The Wasserstein ambiguity sets are first introduced for distributionally robust MDPs in ( (Yang. 2017), and a convex formulation is proposed based on Kantorovich duality. Chen et al. (2019) propose a general formulation to combine both moment- based and cersticcll-disannce Mannor (2020) study the connection between Wasserstein distributionally robust MDPs and regu- larization. Abdullah et al. (2019) propose a reinforcement learning framework with distributionally robustiness. As mentioned before, however, Brann-Cléémett and Kroer, 2021a) is the only work that focuses on the algorithmic development of solving Wasserstein distributionally robust MDPs, which also the focus of this paper.
ambiguity sets together for distributionally robust MDPs. Derman and
3

---

3 Preliminaries
A distributionally robust MDP is a tuple (S, A. Po. r, X, M)), where S and A are the sets of states and actions, respectively. We assume that the state space S (I...., ......,,,, 5} and the action space A= (I.... A} are finite. The initial state 80 follows given initial distribution Po E As. When the decision maker takes an action
in state s E S, the MDP will transit to the next state randomly according to the probability distribution Psa E As, which is assumed to be unknown; the transition from state s E S under action
to the next state s' e S will induce the reward T'as'. We use the vector Tsa E IR to represent the connection of all rewards given state and action a. We denote by À E (0. 1) the discount factor. In distributionally robust MDPs, the transition kernel is assumed to be uncertain, and it is governed by an unknown probability
II Pan E (As) SxA *ES, distribution which is assumed to reside in a prescribed ambiguity set MI that is calibrated from historical data. Note that we also use the notation p. II a Pso (As)² for the transition kernel
given the state S.
The goal of a distributionally robust MDP is to maximize the worst-case expected return; that is,
max min E T) PEM
E proM
MTnaastt 50 Pa
(I)
Here, the decision variable ग E (AA)* called a (randomized) policy where for any given 8 E S, its sub-vector E A represents the probability distribution of taking action AA in state s.
It is worth mentioning that the above formulation of distributionally robust MDPs is generalization of both robust MDPs and standard MDPs. By specifying the anbiguity set M to be P(2) where 2/ is the ambiguity set of robust MDP, the problem (I) is equivalent the robust MDP with ambiguity set U. If the set U is a singleton, then problem (1) with Mi P(2) is equivalent to the standard MDP with transition kernel from the singleton 21
Similar to the case of robust MDPs, we consider the common assumption of rectangularity for the sake of tractability, as otherwise problem (1) is NP-hard (Wiesemann al., 2013). In particular, consider the following s-rectangular Wasserstein ambiguity set (Grand-Clémen) and Kroer, 2021a)
X Mas Hss E My C P ((As)*) Vs $. eES
M n E P ((4s) μ
where My is the marginal ambiguity set for the distribution μ., that governs the transition kernel at the state s S. The superscript g indicates that My is defined by q- W-sssesstein distance; that is,
M3 -
[1s E P ((As)*) Wg (Me: v.) 50} for 1 1 oc,
where u is the reference distribution at the center of the Wasserstein ball and
{ ((E (dg(a, (d<(a.w)))) 1/4
: : K T(Hs us) Pe
Wg (Ja, Uz) =
min sEr((As)*x(Ae))
for S q 00, where I'M vs) the set of all the couplings of μ, and vs (Ambrosio and Nicola Gigli, 2005), and dy is induced by the q-norm dy (z, v) = |I vll, throughout this work. We define the co- Wasserstein distance W~ (Ms, vx)) by taking q + xc.
We adopt the usual setting in distributionally robust data-driven optimization where v, = * EN 6pp
is set to be the empirical distribution and estimated by the samples (Ø) ( where P' E (As) for
every i E [N]. Then, the optimal value function of (1), v*, satisfies the following distributionally robust Bellman equation (Yu and Xu, 2015; Yang, 2017; Ciand-Clément and Kroer, 2021a)
0g max min
TsaPra (rea + X00)) NaEA
Σ
Vs E S,
(2)
R. R.CAA P.EBY

---

where B3 is the set of expected kernels (Yang, 2017; Bertsimas et al., 2018; Xie, 2020) of the following forms
1
B
|p's B'llyy S 09, p (As),, W for 1 15 q < oo,
N
N
i=1
i=1
1
P
|p'
p:11 0, p¹ e (As) wie IMI
N
(3)
4
The Decomposition Algorithm for Distributionally Robust Bellman Update
In this section, we focus on computing the distributionally robust Bellman update, which the most fundamental operator for solving the distributionally robust Bellman equation (2); define the distributionally robust Bellman operator I(-) where for any v E IRS
[I())]+
T max min Σ Naa Psa (r + Xv) E
(4)
T,EAA E.
Given the above definition of ((),, it is well-known that (2) can be solved by iteratively applying T, which is a variant of standard value iteration for solving distributionally robust MDPs (Yu and Xu, 2015;
and Kroer, 2021a). For any initial guess v" IRS, we have lim; v² + v* where of I(v(-1) for 1, 2, ..... and ex satisfies (2). Therefore, the efficiency of evaluating [I()))s is crucial to the computation of solving distributionally robust MDPs.
However, computing [E()) using generic convex optimization solvers is much more computationally demanding compared to the case classical MDPs, which only have time complexity O(S.A for computing their Bellman updates. In this section, exploit the specific problem structure in (4) and reformulate the optimization problem to a form that could be decomposed into smaller problems. We first focus the case of q [1, oe) in Section 4.1 and then consider the case of g = in Section 4.2, where indicates the type of Wasserstein distance used in the ambiguity set. As we will show later, combined with the customized fast algorithms in Section for solving the subproblems, one can compute [3(]],, in time complexity that is quasi-linear in S and linear in A and N. All the proofs of propositions and theorems this section are provided in Appendix A.1.
4.1 Nested bisection method for q- --Wasserstein ambiguity set with E [1,
We consider the optimization problem (4) for the case where E [1, oo). By applying the minimax theorem, we obtain the following result.
Proposition 4.1. Consider the Bellman updates (4) with q (o) Then, minimize
7
subject
Σ (rsa + Ae) Psa St, Va EA
N
[[()]
E
(5)
1
Pb P'aa 11 S 89
E R. P'a E As, Vi E Wa E A.
The above proposition reformulates the maximin optimization problem in (4) into a convex min- imization problem, which can be solved by commercial convex optimization solvers. Moreover, problem (5) can be solved via bisection on 7; that is, we seek for the lowest possible y such that N DN (rea + Xv)T P'm. < y for each a E AA while pi, satisfies the other constraints in (5), for every a E A. To this end, we introduce the following subproblem
N
minimize
D'r. -D'allg
IN bsa. B ((Dia)N
(6)
subject to
DraP's ST
N
1=1
pi
As, Vi [N]
5

---

By setting Dsa = Tsa + λv, the above problem B 1:: allows us to seek for decision variable (PsaJi=) /i=1 that satisfies the first and the last line of constraints in (5) while minimizing the "budget" used in the second line constraint in (5). particular, for any fixed y E IR, we distinguish among the following two cases:
(a) If EnEA & ((pra) aay r¹) - 04, then 7 feasible for (5) and is an upper bound of the optimal objective value in (5).
(b) If DaEA P is lower bound of the optimal objective value in (5).
sa: ?') > 04, then y' must be infeasible for problem (5) and
To apply the bisection method, derive the initial upper and lower bounds of the optimal y" in (5). Proposition 4.2. Consider the distributionally robust Bellman update (5). The optimal objective value 7" bounded by
1 N
max (min (bsa)) S S aEA
max aEA
As opposed to (5), the size of problem (6) independent of A. By using the aforementioned bisection method, the overall time complexity is only linear in A since only need solve (6) O log 61) times, where 6 is the tolerance for the bisection method.
While one can solve problem (6) using off-the-shed solvers and enjoy the reduced complexity by using the above bisection method, this subproblem itself turns out to another structural optimization problem. As we will show in the following proposition, by applying duality (6), the reformulation could be further decomposed via another bisection method on the dual variable.
Proposition 4.3. Consider the problem Q
((D J.V. Dear 7) in (6). If > min {beah, then
N
& ((Psa) ((Ps)) Dsa bi = max -ay + USasa
2 (P Uwa;
(8)
where ā = max ||e; D'a ( ((y- min (bea)) and
Do Do Dsas a) = min
Ip'a - P'alllg + a -abbb.a baaPea
(9)
Pha EAg
where j arg min Dsas'-
Notice that focus on the case where 7 min (bea) since min {bsa) is not larger than the lower bound of y in (7), and the case 7 = min { {b will not occur because y is always taken to be the average of the upper and lower bounds in our bisection method. We refer interested readers to Appendix for more details.
Proposition 4.3 indicates that for any fixed feasible a in (8), the subproblems (9) be solved separately. Therefore, by applying the bisection method on a in problem (8), one can naturally decompose the problem (8) into N smaller problems that have variables and S + constraints. As a consequence, the overall time complexity is linear in the number of kernels N.
By combining both decomposition strategies above, we obtain the proposed nested bisection method, whose pseudocode could be found in Algorithm 1. More details the algorithm is provided in Appendix A.1.
Theorem 4.4. Suppose 7 is the value returned by Algorithm 1, and 7* be the optimal value of (5). have
With the inputs provided in Algorithm and user-specified tolerances e1, 62
Acz
max (ba) -2
max max (aaaabbb b.a +
- -1-7+1 E1
aEA
aEAA
Theorem 4.5. Algorithm computes (5) time O Ohg (S, loge log ","" AS), where (S)
is the time complexity for solving (9).

---

Algorithm 1: Nested bisection method to compute distributionally robust Bellman update (5)
Input: Tolerance 61 for outer bisection method and 62 for inner bisection method; Initialization; Set lower bound y and upper bound 7y that are specified in (7);
while 5 - > do
Compute 7 = (5 + 2)/2 and ro = 0
for a = ,,,,,,,, A do
Set lower bound a = 0 and compute the upper bound da in Proposition 4.3; while â. - a C2 do
Compute 0g (āa + a,)/2;
Solve (P bsa; da), for every E [N] Compute the slope of (8): - + N Li=1 b Dep
if S 0 then set in = da else set an = iaa end
Set a = (ā, + +2a//2:
Solve (p(P'n bra, 0g),, for every i [N];
Set Ta-1 0 + N N Dg bsas ag);
end
rA > 89 then set 2 ÿ else set 7 = 9;
end
Result:
Optimal objective value of (5): y = (5 + 2)/2:
As shown in the above theorem, the proposed nested bisection method has time complexity that is linear both A and N, but the overall complexity depends on the subproblem (9) where its complexity depends h (5) and h2(S), which are the two most common cases. The complexities associated with solving equations (5) and (8) using general convex optimization methods are also provided in the Appendix B.
the choice of q and the number of states S. In Section 5, will derive
4.2 Decomposition scheme for oc-Wasserstein ambiguity set
As opposed to the case where E oo), for oo- Wasserstein ambiguity set, the corresponding Bellman update can be naturally decomposed without using any bisection method.
Proposition 4.6. Consider the Bellman update (4) with q oo. Then,
N
[[()]] max aEA
min (D Pi. EAs
I| - Pia 11 SO} Vs
(10)
The above reformulation requires solving A inner minimization problems that have S variables and S 2 constraints. Hence, the time complexity of Bellman update (10) is linear in A and N. Similar Section 4.1, we will discuss the complexity of solving the inner minimization the next section.
Efficient Algorithms for Subproblems
Section 4 offers decomposition schemes for distributionally robust MDPs with q- Wasserstein am- biguity set. As shown in Theorem 4.5 and Proposition 4.6, the overall complexity of computing distributionally robust Bellman update linear in both A and N but it depends on complexities of solving the subproblem (9) and the inner minimization problem in (10). In this section, we consider the common cases where g E (1,2, x.) and discuss the time complexity of solving the subproblem in each case.
As we will show below, the overall time complexities of our decomposition algorithms only O (SAlog S + NAS ci log 1) for q 1, ONNAS log S o ( for 2, and O(AS(logS + N)) when y xc. These complexities are much lower than ((N3.5 43.3 53.5 log(e -)) and O 22.5522.. log()) -1.5) ods in (Xu and Mannor, 2010) and Grand-CCément and Kroer, 2021a), respectively. The proofs of
from the state-of-thh-aar solution meth-
7

---

the following propositions, theorems, and corollaries and the details of the proposed algorithms are relegated to Appendix A.2.
5.1 1- 1-Wasserstein ambiguity set
When 1, the corresponding subproblem (9) has an equivalent form as follows.
Proposition 5.1. Suppose q 1. The minimization problem (9) is equivalent to
min + min bapP' v0
Ip'.. Psalli S Su}
(11)
Pia Edg
While problem (11) does not appear be trivial at first glance, it has nice mathematical properties which allow us to solve this problem efficiently.
Theorem 5.2. The objective f in (11) is piecewise-linear in v, and it can be solved in time (Slog S).
Therefore, the proposed Algorithm can compute the Bellman update with the time complexity linear in A and N and quali-linear in S. We obtain the following result by combining Theorem 4.5 and Theorem 5.2.
Corollary 5.2.1. The distributionally robust Bellman update with 1- I-Wasserstein ambiguity set can be computed in time O (SAlog NAS NASlog Ci log Cg ) where 6 and 62 the user-specified tolerances in Algorithm 1.
5.2 2- Wasserstein ambiguity set
When q 2, the following result shows that (9) could be transformed to Euclidean projection problem onto the probability simplex.
Theorem 5.3. Suppose
2. The minimization problem (9) is equivalent to
b.a l2 + abTAp' + min
2p
absa
Psa
(12)
EAs
2
2
which can be solved time 0S log S).
Euclidean projection problem the probability simplex can be solved via the existing algorithm (Wang and Carreira-Perpinán, 2013). By combining the results in Theorem 4.5 and Theorem 5.3, we obtain the following result.
Corollary 5.3.1. The distributionally robust Bellman update with Wasserstein ambiguity set can be computed in time O AS log S log 6 log +))), where and e2 are the user-specified tolerances Algorithm 1.
5.3
Wasserstein ambiguity set
When 0o, the corresponding subproblem
min ( (b..Pi p EAs
||p'.
p Ios S
(13)
can be reformulated to a linear program with box constraints and a linear equality constraint which can be solved efficiently.
Theorem 5.4. The minimization problem (13) is solved in time O(S log S).
Therefore, we obtain the overall complexity as follows.
Corollary 5.4.1. The distributionally robust Bellman update with do- Wasserstein ambiguity set can be computed in time O (AS (log S + N)).

---

Runtimes (100s)
-X- Fast *der Garebi FOMID
Runtimes (100s)
**- Fast .le Gurobi
Runtimen (100s)
8
+ Fast
Gurobi
6
6
FOMIB
4
4
2
2
2
0
0
50
70
90
110
50
70
90
110
50
70
90
110
5 =
A
5
N
Figure I: Comparisons of all algorithms for q= where (left)
50, (middle) 70, N 50,
and (right) S A 70.
Runtimes (10s)
8
+ Past
Runtimes (10s)
8
+ Post
Runtimes (10s)
8
+ Fast
.e. Garebi
.xe Gurabi
Gurcti FOMIB
6
FOMIB
FOMB
6
4
4
2
2
0
0
50
70
90
110
50
70
90
110
50
70
90
110
5 = A
5
N
Figure 2: Comparisons of all algorithms for q = where (left) N 50, (middle) 70, N 50,
and (right) S A = 70.
+
Fost
+ Feet
8
+ Fast
Runtimes (5)
Gurabi
Runtimes (s)
. Gurobi
Runtimes (s)
.ee Gurobl
6
FOMIB
FOMI3
5
FOMIB
4
4
4
2
0
0
50
70
90
110
50
70
90
110
50
70
90
110
S = A
S
N
Figure 3: Comparisons of all algorithms for q = co, where (left) N 50, (middle) 70, N 50,
and (right) S A = 70.
Numerical Results
We experiment on the performance of the proposed algorithms (Fast), Gurobi (Gurobi Optimization, LLC, 2023), and the first-order method (FOM) proposed in (Grand-Clément and Kroer, 2021a) with different sizes of the randomly generated distributionally robust MDPs. Due page limit, only re- port the results of the experiments on a single Bellman update, and we provide additional experimental results in the Appendix C, which also contains the detailed settings of all the experiments.
Figure reports the average computation times Bellman updates with L-Wasserstein ambiguity sets. For the first-order method (Grand-Cléman and Kroer, 2021a), we report the computation times of the third Bellman iteration, and these computation times increase at every Bellman iteration of the first-order method. One can see that all three algorithms perform similarly when problem size is small. However, as S, A, or N increases, the runtimes of both Gurobi and first-order method increase rapidly, while the proposed algorithm remains scalable.
The results with Wasserstein ambiguity sets and oo- oasssssrsteein ambiguity sets are shown in Figure 2 and Figure 3, respectively. These results are similar to the case where g 1, which are consistent to our theoretical results on complexities. As expected, the proposed algorithms are several orders of magnitude faster than the existing state-of-the-an the solution methods.

---

7 Conclusion
This paper studies distributionally robust MDPs with Wasserstein ambiguity sets. In particular, we propose fast algorithms to compute distributionally robust Bellman updates. We show that the proposed algorithms outperform other existing methods in both theory and experiments. Future work could address the extensions to approximate dynamic programming for distributionally robust MDPs as well as model-free settings.
8 Acknowledgement
We thank the anonymous reviewers for their supportive comments. This work supported, in part, by the CityU Up Grant (Project No. 9610481), the National Natural Science Foundation of China (Project No. 72032005), and Chow Sang Sang Group Research Fund sponsored by Chow Sang Sang Holdings International Limited (Project No. 9229076).
References
Mohammed Amin Abdullah, Hang Ren, Haitham Bou Ammar, Vladimir Milenkovie, Rui Luo, Mingtian Zhang, and Jun Wang. Wasserstein robust reinforcement learning. arXiv preprint 1rXiv:1907.1319,, 2019.
Luigi Ambrosio and Giuseppe Savaré Nicola Gigli. Gradient flows: in metric spaces and in the space of probability measures. Springer Science & Business Media, 2005.
TW Archibald, KIM McKinnon, and LC Thomas. On the generation of markov decision processes. Journal of the Operational Research Society, 46(3):354-361, 1995.
Kishan Panaganti Badrinath and Dileep Kalathil. Robust reinforcement learning using least squares policy iteration with provable performance guarantees. In International Conference on Machine Learning, pages 51 511--5 1200. PMLR, 2021.
Bahram Behzadian, Marek Petrik, and Chin Pang Ho. Fast algorithms for sren srectangala robust mdps. Advances in Neural Information Processing Systems, 34:25982-25992, 2021.
Dimitri P. Bertsekas. Nonlinear programming. Athena Scientific Belmont. Massachusets, USA, 1999.
Dimitris Bertsimas, Shimrit Shtern, and Bradley Sturt. A data-driven approach for multi-stage linear optimization. Available at Optimization Online, 2018.
Shalabh Bhatnagar, Richard S Sutton, Mohammad Ghavamzadeh, and Mark Lee. Natural actor-critic algorithms. Automatica, 45(1 (5:11)224712482 2009.
Daniel Brown, Scott Niekum, and Marek Petrik. Bayesian robust optimization for imitation learning. Advances in Neural Information Processing Systems, 33:2479-2491, 2020.
Zhi Chen, Penggian Yu, and William B Haskell. Distributionally robust optimization for sequential decision-making. Optimization, 68(12):397-2426, 2019.
Erick Delage and Yinyu Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations research, 58(3):595-611 2010.
Karina V Delgado, Leliane N De Barros, Daniel B Dias, and Scott Sanner. Real-time dynamic programming for markov decision processes with imprecise probabilities. Artificial Intelligence, 230: 192-223, 2016.
Esther Derman and Shie Mannor. Distributional robustness and regularization in reinforcement learning. arXiv preprint arXiv:2003. 022844 2020.
Esther Derman, Matthieu Geist, and Shie Mannor. Twice regularized mdps and the equivalence between robustness and regularization. Advances Neural Information Processing Systems, 34: 22274-22288 2021.
10

---

Rui Gao and Anton Kleywegt. Distributionally robust stochastic optimization with wasserstein distance. Mathematics of Operations Research. 2022.
Mohammad Ghavamzadeh, Marek Petrik, and Yinlam Chow. Safe policy improvement by minimizing robust baseline regret. Advances Neural Information Processing Systems, 29, 2016.
Robert Givan, Sonia Leach, and Thomas Dean. Bounded-parameter markov decision processes.
Artificial Intelligence, 122(1-2):71-199, 2000.
Vineet Goyal and Julien Grand-Clément Robust markov decision processes: Beyond rectangularity. Mathematics of Operations Research, 48(1):203-223 2023.
Julien Grand-Clément and Christian Kroer. First-order methods for wasserstein distributionally robust mdp. In International Conference Machine Learning, pages 2010-2011 PMLR, 2021a.
Julien Grand-Clément and Christian Kroer. Scalable first-order methods for robust mdps. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pages 12086-12094, 2021b.
Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2023. URL https://wwz gurobi com.
Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Fast bellman updates for robust mdps. In International Conference on Machine Learning, pages 1979-1988. PMLR, 2018.
Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Partial policy iteration for 11-robust markov decision processes. The Journal of Machine Learning Research, 221):172622 2021.
Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Robust phi-divergence mdps. arXiv preprint arXiv:2205. 14202, 2022.
Garud N Iyengar. Robust dynamic programming. Mathematics of Operations Research, 30(2): 257-280, 2005.
Narendra Karmarkar, A new polynomial-tin algorithm for linear programming. Proceedings of the sixteenth annual ACM symposium Theory of computing, pages 302-311, 1984.
David I. Kaufman and Andrew Schaefer. Robust modified policy iteration. INFORMS Journal on Computing, 25(3):396 2013.
Navdeep Kumar, Kfir Levy, Kaixin Wang, and Shie Mannor. Efficient policy iteration for robust
markov decision processes via regularization. arXiv preprint arXiv: /4327, 2022.
Yann Le Tallec. Robust, risk-sensitive, and data-driven control of Markov decision processes. PhD thesis, Massachusetts Institute of Technology, 2007.
Yan Li, Tuo Zhao, and Guanghui Lan. First-order policy optimization for robust markov decision process. arXiv preprint a2Xiv:2209. 10579, 2022.
Zijian Liu, Qinxun Bai, Jose Blanchet, Perry Dong, Wei Xu, Zhengqing Zhou, and Zhengyuan Zhou. Distributionally robust q-learning. In International Conference on Machine Learning, pages 13623-13643. PMLR, 2022.
Elita A Lobo, Mohammad Ghavamzadeh, and Marek Petrik. Soft-robust algorithms for batch reinforcement learning. arXiv preprint arXiv:2011.14499 2020.
Daniel Mankowitz, Nir Levine, Rae Jeong, Yuanyuan Shi, Jackie Kay, Abbas Abdolmaleki, Jost To- bias Springenberg, Timothy Mann, Todd Hester, and Martin Riedmiller. Robust reinforcement learning for continuous control with model misspecification. arXiv preprint 110.:v1116.075550 2019.
Shie Mannor, Ofr Mebel, and Huan Xu. Robust mdps with k-rectangular uncertainty. Mathematics of Operations Research, 31(4):1444-150) 2016.

---

Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Programming, 171(1-2):115-16 2018.
Arnb Nilim and Laurent EI Ghaoui. Robust control of markov decision processes with uncertain transition matrices. Operations Research, 53(5):780-798, 2005.
Kishan Panaganti and Dileep Kalathil. Sample complexity of robust reinforcement learning with generative model. In International Conference on Artificial Intelligence and Statistics, pages 9582-9602. PMLR. 2022.
Kishan Panaganti, Zaiyan Xu, Dileep Kalathil, and Mohammad Ghavamzadeh. Robust reinforcement learning using offline data. arXiv preprint arXiv:2208. 05/29, 2022.
Marek Petrik and Dharmashankar Subramanian. Raam: The benefits of robustness in approximating aggregated mdps in reinforcement learning. Advances in Neural Information Processing Systems, 27, 2014.
Lerrel Pinto, James Davidson, Rahul Sukthankar, and Abhinav Gupta. Robust adversarial reinforce- ment learning. In International Conference on Machine Learning, pages 2817-2826. PMLR, 2017.
Martin Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Aurko Roy, Huan Xu, and Sebastian Pokutta. Reinforcement learning under model mismatch.
Advances in neural information processing systems, 30, 2017.
Jay K Satia and Roy E Lave Markovian decision processes with uncertain transition probabilities. Operations Research, (1(3):72 -702, 1973.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Aviv Tamar, Shie Mannor, and Huan Xu. Scaling up robust mdps using function approximation. In International conference machine learning, pages 181-189. PMLR, 2014.
Kaixin Wang, Navdeep Kumar, Kuangqi Zhou, Bryan Hooi, Jiashi Feng, and Shie Mannor. The geometry of robust value functions. In International Conference on Machine Learning, pages 22727-22751. PMLR, 2022.
Weiran Wang and Miguel A algorithm with a simple proof, and an application. arXiv preprint arXiv:1309.1541 2013.
Projection onto the probability simplex: An efficient
Yue Wang and Shaofeng Zou. Online robust reinforcement learning with model uncertainty. Advances in Neural Information Processing Systems, 34:7193-7206, 2021.
Yue Wang and Shaofeng Zou. Policy gradient method for robust reinforcement learning. In Interna- tional Conference on Machine Learning, pages 33484-23526 PMLR, 2022.
Wolfram Wiesemann, Daniel Kuhn, and Berç Rustem. Robust markov decision processes. Mathe- matics of Operations Research, 38(1):153-183, 2013.
Wolfram Wiesemann, Daniel Kuhn, and Melvyn Sim. Distributionally robust convex optimization. Operations Research, 6256:1135 2014.
Weijun Xie. Tractable reformulations of two-stage distributionally robust linear programs over the type-co wasserstein ball. Operations Research Letters, 48(4):513-52) 2020.
Huan Xu and Shie Mannor. The robustness- perfommace tradeoff in markov decision processes. Advances in Neural Information Processing Systems, 19, 2006.
Huan Xu and Shie Mannor. Distributionally robust markov decision processes. Advances Neural
Information Processing Systems, 23, 2010.
12

---

Insoon Yang. A convex optimization approach to distributionally robust markov decision processes with wasserstein distance. IEEE control systems letters, 111):664-16, 2017.
Pengqian Yu and Huan Xu. Distributionally robust counterpart in markov decision processes. JEEE Transactions Automatic Control, 61(9):2538-2543 2015.
Huan Zhang, Hongge Chen, Chaowei Xiao, Bo Li, Mingyan Liu, Duane Boning, and Cho-Jui Hsieh. Robust deep reinforcement learning against adversarial perturbations on state observations. Advances in Neural Information Processing Systems, 23:2102--21037, 2020.
13

---

A Appendix: Proofs and Algorithms
A.1 Proofs of results in Section
Proof of Proposition 4.1. Plug B in (3) into (4), and apply the minimax theorem, the original problem [I(v)]e, Vs E S is given by:
N
max n. EAA
min
Σ
Tea
Σ
Pia
(rea + X2)
P
N
aEA
i=1
N
s.t.
p
P:11' - 04
N
t=
P'a
As,
Vi E [M], Va E A
min
max R.EAA
Σ
Tea Σ
(rea + 2v) Psa
N
REA
N
s.t.
p
1:ll%y S 04
N
i=
D'a
Vi E [N], Va E A As,
1
min
max aEA
(rsa Xv) Dea
N
pp. - 1:,1% S 09 Σ
s.L.
aEA
P'a As, Vie [N], Va E A.
Hence (5) is a direct consequence of above formulation by introducing the following epigraph variable
N
y which satisfies y 2 maX.ed N E (rea + xv) pi..
im
Proof of Proposition 4.2. To prove the upper bound of 7", consider 7
max 1. EN biap'e. which equals to 7 2 E D Va E A. This implies (Paa) J satisfies every constraint problem (6) with the lowest possible objective value 0, for every A. Therefore, EaEA & (fri ii Tsa + Av, 7) = 8", this time, 7 is feasible for (5). Hence, we provide an upper bound
max N Σn b D
To prove the lower bound o", wC assume the contrary max,EA (min [b,a}} 7-, So there exists à E A with y* min (b }. Then there is no Pha satisfies the first constraint in (5) for å. This contradiction verifies the lower bound y = max {min }b
Proof of Proposition 4.3. By definition, for any fixed E S and E A, we have
N
B (IPia MM bsa,
min
llp
P
N
N
1
T P'n S ?:
s.t.
N
pia
E
As.
s: Vi [N].
Then, we introduce the dual variable a for the constraint
E
bTaD' S 7 and obtain
1
Σ
N
Ilpi - Dia 19 +
PapPoo s0 b
}
max a20
min
pr.,Ag, ViE(N)
N
N
1=1
N
max 20
ary
D (D bea, at),
14

---

where Dy bea: a) is defined in the Proposition 4.3.
To show ā defined in the proposition is indeed an upper bound of the optimal a², we denote
N
1 PP b bear a..
f(a) -ay
N
1 Pieele 11 7-min bra) where j E arg min bans' : s'ES
Notice that Vo 2 maX4e[N]
f(a) S -ay+ +
Σ
ej Pially + abnne ej
N
i=1
N
N c
Σ
p'a , + (min {b 7) a
211
S 0 f(0),
where the first inequality is from the definition of D, (P's Doa: a), the second equality is due the selection of j. and the last inequality is due to the selection of a. Hence a E [0, āj by f(a) is concave w.r.t a (otherwise 3a such that f(a*) > f(0) f(ā), where nonconcavity of follows), and (8) is true. One can further compute the subdifferential of f(a) by Danskin's 's theorem (Bertsekas, 1999). Typically,
N E PI. E @f(a)).
N
where pha is any minimizer of the inner minimization problem corresponding with given a 2
Proof of Theorem 4.4. For simplicity of notations, throughout this proof we use
Ac2 2
A
max max +7)
and f(7) 2
Σ B (I by E by 7]. aEA
aeAA
By (6), we can get that P ((Dt sa ) Dwn, 7) and ((7) are non-increasing in l2 7]. For each given Y 7 and a E A. we denote & (fPt. bsa. 7) the corresponding value calculated by the Algorithm 1. Furthermore, we call
((7)
Σ $ aEA
Vy E h.
We can show that
56 by l).
(14)
We fix 7 E 12, 7). and consider (8). Algorithm provides the optimal solution for (8) with tolerance c2/2, so we have |arbe, ása,zl S <2/ where a'a,- the true optimal solution (8) and &sa,,y is the solution computed in the inner bisection in Algorithm 1. Thus we get
B
bea. $ ((Pia)
1N bea: n)
N
N
= = -
+
Σ
Dr (P'ar bsa: - (-aeur? R +
Σ Dy(P'an bsa: aaa.))
2san x
Dg(Pin Dsas asa,7) - D, Dr Dsa, 2sa,))
N
1=1
7 max max (bsa)) aEA
2
15

---

where the last step is due to that
Dg bsnn Noart) Dj (Ppoas Dsas asa,))
|p'aa - Peall' + Yearr b DeaP'n min
min Piu EAs
llp'o
Psn 11 &say D
Pla Eds
(Qaa, asa;)) boapsa in
62 2 aEA
max max
and
y (Psa: beer @sart) D D Doas Psa,y)
min Pi, EAs
||'aa Poally + deair b DaaP'o
min pin Eds
llp'a D'aally + b
< (āea., - ('a,)) Pa inn"
2 max max(bbal,
2 aEA
here pba and p.o are the optimal solutions to (9) for a = âsa,y and a = d,a,y respectively.
Hence, (14) is the direct consequence of above estimation, together with the definitions of f and 6. And we have y- in [xs 7) f(>) 04). We further define 7 - inf(> 7] f(7) 09}. The outer bisection of Algorithm I implies that h - + S 7 to prove the claimed result in the theorem, suffices to show that
28 max max (baa) -
lr' -
(15)
saaliii Dsan 7) Vy E [g, by definitions, we get f(>) s f(7). so fy E h 71 f()) < 50 s E h f(7) I 8²), hence we get y S We assume that y - 7 2 (15), we claim the following statement:
Notice that $3 (Lo) 1N bea, 7) < 5 (1iba
(PaaJis
25 (maxuea max (b.a) otherwise (15) trivially satisfied. To achieve
09
max max (bma) -2)
25
aEA
09 + S.
(16)
09
If the statement is true, have Myy 12,7" (max max(bra) -2) (25 (8)/
28
max max Dsa) -2)
16)) ((7) - 5225
7"
aeA
- 8 89
09
where the first inequality is from (14) and f(r) f()), the second inequality is due that f(y) is non-increasing, and the third inequality is from the above statement.
So (> E [7, 7) f(7) < 04) C (7* - (maxaeA max baxxbend bba) -7) (25/09), 7]. thus
28 (max.e max(beat
S 5 39'
09
which implies the desired (15).
For the proof of the statement (16). We argue that
26 (max,e.A max [ba) - 7)
02 (max.eA max(baa) 2) (17)
Vy E b--
xT)): Σ asan 2
09
16

---

26 28(aaxxed max(b ,7) with that
by contradiction.
Assume there is some 7" E h-" 1
An
EnEA "sa,""
Clearly f(>") 0 since y" <. Then
maNucc max(bba)-1]
09
Σ asa,y"
max 6ea -
aEA
aeA
N
asaaz"
Σ
b
buaP's -7)
N
aEA
N
Gsa,y" (b (B -7)
N
N
|P Psall 1 + aang"" ( (D -7)
min Piu EAg
N
aEA
Σ
Dy Dsay 2,0
N
Σn F ((p ii
aeA
=f(?"),
which is contradicted with f(7") > 89. This contradiction implies that (17) is true. For simplicity of notations, we call 7" then
26( maxx a max (b..) -22 and we fix any ge (,"","*).
0
25 (maxaed max (bsa) - 7)
09 8 7-
09
<{(?) - -(("")
Σ (H (1i'a I.N (i 7²) 8 (IDma) Dwa, baan)"")) 7"
N
N
Σ
- C
Σ i=1
Dr Dsan )snn)) - a
+
Σ D
(i Dany maan))
N
N
aEd
1=1
Σ
Ox 0ra,zel "" -79)
aEd
<(>"" -7
| max bsa ?)
where the first step is due to ndh < <n" and definition of 7 " the second step is from the definition of function f, the third step is from the definition of ((fria) 1. bea, 7) and ata.. and the last step is due to (17) and 7 E (7" >)) Notice that above inequality is true for all 7f (""","") we could let xt + 7* which leads to
(x" ( 25 (maxae mnax(bsa) - 2) 2)
09
-2 -8.
0g
Then (16) is the direct consequence of above inequality. This finishes the proof of statement, hence finishes the proof of the theorem.
Proof of Theorem 4.5. The Algorithm 1 is the direct consequence of the procedures in the content, except computing the slope, which has been explained at the end of the proof of Proposition 4.3. For the time complexity, O e log E ). For the subproblem (9), which costs time complexity h₂(S), need to solve
that the bisection method on y and uses complexity
17

---

Algorithm 2: Fast algorithm to solve (9) with g =
Input: Sorted baa with Usann Uyang
Usang:
Initialization: f- b b and Psaa f Psaa
for k = 1....., saan + aabbans Deans Peans 2 then break
do
else
r+ 2P + albsang Dsan,, PFass Psans + Psan; and D'sane O.
end
Result: Optimal objective value and optimal solution pia. of (9) with 1.
it N times. Besides, computing the upper bound claimed in Proposition 4.2 requires finding min [b for each A. which in time complexity ((AS). So we get the time complexity of Algorithm is O (hg (S)NA cog logez" + AS).
Proof of Proposition 4.6. Plug B in (3) into (4), we get the Bellman update is given by
UsaPso (rsa Xv) I| Dealloo 50, p'a Es,, Vi [N], Va E E
max N T,AAA
Σ Tsa
min Dea T (rsa Av) I|pea Doalloo So p'a E {ss}
max
TAEA
aEA
=
max
min
(b : I|peo - Pea 11.00 S $0}.
N AEA
Pi. EAg
The first equality is due that the decision varaibles Pi.. Vi [N] and Va E A, are independent from each other. Then we can divide the orignal problem into N A subproblems. The second equality from the fact that the objective function is affine w.r.t. waa and maximum simply the greatest coefficient of Tea-
A.2 Proofs results in Section
A.2.1 Proof of results in Section 5.1
Proof of Proposition 5.1. We introduce the variable IIp`s 1 so
D1(p' bsa, a.) = min Ilp'a - P'a /1, + biaP's
Ple EAg
min min {o + ab aa Ilp's p'a pi S 5v}
v20 pia EDs
(18)
min
U min b b llp' p'a S 5o v20
pia EAg
Proof of Theorem 5.2. We denote the objective function in (11) as
F(v) v v
min
II
P'all1 Su suy.
p:. EAy
18

---

W.L.O.G., we assume that deal > Dou2 > ... > beas, and P'a > 0. We claim that
(bsal
beas)
1. + a
+ b.a Psn
if E [0, 2f
2
K
K+1
F(v)
8 + a
baba
P
2
Pkak
for some K E [S - 2].
km K
km
1. + cxbsas
if fee
Paak: t00
k=1
where
K 1/
Σ Poak / if +2 k
1 Poa(K+1)
if = K +
Tk
Prak Peas +
k=1
k = S.
To prove the first case of the claim, it suffices show that the optimal solution for the minimization problem in F()) is given by p". whose components are p = -Pse1 P = Paak. V2 S k S S
and P's P + 2. It can be easily verified that p" defined in this way satisfies the constraints of the minimization problem in F(v).
To see the optimality, we consider any optimal solution p. We first notice that > 0 and ps < 1.
Actually, Let
N (k E [S]
Du Feahh
p (k [S]
P& p }
(k E [S]
Pk p 3
Then by e TTp p=pp- e pi.. and |p Piall1 S U. we get
Σ
Deak Đk Σ
Dk
Paak
KEN
kEP
Σ Weec Dw)) + Σ KEN
(px:
Deak) Sv
kEP
Hence EKEN - Pt =kEP Pk - -Paak S " Piat so we get 11 > and ps 1.
Next we show that Pk = Plakk 42 S k S - 1. Otherwise we have some 2 < h: < 5 - such that 17
O. If Px p Peak. we define P with that P = P and Ps = Ps + here
19 D 1=0=), while keeping the other components of p same as p. We can that P achieves smaller objective value than p does due bakk beas, which contradicts the optimality of P. If Dr pe then we define with that P.- Pk + E and pt Pt s, here
O < < min(
0<E< < minf Pia-l 2 }. while keeping the other components of p same as p. Similarly, achieves smaller objective value and this implies the contradiction.
Finally, verify the rest components. We introduce the variables di P1 Fea and ds = Ps - Plas. then we get the equivalent reformulation of the inner minimization in F(u):
min diddg s.1.
bsai + beasds
di + ds 0, |dil + dsl S.. v.
The optimal dj and ds are given by - u/2 and v/2 respectively. Hence we proved pz is indeed an optimal solution.
To prove the second case of the claim, the optimal solution for the minimization problem in F(v) given by p", whose components are P 0, vk E [K]. and pi = rk, AK + 1 < < The decision variables p* defined in this way satisfies the constraints of the minimization problem in F(w).
To the optimality, we consider any optimal solution P. We first notice that P > 0 and ōs 1.
19

---

To argue this by contradiction, we suppose the contrary PK+1 = 0. Define the notations NN P and 3 same as before. Then there exists & < K with P3z 0, otherwise we assume that Ps. 0, vk E [K], which implies the contradiction as follows.
K-1 Σ
Psak
Σ KEN
Prak-
Pk
Poak
Here the first inequality is due to the selection of v. the second inequality has been deduced in the first case and the third inequality is from Pk O, Vk. E [K + 1]. So we are able to find À < K with
min Pg 2: P from Pg to PK+I, we can achieve smaller objective value while keeping the feasibility, hence we get the contradiction, which implies that PK+1 > 0 and ps 1. By applying the similar procedures in the first case (moving some probability -
Pg > O. By moving probability
from PK+1 or to ps). we can show that Ps. = Paak = ты for K + 2
Next we prove that Pk тк 0 for k [K] Suppose the contrary is true; that is, Pk > 0 for some k K. Then we are able to apply the same procedures as before which illustrate that
Pk = -Pakk Vi < 1: < S. Provided this, one can verified that the optimal strategy is putting all the rest probability 1 - - Psak k to Ps since beal bau2
beas, and v 2 2 > k= Psak 2
2 > Peak. which implies Paak = and we get contradiction. Hence Pk Tx = 0, Ak E [K]. Finally, we verify the rest two components. We introduce the variables dx+1 = pr+1 1' (k))))
and ds Ps Psas: then we get the equivalent reformulation of the inner minimization F F(u):
min dx-1.d
1sa(K+1)dx+1 + bsasds
dk+1 + ds = Σ Penker |dk+1l + ldsl S v
Σ P
s.t.
sak
The optimal dx+1 and ds are given by 2-2 Plas and z respectively. Hence we proved p* is
an optimal solution.
To prove the third case of the claim, we notice that the optimal solution to minpi eAg b b.P is also feasible for the inner minimization problem in F(v). hence it becomes the optimal solution we desire, which implies F Fw)
obsas at this time. This finishes the proof of our claim.
Our claim directly illustrates that F(v) piecewise-linear function in v with breakpoints (0) - 1]}. Furthermore, based on the provided formulation of F(u), we can compute the difference of value for F(-) between any two adjacent breakpoints, given by
[K Peak : AK E [
F(UK) F(Ux-1) = 2fp + abbas bsak Daak 4K E [S-1] - 1].
where UK = _ Pooks AK E { [S- 1] and uo = 0.
k=
Hence provide Algorithm 2 to compute (11), whose time complexity is O(S log 8) generally and can be reduced to O(S) if the sorted bea is provided.
Proof of Corollary 5.2.1. We can see that Algorithm is in time complexity O(S) if Uge Absa are sorted, which can be done at the Initialization step in Algorithm I with time complexity O(AS log S). So by Theorem 4.5, we get the overall complexity is O (N AS log -1 log 5 AS ASlog S).
20

---

Algorithm 3: Fast algorithm to solve the inner minimization problem in (10)
Input:e, Tsa. and P'a E Ag. Initialization:paa. 0. Sort ba as bsan: S .... S b Deans-
Find the smallest & such that Σ/=1 (Ban,, + +0) 2 1.
Set it Psan:
= Psany + for j - and Psann = Σj= (Pean, +0)
T =b Pa Result: Optimal objective value and optimal solution p.a of the inner minimization problem in (10).
A.2.2 Proof of results in Section 5.2
Proof of Theorem 5.3. From problem (9), we can get the minimization problem for = 2
min Eis
I|p'a.
p'aa 1 ab abaap's
min ||pbal12 2p' p'a +
| 12 ab ab salP
Pia EDs
|| 11 (2pi.. aaba) P'a + | l|
min
Pin EDs
a² beal
2p
abea
ab
PaPs
+ min Eds
Pso
4
2
2
28
ab sa
So suffices to solve minp;, <As
D'a
which can done by Euclidean projec-
tion algorithm (Wang and Carreira Perpinén, 2013) with time complexity O(S S).
Proof of Conollary 5.3.1. The result is the direct consequence of Theorem 4.5 with h2(S) is
O log S), provided by Theorem 5.3.
A.2.3 Proof of results in Section 5.3
Proof of Theorem 5.4. We claim that the Algorithm solves problem (13) with time complexity OSS log S). By expanding the oc-norm, we formulate (13) as the following box constraints problem.
min b P
sa s.t. max 0. O.P -0ee S <p', S min (e. + 0e).
e p' 1, P'_ R
To get the optimal solution, we put the probability on the index where bsa is small as much as possible. Specifically, we assume baal bsa2 <... bsas w.L.o.g,, and assume & is the smallest index such that Σラ (Pins 0) 2 1. We claim that
+0
ifl $iijske 5-
R-
Psaj
1-
(p
+
if
8=1
otherwise
is the optimal solution to the above formulation. To this, suppose is optimal and there is some + 0, so there exists 3 Z k such that ūÚ > S saj. By moving the probability from f7 Pj. we can achieve strictly smaller objective value, which is contradiction with that p is optimal. This gives us an optimal solution with the first & - components coincides p Then get p is indeed optimal by putting the
j E Ik with Pj * Peoy + 0. By above box constraint, get Pj < p eaj
21

---

extra 1 - Et (Peet + 6) probability on P since book S.. ..... beas. The major time complexity of Algorithm 3 is sorting the vector bea E 12,, which is O log S).
Proof of Corollary 5.4.1. 1. As for each E A, we need to sort b.a. which costs (o4.logg S), then we need to solve N A subproblems, which costs ((V.A))) so the whole problem computed in time (AS log S + AS).
B Appendix: Computational Complexity for General Convex Optimzation
To compare our algorithm with general convex optimization algorithm, we use general convex optimization compute the problem (5) and problem (8). The time complexities will be discussed in different situations:
Suppose 1: For general convex optimization problem, problem (5) is equivalent with the following problem:
minimize
subject to
(rea + TT)TP'n ST, Va EA
N
[[()]]
4s S.
#EΣΣ Psas' Paae 5
N
1= dEA'ESS
Y R, D'o E As, ViE [N], Va E A.
By introducing the variables Elass above problem is equivalent with
P'ass Peas Vi E [N]. Va E A, As' E S, the
minimize
N
subject to
Σ
(Tea + Ae) P'a Sn, Va A
N
N
+EΣ taass 5
N 1= HEANES tlass 2 Peas' - Pans' Paas 2 Peos - Peas' Vi [], V A, Vs' E 7 R, D.a E As, to taas E R, Vi [N], Va A. Vs E
There are + NSA + NSA O(NSA) decision variables, and the number of bits in the input is O(1) + + O(NAS) + O(NAS)+ + O(N + O(NAS) O(NAS) = O(NAS). So by (Karmarkar, 1984), the complexity of solving this LP ON A1.5).
We utilize outer bisection, solving (6) directly using convex optimization. Typically, for each fixed a E A and 7, (6) is equivalent with the I.P that
N
minimize
ti
Eas'
N
i=
SES
N1
1
subject to
bT ss Sg
b
N
eea'' 2 Peas' Peas yas' 2 Brag -Pess Vi [N], Vs E Pia E As. toas E R, Vi [N], Vs' S.
There are 2NS decision variables, and the number of bits in the input is O(NS + NS + NS + NS) O(NS). So the complexity of solving this LP is ((N455445) Together with the outer bisection, the total complexity for each Bellman update with é tolerance is 5.555.51..5 log 1t)
22

---

We utilize our nested bisection scheme. solving (9) using the general convex optimization algorithm. Typically, for each fixed a E A. E [N] and a, (9) is equivalent with
minimize
Etaar a DaaP'a
s'ES
subject to E'ags' 2 Psass Psnss Vang Pans' Prass A9 E Pia As. saas' E R, vs' E E.S.
There are 25 decision variables, and the number of bits in the input is -($+++5)= So the complexity of solving this LP OSS4-)). Together with the nested bisection, the total complexity for each Bellman update where the tolerances of bisections are and is ( A log i log 62
Suppose q 2: For general convex optimization problem, problem (5) is equivalent with the following SOCP:
minimize
N TTea + (w) Ax)) " Sn, Va E
subject to
[()]]
As E S.
p 80
Psa
02
aEA E As, [N], Va A.
I, P
There are 1 NSA = O(NSA) decision variables, and the number of constraints are A+1+ D(NAS) So the complexity of solving the SOCP with e-accuracy is
(NSA NA A(S+2))) O(N3.5 43.5 loge -1).
We utilize our outer bisection, solving (6) directly using convex optimization. Typically, for each fixed AA and 7, (6) equivalent with the SOCP that
minimize
8
1
subject to
b
PaPPo Sy
N
1
Σ N
I' P/ < 5
i=1
E R, p'.. E As, E
There are 1 + NS = O(NS) decision variables, and the number of constraints are (NNS). So the complexity of solving the above SOCP with e-accuracy is O(VNS logeloge---.. s2 + + 2(1 N(S + 2)))) 0(N3.5 loge-") ")), hence the total complexity of the Bellman update is O(N3.5 Aloge loge
We utilize our nested bisection scheme, solving problem (9) using the general convex optimization algorithm. For each fixed a E A, E [N] and a, problem (9) equivalent with
minimize
8 a bT P
subject to Ip'o P < P/. E As, deR TR.
There are + 5 = O(S) decision variables, and the number of constraints are OSS). So the complexity of solving the above SOCP with e-accuracy is O(V5 OovSloge-+ 5 + + 2(1 + S))) = O loge hence the total com- plexity of the Bellman update is ONN53.5 log - log loge,)"
Suppose 9 = 0o: convex optimization, which is equivalent with
0o: We can also consider solving the inner problem in (10) using general
minimize subject to
b bap'
P Paas Pras' S 0, Pend P'as 50, Vs' E
P'm. E As,
23

---

Table 1: Comparisons of runtime (average and standard deviation)(sccond) of Bellman updates for all algorithms in L1 norm.
N ==
10
20
30
40
Algorithm
Fast
0.0253 (0.02)
0.1243 (0.02)
0.2794 (0.02)
0.5112 (0.02)
Gurobi
0.039 (0.02)
0.7084 (0.03)
4.3617 (0.05)
22.6013 (0.92)
FOM(3 its)
0.644 (1.0193)
4.0821 (1.05)
12.6608 (1.00)
29.22 (1.18)
Table 2: Comparisons of runtime (average and standard deviation)(seconnd) of Bellman updates for all algorithms in Lz2 norm.
= S
50
60
70
80
Algorithm
Fast
7.3717 (0.11)
10.8293 (0.12)
16.8888 (0.12)
24.6080 (0.07)
Gurobi
5.1669 (0.09)
11.6760 (0.14)
23.6852 (0.34)
44.3305 (0.18)
FOM(3 its)
9.1936 (1.83)
14.5653 (2.01)
21.3651 (1.97)
31.1326 (1.79)
where [N] and A are fixed.
There are S decision variables, and the number of bits in the input is +5)+++5)=000 So the complexity of solving this LP (($1.5). Then the total complexity for each Bellman update ((.5 A)).
C Appendix: Details for Numerical Experiments
We compare our fast algorithm with the state-of-the-a art solver Gurobi with version IreO (Gurobi Optimization, LLC, 2023) and the first-order method of (Grand-Clen Clément and Kroer, 2021a). All experiments are implemented in Python 3.8, and they are run on a 2.3 GHz 4-Core Intel Core i7 CPU with 32 GB 3733 MHz DDR4 main memory. We will release our code to ensure reproducibility. https://gi //gttbb.. com/Chill ChhhCChhllllll (all1ldd/Fast-.. Pat/FaasttBBllaann..e - DRMDP
Our algorithms are tested on some random instances generated by the Generalized Average Reward Non-stationary Environment Test-bench (Garnet MDPs) (Archibald et al., 1995; Bhatnagar al., 2009). The Garnet MDPs are a collection of assessment problems designed to assess the performance of reinforcement learning algorithms in non-stationary environments. It is convenient to construct and implement these problems. We utilize the parameter to regulate the proportion of the next states accessible for each state-action pair (s, a). Following the same setting as Serand-Clément and Kroer, 2021a), we set to be 0.2 and random uniform rewards to be in [0, 10). The discount factor is fixed at 0.8, and parameter p' p!....,p' p', considering N small random (Garnet) perturbations around the nominal kernel p". We set parameter in Proposition 4.1 to be Vr,.A.
0.1. For each MDP instance, we generate the sampled kernels
To test the speed of Bellman update, we run the random instances 50 times for all the algorithms, and show the average time of them in tables 1, and 3. We see that algorithm performs better than Gurobi and the first-order method. When the states number increases, the running time of our algorithm keeps small standard deviation.
We also compare the speed value iteration for all algorithms using the same convergence criteria: || - 0" ae < 2Ac(1 A)-1, which follows Chann-Clémmet and Kroer, 2021a). The results shown in tables 4, and 6. The runtimes that exceed 4000s for L₁ and Loc case or exceed 10000s for
Table 3: Comparisons of runtime (average and standard deviation) (millisecond) of Bellman updates for all algorithms in Tog
- A
10
20
30
40
Algorithm
Fast
0.06 (0.1)
0.21 (0.1)
0.47 (0.2)
0.80 (0.4)
Gurobi
0.92 (0.1)
7.33 (0.8)
26.50 (3.8)
65.22 (7.9)
FOM
144.92 (981.1)
166.19 (966.5)
237.24 (978.2)
366.74 (976.6)
24

---

L2 case will be shown as "-" We can see from tables 4, 5 and 6 that our algorithm always performs better than Gurobi and the first-order method. We point out that our algorithm generally becomes better as the state number increases.
Table 4: Runtime (second) of value iteration for all algorithms in L norm.
N = A
10
20
30
40
Algorithm
Fast
3.6164
34.0930
116.4698
281.9007
Gurobi FOM
5.7131
202.0723
1967.9552
202.4039
3578.1377
Table Runtime (second) of value iteration for all algorithms L2 norm.
SS=
40
50
60
70
Algorithm
Fast Gurobi FOM
2198.5413
4847.4994
7223.7663
12387.5601
2543.9588
4861.4114
7949.4562
22096.7578
Table 6: Runtime (second) of value iteration for all algorithms in Lo norm.
N = A
10
20
30
40
Algorithm
Fast
0.0045
0,0348
0.1371
0.3181
Gurobi
0.0907
0.8746
3,9243
11,5589
FOM
16.4344
97.44293
1131.5260
3933.8585
25


Reviews:
[{"review_id": "cYJJzCmu2d", "reviewer": "Reviewer_NtRj", "summary": "The paper focuses on the computational complexity of Wasserstein Distributionally Robust MDPs with Lp norm. By decomposing the calculation of Bellman updates to smaller subproblems, the algorithm can achieve linear complexity in the number of actions and kernels, and quasi-linear complexity in the number of states. ", "questions": "* The paper follows the common rectangularity assumption. However, it is known that the rectangularity assumption would lead to overly conservative policies. Can the proposed method easily extend to non-rectangular settings?\n* How will the problem decomposition in section 4 affect the resulting value of the Bellman update? In other words, will there be tradeoffs between the computational complexity and the accuracy of the optimal value introduced by the proposed decomposition?\n* Considering that the following paper also focuses on the model-based distributionally robust RL setting. It would be interesting to also compare with it.\n    * Shi, Laixi, and Yuejie Chi. \"Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity.\" arXiv preprint arXiv:2208.05767 (2022).\n", "limitations": "The limitation of the proposed method is not fully discussed in the paper. ", "rating": 5, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The proposed method shows state-of-the-art complexity results. \n\n**DISCLAIMER:**\n I have not checked the proof thoroughly and cannot verify the correctness of the theorems.", "weaknesses": "* The experiments are carried out on very simple, randomly generated distributionally robust MDP. While this approach provides a controlled setting for their work, it risks oversimplifying the problem and limiting the potential real-world applicability of the proposed methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your comments and your time to review our paper!\n\n**Weaknesses**\n\n1. The experiments are carried out on very simple, randomly generated distributionally robust MDP. While this approach provides a controlled setting for their work, it risks oversimplifying the problem and limiting the potential real-world applicability of the proposed methods.\n\nWe are sorry for being unclear on the purpose of our experimental setup. This paper focuses on algorithmic development and computational complexity for solving distributionally robust MDPs. Therefore, randomly generated distributionally robust MDPs are perfect for our purpose, as additional problem-specific structures should not be considered for the sake of generality. Moreover, these randomly generated distributionally robust MDPs are more computationally challenging than stylized problems as their transitional kernels are dense, and so they are used to test the performance of various algorithms in our experiments. \n\n**Questions**\n\n1. The paper follows the common rectangularity assumption. However, it is known that the rectangularity assumption would lead to overly conservative policies. Can the proposed method easily extend to non-rectangular settings?\n\nThank you for raising this interesting question. Solving the general robust MDP problem with nonrectangular ambiguity sets is strongly NP-hard and intractable [ii]. Moreover, Bellman equations for nonrectangular robust MDPs do not exist in general, so our approach could not be extended to non-rectangular settings. However, the rectangularity assumption can always be satisfied via outer approximation [i]. Therefore, most research focuses on rectangular robust MDPs [i-iii], such as $s$-rectangular and $s,a$-rectangular robust MDPs, and so does the research in distributionally robust MDPs [iv-v]. In this paper, we adopt the $s$-rectangular setting, which is less conservative than the $s,a$-rectangular case. \n\n2. How will the problem decomposition in section 4 affect the resulting value of the Bellman update? In other words, will there be tradeoffs between the computational complexity and the accuracy of the optimal value introduced by the proposed decomposition?\n\nWe are sorry for the confusion caused. By exploiting the specific problem structure of the optimization problem for the distributionally robust Bellman update, the proposed algorithm is exact (up to user-specified tolerances on the accuracy). Therefore, we can compute the Bellman update exactly up to an arbitrarily small tolerance. We will clarify this point in our next version of manuscript.\n\n3. Considering that the following paper also focuses on the model-based distributionally robust RL setting. It would be interesting to also compare with it.\nShi, Laixi, and Yuejie Chi. \"Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity.\" arXiv preprint arXiv:2208.05767 (2022).\n\nThank you for bringing up this interesting reference paper. We would like to clarify that our paper focuses on the computations and time complexity of distributionally robust MDPs, but the reference paper focuses on learning and sample complexity. Therefore, the purposes of both papers are complementary to each other. Moreover, by our definition in Section 1 (which is also adopted in [i-v]), this reference paper considers the settings of robust MDPs rather than distributionally robust MDPs; thus, the problem settings are different in the two papers. We are sorry for the confusion caused, and we will cite this paper and clarify this point in our next version of manuscript.\n\n**Limitations**\n\n1. The limitation of the proposed method is not fully discussed in the paper.\n\nWe apologize for missing this discussion. Our method does not consider the settings of continuous state and action spaces. Besides, we focus on discounted MDPs, and it is not obvious that our proposed algorithms could be extended to the case of optimizing average expected return. We will provide the discussion in the next version of manuscript.\n\n\n\n\n\n[i] A. Nilim and L. El Ghaoui, Robust control of Markov decision processes with uncertain transition matrices, 2005.\n\n[ii] W. Wiesemann, D. Kuhn, and B. Rustem, Robust Markov decision processes, 2013.\n\n[iii] G. N. Iyengar, Robust dynamic programming, 2005.\n\n[iv] I. Yang, A convex optimization approach to distributionally robust Markov decision processes with Wasserstein distance, 2017.\n\n[v] J. Grand-Clement and C. Kroer, First-order methods for Wasserstein distributionally robust MDP, 2021.\n\n[vi] L. Shi and Y. Chi, Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity, 2022."}, {"review_id": "YDj6Ptmhg7", "reviewer": "Reviewer_SCaM", "summary": "The paper studies Wasserstein Distributionally robust MDPs (WDRMDP) problem when the ambiguity set is defined based on Wasserstein distance and rectangular. It is then well known that the optimal policy can be computed by solving Bellman equations, which have the form of distributionally robust linear programs. The authors then focus on solving the Bellman updates. The main idea is to transform the robust linear problem into parameterized problems so that the Bellman updates can be done via bisection. The paper then provides some new algorithms running in polynomial time complexity and claims that their complexities are smaller than those from prior works. Experiments show that their algorithm performs better than some baselines. ", "questions": "- What would happen if you solve (5) and (8) by convex optimization? How does this affect the performance,  in terms of both theoretical complexities and numerical experiments?\n- In Theorem 4.4 and Corollary 5.2.1, what is the quality of the returned solutions? Does the bisection return an optimal solution or just a near-optimal one? Does the quality of the returned solution depend on $\\epsilon_1$ and $\\epsilon_2$ \n- What are the technical distinctions between the proposed algorithm and the bisection method presented in [1]\n- Can the results be extended to the robust MDP with (s)-rectangular ambiguity sets?\n\n[1] Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Fast bellman updates for robust mdps. In International Conference on Machine Learning, pages 1979–1988. PMLR, 2018.", "limitations": "I do not see any negative societal impact from this work.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The problem under investigation holds significant importance. Over the past few decades, robust and distributionally robust Markov Decision Processes (MDPs) have garnered considerable attention. Resolving robust MDPs typically entails addressing computationally expensive minimax problems, necessitating the development of efficient algorithms. In this regard, the paper appears to succeed by presenting novel and efficient algorithms for solving the problem. The proposed algorithms exhibit sound and good time complexities. The experimental results seem convincing, effectively showcasing the efficiency of the new algorithm in comparison to prior algorithms.", "weaknesses": "The WDRMDP problem itself is not novel, and the main focus of the paper centers around solving the Bellman equation, which is a well-studied distributionally robust linear program extensively explored in the fields of optimization and mathematical programming. As a result, certain results claimed in the paper, such as Propositions 4.1 and 4.3, appear rather obvious. Additionally, it seems that some of the problems, such as those presented in Equations (11) and (13), bear a strong resemblance to those examined in a previous work cited as [1]. It is possible that similar techniques have been employed to derive the bisection algorithm and determine the running time complexities. In this regard, the technical contributions seem quite incremental.\n\nSome theorems and corollaries are not properly stated. Please see the Question section. \n\nThe experiments conducted in the paper do not achieve a satisfactory level of validation. The comparisons made are limited to just two basic baselines, which fail to provide sufficient justification for the introduction of the bisection algorithm. For example, the paper suggests that (5)  and (8) can be solved directly using convex optimization, rather than employing bisection. Therefore, it would be necessary to include a comparison between solving (5) and (8) through convex optimization and the proposed approach.\n\n[1] Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Fast bellman updates for robust mdps. In International Conference on Machine Learning, pages 1979–1988. PMLR, 2018.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks a lot for your encouraging comments and your time to review our paper!\n\n**Weaknesses**\n\n1. The WDRMDP problem itself is not novel...well-studied distributionally robust linear program extensively explored in the fields of optimization and mathematical programming...Propositions 4.1 and 4.3, appear rather obvious...Equations (11) and (13), bear a strong resemblance to those examined in a previous work cited as [i]. It is possible that similar techniques have been employed...the technical contributions seem quite incremental. \n\nWe would like to clarify that, although distributionally robust linear programs are extensively explored in the fields of optimization and mathematical programming, our contributions focus on the exploitation of the specific problem structure of distributionally robust Bellman updates. In particular, the proposed decomposition scheme does not work for generic distributionally robust linear programs. \n\nIn the same spirit of using first order methods for solving both the robust MDPs (RMDPs) and the distributionally robust MDPs (DRMDPs) [ii-iii], while we agree that the proposed algorithm belongs to the class of bisection methods that solve RMDPs (such as [i]), we would like to highlight the differences between our work and [i]. (a) Firstly, we consider DRMDPs, which is a generalization of RMDPs, and so we are different from [i] in terms of problem settings. (b) Secondly, the resemblance of our bisection step and [i] only occurs in lines 200 - 213, which is less than 1/2 page in this paper. In particular, the second bisection step in the proposed nested bisection method and Proposition 4.3 are novel in this paper; without them, it is not easy to design algorithm(s) that has a time complexity that is almost linear in the number of expected kernels, which do not exist in RMDPs and [i]. In our opinion, this is one of the key factors that our proposed algorithm outperforms state-of-the-art algorithms, such as [iii]. (c) Finally, the subproblems, such as (11), are different from the case of RMDPs, and so it requires new theoretical results and algorithms for fast computations. Thank you for your question and we will clarify this point in our next version of the manuscript.\n\n2. Some theorems and corollaries are not properly stated. Please see the Question section. Also, The experiments conducted in the paper do not achieve a satisfactory level of validation. \n\nPlease find our answers to your questions below.\n\n**Questions**\n\n1. What would happen if you solve (5) and (8) by convex optimization? How does this affect the performance, in terms of both theoretical complexities and numerical experiments?\n\nSorry for the confusion caused. In our experiments, one of our baselines is Gurobi, which is a state-of-the-art commercial solver for solving convex programs, such as (5) and (8). In Section 6, we provide the comparisons between our proposed algorithms and Gurobi (i.e. convex optimization solver) on the runtimes of solving (5), and our algorithm outperforms Gurobi. In terms of solving (8), below please find our additional results for this rebuttal. As we can see, the results are consistent with our conclusion in Section 6. \n\nS=A=N=50, Gurobi (q=1): 722ms, fast (q=1): 6.10ms, Gurobi (q=2): 2959ms, fast (q=2): 32.61ms \n\nS=A=N=70, Gurobi (q=1): 998ms , fast (q=1): 6.26ms, Gurobi (q=2): 5372ms, fast (q=2): 35.21ms\n\nS=A=N=90, Gurobi (q=1): 1292ms, fast (q=1): 6.91ms, Gurobi (q=2): 8861ms, fast (q=2): 35.82ms\n\nIn terms of theoretical complexities, from [iv-v], the time complexities for solving (5) using convex optimization solver are $\\mathcal{O}(N^{4.5}S^{4.5}A^{4.5})$ and $\\mathcal{O}(N^{3.5}S^{3.5}A^{3.5}\\log \\epsilon^{-1})$ for $q=1$ and $q=2$, respectively. The time complexities for solving (8) using convex optimization solver are $\\mathcal{O}(N^{4.5}S^{4.5})$ and $\\mathcal{O}(N^{3.5}S^{3.5}\\log \\epsilon^{-1})$ for $q=1$ and $q=2$, respectively. Therefore, the theoretical complexities of the proposed algorithms, as stated in Section 5, are orders of magnitudes lower than using generic convex optimization solvers.\n\n2.  In Theorem 4.4 and Corollary 5.2.1, what is the quality of the returned solutions? Does the bisection return an optimal solution or just a near-optimal one? Does the quality of the returned solution depend on $\\epsilon_1$ and $\\epsilon_2$?\n\nThe proposed algorithm is exact (up to tolerances $\\epsilon_1$ and $\\epsilon_2$). In particular, we have proved a new result on the error bound for the Bellman update, which is $\\mathcal{O}(\\epsilon_1+\\epsilon_2)$. Since the complexities of the proposed algorithms are in $\\mathcal{O}(\\log\\epsilon_1^{-1}\\log\\epsilon_2^{-1})$, one can compute a highly accurate solution with very small $\\epsilon_1$ and $\\epsilon_2$. We will provide this result in our next version of manuscript.\n\n3. What are the technical distinctions between the proposed algorithm and the bisection method presented in [i].\n\nPlease refer to the **Weaknesses**\n\n4.  Can the results be extended to the robust MDP with (s)-rectangular ambiguity sets?\n\nThank you for your insightful question.  The problem of interest, s-rectangular distributionally robust MDPs, is in fact a generalization of s-rectangular robust MDPs. Thus, distributionally robust MDPs, can recover (instead of extend) to s-rectangular robust MDPs. By setting the radius of the Wasserstein ball to be $\\infty$ (or a large enough number, such as the diameter of $(\\Delta_S)^A$), the ambiguity set contains all possible probability distributions. It is known that the worst-case distribution is a Dirac distribution, hence the distributionally robust MDP degenerates into robust MDP.\n\n[i] Fast Bellman updates for robust MDPs. 2018.\n\n[ii] Scalable first-order methods for robust MDPs. 2021.\n\n[iii] First-order methods for Wasserstein distributionally robust MDP. 2021.\n\n[iv] A new polynomial-time algorithm for linear programming. 1984.\n\n[v] Applications of second-order cone programming. 1998."}, {"review_id": "YO9aeirn5m", "reviewer": "Reviewer_JxmA", "summary": " - The paper proposes a computationally efficient solution framework to solve the distributionally robust Bellman operator induced by Wasserstein ambiguity sets, which is critical in performing distributionally robust value iteration algorithms.\n- The proposed framework features a novel decomposition of the optimization problem involved in DR Bellman updates, based on which the updates are reduced to solving small subproblems.\n- The paper shows the overall complexity of the proposed framework is quasi-linear in $S$ and $A$ when considering Wasserstein distance with $L_1$, $L_2$, or $L_\\infty$ norm, proving the advantage of the proposed framework. The theory is further supported by numerical experiments.\n", "questions": "- Regarding the formulation of distributionally robust MDPs, the reference distribution of the ambiguity set $\\nu_s$ is defined as $\\mu_s(\\cdot) = \\frac{1}{N}\\sum_{i=1}^N\\delta_{\\hat{\\boldsymbol{p}}_s^i}(\\cdot)$ with $\\hat{\\boldsymbol{p}}^i$ being $N$ empirical transition kernels. What do these empirical distributions corresponds to in the real-world application of reinforcement learning algorithms? Typically, we have a single empirical transition kernel estimated from historical data. How can we interpret these $N$ empirical transition kernels?", "limitations": "The limitation is stated in previous parts.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The problem of solving distributionally robust Bellman operator in a computationally efficient manner is a critical problem in DRMDP literatures. The presentation of the paper is also clear.", "weaknesses": "-  The framework seems to rely on the specific structure of the Wasserstein ambiguity set with specific reference distribution of the Wasserstein ball, which makes the application of the framework limited. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your positive comments and your time to review our paper!\n\n**Weaknesses**\n\n1. The framework seems to rely on the specific structure of the Wasserstein ambiguity set with the specific reference distribution of the Wasserstein ball, which makes the application of the framework limited.\n\nWe would like to clarify that Wasserstein ambiguity sets have many nice properties such as consistency in optimality and finite-sample bounds [i], and thus provides a general and powerful metric for distributionally robust MDPs. In recent years, Wasserstein ambiguity sets are widely used in stochastic optimization [iv] and multistage stochastic optimization [v]. Compared to the $\\phi$-divergence, the Wasserstein structure does not suffer from the absolutely continuous restriction and addresses how close two points in the support are to each other. More advantages of the Wasserstein structure are provided in [iii]. \n\n**Questions**\n\n1. Regarding the formulation of distributionally robust MDPs, the reference distribution of the ambiguity set is defined with $N$ empirical transition kernels. What do these empirical distributions correspond to in the real-world application of reinforcement learning algorithms? Typically, we have a single empirical transition kernel estimated from historical data. How can we interpret these $N$ empirical transition kernels?\n\nThank you for your very insightful question! You are absolutely correct; the common practice is to estimate just one empirical transition kernel. However, in our context, one could/should consider the Bayesian point of view (instead of frequentist), see e.g. [vi]; That is, one could use historical data to estimate the posterior distribution of the transition kernel. Then the $N$ empirical transition kernels could be sampled from the posterior distribution.\n\n\n[i] PM Esfahani and D Kuhn. Data-driven distributionally robust optimization using the Wasserstein metric: Performance guarantees and tractable reformulations. 2018.\n\n[ii] I. Yang. A convex optimization approach to distributionally robust Markov decision processes with Wasserstein distance. 2017.\n\n[iii] R. Gao, A. Kleywegt. Distributionally robust stochastic optimization with Wasserstein distance. 2022.\n\n[iv] D. Wozabal. Robustifying convex risk measures for linear portfolios: A nonparametric approach. 2014.\n\n[v] G.C. Pflug and A. Pichler. Multistage stochastic optimization. 2014.\n\n[vi] R. H. Russel, M Petrik. Beyond confidence regions: tight Bayesian ambiguity sets for robust MDPs. 2019."}, {"review_id": "zcjD0AqQvf", "reviewer": "Reviewer_zZNP", "summary": "* tailored algorithms for solving Wasserstein distributionally robust MDPs and fast implementations with $L_1$, $L_2$, $L_\\infty$-based Wasserstein distance.", "questions": "Model : can your algorithms be extended to handle the case of both transitions and reward uncertainty?\n\nTheorem 4.4: how do the $\\epsilon_1$ and $\\epsilon_2$ error propagate into the overall error of $v_t$, the t-th iterate of value iteration?\n\nWasserstein balls: how is the radius $\\theta$ chosen in practice, given some dataset? ", "limitations": "I do not see any serious limitations in this work. The only downside is the somewhat narrow scope in terms of outreach.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* Fastest algorithms (in terms of dependency on N, S, and A - the number of kernels/states/actions) for solving Wasserstein distributionally robust MDPs.\n* Paper is easy to follow, short and straight-to-the-point.", "weaknesses": "* This is a rather niche application of principles that have been developed in several other papers for robust MDPs [1,2,3].\n\n[1] Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Fast bellman updates for robust mdps. In  International Conference on Machine Learning, pages 1979–1988. PMLR, 2018.\n\n[2] Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Partial policy iteration for l1-robust markov\n decision processes. The Journal of Machine Learning Research, 22(1):12612–12657, 2021.\n\n[3]  Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Robust phi-divergence mdps. arXiv preprint\n arXiv:2205.14202, 2022.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you very much for your positive comments and for taking the time to read our manuscript!\n\n**Weaknesses**\n\n1. This is a rather niche application of principles that have been developed in several other papers for robust MDPs.\n\nIn the same spirit of using first order methods for solving both the robust MDPs (RMDPs) and the distributionally robust MDPs (DRMDPs) [iv-v], while we agree that the proposed algorithm belongs to the class of bisection methods that solve RMDPs [i-iii], we would like to highlight the differences between our work and [i-iii]. (a) Firstly, we consider DRMDPs, which is a generalization of RMDPs, and so we are different from [i-iii] in terms of problem settings. (b) Secondly, the resemblance of our bisection step and [i-iii] only occurs in lines 200 - 213, which is less than 1/2 page in this paper. In particular, the second bisection step in the proposed nested bisection method and Proposition 4.3 are novel in this paper; without them, it is not easy to design algorithm(s) that have time complexity that is almost linear in the number of expected kernels, which do not exist in RMDPs [i-iii]. In our opinion, this is one of the key factors that our proposed algorithm outperforms state-of-the-art algorithms, such as [v]. (c) Finally, the subproblems, such as (11), are different from the case of RMDPs, and so it requires new theoretical results and algorithms for fast computations. Thank you for your question and we will clarify this point in our next version of manuscript.\n\n**Questions**\n\n1. Model : can your algorithms be extended to handle the case of both transitions and reward uncertainty?\n\nThank you for your great question. Yes, our algorithm can be extended to reward uncertainty, because the decomposition technique can be seamlessly applied to this extension. We will clarify this point in our next version of the manuscript.\n\n2. Theorem 4.4: how do the $\\epsilon_1$ and $\\epsilon_2$ error propagate into the overall error of $v_t$, the t-th iterate of value iteration? \n\nThanks again for your insightful question. We have proved a new result on the error bound for the Bellman update, which is $\\mathcal{O}(\\epsilon_1+\\epsilon_2)$. Since the complexities of the proposed algorithms are in $\\mathcal{O}(\\log\\epsilon_1^{-1}\\log\\epsilon_2^{-1})$, one can compute a highly accurate solution with very small $\\epsilon_1$ and $\\epsilon_2$. We will provide this result in our next version of the manuscript.\n\n3. Wasserstein balls: how is the radius $\\theta$ chosen in practice, given some dataset?\n\nIn practice, $\\theta$ is often chosen via cross-validation, e.g. [vi]. However, one could also derive theoretical bounds such that the ambiguity set contains the unknown true distribution with high confidence [vii]. \n\n[i] Fast Bellman updates for robust MDPs. 2018.\n\n[ii] Partial policy iteration for l1-robust Markov decision processes. 2021.\n\n[iii] Robust phi-divergence MDPs. 2022.\n\n[iv] Scalable first-order methods for robust MDPs. 2021.\n\n[v] First-order methods for Wasserstein distributionally robust MDP. 2021.\n\n[vi] Distributionally robust inverse covariance estimation: the Wasserstein shrinkage estimator, 2020.\n\n[vii] Data-driven distributionally robust optimization using the Wasserstein metric: Performance guarantees and tractable reformulations. 2018."}]

Author rebuttal:

Meta-review: Robust MDP (RMDP) and robust RL are now a popular approach for modeling and solving robust control problems. While the basic formulation is straight forward, it typically hides the complexity of the inner optimization over the uncertainty set. In fact, many robust RL algorithms are not scalable to large state/action spaces due to the complexity of this optimization. This paper presents a computationally efficient solution framework for solving distributionally robust MDPs with Wasserstein ambiguity sets. 

I have read this paper, reviewer comments, and the authors’ rebuttal carefully. The paper is well written and technically solid. The contribution is novel and will be of interest and use to the robust RL community. Reviewers had some questions about the experiments, and I the authors have given clear response those. 

For the final version, please make all the changes you have promised during the rebuttal, and update the paper.

Regards,

AC
Decision: Accept (poster)