Policy Optimization for Robust Average Reward MDPs

Authors: zhongchang  sun
sihong  he
fei  miao
shaofeng  zou
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 15:42:47

Abstract: This paper studies first-order policy optimization for robust average cost Markov decision processes (MDPs). Specifically, we focus on ergodic Markov chains. For robust average cost MDPs, the goal is to optimize the worst-case average cost over an uncertainty set of transition kernels. We first develop a sub-gradient of the robust average cost. Based on the sub-gradient, a robust policy mirror descent approach is further proposed. To characterize its iteration complexity, we develop a lower bound on the difference of robust average cost between two policies and further show that the robust average cost satisfies the PL-condition. We then show that with increasing step size, our robust policy mirror descent achieves a linear convergence rate in the optimality gap, and with constant step size, our algorithm converges to an $\epsilon$-optimal policy with an iteration complexity of $\mathcal{O}(1/\epsilon)$. The convergence rate of our algorithm matches with the best convergence rate of policy-based algorithms for robust MDPs. Moreover, our algorithm is the first algorithm that converges to the global optimum with general uncertainty sets for robust average cost MDPs. We provide simulation results to demonstrate the performance of our algorithm.

Main paper: 
 introduction: 
 Markov decision process (MDPs) [34] has been widely used to model agent-environment interactions in sequential decision-making problems. An MDP consists of a set of states, a set of actions, a transition kernel describing the dynamics of the environment, and a cost function of state-action pairs. The agents aim to minimize the cumulative cost obtained over time under a given transition kernel. However, real-world environments often exhibit uncertainties and non-stationarity that challenge the assumptions of traditional RL approaches. When there is a mismatch between the training environment and the real environment, minimizing the cumulative cost under the training environment may lead to poor performance under the real environment.To address the challenges raised by the model mismatch, the robust MDP was proposed [13, 27] , where the transition kernel of the MDP is not fixed but lies in an uncertainty set. The goal of the robust MDP is to optimize the worst-case performance over the uncertainty set of transition kernels. The obtained policy under the robust setting is thus robust to the model mismatch.Existing works on robust MDPs mainly focus on the discounted cost setting, where the goal is to minimize the worst-case cumulative discounted cost. However, in many real-world applications with long time horizons, such as inventory management in supply chains and applications in communication networks [16] , the optimal policies obtained from the discounted cost setting may have poor long-term performance [15] .To address these challenges, recent research has shifted focus towards robust average cost MDPs, where the objective is to optimize the worst-case average cost obtained per time step. The average cost MDPs offer several advantages over discounted cost MDPs, including better stability and applicability to infinite-horizon tasks. However, achieving robust and efficient learning in average cost settings remains a significant challenge. Since the cost is discounted exponentially with time in discounted cost MDPs, establishing a contraction only requires a discount factor strictly less than one. Compared with discounted cost MDPs, average cost MDPs depend on the long-term performance of the underlying MDPs as it assigns equal weight to both immediate and future costs. Research on robust average cost MDPs is relatively scarce in the existing literature, with only a few notable studies, such as those by [36, 22, 39, 40, 12] . None of the above-mentioned works focuses on the fundamental characterization of gradient-based algorithms. Therefore, in our paper, we present the first theoretical analysis of the global convergence of policy optimization algorithms in the context of robust average cost MDPs with general uncertainty sets.In this paper, for ergodic Markov chains, we propose a policy-based optimization algorithm called robust policy mirror descent to solve the robust average cost MDPs. We further show that with increasing step size, our robust policy mirror descent achieves linear convergence in the optimality gap, and with constant step size, our algorithm converges to an ϵ-optimal policy with iteration complexity of O(1/ϵ). Our algorithm is the first policy-based algorithm with global convergence and finite iteration complexity analysis for robust average cost MDPs with general uncertainty sets. In particular, our main contributions are summarized as follows.We derive a policy (sub)-gradient for the robust average cost MDPs. In robust average cost MDPs, the goal is to optimize the worst-case performance, known as the robust average cost, which considers the worst-case over the uncertainty set of transition kernels. However, the robust average cost function typically involves a "max" operator over transition kernels, making it non-differentiable with respect to the policy. Therefore, in this paper, we first develop the Fréchet sub-gradient of the robust average cost with general uncertainty sets, which serves as a foundation of the policy-based algorithm. The robust policy gradient under the discounted setting was derived in [42, 21] . The work [42] considers the specific R-contamination uncertainty set, and the derivative of the robust policy gradient in [21] relies on the fact that the discount factor is strictly less than 1, which can not be extended to the average cost setting.We propose a robust policy mirror descent algorithm. Based on the derivative of the Fréchet sub-gradient, we propose the robust policy mirror descent algorithm. We apply the dynamically weighted divergence to the policy mirror descent so that the policy can be updated for each state separately, which further ensures global convergence.We show that our algorithm converges to the global optimum, and we further characterize the iteration complexity. To show the global convergence, we first prove that the robust average cost satisfies the Polyak-Łojasiewicz (PL) condition [29, 23] . We then prove the global optimality of our algorithm and characterize its iteration complexity. We show that with increasing step size, our robust policy mirror descent achieves linear convergence in the optimality gap, and with constant step size, our algorithm converges to an ϵ-optimal policy with iteration complexity O(1/ϵ). For increasing step size, the linear convergence of our algorithm matches with the robust discounted cost setting [21] . In [21] , the robust policy mirror descent was shown to converge to the global optimum with iteration complexity O(1/ϵ) for discounted cost MDPs when the step size is sufficiently large. Conversely, for non-robust discounted cost MDPs, the policy mirror descent was demonstrated to achieve a O(1/ϵ) iteration complexity in [45] , representing the state-of-the-art convergence rate for policy mirror descent with a constant step size. Therefore, the convergence rate of our algorithm matches with the performance of the best non-robust counterpart. The convergence analysis of policy-based algorithms for non-robust MDPs relies on the fact that the value function is smooth, which is not the case for the robust average cost function. In this paper, we combine the first-order optimality condition of the policy update and the PL condition to develop a novel proof for the convergence rate of our algorithm.In this section, we discuss works on policy-based approaches for non-robust MDPs, robust discounted cost MDPs and robust average cost MDPs.Policy-based approaches for non-robust MDPs. In the non-robust setting, policy-based algorithms [44, 35, 17, 14, 31, 32] have demonstrated remarkable success across various applications. Recently, the global convergence of the policy-based algorithms were established [4, 5, 1, 24, 20, 19, 49, 9, 48, 45] . For discounted cost MDPs, it was shown in [1] that the projected gradient descent converges to a global optimum with iteration complexity O(1/ϵ 2 ). In both [48] and [45] , the authors show that the projected gradient descent converges to a global optimum with a less iteration complexity O(1/ϵ). For average cost MDPs, [18] presents a sublinear convergence bound for projected gradient descent, where the bound involves the parameter that characterizes the complexity of the underlying MDP. In all the above works, the convergence analysis depends on the smoothness of the value function. However, since the smoothness may not hold for the robust value function, the methodologies applied in [1, 48, 45, 18] can not be extended to our case. The policy mirror descent with increasing step size was shown to achieve linear convergence in [45] . The policy mirror descent with constant step size was also proved to converge with iteration complexity O(1/ϵ) for discounted cost MDPs in [45] . Their proof relies on the performance difference lemma and the fact that the underlying transition kernel doesn't change with time. In this paper, we derive the policy sub-gradient for robust average cost MDPs and design the robust policy mirror descent. We then develop a lower bound on the difference of robust average cost between two policies and further combine it with the first-order optimality condition of the policy update to characterize the iteration complexity of our algorithm.Robust discounted cost MDPs. Robust discounted cost MDPs have been widely studied [13, 27, 41, 3, 11, 47, 21, 46, 28, 33] , where the goal is to minimize the worse-case cumulative discounted cost over the uncertainty set of transition kernels. In this section, we introduce works on policy-based algorithms, which are closely related to our work. In [42] , the robust MDPs are considered under the R-contamination uncertainty set and the robust policy gradient algorithm is designed. It is shown in [38] that the robust policy gradient algorithm converges to the global optimum with iteration complexity O(1/ϵ 3 ). Later in [38] , the double-loop robust policy gradient was proposed for general uncertainty sets and was further proved to converge to the global optimum with iteration complexity O(1/ϵ 4 ). The robust policy mirror descent was designed for discounted cost MDPs in [21] . With increasing step size, the robust policy mirror descent converges linearly to the global optimum, while with constant but sufficiently large step size, the algorithm converges to the global optimum with iteration complexity O(1/ϵ). In this paper, we study the average cost setting and show that with increasing step size, the robust policy mirror descent achieves linear convergence, and with constant step size, our algorithm converges to the global optimum with iteration complexity O(1/ϵ), which matches with the best achievable iteration complexity of policy-based algorithms for robust MDPs. For the case with constant step size, our analysis doesn't require the step size to be sufficiently large.Robust average cost MDPs. Existing literature that focuses on robust average cost MDPs is relatively limited. The robust average cost MDPs were initially explored by [36] , where a specific finite interval uncertainty set was considered and the O(1/ϵ) convergence rate was achieved. In [22] , the robust average cost MDPs were studied under the l 1 uncertainty set. However, the approaches in [36, 22] are not applicable for general uncertainty sets. Later in [39, 40, 12] , the robust value iteration based algorithms were proposed and the global convergence was proved. In [12] , the connection are built between the discounted reward MDPs and average reward MDPs and the existence of Blackwell optimal policies are proved. In [39, 40] , the model-based and model-free robust average reward MDPs are studied and the robust relative value iteration algorithms are proposed. However, finding a stopping criterion and characterizing the iteration complexity for robust value iteration based algorithms remain elusive. In this paper, we propose the first policy-based algorithm for robust average cost MDPs. We show that our algorithm converges to the global optimum and we further characterize its iteration complexity. Therefore, our algorithm is the first algorithm with finite iteration complexity analysis for robust average cost MDPs with general uncertainty sets.Exponential cost robust MDPs. For the robust average cost MDPs, when the uncertainty set is defined by the KL-divergence metric, the problem admits a dual formulation, which is the exponential cost robust MDPs. The exponential cost robust MDPs have also been studied in the literature. In [6] , the Q-learning and the actor-critic method are described and the asymptotic performance are characterized for risk sensitive robust MDPs. In [7] , the value iteration and policy iteration algorithms are also analyzed for risk sensitive MDPs. Recently, in [26] , the modified policy iteration is proved to converge to the global optimum for exponential cost risk sensitive MDPs. The policy gradient algorithm for the risk sensitive exponential cost MDPs is studied and the asymptotic convergence bounds to a stationary point are provided in [25] . In our paper, we study the robust average cost MDPs with general uncertainty sets and characterize the global convergence of our algorithm. 
