Provably Efficient Algorithms for S- and Non-Rectangular Robust MDPs with General Parameterization

TL;DR

This paper develops efficient algorithms for robust Markov decision processes with general policy parameterizations, extending theoretical guarantees and improving sample complexity bounds for both discounted and average reward settings.

Contribution

It introduces novel algorithms with provable sample complexity guarantees for RMDPs beyond tabular policies, including infinite state spaces and non-rectangular uncertainty sets.

Findings

Introduces a multilevel Monte Carlo gradient estimator with $ ilde{O}(rac{1}{ ext{epsilon}^2})$ complexity.

Provides the first sample complexity guarantees for RMDPs with general policy parameterization beyond $(s,a)$-rectangularity.

Achieves significant improvements in sample complexity bounds for both discounted and average reward robust MDPs.

Abstract

We study robust Markov decision processes (RMDPs) with general policy parameterization under s-rectangular and non-rectangular uncertainty sets. Prior work is largely limited to tabular policies, and hence either lacks sample complexity guarantees or incurs high computational cost. Our method reduces the average reward RMDPs to entropy-regularized discounted robust MDPs, restoring strong duality and enabling tractable equilibrium computation. We prove novel Lipschitz and Lipschitz-smoothness properties for general policy parameterizations that extends to infinite state spaces. To address infinite-horizon gradient estimation, we introduce a multilevel Monte Carlo gradient estimator with $\tilde{\mathcal{O}}(\epsilon^{-2})$ sample complexity, a factor of $\mathcal{O}(\epsilon^{-2})$ improvement over prior work. Building on this, we design a projected gradient descent algorithm for s-rectangular uncertainty ($\mathcal{O}(\epsilon^{-5})$) and a Frank--Wolfe algorithm for non-rectangular uncertainty ($\mathcal{O}(\epsilon^{-4})$ discounted, $\mathcal{O}(\epsilon^{-10.5})$ average reward), significantly improving prior results in both the discounted setting and average reward setting. Our work is the first one to provide sample complexity guarantees for RMDPs with general policy parameterization beyond $(s, a)$-rectangularity. It also provides the first such guarantees in the average reward setting and improves existing bounds for discounted robust MDPs.