Title: Fast Bellman Updates for Wasserstein Distributionally Robust MDPs
Venue: NeurIPS 2023
Accepted: yes
Keywords: Markov decision processes

distributionally robust optimization
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.

Fast Bellman Updates for Wasserstein Distributionally Robust MDPs

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

1 Introduction

Markov Decision Processes (MDPs) provide a flexible and powerful modeling framework for sequential 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, it is well-known that the optimal policy of MDP is sensitive to estimation errors because of the dynamic nature of the problem. In particular, small errors in estimating transition kernels could lead to catastrophic failures in practice (Iyengar, 2005; Nilim and El Ghaoui, 2005; Le Tallec, 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 to 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 a 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

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 is 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 or medium size of problems. In recent years, many efficient algorithms are proposed for solving robust MDPs to address this issue (Iyengar, 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é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 (Grand-Clément and Kroer, 2021a) 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, we reformulate the optimization problem of the distributionally 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 of 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; Grand-Clément 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 4, we illustrate our decomposition scheme for computing the distributionally robust Bellman updates when the Wasserstein distance in the ambiguity set is chosen to be L q norm. In Section 5, customized fast algorithms and their time complexities are particularly provided for the three most well-known cases: L 1 , L 2 , and L ∞ . The numerical experiments in Section 6 verify that our methods outperform other state-of-the-art algorithms.

Notations We use boldface lowercase and uppercase letters to denote vectors and matrices, respectively. We denote by [ N ] the set of natural numbers from 1 to N ; that is, [ N ] = { 1 , 2 , . . . , N } . The vector e is denoted as the vector of all ones while the dimension depends on the context. The vector e j represents the vector of all zeros except the j -th component is one. Probability simplex is denoted as ∆ N = { x ∈ R N : e ⊤ x = 1 , x ≥ 0 } . The smallest component of a vector x is denoted by min { x } . Given a set X , we call the set of all the Borel probability measures as P ( X ) . The general L q norm for a vector x ∈ R N is denoted by ∥ x ∥ q = ( ∑ N i =1 | x i | q ) 1 /q for 1 ≤ q < ∞ , and the L ∞ norm is ∥ x ∥ ∞ = max i ∈ [ N ] | x i | . The Dirac distribution at x ∈ R N is denoted by δ x . 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 )-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 s - and ( s, a )-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 -rectangularity (Goyal and Grand-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 et al. (2018) propose fast algorithms for solving s -rectangular robust MDPs with weighted L 1 norm ambiguity sets and propose to compute robust Bellman updates using the combination of bisection method and homotopy method, and Behzadian et al. (2021) further extend the solution scheme for the unweighted L ∞ norm cases (Delgado et al., 2016). A first order method is introduced in (Grand-Clément and Kroer, 2021b) to approximate robust Bellman updates for robust MDPs with ellipsoidal and Kullback-Leibler (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 (Iyengar, 2005) and it has better empirical performance compared to robust value iteration. Kaufman and Schaefer (2013) propose a modified policy iteration for ( s, a )-rectangular robust MDPs. Ho et al. (2021) introduce partial policy iteration for s - and ( s, a )-rectangular robust MDPs with L 1 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 et al. (2014) apply value function approximation and propose robust approximate dynamic programming for robust MDPs.

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 ( s, a )-rectangular robust MDPs. A function approximation approach is proposed in (Panaganti et al., 2022) for ( s, a )-rectangular robust MDPs with offline datasets. The combination of robust RL and deep RL is 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 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 it 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 statistical-distance-based ambiguity sets together for distributionally robust MDPs. Derman and Mannor (2020) study the connection between Wasserstein distributionally robust MDPs and regularization. Abdullah et al. (2019) propose a reinforcement learning framework with distributionally robustiness. As mentioned before, however, (Grand-Clément and Kroer, 2021a) is the only work that focuses on the algorithmic development of solving Wasserstein distributionally robust MDPs, which is also the focus of this paper.

3 Preliminaries

A distributionally robust MDP is a tuple ( S , A , p 0 , r , λ, M ) , where S and A are the sets of states and actions, respectively. We assume that the state space S = { 1 , . . . , S } and the action space A = { 1 , . . . , A } are finite. The initial state s 0 follows a given initial distribution p 0 ∈ ∆ S . When the decision maker takes an action a ∈ A in state s ∈ S , the MDP will transit to the next state randomly according to the probability distribution p sa ∈ ∆ S , which is assumed to be unknown; the transition from state s ∈ S under action a ∈ A to the next state s ′ ∈ S will induce the reward r sas ′ . We use the vector r sa ∈ R S to represent the connection of all rewards given state s and action a . We denote by λ ∈ (0 , 1) the discount factor. In distributionally robust MDPs, the transition kernel Π s ∈S ,a ∈A p sa ∈ (∆ S ) S × A is assumed to be uncertain, and it is governed by an unknown probability distribution µ which is assumed to reside in a prescribed ambiguity set M that is calibrated from historical data. Note that we also use the notation p s = Π a ∈A p sa ∈ (∆ S ) A for the transition kernel given the state s ∈ S .

The goal of a distributionally robust MDP is to maximize the worst-case expected return; that is,

Here, the decision variable π ∈ (∆ A ) S is called a (randomized) policy where for any given s ∈ S , its sub-vector π s ∈ ∆ A represents the probability distribution of taking action a ∈ A in state s .

It is worth mentioning that the above formulation of distributionally robust MDPs is a generalization of both robust MDPs and standard MDPs. By specifying the ambiguity set M to be P ( U ) where U is the ambiguity set of a robust MDP, the problem (1) is equivalent to the robust MDP with ambiguity set U . If the set U is a singleton, then problem (1) with M = P ( U ) is equivalent to the standard MDP with transition kernel from the singleton U .

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 et al., 2013). In particular, we consider the following s -rectangular q -Wasserstein ambiguity set (Grand-Clément and Kroer, 2021a)

where M q s is the marginal ambiguity set for the distribution µ s that governs the transition kernel at the state s ∈ S . The superscript q indicates that M q s is defined by q -Wasserstein distance; that is,

where ν s is the reference distribution at the center of the Wasserstein ball and

for 1 ≤ q < ∞ , where Γ( µ s , ν s ) is the set of all the couplings of µ s and ν s (Ambrosio and Nicola Gigli, 2005), and d q is induced by the q -norm d q ( x, y ) = ∥ x -y ∥ q throughout this work. We define the ∞ -Wasserstein distance W ∞ ( µ s , ν s ) by taking q →∞ .

Weadopt the usual setting in distributionally robust data-driven optimization where ν s = 1 N ∑ N i =1 δ ˆ p i s is set to be the empirical distribution and estimated by the samples { ˆ p i s } N i =1 where ˆ p i s ∈ (∆ S ) A for every i ∈ [ N ] . Then, the optimal value function of (1), v ⋆ , satisfies the following distributionally robust Bellman equation (Yu and Xu, 2015; Yang, 2017; Grand-Clément and Kroer, 2021a)

where B q s is the set of expected kernels (Yang, 2017; Bertsimas et al., 2018; Xie, 2020) of the following forms

4 The Decomposition Algorithm for Distributionally Robust Bellman Update

In this section, we focus on computing the distributionally robust Bellman update, which is the most fundamental operator for solving the distributionally robust Bellman equation (2); we define the distributionally robust Bellman operator T ( · ) where for any v ∈ R S ,

Given the above definition of T ( · ) , 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; Grand-Clément and Kroer, 2021a). For any initial guess v 0 ∈ R S , we have lim t →∞ v t → v ⋆ where v t = T ( v t -1 ) for t = 1 , 2 , . . . and v ⋆ satisfies (2). Therefore, the efficiency of evaluating [ T ( v )] s is crucial to the computation of solving distributionally robust MDPs.

However, computing [ T ( v )] s using generic convex optimization solvers is much more computationally demanding compared to the case of classical MDPs, which only have time complexity O ( SA ) for computing their Bellman updates. In this section, we 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 on the case of q ∈ [1 , ∞ ) in Section 4.1 and then consider the case of q = ∞ in Section 4.2, where q indicates the type of Wasserstein distance used in the ambiguity set. As we will show later, combined with the customized fast algorithms in Section 5 for solving the subproblems, one can compute [ T ( v )] s in time complexity that is quasi-linear in S and linear in A and N . All the proofs of propositions and theorems in this section are provided in Appendix A.1.

4.1 Nested bisection method for q -Wasserstein ambiguity set with q ∈ [1 , ∞ )

We consider the optimization problem in (4) for the case where q ∈ [1 , ∞ ) . By applying the minimax theorem, we obtain the following result.

Proposition 4.1. Consider the Bellman updates (4) with q ∈ [1 , ∞ ) . Then,

The above proposition reformulates the maximin optimization problem in (4) into a convex minimization problem, which can be solved by commercial convex optimization solvers. Moreover, problem (5) can be solved via bisection on γ ; that is, we seek for the lowest possible γ such that 1 N ∑ N i =1 ( r sa + λ v ) ⊤ p i sa ≤ γ for each a ∈ A while p i sa satisfies the other constraints in (5), for every a ∈ A . To this end, we introduce the following subproblem

By setting b sa = r sa + λ v , the above problem P ( { ˆ p i sa } N i =1 ; b sa , γ ) allows us to seek for decision variable { p i sa } N i =1 that satisfies the first and the last line of constraints in (5) while minimizing the 'budget' used in the second line of constraint in (5). In particular, for any fixed γ ′ ∈ R , we distinguish among the following two cases:

(a) If ∑ a ∈A P ( { ˆ p i sa } N i =1 ; b sa , γ ′ ) ≤ θ q , then γ ′ is feasible for (5) and it is an upper bound of the optimal objective value in (5).

(b) If ∑ a ∈A P ( { ˆ p i sa } N i =1 ; b sa , γ ′ ) > θ q , then γ ′ must be infeasible for problem (5) and so it is a lower bound of the optimal objective value in (5).

To apply the bisection method, we derive the initial upper and lower bounds of the optimal γ ⋆ in (5).

Proposition 4.2. Consider the distributionally robust Bellman update (5) . The optimal objective value γ ⋆ is bounded by

As opposed to (5), the size of problem (6) is independent of A . By using the aforementioned bisection method, the overall time complexity is now only linear in A since we only need to solve (6) O ( A log ϵ 1 ) times, where ϵ 1 is the tolerance for the bisection method.

While one can solve problem (6) using off-the-shelf solvers and enjoy the reduced complexity by using the above bisection method, this subproblem itself turns out to be another structural optimization problem. As we will show in the following proposition, by applying duality on (6), the reformulation could be further decomposed via another bisection method on the dual variable.

Proposition 4.3. Consider the problem P ( { ˆ p i sa } N i =1 ; b sa , γ ) in (6) . If γ > min { b sa } , then

where ¯ α = max i ∈ [ N ] ∥ ∥ e j -ˆ p i sa ∥ ∥ q q / ( γ -min { b sa } ) and

where j ∈ arg min s ′ ∈S b sas ′ .

Notice that we focus on the case where γ > min { b sa } since min { b sa } is not larger than the lower bound of γ in (7), and the case γ = min { b sa } will not occur because γ 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 α in (8), the subproblems (9) can be solved separately. Therefore, by applying the bisection method on α in problem (8), one can naturally decompose the problem (8) into N smaller problems that have S variables and S +1 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 of the algorithm is provided in the Appendix A.1.

Theorem 4.4. Suppose γ ′ is the value returned by Algorithm 1, and γ ⋆ be the optimal value of (5) . With the inputs provided in Algorithm 1 and user-specified tolerances ϵ 1 , ϵ 2 > 0 , we have

Theorem 4.5. Algorithm 1 computes (5) in time O ( h q ( S ) NA log ϵ -1 1 log ϵ -1 2 + AS ) , where h q ( S ) is the time complexity for solving (9) .

As shown in the above theorem, the proposed nested bisection method has a time complexity that is linear in both A and N , but the overall complexity depends on the subproblem (9) where its complexity depends on the choice of q and the number of states S . In Section 5, we will derive h 1 ( S ) and h 2 ( 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.

4.2 Decomposition scheme for ∞ -Wasserstein ambiguity set

As opposed to the case where q ∈ [1 , ∞ ) , for ∞ -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 = ∞ . Then,

The above reformulation requires solving NA 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 to Section 4.1, we will discuss the complexity of solving the inner minimization in the next section.

5 Efficient Algorithms for Subproblems

Section 4 offers decomposition schemes for distributionally robust MDPs with q -Wasserstein ambiguity set. As shown in Theorem 4.5 and Proposition 4.6, the overall complexity of computing distributionally robust Bellman update is 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 q ∈ { 1 , 2 , ∞} 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 are only O ( SA log S + NAS log ϵ -1 1 log ϵ -1 2 ) for q = 1 , O ( NAS log S log ϵ -1 1 log ϵ -1 2 ) for q = 2 , and O ( AS (log S + N )) when q = ∞ . These complexities are much lower than O ( N 3 . 5 A 3 . 5 S 3 . 5 log( ϵ -1 )) and O ( NA 2 . 5 S 2 . 5 log( S ) ϵ -1 . 5 ) from the state-of-the-art solution methods in (Xu and Mannor, 2010) and (Grand-Clément and Kroer, 2021a), respectively. The proofs of

the following propositions, theorems, and corollaries and the details of the proposed algorithms are relegated to Appendix A.2.

5.1 1 -Wasserstein ambiguity set

When q = 1 , the corresponding subproblem (9) has an equivalent form as follows.

Proposition 5.1. Suppose q = 1 . The minimization problem (9) is equivalent to

While problem (11) does not appear to be trivial at first glance, it has nice mathematical properties which allow us to solve this problem efficiently.

Theorem 5.2. The objective function in (11) is piecewise-linear in v , and it can be solved in time O ( S log S ) .

Therefore, the proposed Algorithm 1 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 -Wasserstein ambiguity set can be computed in time O ( SA log S + NAS log ϵ -1 1 log ϵ -1 2 ) , where ϵ 1 and ϵ 2 are 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 an Euclidean projection problem onto the probability simplex.

Theorem 5.3. Suppose q = 2 . The minimization problem (9) is equivalent to

which can be solved in time O ( S log S ) .

Euclidean projection problem onto 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 2 -Wasserstein ambiguity set can be computed in time O ( NAS log S log ϵ -1 1 log ϵ -1 2 ) , where ϵ 1 and ϵ 2 are the user-specified tolerances in Algorithm 1.

5.3 ∞ -Wasserstein ambiguity set

When q = ∞ , the corresponding subproblem

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 ∞ -Wasserstein ambiguity set can be computed in time O ( AS (log S + N )) .

50

70

90

110

S = A

0

2

4

6

8

Runtimes (100s)

Fast

Gurobi

FOM(3 its)

50

70

90

110

S

0

2

4

6

8

Runtimes (100s)

Fast

Gurobi

FOM(3 its)

50

70

90

110

N

0

2

4

6

8

Runtimes (100s)

Fast

Gurobi

FOM(3 its)

50

70

90

110

S = A

0

2

4

6

8

Runtimes (10s)

Fast

Gurobi

FOM(3 its)

50

70

90

110

S

0

2

4

6

8

Runtimes (10s)

Fast

Gurobi

FOM(3 its)

50

70

90

110

N

0

2

4

6

8

Runtimes (10s)

Fast

Gurobi

FOM(3 its)

50

70

90

110

S = A

0

2

4

6

8

Runtimes (s)

Fast

Gurobi

FOM(3 its)

50

70

90

110

S

0

2

4

6

8

Runtimes (s)

Fast

Gurobi

FOM(3 its)

50

70

90

110

N

0

2

4

6

8

Runtimes (s)

Fast

Gurobi

FOM(3 its)

6 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 to page limit, we only report 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 1 reports the average computation times of Bellman updates with 1 -Wasserstein ambiguity sets. For the first-order method (Grand-Clément 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 2 -Wasserstein ambiguity sets and ∞ -Wasserstein ambiguity sets are shown in Figure 2 and Figure 3, respectively. These results are similar to the case where q = 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-art 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.

References

Mohammed Amin Abdullah, Hang Ren, Haitham Bou Ammar, Vladimir Milenkovic, Rui Luo, Mingtian Zhang, and Jun Wang. Wasserstein robust reinforcement learning. arXiv preprint arXiv:1907.13196 , 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.

TWArchibald, 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 511-520. PMLR, 2021.

Bahram Behzadian, Marek Petrik, and Chin Pang Ho. Fast algorithms for l ∞ -constrained s-rectangular 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(11):2471-2482, 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, Pengqian Yu, and William B Haskell. Distributionally robust optimization for sequential decision-making. Optimization , 68(12):2397-2426, 2019.

Erick Delage and Yinyu Ye. Distributionally robust optimization under moment uncertainty with application to data-driven problems. Operations research , 58(3):595-612, 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.02894 , 2020.

Esther Derman, Matthieu Geist, and Shie Mannor. Twice regularized mdps and the equivalence between robustness and regularization. Advances in Neural Information Processing Systems , 34: 22274-22287, 2021.

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 in Neural Information Processing Systems , 29, 2016.

Robert Givan, Sonia Leach, and Thomas Dean. Bounded-parameter markov decision processes. Artificial Intelligence , 122(1-2):71-109, 2000.

Vineet Goyal and Julien Grand-Clément. Robust markov decision processes: Beyond rectangularity. Mathematics of Operations Research , 48(1):203-226, 2023.

Julien Grand-Clément and Christian Kroer. First-order methods for wasserstein distributionally robust mdp. In International Conference on Machine Learning , pages 2010-2019. 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://www. 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 l1-robust markov decision processes. The Journal of Machine Learning Research , 22(1):12612-12657, 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-time algorithm for linear programming. Proceedings of the sixteenth annual ACM symposium on Theory of computing , pages 302-311, 1984.

David L Kaufman and Andrew J Schaefer. Robust modified policy iteration. INFORMS Journal on Computing , 25(3):396-410, 2013.

Navdeep Kumar, Kfir Levy, Kaixin Wang, and Shie Mannor. Efficient policy iteration for robust markov decision processes via regularization. arXiv preprint arXiv:2205.14327 , 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 arXiv: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.14495 , 2020.

Daniel J Mankowitz, Nir Levine, Rae Jeong, Yuanyuan Shi, Jackie Kay, Abbas Abdolmaleki, Jost Tobias Springenberg, Timothy Mann, Todd Hester, and Martin Riedmiller. Robust reinforcement learning for continuous control with model misspecification. arXiv preprint arXiv:1906.07516 , 2019.

Shie Mannor, Ofir Mebel, and Huan Xu. Robust mdps with k-rectangular uncertainty. Mathematics of Operations Research , 41(4):1484-1509, 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-166, 2018.

Arnab Nilim and Laurent El 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 a 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.05129 , 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 reinforcement learning. In International Conference on Machine Learning , pages 2817-2826. PMLR, 2017.

Martin L 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 Jr. Markovian decision processes with uncertain transition probabilities. Operations Research , 21(3):728-740, 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 on 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 Carreira-Perpinán. Projection onto the probability simplex: An efficient algorithm with a simple proof, and an application. arXiv preprint arXiv:1309.1541 , 2013.

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 International Conference on Machine Learning , pages 23484-23526. PMLR, 2022.

Wolfram Wiesemann, Daniel Kuhn, and Berç Rustem. Robust markov decision processes. Mathematics of Operations Research , 38(1):153-183, 2013.

Wolfram Wiesemann, Daniel Kuhn, and Melvyn Sim. Distributionally robust convex optimization. Operations Research , 62(6):1358-1376, 2014.

Weijun Xie. Tractable reformulations of two-stage distributionally robust linear programs over the type-∞ wasserstein ball. Operations Research Letters , 48(4):513-523, 2020.

Huan Xu and Shie Mannor. The robustness-performance tradeoff in markov decision processes. Advances in Neural Information Processing Systems , 19, 2006.

Huan Xu and Shie Mannor. Distributionally robust markov decision processes. Advances in Neural Information Processing Systems , 23, 2010.

Insoon Yang. A convex optimization approach to distributionally robust markov decision processes with wasserstein distance. IEEE control systems letters , 1(1):164-169, 2017.

Pengqian Yu and Huan Xu. Distributionally robust counterpart in markov decision processes. IEEE Transactions on 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 , 33:21024-21037, 2020.

A Appendix: Proofs and Algorithms

A.1 Proofs of results in Section 4

Proof of Proposition 4.1. Plug B q s in (3) into (4), and apply the minimax theorem, the original problem [ T ( v )] s , ∀ s ∈ S is given by:

Hence (5) is a direct consequence of above formulation by introducing the following epigraph variable

Proof of Proposition 4.2. To prove the upper bound of γ ⋆ , we consider ¯ γ = max a ∈A 1 N ∑ N i =1 b ⊤ sa ˆ p i sa , which equals to ¯ γ ≥ 1 N ∑ N i =1 b ⊤ sa ˆ p i sa , ∀ a ∈ A . This implies { ˆ p i sa } N i =1 satisfies every constraint in problem (6) with the lowest possible objective value 0 , for every a ∈ A . Therefore, ∑ a ∈A P ( { ˆ p i sa } N i =1 ; r sa + λ v , ¯ γ ) = 0 < θ q . At this time, ¯ γ is feasible for (5). Hence, we provide an upper bound ¯ γ = max a ∈A 1 N ∑ N i =1 b ⊤ sa ˆ p i sa .

To prove the lower bound of γ ⋆ , we assume the contrary max a ∈A { min { b sa }} > γ ⋆ . So there exists ˆ a ∈ A with γ ⋆ < min { b s ˆ a } . Then there is no p i s ˆ a satisfies the first constraint in (5) for a = ˆ a . This contradiction verifies the lower bound γ = max a ∈A { min { b sa }} .

Proof of Proposition 4.3 . By definition, for any fixed s ∈ S and a ∈ A , we have

Then, we introduce the dual variable α ≥ 0 for the constraint 1 N ∑ N i =1 b ⊤ sa p i sa ≤ γ and obtain

where D q ( ˆ p i sa , b sa , α ) is defined in the Proposition 4.3.

To show ¯ α defined in the proposition is indeed an upper bound of the optimal α ⋆ , we denote

Notice that ∀ α ≥ max i ∈ [ N ] ∥ e j -ˆ p i sa ∥ q q γ -min { b sa } , where j ∈ arg min s ′ ∈S b sas ′ :

where the first inequality is from the definition of D q ( ˆ p i sa , b sa , α ) , the second equality is due to the selection of j , and the last inequality is due to the selection of α . Hence α ⋆ ∈ [0 , ¯ α ] by f ( α ) is concave w.r.t. α (otherwise ∃ α ⋆ > ¯ α such that f ( α ⋆ ) > f (0) ≥ f (¯ α ) , where nonconcavity of f follows), and (8) is true. One can further compute the subdifferential of f ( α ) by Danskin's theorem (Bertsekas, 1999). Typically,

where p i,⋆ sa is any minimizer of the inner minimization problem corresponding with given α ≥ 0 .

Proof of Theorem 4.4. For simplicity of notations, throughout this proof we use

By (6), we can get that P ( { ˆ p i sa } N i =1 ; b sa , γ ) and f ( γ ) are non-increasing in [ γ, ¯ γ ] . For each given γ ∈ [ γ, ¯ γ ] and a ∈ A , we denote ˆ P ( { ˆ p i sa } N i =1 ; b sa , γ ) the corresponding value calculated by the Algorithm 1. Furthermore, we call

We can show that

We fix γ ∈ [ γ, ¯ γ ] , and consider (8). Algorithm 1 provides the optimal solution for (8) with tolerance ϵ 2 / 2 , so we have | α ⋆ sa,γ -ˆ α sa,γ | ≤ ϵ 2 / 2 , where α ⋆ sa,γ is the true optimal solution of (8) and ˆ α sa,γ is the solution computed in the inner bisection in Algorithm 1. Thus we get

where the last step is due to that

and

here p i, ˆ α sa and p i,α ⋆ sa are the optimal solutions to (9) for α = ˆ α sa,γ and α = α ⋆ sa,γ respectively.

Hence, (14) is the direct consequence of above estimation, together with the definitions of f, ˆ f and δ . And we have γ ⋆ = inf { γ ∈ [ γ, ¯ γ ] : f ( γ ) ≤ θ q } . We further define ˆ γ ≜ inf { γ ∈ [ γ, ¯ γ ] : ˆ f ( γ ) ≤ θ q } . The outer bisection of Algorithm 1 implies that | γ ′ -ˆ γ | ≤ ϵ 1 2 , so to prove the claimed result in the theorem, it suffices to show that

Notice that ˆ P ( { ˆ p i sa } N i =1 ; b sa , γ ) ≤ P ( { ˆ p i sa } N i =1 ; b sa , γ ) , ∀ γ ∈ [ γ, ¯ γ ] by definitions, we get ˆ f ( γ ) ≤ f ( γ ) , so { γ ∈ [ γ, ¯ γ ] : f ( γ ) ≤ θ q } ⊆ { γ ∈ [ γ, ¯ γ ] : ˆ f ( γ ) ≤ θ q } , hence we get ˆ γ ≤ γ ⋆ .

We assume that γ ⋆ -γ ≥ 2 δ ( max a ∈A max { b sa }-γ ) θ q , otherwise (15) is trivially satisfied. To achieve (15), we claim the following statement:

If the statement is true, we have ∀ γ ∈ [ γ, γ ⋆ -( max a ∈A max { b sa } -γ ) (2 δ/θ q )] :

where the first inequality is from (14) and ˆ f ( γ ) ≤ f ( γ ) , the second inequality is due to that f ( γ ) is non-increasing, and the third inequality is from the above statement.

So { γ ∈ [ γ, ¯ γ ] : ˆ f ( γ ) ≤ θ q } ⊆ ( γ ⋆ -( max a ∈A max { b sa } -γ ) (2 δ/θ q ) , ¯ γ ] , thus

which implies the desired (15).

For the proof of the statement (16). We argue that

by contradiction. Assume there is some γ ′′ ∈ [ γ ⋆ -2 δ ( max a ∈A max { b sa }-γ ) θ q , γ ⋆ ) with that ∑ a ∈A α ⋆ sa,γ ′′ < θ q ( max a ∈A max { b sa }-γ ) . Clearly f ( γ ′′ ) > θ q since γ ′′ < γ ⋆ . Then

which is contradicted with f ( γ ′′ ) > θ q . This contradiction implies that (17) is true.

For simplicity of notations, we call γ ′′′ = γ ⋆ -2 δ ( max a ∈A max { b sa }-γ ) θ q , and we fix any γ ℓ ∈ ( γ ′′′ , γ ⋆ ) , then

where the first step is due to γ ℓ < γ ⋆ and definition of γ ′′′ , the second step is from the definition of function f , the third step is from the definition of P ( { ˆ p i sa } N i =1 ; b sa , γ ) and α ⋆ sa,γ , and the last step is due to (17) and γ ℓ ∈ ( γ ′′′ , γ ⋆ ) . Notice that above inequality is true for all γ ℓ ∈ ( γ ′′′ , γ ⋆ ) , we could let γ ℓ → γ ⋆ , which leads to

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, we can see that the bisection method on γ and α uses complexity O (log ϵ -1 1 log ϵ -1 2 ) . For the subproblem (9), which costs time complexity h q ( S ) , we need to solve

Input: Sorted b sa with b san 1 ≥ b san 2 ≥ · · · ≥ b san S .

Initialization:

for

k

= 1

r

, . . . , S

-

b

1

⊤

sa

ˆ

p

do

end

Result: Optimal objective value r and optimal solution p i sa of (9) with q = 1 .

it NA times. Besides, computing the upper bound claimed in Proposition 4.2 requires finding min { b sa } for each a ∈ A , which is in time complexity O ( AS ) . So we get the time complexity of Algorithm 1 is O ( h q ( S ) NA log ϵ -1 1 log ϵ -1 2 + AS ) .

Proof of Proposition 4.6. Plug B ∞ s in (3) into (4), we get the Bellman update is given by

The first equality is due to that the decision varaibles p i sa , ∀ i ∈ [ N ] and ∀ a ∈ A , are independent from each other. Then we can divide the orignal problem into NA subproblems. The second equality is from the fact that the objective function is affine w.r.t. π sa , and the maximum is simply the greatest coefficient of π sa .

A.2 Proofs of results in Section 5

A.2.1 Proof of results in Section 5.1

Proof of Proposition 5.1. We introduce the variable v ≥ ∥ p i sa -ˆ p i sa ∥ 1 , so

Proof of Theorem 5.2. We denote the objective function in (11) as

i

and

p

i

i

ˆ

p

.

sa sa

sa

W.L.O.G., we assume that b sa 1 > b sa 2 > · · · > b saS , and ˆ p i sa > 0 . We claim that

where

To prove the first case of the claim, it suffices to show that the optimal solution for the minimization problem in F ( v ) is given by p ⋆ , whose components are p ⋆ 1 = ˆ p i sa 1 -v 2 , p ⋆ k = ˆ p i sak , ∀ 2 ≤ k ≤ S -1 and p ⋆ S = ˆ p i saS + v 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 ¯ p 1 > 0 and ¯ p S < 1 . Actually, Let

Then by e ⊤ ¯ p = e ⊤ ˆ p i sa = 1 and ∥ ¯ p -ˆ p i sa ∥ 1 ≤ v , we get

Hence ∑ k ∈N ˆ p i sak -¯ p k = ∑ k ∈P ¯ p k -ˆ p i sak ≤ v 2 < ˆ p i sa 1 , so we get ¯ p 1 > 0 and ¯ p S < 1 . Next we show that ¯ p k = ˆ p i sak , ∀ 2 ≤ k ≤ S -1 . Otherwise we have some 2 ≤ ˆ k ≤ S -1 that | ¯ p ˆ k -ˆ p i sa ˆ k | > 0 . If ¯ p ˆ k > ˆ p i sa ˆ k , we define ˜ p with that ˜ p ˆ k = ¯ p ˆ k -ε and ˜ p S = ¯ p S + ε

such , here

0 < ε < min { | ¯ p ˆ k -ˆ p i sa ˆ k | 2 , 1 -¯ p S 2 } , while keeping the other components of ˜ p same as ¯ p . We can see that ˜ p achieves smaller objective value than ¯ p does due to b sa ˆ k > b saS , which contradicts the optimality of ¯ p . If ¯ p ˆ k < ˆ p i sa ˆ k , then we define ˜ p with that ˜ p ˆ k = ¯ p ˆ k + ε and ˜ p 1 = ¯ p 1 -ε , here

0 < ε < min { | ¯ p ˆ k -ˆ p i sa ˆ k | 2 , ¯ p 1 2 } , while keeping the other components of ˜ p same as ¯ p . Similarly, ˜ p achieves smaller objective value and this implies the contradiction.

Finally, we verify the rest two components. We introduce the variables d 1 = ¯ p 1 -ˆ p i sa 1 and d S = ¯ p S -ˆ p i saS , then we get the equivalent reformulation of the inner minimization in F ( v ) :

The optimal d 1 and d S are given by -v/ 2 and v/ 2 respectively. Hence we proved p ⋆ is indeed an optimal solution.

To prove the second case of the claim, the optimal solution for the minimization problem in F ( v ) is given by p ⋆ , whose components are p ⋆ k = 0 , ∀ k ∈ [ K ] , and p ⋆ k = r k , ∀ K +1 ≤ k ≤ S . The decision variables 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 ¯ p K +1 > 0 and ¯ p S < 1 .

To argue this by contradiction, we suppose the contrary ¯ p K +1 = 0 . Define the notations N , P and E same as before. Then there exists ˆ k ≤ K with ¯ p ˆ k > 0 , otherwise we assume that ¯ p k = 0 , ∀ k ∈ [ K ] , which implies the contradiction as follows.

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 ¯ p k = 0 , ∀ k ∈ [ K +1] . So we are able to find ˆ k ≤ K with ¯ p ˆ k > 0 . By moving probability ε = min { ¯ p ˆ k 2 , ˆ p i sa ( K +1) 2 } from ¯ p ˆ k to ¯ p K +1 , we can achieve smaller objective value while keeping the feasibility, hence we get the contradiction, which implies that ¯ p K +1 > 0 and ¯ p S < 1 . By applying the similar procedures in the first case (moving some probability from ¯ p K +1 or to ¯ p S ), we can show that ¯ p k = ˆ p i sak = r k for K +2 ≤ k ≤ S -1 . Next we prove that ¯ p k = r k = 0 for k ∈ [ K ] . Suppose the contrary is true; that is, ¯ p ˆ k > 0 for some ˆ k ≤ K . Then we are able to apply the same procedures as before which illustrate that ¯ p k = ˆ p i sak , ∀ ˆ k < k < S . Provided this, one can verified that the optimal strategy is putting all the rest probability 1 -∑ S -1 k =ˆ k +1 ˆ p i sak to ¯ p S since b sa 1 > b sa 2 > · · · > b saS , and v ≥ 2 ∑ K k =1 ˆ p i sak ≥ 2 ∑ ˆ k k =1 ˆ p i sak , which implies ¯ p sa ˆ k = 0 and we get a contradiction. Hence ¯ p k = r k = 0 , ∀ k ∈ [ K ] . Finally, we verify the rest two components. We introduce the variables d K +1 = ¯ p K +1 -ˆ p i sa ( K +1) and d S = ¯ p S -ˆ p i saS , then we get the equivalent reformulation of the inner minimization in F ( v ) :

The optimal d K +1 and d S are given by -v -2 ∑ K k =1 ˆ p i sak 2 and v 2 respectively. Hence we proved p ⋆ is an optimal solution.

To prove the third case of the claim, we notice that the optimal solution e S to min p i sa ∈ ∆ S b ⊤ sa p i sa is also feasible for the inner minimization problem in F ( v ) , hence it becomes the optimal solution we desire, which implies F ( v ) = v + αb saS at this time. This finishes the proof of our claim.

Our claim directly illustrates that F ( v ) is a piecewise-linear function in v with breakpoints { 0 } ∪ { 2 ∑ K k =1 ˆ p i sak : ∀ K ∈ [ S -1] } . Furthermore, based on the provided formulation of F ( v ) , we can compute the difference of value for F ( · ) between any two adjacent breakpoints, given by

where v K = 2 ∑ K k =1 ˆ p i sak , ∀ K ∈ [ S -1] and v 0 = 0 . Hence we provide Algorithm 2 to compute (11), whose time complexity is O ( S log S ) generally and can be reduced to O ( S ) if the sorted b sa is provided.

Proof of Corollary 5.2.1. We can see that Algorithm 2 is in time complexity O ( S ) if ∪ a ∈A b sa are sorted, which can be done at the Initialization step in Algorithm 1 with time complexity O ( AS log S ) . So by Theorem 4.5, we get the overall complexity is O ( NAS log ϵ -1 1 log ϵ -1 2 + AS log S ) .

Input: v , r sa , and ˆ p i sa ∈ ∆ S .

Initialization: p i,⋆ sa = 0 .

Sort b sa as b san 1 ≤ · · · ≤ b san S .

Find the smallest k such that ∑ k j =1 ( ˆ p i san j + θ ) ≥ 1 .

Result: Optimal objective value r and optimal solution p ⋆ sa 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 q = 2

So it suffices to solve min p i sa ∈ ∆ S ∥ ∥ ∥ ∥ p i sa -2ˆ p i sa -α b sa 2 ∥ ∥ ∥ ∥ 2 2 , which can be done by Euclidean projection algorithm (Wang and Carreira-Perpinán, 2013) with time complexity O ( S log S ) .

Proof of Corollary 5.3.1. The result is the direct consequence of Theorem 4.5 with h 2 ( S ) is O ( S 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 3 solves problem (13) with time complexity O ( S log S ) . By expanding the ∞ -norm, we formulate (13) as the following box constraints problem.

To get the optimal solution, we put the probability on the index where b sa is small as much as possible. Specifically, we assume b sa 1 < b sa 2 < · · · < b saS w.l.o.g., and assume k is the smallest index such that ∑ k j =1 ( ˆ p i saj + θ ) ≥ 1 . We claim that

is the optimal solution to the above formulation. To see this, suppose ¯ p is optimal and there is some ˆ j ∈ [ k -1] with ¯ p ˆ j = ˆ p i saj + θ . By above box constraint, we get ¯ p ˆ j < ˆ p i sa ˆ j + θ , so there exists ¯ j ≥ k such that ¯ p ¯ j > p i,⋆ sa ¯ j . By moving the probability from ¯ p ¯ j to ¯ p ˆ j , we can achieve strictly smaller objective value, which is a contradiction with that ¯ p is optimal. This gives us an optimal solution with the first k -1 components coincides p i,⋆ sa . Then we can get p i,⋆ sa is indeed optimal by putting the

̸

Proof of Corollary 5.4.1. As for each a ∈ A , we need to sort b sa , which costs O ( AS log S ) , then we need to solve NA subproblems, which costs O ( NAS ) , so the whole problem is computed in time O ( AS log S + NAS ) .

B Appendix: Computational Complexity for General Convex Optimzation

To compare our algorithm with general convex optimization algorithm, we use general convex optimization to compute the problem (5) and problem (8). The time complexities will be discussed in different situations:

Suppose q = 1 : For general convex optimization problem, problem (5) is equivalent with the following problem:

By introducing the variables t i sas ′ = ∣ ∣ p i sas ′ -ˆ p i sas ′ ∣ ∣ , ∀ i ∈ [ N ] , ∀ a ∈ A , ∀ s ′ ∈ S , the above problem is equivalent with

There are 1 + NSA + NSA = O ( NSA ) decision variables, and the number of bits in the input is O (1) + O ( NAS ) + O ( NAS ) + O ( NAS ) + O ( NAS ) + O ( NAS ) = O ( NAS ) . So by (Karmarkar, 1984), the complexity of solving this LP is O ( N 4 . 5 S 4 . 5 A 4 . 5 ) .

We utilize our outer bisection, solving (6) directly using convex optimization. Typically, for each fixed a ∈ A and γ , (6) is equivalent with the LP that

There are 2 NS 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 O ( N 4 . 5 S 4 . 5 ) . Together with the outer bisection, the total complexity for each Bellman update with ϵ 1 tolerance is O ( N 4 . 5 S 4 . 5 A log ϵ -1 1 ) .

We utilize our nested bisection scheme, solving (9) using the general convex optimization algorithm. Typically, for each fixed a ∈ A , i ∈ [ N ] and α , (9) is equivalent with

There are 2 S decision variables, and the number of bits in the input is O ( S + S + S ) = O ( S ) . So the complexity of solving this LP is O ( S 4 . 5 ) . Together with the nested bisection, the total complexity for each Bellman update where the tolerances of bisections are ϵ 1 and ϵ 2 , is O ( NS 4 . 5 A log ϵ -1 1 log ϵ -1 2 ) .

Suppose q = 2 : For general convex optimization problem, problem (5) is equivalent with the following SOCP:

There are 1 + NSA = O ( NSA ) decision variables, and the number of constraints are A +1+ NA ( S +2) = O ( NAS ) . So the complexity of solving the SOCP with ϵ -accuracy is O ( √ NAS log ϵ -1 · ( NSA ) 2 ( NAS +1+2( A + NA ( S +2)))) = O ( N 3 . 5 S 3 . 5 A 3 . 5 log ϵ -1 ) .

We utilize our outer bisection, solving (6) directly using convex optimization. Typically, for each fixed a ∈ A and γ , (6) is equivalent with the SOCP that

There are 1 + NS = O ( NS ) decision variables, and the number of constraints are O ( NS ) . So the complexity of solving the above SOCP with ϵ -accuracy is O ( √ NS log ϵ -1 · N 2 S 2 ( NS + 1 + 2(1 + N ( S + 2)))) = O ( N 3 . 5 S 3 . 5 log ϵ -1 ) , hence the total complexity of the Bellman update is O ( N 3 . 5 S 3 . 5 A log ϵ -1 log ϵ -1 1 ) .

We utilize our nested bisection scheme, solving problem (9) using the general convex optimization algorithm. For each fixed a ∈ A , i ∈ [ N ] and α , problem (9) is equivalent with

There are 1 + S = O ( S ) decision variables, and the number of constraints are O ( S ) . So the complexity of solving the above SOCP with ϵ -accuracy is O ( √ S log ϵ -1 · S 2 ( S + 1 + 2(1 + S ))) = O ( S 3 . 5 log ϵ -1 ) , hence the total complexity of the Bellman update is O ( NS 3 . 5 A log ϵ -1 log ϵ -1 1 log ϵ -1 2 ) .

Suppose q = ∞ : We can also consider solving the inner problem in (10) using general convex optimization, which is equivalent with

where i ∈ [ N ] and a ∈ A are fixed.

There are S decision variables, and the number of bits in the input is O ( S + S + S ) = O ( S ) . So the complexity of solving this LP is O ( S 4 . 5 ) . Then the total complexity for each Bellman update is O ( NS 4 . 5 A ) .

C Appendix: Details for Numerical Experiments

We compare our fast algorithm with the state-of-the-art solver Gurobi with version v10.0.1rc0 (Gurobi Optimization, LLC, 2023) and the first-order method of (Grand-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 i 7 CPU with 32 GB 3733 MHz DDR 4 main memory. We will release our code to ensure reproducibility. https://github.com/Chill-zd/Fast-Bellman-Updates-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 et 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 n b to regulate the proportion of the next states accessible for each state-action pair ( s, a ) . Following the same setting as (Grand-Clément and Kroer, 2021a), we set n b to be 0 . 2 and random uniform rewards to be in [0 , 10] . The discount factor λ is fixed at 0 . 8 , and parameter ϵ = 0 . 1 . For each MDP instance, we generate the sampled kernels p 1 , . . . , p N , considering N small random (Garnet) perturbations around the nominal kernel p 0 . We set parameter θ in Proposition 4.1 to be √ n b A .

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, 2 and 3. We can see that our algorithm performs better than Gurobi and the first-order method. When the states number increases, the running time of our algorithm keeps a small standard deviation.

We also compare the speed of value iteration for all algorithms using the same convergence criteria: ∥ v -v ⋆ ∥ ∞ ≤ 2 λϵ (1 -λ ) -1 , which follows (Grand-Clément and Kroer, 2021a). The results are shown in tables 4, 5 and 6. The runtimes that exceed 4000 s for L 1 and L ∞ case or exceed 10000 s for

L 2 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.

Peer reviews:
comments: SUMMARY:
* tailored algorithms for solving Wasserstein distributionally robust MDPs and fast implementations with $L_1$, $L_2$, $L_\infty$-based Wasserstein distance.

STRENGTHS:
* Fastest algorithms (in terms of dependency on N, S, and A - the number of kernels/states/actions) for solving Wasserstein distributionally robust MDPs.
* 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].

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

[2] Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Partial policy iteration for l1-robust markov
 decision processes. The Journal of Machine Learning Research, 22(1):12612–12657, 2021.

[3]  Chin Pang Ho, Marek Petrik, and Wolfram Wiesemann. Robust phi-divergence mdps. arXiv preprint
 arXiv:2205.14202, 2022.

QUESTIONS:
Model : can your algorithms be extended to handle the case of both transitions and reward uncertainty?

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?

Wasserstein 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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-08

comments: 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.
- 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.
- 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.

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.

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.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-04

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nI do not see any negative societal impact from this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

comments: 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.

STRENGTHS:
The proposed method shows state-of-the-art complexity results. 

**DISCLAIMER:**
 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.

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?
* 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?
* Considering that the following paper also focuses on the model-based distributionally robust RL setting. It would be interesting to also compare with it.
    * Shi, Laixi, and Yuejie Chi. "Distributionally robust model-based offline reinforcement learning with near-optimal sample complexity." arXiv preprint arXiv:2208.05767 (2022).

LIMITATIONS:
The limitation of the proposed method is not fully discussed in the paper.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-07