preliminaries and problem formulation: 
 In this section, we introduce some preliminaries on discounted cost MDPs, average cost MDPs and present our problem formulation.A discounted cost MDP is defined by the tuple (S, A, P, r, γ), where S denotes the finite state space, A denotes the finite action space, P = {P a s ∈ ∆(S), s ∈ S, a ∈ A}1 is the transition kernel, r : S × A → [0, 1] denotes the cost function, and γ ∈ [0, 1) is the discount factor. Denote by S the number of states and A the number of actions, respectively.We consider the set of all stationary and randomized policies Π = {π : S → ∆(A)}. For each policy π, it maps from state s ∈ S to a distribution over action a ∈ A. At each state s, the agent takes action a with probability π(a|s), and the environment transits from state s to state s ′ according to P a s . The discounted value function of a policy π starting from an initial state s is defined asEQUATIONwhere E P denotes the expectation with respect to the distribution induced by the transition kernel P. To align with conventions in the optimization literature, in this paper, we adopt a minimization formulation. For discounted cost MDPs, the goal is to find a policy π that minimizes the discounted value function V π P,γ (s) for any initial state s.Average cost is another fundamental criterion for MDPs. For discounted cost MDPs, the agent penalizes the future cost with discount factor γ to demonstrate the preference for the current cost.The average cost MDPs focus on the long-term performance of the underlying MDPs under the steady-state distribution. The average cost MDP can be defined by the tuple (S, A, P, r). For a policy π, define the average cost under transition kernel P starting from an initial state s as followsEQUATIONWe also define the relative value function V π P and the relative state-action value function Q π P for average cost MDPs as followsEQUATIONThe relative value function V π P and the average cost g π P satisfy the following Bellman equation [30]EQUATIONwhere P a s,s ′ denotes the probability of transiting to state s ′ when choosing action a at state s. Let d π P denote the stationary probability induced by the policy π and transition kernel P, and it satisfies that d π P P = d π P . Similar as [18] , we consider the projection of the value function onto the subspace orthogonal to the 1 vector so that V π P and Q π P are unique. For average cost MDPs, the goal is to find a policy π that minimizes the average cost g π P (s) for any initial state s.For robust MDPs, the transition kernel P is not fixed but lies in some uncertainty set P. Define the robust average cost MDP by the tuple (S, A, P, r). In this work, we consider (s, a)-rectangular uncertainty set:EQUATIONwhere P 0 is a known nominal transition kernel, D measures the difference between two distributions, e.g., KL divergence, and R is the pre-specified radius of the uncertainty set.For robust average cost MDPs, the agent aims to optimize the worst-case performance over the uncertainty set P. Define the worst-case average cost as followsEQUATIONSimilarly, we denote by V π P and Q π P the robust relative value function and the robust relative state action value function, respectively. The robust relative value function V π P and the robust average cost g π P satisfy the following Bellman equation [39] .EQUATIONFor any policy π, denote by d π P the stationary distribution of the state under the worst-case transition kernel of π. The goal is to find a policy π such that the worst-case average cost g π P is minimized, i.e., min π∈Π g π P (s), for any s ∈ S.Denote the optimal policy by π * and the robust average cost of π * by g * P . In this paper, we state the following assumption to guarantee that the average cost is independent of the initial state, which is widely used in the studies of average cost MDPs [43, 50, 10, 37, 39, 40] . Assumption 2.1. For any π ∈ Π and P ∈ P, the induced Markov chain is ergodic. 
