Policy Regularized Distributionally Robust Markov Decision Processes with Linear Function Approximation
Jingwen Gu, Yiting He, Zhishuai Liu, Pan Xu
TL;DR
This work addresses distribution shift in online reinforcement learning by formulating policy-regularized, distributionally robust MDPs with linear function approximation. It introduces the model-free DR-RPO algorithm, which combines softmax policy updates near a reference policy with UCB-based exploration to achieve sublinear average regret in both $d$-rectangular DRMDP and RRMDP settings. Theoretical results establish robust Bellman properties, linearity of robust Q-functions, and a regret bound that matches greedy online methods up to logarithmic factors, while preserving stochastic policies. Empirical evaluations on simulated off-dynamics linear MDPs and American Put Option problems demonstrate DR-RPO’s robustness to distribution shifts and its practical viability. Overall, the paper advances efficient policy optimization for online robust RL, offering concrete tools for robust control under model uncertainty with scalable function approximation.
Abstract
Decision-making under distribution shift is a central challenge in reinforcement learning (RL), where training and deployment environments differ. We study this problem through the lens of robust Markov decision processes (RMDPs), which optimize performance against adversarial transition dynamics. Our focus is the online setting, where the agent has only limited interaction with the environment, making sample efficiency and exploration especially critical. Policy optimization, despite its success in standard RL, remains theoretically and empirically underexplored in robust RL. To bridge this gap, we propose \textbf{D}istributionally \textbf{R}obust \textbf{R}egularized \textbf{P}olicy \textbf{O}ptimization algorithm (DR-RPO), a model-free online policy optimization method that learns robust policies with sublinear regret. To enable tractable optimization within the softmax policy class, DR-RPO incorporates reference-policy regularization, yielding RMDP variants that are doubly constrained in both transitions and policies. To scale to large state-action spaces, we adopt the $d$-rectangular linear MDP formulation and combine linear function approximation with an upper confidence bonus for optimistic exploration. We provide theoretical guarantees showing that policy optimization can achieve polynomial suboptimality bounds and sample efficiency in robust RL, matching the performance of value-based approaches. Finally, empirical results across diverse domains corroborate our theory and demonstrate the robustness of DR-RPO.