First-Order Methods for Wasserstein Distributionally Robust MDP

TL;DR

This paper introduces a first-order method framework for solving Wasserstein distributionally robust MDPs, achieving faster convergence and better scalability than existing approaches, with strong empirical performance.

Contribution

The paper develops a novel first-order optimization framework for Wasserstein distributionally robust MDPs, improving convergence rates and computational efficiency over prior methods.

Findings

Proposed algorithms achieve a convergence rate of $O(NA^{2.5}S^{3.5} ext{log}(S) ext{log}(rac{1}{ extepsilon}) extepsilon^{-1.5}$.

Dependence on $N$, $A$, and $S$ is significantly better than existing methods.

Numerical experiments demonstrate superior scalability and performance across multiple domains.

Abstract

Markov decision processes (MDPs) are known to be sensitive to parameter specification. Distributionally robust MDPs alleviate this issue by allowing for \emph{ambiguity sets} which give a set of possible distributions over parameter sets. The goal is to find an optimal policy with respect to the worst-case parameter distribution. We propose a framework for solving Distributionally robust MDPs via first-order methods, and instantiate it for several types of Wasserstein ambiguity sets. By developing efficient proximal updates, our algorithms achieve a convergence rate of $O\left(NA^{2.5}S^{3.5}\log(S)\log(\epsilon^{-1})\epsilon^{-1.5} \right)$ for the number of kernels $N$ in the support of the nominal distribution, states $S$, and actions $A$; this rate varies slightly based on the Wasserstein setup. Our dependence on $N,A$ and $S$ is significantly better than existing methods, which have a complexity of $O\left(N^{3.5}A^{3.5}S^{4.5}\log^{2}(\epsilon^{-1}) \right)$. Numerical experiments show that our algorithm is significantly more scalable than state-of-the-art approaches across several domains.