robust policy mirror descent: 
 In this section, we first derive the robust average cost policy gradient. We then propose the robust policy mirror descent algorithm.Since the worst-case average cost g π P takes "max" over all P ∈ P, g π P might not be differentiable. To address this issue, we introduce the concept of Fréchet sub-gradient. Let ∥ • ∥ denote the L 2 norm of a vector. Definition 3.1. For any function f :EQUATION)When f is differentiable at x, the Fréchet sub-gradient u is the gradient of f . In this paper, we consider the direct policy parameterization. We derive the sub-gradient for the robust average cost g π P in the following lemma. Lemma 3.2. Let ∇g π P (s, a) = d π P (s)Q π P (s, a). Then ∇g π P is the Fréchet sub-gradient of g π P . Note that the Fréchet sub-gradient has been derived for robust discounted cost MDPs [21] , of which the (s, a) entry takes the form 1 1-γ d π P (s)Q π P (s, a). Here, with a little abuse of notation, we use d π P to denote the visitation distribution of policy π under the worst-case transition kernel and use Q π P to denote the worst-case action value function of policy π. The Fréchet sub-gradient for robust discounted cost MDPs in [21] can not be extended to the average setting since in [21] , the discounted factor γ is required to be strictly less than 1. In this paper, we derive the Fréchet sub-gradient of robust average cost MDPs by applying the performance difference lemma for average cost MDPs [8] and the Lipschitz property of the relative action value function [18] .Input: step size η k , initial policy π 0 for k = 0, 1, • • • , K -1 do for s ∈ S do Update policy: π k+1 (•|s) = arg min p∈∆(A) η k ⟨Q π k P (s, •), p⟩ + D(p, π k (•|s)). end for end for Output: π KWith Lemma 3.2, we are ready to present our robust policy mirror descent algorithm for average cost MDPs. We assume that for a given policy π, there exists an oracle that outputs the robust relative state action value function Q π P . We denote by D(π(•|s), π ′ (•|s)) the Bregman divergence between two policies π(•|s) and π ′ (•|s). We further define the weighted Bregman divergence functionD d (π, π ′ ) = s∈S d(s)D(π(•|s), π ′ (•|s))for any d ∈ ∆(S). We define the following robust policy mirror descent with dynamically weighted divergenceEQUATIONwhere η k is the step size. Note that g π P might not be differentiable, thus in our algorithm we replace the gradient of g π P by its Fréchet sub-gradient ∇g π P . By plugging in the sub-gradient formula of g π P in Lemma 3.2, we have thatEQUATIONwhich is equivalent toEQUATIONWe summarize our algorithm in Algorithm 1.Note that for the projected policy (sub)-gradient algorithm, the policy is updated as followsEQUATIONIn our paper, we set the Bregman divergence D(•, •) to be the squared Euclidean distance. In this case, the difference between our robust policy mirror descent and the projected policy gradient lies in that we replace the policy (sub)-gradient ∇g π k P by Q π k P . In the next section, we show that though the robust average cost g π P might not be differentiable, our robust policy mirror descent achieves linear convergence in the optimality gap with increasing step size, and converges to an ϵ-optimal policy with iteration complexity O(1/ϵ) with constant step size. 
theoretical results: 
 Before we show the global optimality of the robust policy mirror descent, we first provide some important properties of the robust average cost MDPs.We first provide a lower bound on the difference of robust average cost between two policies, which is a key step to derive the global optimality. Lemma 4.1. For any two policies π, π ′ , we have thatg π P -g π ′ P ≥ E s∼d π ′ P ⟨Q π P (s, •), π(•|s)-π ′ (•|s)⟩ .Lemma 4.1 is not a straightforward extension of performance difference lemma [8] to the robust setting since the worst-case transition kernels are different for different policies. A similar bound was derived in [21] for robust discounted value function by applying the Bellman equation of robust value function, which is not applicable in our case as the average cost itself does not satisfy the Bellman equation. Therefore, we apply the Bellman equation of the robust relative value function V π P and the robust average cost g π P in (7) and the observation that g π P is independent of the initial state s to obtain Lemma 4.1.We then show that the robust average cost g π P satisfies the PL-condition in the following lemma. Lemma 4.2. The suboptimality of any π satisfies g π P -g * P ≤ C P L max π ⟨∇g π P , π -π⟩, whereC P L = max π,s d π * P (s) d π P (s) .The PL-condition implies that when the subgradient ∇g π P is small, the policy π lies in the small neighborhood of the global optimum.Our convergence analysis also leverages the following Lipschitz property of the non-robust relative value function V π P [18] . Lemma 4.3. The relative value function V π P is Lipschitz in π, i.e., there exists a constant L π such that|V π P -V π ′ P | ≤ L π ∥π -π ′ ∥.In this section, we show that our algorithm achieves linear convergence rate with increasing step size. We first characterize some properties of our robust policy mirror descent algorithm. Lemma 4.4. For any p ∈ Π and s ∈ S, we have thatEQUATIONIn the following lemma, we establish the convergence property for each iteration of our algorithm. Lemma 4.5. At each iteration of our algorithm, we have thatEQUATIONwhere M = sup π,P∈Pd π * Pπ d π P ∞ .We show that with increasing step size, our robust policy mirror descent converges linearly.Theorem 4.6. Under Assumption 2.1, set the step sizeη k ≥ η k-1 1 -1 M -1M . The robust policy mirror descent satisfiesEQUATION)Our analysis in this section mainly leverages the performance difference lemma [8] and the Bregman divergence three-point lemma. The convergence rate for our robust policy mirror descent with increasing step size matches with the best convergence rate of policy-based algorithm for robust MDPs [21] . Our algorithm is the first algorithm that converges to the global optimum with finite iteration complexity for robust average cost MDPs with general uncertainty sets.We proceed to show that with constant step size, our algorithm achieves the global optimum with iteration complexity O(1/ϵ).Theorem 4.7. Under Assumption 2.1, let step size η = 1 Lπ for all k ≥ 1. We have that the each iteration of the robust policy mirror descent satisfiesEQUATIONwhere ω = (2√ 2SC P L ) -2 .For Algorithm 1, to find an ϵ-optimal policy, Theorem 4.7 shows that the iteration complexity is upper bounded by O(1/ϵ). In [21] , the robust policy mirror descent was studied for discounted cost MDPs. The iteration complexity O(1/ϵ) in [21] can only be achieved with a sufficient large step size η k = 1/ϵ. For the non-robust discounted cost MDPs, the policy mirror descent was shown to enjoy a O(1/ϵ) iteration complexity in [45] , which is the state-of-the-art convergence rate for policy mirror descent with constant step size. Our robust policy mirror descent with constant step size converges to the global optimum with iteration complexity O(1/ϵ), which matches with the best non-robust counterpart [45] .Define the gradient mappingG 1 η (π k (•|s)) = 1 η (π k (•|s) -π k+1 (•|s)). Note that if π k is updated exactly by the (sub)-gradient descent, then G 1 η (π k (•|s)) = ∇g π P (s, •). The norm ∥G 1 η (π k (•|s))∥ measures the closeness of the current step to the first-order stationary point. Our proof relies on the following key ingredient.EQUATION)In [1, 45] , similar steps are adopted to derive the convergence rate of the policy gradient descent for non-robust discounted cost MDPs. In their analyses, to obtain (18) , the smoothness of the value function is required. However, since the worst-case transition kernel is a function of the policy π, the robust average cost g π P might not be smooth. Therefore, the approaches applied in [1, 45] can not be directly extended to robust average cost MDPs. Without the smoothness of the value function, we still prove that (18) holds by leveraging Lemma 4.1 and the following first-order optimality, which is from the optimality condition of the robust policy mirror descent update in ( 12)EQUATIONRemark 4.8. Since for robust discounted cost MDPs, Lemma 4.1, Lemma 4.2 and Lemma 4.3 also hold, the result in Theorem 4.7 also holds for robust discounted cost MDPs. Therefore, our robust policy mirror descent with constant step size finds an ϵ-optimal policy with iteration complexity O(1/ϵ) for robust discounted cost MDPs. Compared with [21] , our step size doesn't need to be sufficiently large. 
simulation results: 
 In this section, we provide some simulation results to demonstrate the performance of our algorithm. We verify our method on one classical problem: the Garnet problem, and a robotic application problem: the recycling robot problem.Garnet environments are synthetic benchmarks designed to be used for studying the performance of RL algorithms. The Garnet framework provides a way to create randomly generated MDPs with specified properties, such as the number of states, actions, and transition probabilities. More details can be found in [2] .In the recycling robot problem, a mobile robot powered by a rechargeable battery is tasked with collecting empty soda cans. The robot operates with two battery levels: low and high. It has three possible actions: (1) search for empty cans; (2) remain stationary and wait for someone to bring it a can; or (3) return to its home base to recharge. When the robot's battery is low (high), it has a probability of α (β) of finding an empty can and maintaining its current battery level. If the robot searches for cans but does not find any, it will deplete its battery completely and must be carried back by humans. For more details, refer to [34] . In this paper, we set α = 0.9, β = 0.9.We compare our robust policy mirror descent with the non-robust method [45] . We consider Garnet(3, 2), Garnet(5, 2), Garnet (10, 5) and recycling robot problems. Both methods use a uniform random Figure 1 showcases the mean values of the optimal robust average cost values and standard deviations for our robust method and the non-robust baseline over the training episode in Garnet(3, 2). Our robust method converges faster than the non-robust baseline, though these two methods converge to the same value. We speculate that the reason why the two methods converge to similar numbers is because this MDP environment has a small size. Therefore, we further compare our robust method with the non-robust baseline in two larger Garnet problems Garnet(5, 2) and Garnet (10, 5) , and the recycling robot problem. Figure 2 , 3, and 4 respectively show the mean optimal robust average cost and standard deviations in these three problems. From Figure 2 , 3 and 4, we find that our robust method outperforms the non-robust baseline in terms of mean costs in all three problems. The simulation results demonstrate the robustness of our algorithm. 
conclusion: 
 In this paper, we investigated the policy-based algorithm for robust average cost MDPs. We first introduced the sub-gradient of the robust average cost. Based on the sub-gradient, we proposed the robust policy mirror descent. Our theoretical analysis demonstrates that the proposed algorithm with increasing step size achieves linear convergence rate, and with constant step size, the proposed algorithm finds an ϵ-optimal policy with iteration complexity O(1/ϵ), matching the best convergence rate observed in policy mirror descent algorithms for robust MDPs. Moreover, our algorithm is the first algorithm that converges to the global optimum with finite iteration complexity for robust average cost MDPs with general uncertainty sets. Our paper focuses on the model-based setting. In the future, it is of interest to design policy-based model-free algorithms for robust average cost MDPs.A Proof of Lemma 3.2In this section, we provide the proof for Lemma 3.2.Consider polices π, π ′ ∈ Π. For the non-robust average cost g π P , we have thatEQUATION)where (a) is from the performance difference lemma [8] .Since the relative value functionV π P (s) is L π -Lipschitz in π (Lemma 4.3), we have thatQ π P (s, •) is L π -Lipschitz in π. Therefore, we have that Q π ′ P (s, •) -Q π P (s, •) = O(∥π -π ′ ∥), which further implies that s d π P (s)⟨Q π ′ P (s, •) -Q π P (s, •), π ′ (•|s) -π(•|s)⟩ = o(∥π -π ′ ∥). Let P π denote the worst-case transition kernel for policy π. For the robust average cost, we have thatEQUATION)We then have thatEQUATIONTherefore, the Fréchet sub-gradient of g π P is ∇g π P with each (s, a)-entry ∇g π P (s, a) = d π P (s)Q π P (s, a). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: The paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: We do not use any existing assets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: For the robust relative value function V π P (s), we have thatV π ′ P (s) -V π P (s) = a π ′ (a|s)Q π ′ P (s, a) - a π(a|s)Q π P (s, a) = a π ′ (a|s) -π(a|s) + π(a|s) Q π ′ P (s, a) - a π(a|s)Q π P (s, a) = a π ′ (a|s) -π(a|s) Q π ′ P (s, a) - a π(a|s) Q π ′ P (s, a) -Q π P (s, a) (a) = a π ′ (a|s) -π(a|s) Q π ′ P (s, a) + π(a|s) s ′ P π ′ (s ′ |s, a) r(s, a) -g π ′ P (s ′ ) + V π ′ P (s ′ )-π(a|s)EQUATIONwhere (a) is from the Bellman equation ( 4). By re-arranging terms and from the fact that g π P is independent of s, we have thatEQUATIONwhere (a) is due to the fact that P π is the worst-case transition kernel for π. Since the LHS of ( 24) is independent of s, the RHS of ( 24) doesn't depend on s either. We then take the weighted sum of the RHS with respect to d π Pπ :EQUATION)From the definition of the stationary distribution, we have that s a d π Pπ (s)π(a|s)P π (s ′ |s, a) = d πPπ (s ′ ). Therefore,EQUATION)This completes the proof.From the performance difference lemma [8] , we have thatEQUATIONNote thatd π Pπ (s) max π ⟨Q π Pπ (s, •), π(•|s) -π(•|s)⟩ ≥ 0 since ⟨Q π Pπ (s, •), π(•|s) -π(•|s)⟩ = 0.We then have thatEQUATIONwhereC P L = max π,s d π * P (s) d π P (s)and the last step is from Lemma 3.2.For the update of our algorithmEQUATIONthe following first-order optimality condition holds for any p ∈ Π and sEQUATIONFrom the three-point Lemma of the Bregman divergence, we have thatEQUATIONCombining ( 30) and ( 31), we have that Lemma 4.4 holds.Plug p = π k in ( 14), we have that 14), we have thatη k ⟨Q π k P (s, •), π k+1 (•|s) -π k (•|s)⟩ ≤ -∥π k+1 (•|s) -π k (•|s)∥ 2 -∥π k (•|s) -π k+1 (•|s)∥ 2 . (32) Moreover, plug p = π * in (EQUATION)For the term ⟨Q π k P (s, •), π k+1 (•|s) -π k (•|s)⟩, we have thatEQUATIONwhere (a) is from the performance difference lemma [8] and (b) is from the definition of M and (32).For the term ⟨QEQUATIONTherefore, we have thatEQUATION)By re-arranging the terms in the above equation, we have thatEQUATIONF Proof of Theorem 4.6By recursively applying Lemma 4.5, we have thatEQUATION)Since the step size satisfiesη k ≥ η k-1 1 -1 M -1M , we have thatEQUATIONTherefore, we have thatEQUATION)G Proof of Theorem 4.7Set the step size η k = 1 L and define the gradient mappingG L (π k (•|s)) = L(π k (•|s) -π k+1 (•|s)). Note that if π k is updated exactly by the (sub)-gradient descent, then G L (π k (•|s)) = ∇g π P (s, •). The norm ∥G L (π k (•|s))∥ measures the closeness of the current step to the first-order stationary point. To prove Theorem 4.7, we first show the following lemma. Lemma G.1. For the robust average cost, setting L = L π we have the following inequalities.EQUATIONProof. For the update of our algorithmEQUATIONthe following first-order optimality condition holds for any p ∈ Π and sEQUATIONBy letting p = π k , we have thatEQUATIONFrom Lemma 4.1, we further have thatEQUATIONTherefore, we have thatEQUATION)On the other hand, from Lemma 4.2, we have thatEQUATIONFor Qπ k+1 P (s, •), π k+1 (•|s) -π(•|s), from the first-order optimality condition (43) , we have thatEQUATIONwhere (a) is due to the fact that the robust relative action value functionQ π P is L π -Lipschitz (Lemma 4.3) and (b) is due to the fact that ∥π k+1 (•|s) -π(•|s)∥ ≤ 2|S|.Therefore, we have thatEQUATIONBy setting L = L π , we have thatEQUATIONwhere the second inequality is from Jensen's inequality.We are now ready to prove Theorem 4.7.Proof. Let ω = 1 (C P L 2 √ 2|S|) 2 .We have thatEQUATIONLet δ k = g π k P -g * P . We then have that ω L π δ 2 k+1 ≤ (g π k P -g * P ) + (g * P -gEQUATIONWe then divide both side by δ k δ k+1 and have thatEQUATION)We then take the sum of the above inequality over iterations 0, 1, • • • , k -1. It then follows thatEQUATIONFrom (46) , we have that g π k P -gEQUATIONLπ ∥G Lπ (π k (•|s))∥ 2 ≥ 0. Therefore, we have that g π k P ≥ gπ k+1 Pand thus δ i+1 ≤ δ i .For any two constant l, m ∈ (0, 1), let n(k, l) be the number of steps that δi+1 δi is at least l over the first k iterations. If n(k, l) ≥ mk, we then have that δi+1 δi ≥ l at least ⌈mk⌉ times. Therefore,EQUATIONWe then have thatEQUATIONIf n(k, l) < mk, we have that δi+1 δi < l at least ⌈(1 -m)k⌉ times. Since δ i+1 ≤ δ i , we have thatEQUATIONSince ( 56) and (57) hold for any l, m ∈ (0, 1), we have thatEQUATION)By letting l = m = 1 2 , we have thatEQUATION)This completes the proof.Answer Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).Justification: There is no potential societal impact that needs to be specified.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).38th Conference on Neural Information Processing Systems (NeurIPS 2024).∆(S) denotes the probability simplex defined on S. 


Decision: The paper analyzed policy gradient methods for average reward robust MDPs.
The contribution and challenge here are mainly technical, since the results for the discounted case already appear in previous work.
Nevertheless, all reviewers (some of them experts in the field) unanimously agree that the contribution is important and relevant to the neurips community and that the paper is well written.
I will therefore accept.

Reviews: review: paper_summary: paper_summary: The authors consider the mirror descent algorithm in the context of robust average cost MDPs. The consider $(s,a)$-rectangular uncertainty sets across a general distance metric. The Bregman divergence chosen for the purpose of analysis is the Euclidean 2-norm distance. The authors leverage on a prior result proving the existence of a corresponding Bellman equation for the robust average cost and another prior result proving Lipschitzness of the risk neutral average cost value function to derive the results in the paper. Since the robust average cost is the maximum average cost across the uncertainty set it is not differentiable and hence the authors resort to using the sub-gradient as a substitute. This sub gradient is shown to be the robust $Q$ function. Subsequently, the performance difference inequality in conjunction with a PL result is proven and leveraged to obtain the final convergence bounds which are of $O(1/\epsilon)$. They characterize these bounds for both a constant and increasing step size.
strength_weakness: strength_weakness: 1. The authors seem to have missed a line of work pertaining to the robust average cost. Although they consider general $(s,a)$-rectangular sets with any distance metric, much of the prior work has considered the KL (Kullback Leibler) distance metric whose dual formulation yields the exponential cost robust MDPs. Some of the earlier works on this include [Borkar 2010], [BS02] where they characterize asymptotic behavior of dynamic programming algorithms for this robust formulation. Some of the more recent work in this domain can be found in [MMS23] and [MMS24]. [MMS23] considers a policy gradient algorithm for this robust setting, but provides asymptotic convergence bounds to a stationary point. [MMS24] considers modified policy iteration, an algorithm that generalizes both policy and value iteration and provides with finite time global convergence guarantees. Hence the repeated claim in the paper to have characterized the first policy based global convergence bounds in the context of robust average cost MDPs has to be modified accordingly. 

2. The authors assume an oracle provides with $Q^\pi_\mathcal{P}$, that is the value function corresponding to the argmax kernel associated with policy $\pi$. In the risk neutral case, one can solve for the state action value function provided corresponding to a policy $\pi$ by solving for the average cost Bellman equation (which typically involves a matrix inversion). It is not clear as to how to evaluate this robust value function even when you have access to the model, that is the nominal transition matrix and the single step costs. 

3. The authors assume that the underlying Markov chain is unichain. However, unichain Markov chain implies the existence of some transient states whose associated stationary measure is zero. In that case, the concentrability coefficient $C_{PL}$ will be infinity. Hence it seems like ergodicity is a requirement. 

4. The increasing step size requires knowledge of $M$, which is not apparent even in the full information setting. For instance, reference [16] in the paper characterizes the Lipschitz constant in terms of the model considered and hence can be evaluated when given access to the complete model information such as the probability transition matrix, etc. However, it is not apparent as to how to determine M based on its definition in line 254. 



[Borkar 2010] Borkar, V. S. (2010, July). Learning algorithms for risk-sensitive control. In Proceedings of the 19th International Symposium on Mathematical Theory of Networks and Systems–MTNS (Vol. 5, No. 9).

[BS02] Borkar, V. S., & Meyn, S. P. (2002). Risk-sensitive optimal control for Markov decision processes with monotone cost. Mathematics of Operations Research, 27(1), 192-209.

[MMS24] Moharrami, M., Murthy, Y., Roy, A., & Srikant, R. (2024). A policy gradient algorithm for the risk-sensitive exponential cost MDP. Mathematics of Operations Research.

[MMS23] Murthy, Y., Moharrami, M., & Srikant, R. (2023, June). Modified Policy Iteration for Exponential Cost Risk Sensitive MDPs. In Learning for Dynamics and Control Conference (pp. 395-406). PMLR.
questions: questions: 1. Is there a closed form expression for the value function corresponding to the solution of Equation 7? In exponential cost MDPs, these value functions do indeed have a closed form expression and I was curious about their generalizations to other distance metrics.

2. Since the Bregman divergence considered is the Euclidean 2 norm, it feels more accurate to refer to the algorithm as policy gradient rather than mirror descent. Mirror descent is more commonly employed when the KL divergence is used as the Bregman divergence. Is there any specific reason for not referring to this algorithm as policy gradient?

3. Since the minimization problem is considered, it is more relevant to refer to the setting as average cost rather than average reward as is currently done. 


Some typos: 

1. Line 463.5 in the appendix: should be $+\sum_a\pi(a|s)$ instead of $-\sum_a\pi(a|s)$

2. In equation 10: should be $\nabla g^{\pi_k}_\mathcal{P}$ instead of $\nabla g^{\pi}_\mathcal{P}$
limitations: limitations: Even though the limitations are not explicitly addressed, the paper has been presented clearly and hence the limitations are apparent.
paper_summary: paper_summary: The authors present a gradient-based algorithm for average reward robust MDP (finite MDP). This setting has been studied in prior works, however, this paper proposes a policy optimization-based algorithm which is not yet done. They do a theoretical analysis of this setting and show linear convergence (by increasing step size), finite iteration complexity and quite nicely show satisfying PL condition and therefore perhaps trivially global convergence.
strength_weakness: strength_weakness: - Presentation of theoretical results is not sharp or rather quite bad. They do not mention the assumptions for the lemma/theorem statement holds. I am uncertain if there are any further assumptions.
-  Much of the used theory is already exists in the literature of robust MDP and infinite horizon case (which is not bad since authors make sure to explain that they not trivially using but extending it, e.g., 197-204 lines). This is currently described in pieces throughout the paper, I would appreciate it if authors could concretely explain at one place (perhaps in related works) the tools used from prior works.
questions: questions: Questions:

43-44 -> What are these algorithms based upon in contrast to gradient-based? 

Is there an assumption on eqn 3 about boundedness? Why does V and Q are necessarily finite?

What will be a sufficient condition for assumption 2.1? Also motivating beforehand on why you need this assumption will be good.

Is there a bound/relation between C_{PL} and M? 

What does the non-robust method line 307 optimize for? what is it's objective and used transition?
limitations: limitations: None
paper_summary: paper_summary: The paper studies gradient-based methods for robust average reward MDPs. The paper first derives a sub-gradient for the robust average reward (which is nonsmooth), and then uses it to define a mirror descent algorithm. They prove a few structural properties of the setting and then use them to provide a convergence guarantee for their algorithm. They also provide some experimental results.
strength_weakness: strength_weakness: The main weakness of the paper is what I perceive to be several technical concerns (listed in the questions section). I would raise my soundness score and overall evaluation if they were all adequately addressed.
questions: questions: Maybe the unichain assumption 2.1 should be placed earlier, since the stationary distributions $d^\pi$ are not generally well-defined in the absence of this assumption.

In line 177, is the definition of $d_{\mathcal{P}}^\pi$ unique?

Lemma 3.2: similar question, I think the worst-case transition kernel for a policy $\pi$ is not unique. In this case I wonder if $d_{\mathcal{P}}^\pi$ and $Q_{\mathcal{P}}^\pi$ are uniquely defined? And thus my main question is, for the sub-gradient, can we use different worst-case transition kernels in $d_{\mathcal{P}}^\pi$ and $Q_{\mathcal{P}}^\pi$, or do they need to use the same one? (The proof makes it seem like they are using the same one?) Thus maybe the lemma statement would need to be amended to discuss this.

Lemma 4.2: Is $C_{PL}$ guaranteed to be finite? The usual definition of unichain (e.g. in Puterman) is that the transition matrix has a single recurrent class plus a (possibly empty) set of transient states. The transient states would have stationary measure 0 which could cause $C_{PL}$ to blow up?

I have a similar concern about the quantity $M$ in Lemma 4.5. I am also confused about the denominator (is there a missing subscript $\pi$?)

Where is the proof of Lemma 4.3? It is claimed that a result from [16] can be easily extended to a more general (unichain) setting; however, because such details can be very thorny in the average reward setting, I think a proof should be provided.
limitations: limitations: No major limitations
paper_summary: paper_summary: The authors introduce study a policy gradient algorithm for solving unichain average reward robust MDPs. They show a linear convergence rate for increasing step sizes and a $O(1/k)$ convergence rate for fixed step size, where $k$ is the number of iterations of the algorithms.
strength_weakness: strength_weakness: I only have three minor comments as regards the claims of the authors, which I find a bit too far from the exact results proved in their theorems. I think that the paper requires some clarifications. 

* The authors should mention the unichain assumption much earlier in the paper. The word *unichain* isn’t even present in the abstract, in the introduction, it is only mentioned in the very end of section 2. This assumption is crucial to establish several important properties of average value function (such as uniformity across states) and relative value function (such as continuity). This needs to be emphasized much earlier, and highlighted when the authors give their main contributions. Given the connection between mean-payoff perfect-information stochastic games and robust MDPs [0], the sentence from l13-14 of the abstract: *Moreover, our algorithm is the first algorithm that converges to the global optimum with polynomial time for robust average reward MDPs* would mean that the authors have solved a long-standing open problem, which is not the case since they solely focus on unichain instances (although I appreciate the results of the authors). Please clarify. 

* I am a bit confused with the claim about the polynomial time complexity of the algorithms. Usually, a (weakly) polynomial time complexity refers to a number of arithmetic iterations polynomial in $\log(1/\epsilon)$, in the dimension of the problems and in the logarithm of the entries of the problem (here, the number of bits necessary to represent the rewards and the transition probabilities). So why is $O(1/\epsilon)$ a polynomial-time complexity? Similarly, the $O(.)$ notations seem to be hiding the $C_{PL}$ term from (15), and it’s not clear how this is bounded in terms of $S,A$ and the parameter of the problem. Same with the constant $L_{\pi}$ from (16), can you provide a bound on this? The closed statement from polynomial-time complexity appears to be Th 4.6 but there again we have terms like $M$ appearing in the bound, and it’s not clear how to bound them. Please clarify or reformulate your claims about *polynomial-time* algorithms.

* Some of the definitions appear incorrect. In all generality the definition of the bias should involve some Cesaro averages or the sum may diverge, the definition given by the authors is correct only for aperiodic Markov chains (see my other comments below).

[0] Chatterjee, K., Goharshady, E. K., Karrabi, M., Novotný, P., & Žikelić, Đ. (2023). Solving long-run average reward robust mdps via stochastic games. arXiv preprint arXiv:2312.13912.
questions: questions: * The authors in [1] mention a $O(1/\epsilon)$ convergence rate for their algorithm in their conclusion. How does that compare to your $O(1/\epsilon)$ result?

* Under the unichain assumption, it is well-known that discounted MDP with a discount factor sufficiently large can solve average reward MDPs (typically with a discount factor larger than $1-\epsilon/D$ where $D$ is the diameter of the MDP). There is a vast number of recent papers on policy gradient methods for nominal discounted MDPs and discounted robust MDPs. Can the authors compare the complexity of their algorithms with the complexity of the best policy gradient methods for discounted MDPs combined with the appropriate (large) discount factor for unichain MDPs?

* I understand that the robust value function may be non-smooth for robust MDPs. But isn’t it also the case in the discounted setting? How do authors in the discounted setting bypass that for designing policy gradient algorithms? Is it the same `trick' in your paper? Also, for what geometry of the uncertainty set is the robust value function smooth? 

* l159-160, definition of the bias: in all generality the definition of the bias/relative value function should involve the Cesaro average, see p 338 in [2]. In my understanding the definition given by the authors only holds fors for *aperiodic* Markov chains. Please clarify.

* A minor comment: Usually, rewards are maximized, and costs are minimized. In this paper, rewards are minimized, which is a bit weird, even though the maths don’t change.

* Please introduce the equations for $d^{\pi}$ in the main body.

l287: typo, the first two sentences of rmk 4.8 should be one sentence.

[1] A. Tewari and P. L. Bartlett. Bounded parameter markov decision processes with average reward criterion. In International Conference on Computational Learning Theory, pages 263–277. Springer, 2007.

[2] M. L. Puterman. Markov decision processes: Discrete stochastic dynamic programming, 1994
limitations: limitations: No problem here.

Comments: title: Response to the authors
comment: Thanks for addressing my comments! 

I am satisfied with the responses and am increasing my score. I am still not convinced that the algorithm should be referred to as policy mirror descent since the Euclidean norm as Bregman divergence yields policy gradient. I understand the need to use the sub gradient (due the maximization objective), but the divergence considered is however the Euclidean norm and hence policy gradient (or sub-gradient)  is perhaps more accurate than mirror descent. 

Once again, thank you for your response!
We thank the reviewer again for the response and helpful comments, which greatly helped to improve the quality of the
paper.
We thank the reviewer for the response. After reading the reviewer's comments, we found that assuming ergodicity instead of unchain is more reasonable. In this case, we don't need to make additional assumptions on $C_{PL}$ and $M$. We will present the assumption earlier in the paper and make it clear.
We thank the reviewer for the response. We will revise our paper based on the comments.
Thank you for your response. From this response it seems like therefore the results are based on the ergodic (one recurrent class with no transient states) assumption, rather than unichain? If so, I would be happy to increase my score if this is presented earlier and more centrally in the paper.
I thank the authors for their replies. I went over them and do not have any further questions. Please include the changes as promised in your revised version, particularly regarding assumptions.
I would like to thank the authors for their (very!) detailed response to my comments. Provided that they indeed make all the changes listed above, the paper will read better and contribute to the literature on gradient methods for robust MDPs. I increased my score in consequence.