Provable Policy Gradient for Robust Average-Reward MDPs Beyond Rectangularity

Thank you for your insightful questions and taking the time to assess our manuscript.

1. Comment 1: Additional clarification in proofs could be helpful .

Thank you very much for providing helpful suggestions on writing style and missing clarification! We will clarify the wording in the next version of our manuscript, including more detailed explanations, to improve readability.

2. Comment 2: Technical novelty need to be further elaborated.

We'd like to emphasis that the extension from the discounted setting to the average-reward setting is non-trivial, even the average-reward MDP can be seen as a special case of discounted MDPs ($\gamma=1$). Please see our reply to Reviewer iYjz, Comment 2 and the introduction of [2] for a detailed explanation.

For the outer loop, we establish Lipschitz continuity via a novel sensitivity analysis and introduce a gradient dominance condition specifically tailored to the average-reward setting, leading to an outer-loop convergence guarantee. Compared to [2], we provide improved smoothness coefficients. While our approach adopts a double-loop structure, which is widely used in various fields such as game theory [3], min-max optimization [4,5], and robust MDPs [1], our analysis relies on standard non-convex minimax optimization techniques, which are significantly different from the mirror descent analysis used in [1].

For the inner loop, we propose two tailored algorithms for worst-case transition evaluation under both rectangular and non-rectangular settings, whereas [1] only addresses the rectangular case. To establish the inner-loop convergence guarantee, we derive the first form of the adversarial transition gradient (Lemma 4.1), establish first adversarial smoothness conditions using sensitivity analysis techniques (Lemma 4.3), and prove the first adversarial gradient dominance condition (Theorem 4.4), along with corresponding inner convergence guarantee.

Therefore, our work differs substantially from both references [1,2]. Building on these theoretical advancements, our approach offers a comprehensive framework for addressing robust average-reward MDPs beyond existing methods. We will further clarify these contributions in the next version of our paper.

[1] Wang, Q., Xu, S., Ho, C. P., and Petrick, M. Policy gradient for robust markov decision processes.

[2] Kumar, N., Murthy, Y., Shufaro, I., Levy, K. Y., Srikant, R., and Mannor, S. On the global convergence of policy gradient in average reward markov decision processes.

[3] Ding D., Wei C. Y., Zhang K., & Jovanovic M. 2022. Independent policy gradient for large-scale markov potential games: Sharper rates, function approximation, and game-agnostic convergence.

[4] Jin C., Netrapalli P., & Jordan M. 2020. What is local optimality in nonconvex-nonconcave minimax optimization?

[5] Davis D., & Drusvyatskiy D. 2019. Stochastic model-based minimization of weakly convex functions.

Thank you very much for taking the time to review our paper and your insightful comments as well as suggestions.

Comment 1: Technical details in Section 4.3 should be added.

Thank you so much for pointing out the missing details in this section! That would be definitely helpful if full technical details are provided. We will include a detailed discussion of the proof and analysis in the next version of our paper.

Comment 2: Statement with "... is the first..." could be misleading.

Thank you for your suggestion. We would like to emphasize that while the technical tools for solving generic non-convex problems are well established [1], no existing literature has applied such analysis to robust average-reward MDPs. To the best of our knowledge, the most recent policy gradient method for robust average-reward MDPs [2] is limited to the more conservative $(s,a)$-rectangular case. We will clarify the wording in the manuscript in our revision to improve readability.

Comment 3: Considering additional benchmarks in the experiment could be helpful.

Thank you for this insightful suggestion. We do totally agree that incorporating additional benchmarks would strengthen our experimental evaluation. As suggested by Reviewer sWdH, we conducted an experiment in an inventory control setting to demonstrate RP2G's superior performance (see https://drive.google.com/file/d/1VKnmT5_Wzpj6PwH_UbHImimhrhKBpVgq/view?usp=sharing). Our results show that the policy obtained by solving the non-rectangular RAMDP is less conservative (see our reply to Reviewer sWdH, Comment 1 for a detailed explanation). We will evaluate our algorithms on other benchmarks in the next version of this paper.

Comment 4: Typos.

Thank you for pointing out the typos. We will update our paper according to these suggestions!

Question 5: Could you incorporate the estimation error of Algorithm 3 into the overall convergence bound in Theorem 3.5?

Thank you for raising this important question. We would like to emphasize that the theoretical guarantee for RP2G (Theorem 3.5) accounts for the inner output error ($\delta_{t}$ at $t$-th iteration). However, unlike Algorithm 2 for the inner problem, which features a simple and effective structure for updating the inner transition kernel, Algorithm 3 employs Monte Carlo sampling to achieve a probabilistic convergence guarantee, which introduces additional estimation error. Therefore, incorporating this estimation error from Algorithm 3 (as detailed in Theorem 4.6) into a deterministic convergence guarantee (shown in Theorem 3.5) remains challenging. We believe that investigating the overall convergence bound while accounting for the inner worst-case evaluation error is an interesting direction for future work that could further enhance the significance of our results.

Question 6: What makes average-reward MDPs so different from discounted MDPs?

This difference arises in both theory and practice. While the RAMDP be seen as a special case of the RMDP ($\gamma=1$), extending the theoretical framework from the discounted setting to the average-reward setting is non-trivial (see our reply to Reviewer iYjz, Comment 2 and [3] for a detailed explanation). In practice, many real-world systems prioritize steady-state or long-term behaviour, where policies derived from discounted MDPs may perform poorly (see Introduction line 37-48).

Question 7: Is the mixing time important to the theoretical analysis?

At this stage, the mixing time is crucial for bounding the differential (action) value function, which plays a key role in establishing Lipschitz continuity and the gradient dominance condition.

Question 8: Is this approach potentially adaptable to discounted MDPs or episodic MDPs?

Thank you for this insightful question. Regarding the discounted MDPs, [4] already proposed methods with a double-loop structure for solving RMDPs. However, extending our approach to episodic MDPs seems challenging, as the ergodicity assumption does not hold in this setting, preventing a direct application of our theoretical framework. Developing a suitable policy gradient framework for episodic MDPs remains an important direction for future research.

[1] Lamperski, A. 2021. Projected stochastic gradient langevin algorithms for constrained sampling and non-convex learning.

[2] Sun, Z., He, S., Miao, F., & Zou, S. 2024. Policy optimization for robust average reward mdps.

[3] Kumar N, Murthy Y, Shufaro I, et al. 2024. On the global convergence of policy gradient in average reward markov decision processes.

[4] Li M., Kuhn, D., & Sutter T. 2023. Policy gradient algorithms for robust mdps with non-rectangular uncertainty sets.

We are sincerely grateful to the reviewer for the insightful comments and valuable questions.

Comment 1. Additional clarification on the difference between our work and Li et al. (2023) is needed.

We apologize for the insufficient explanation of the contributions of this work. It is important to clarify that the extension from the discounted-reward setting to the average-reward setting is non-trivial (see our reply to Comment 2 for a detailed discussion). While both our work and that of Li et al. (2023) rely on a double-loop structure, along with similar properties such as the smoothness of the objective function $J$ and the gradient-dominance condition, the technical foundations used to establish these properties are quite different. Compared to their work, our contributions span both loops. For the outer loop, we establish Lipschitz continuity via a novel sensitivity analysis and introduce a gradient dominance condition specifically tailored to the average-reward setting, leading to the convergence guarantee. For the inner loop, We propose a tailored projected gradient ascent algorithm for worst-case transition evaluation under rectangularity, with a convergence guarantee. This approach is structurally different from Algorithm 3.2 in Li et al. (2023). Although our inner convergence results may appear similar, the underlying analytical tools are fundamentally different.

We hope this clarifies the key differences between our work and that of Li et al. (2023), both in theoretical foundations and algorithmic contributions.

Comment 2: Recovering average-reward results from the discounted setting could be discussed.

We thank the reviewer for raising this point. While it is theoretically possible to recover average-reward results by taking the limit as $\gamma \to 1$ in the discounted-reward setting, this approach is often infeasible in practice. As shown in Li et al. (2023), the number of iterations is $\mathcal{O}(\frac{1}{(1 - \gamma)^4})$ (hidden within their constants), which increases rapidly as $\gamma\to 1$. This makes the computational cost prohibitively high for large $\gamma$. We also provide empirical evidence in Appendix E.5, where we observe a significant increase in iteration count as $\gamma$ increases. These results highlight the practical limitations of using discounted-reward methods for solving robust average-reward MDPs, and motivate the need for a tailored algorithm, as proposed in our work.

Comment 3: Additional explanation for Figure 2 could be helpful.

Thank you for pointing this out. Regarding the experimental setup, we obtain and record the policies from both robust and non-robust AMDPs at each iteration. To assess policy robustness, we evaluate their performance under the worst-case transition scenario, i.e., computing $\max_{\boldsymbol{ p} \in \mathcal{P}} J(\boldsymbol{\pi}, \boldsymbol{ p})$ with given $\boldsymbol{\pi}$, and then record these values for plotting. This comparison is wildly adopted in robust MDPs literature [8,9] as a standard approach to demonstrate robustness.

In terms of result interpretation, the RAMDP policies consistently achieve lower costs under worst-case transitions, highlighting their effectiveness against adversarial kernels. Moreover, as the number of iterations increases, the worst-case evaluation cost stabilizes, indicating convergence. These findings are in line with our theoretical results in Section 4.3.

We hope these additional explanations could help clarify our numerical experiment.

[1] Lamperski, A. 2021. Projected stochastic gradient langevin algorithms for constrained sampling and non-convex learning.

[2] Wang, Y., Velasquez, A., Atia, G., Prater-Bennette, A., & Zou, S. 2024. Robust Average-Reward Reinforcement Learning.

[3] Riemer, M., Khetarpal, K., Rajendran, J., & Chandar, S. 2024. Balancing context length and mixing times for reinforcement learning at scale.

[4] Kearns, M., Mansour, Y., & Ng, A. 1999. Approximate planning in large POMDPs via reusable trajectories.

[5] Jin, Y., & Sidford, A. 2020. Efficiently solving MDPs with stochastic mirror descent.

[6] Puterman, M. L. 2014. Markov decision processes: discrete stochastic dynamic programming.

[7] Xiao, L. 2022. On the convergence rates of policy gradient methods.

[8] Sun, Z., He, S., Miao, F., & Zou, S. 2024. Policy optimization for robust average reward mdps.

[9] Tamar, A., Mannor, S., & Xu, H. 2014. Scaling up robust MDPs using function approximation.

Thank you for your insightful questions and constructive suggestions!

Comment 1: Additional comparison between rectangular and non-rectangular RAMDPs would be helpful.

Thanks for the suggestion! As the (non-)rectangularity only affects the ambiguity structure and appears to be independent of the chosen reward criterion [1], we adopt the non-rectangularity to the average-reward case without further justification; thus we did not design an experiment specifically to show the advantages of non-rectangularity over rectangularity under this criterion. We do agree that it would be helpful if the superiority of the non-rectangularity could be better demonstrated in this setting. To address this, we compare the performance of $(s,a)$-rectangular and non-rectangular RAMDPs in an inventory control setting with the same set size (see https://drive.google.com/file/d/1VKnmT5_Wzpj6PwH_UbHImimhrhKBpVgq/view?usp=sharing). The results, presented in the table, show that the policy obtained from the non-rectangular RAMDP is less conservative, as evidenced by its lower average cost. We will provide a more detailed discussion in the next version of our paper and highlight the advantage and superiority of non-rectangularity.

Comment 2: Further clarification of the technical novelty could be useful.

We appreciate the reviewer's concerns regarding our technical novelty and would like to emphasize that the extension from the discounted setting to the average-reward setting is non-trivial (see the example in Reviewer iYjz, Comment 2)

As our proposed algorithm follows a double-loop structure, widely used in various fields, including game theory [2], min-max optimization [3,4], and robust MDPs [5], our analysis builds on standard non-convex optimization techniques, which were also adopted by [5]. Compared to [5], our novel contributions are twofold: (1) for the outer loop, we establish Lipschitz continuity via a novel sensitivity analysis and introduce a gradient dominance condition with an convergence guarantee tailored to the average-reward setting; (2) for the inner loop, we provide an effective algorithm for worst-case transition evaluation under rectangularity, along with which is structurally different from Algorithm 3.2 in [5]. While achieving a similar convergence rate, our analysis differs from [5], relying on standard non-convex optimization techniques [6].

Comment 3. Essentiality of using average-reward instead of discounted reward should be highlighted.

Thanks for your suggestion. As our work is theoretically oriented, our contribution focus on developing efficient algorithms with theoretical convergence guarantees for RAMDPs. We do understand average-reward setting is well-suited for agents concerned with long-term or steady-state policy behaviour, such as resource allocation, portfolio management, and healthcare [7,8]. In this sense, exploring the practical advantages of RAMDPs in these applications represents a promising direction for our future research.

Moreover, as shown in Appendix E.5, when $\gamma$ approaches $1$, the computational cost increases significantly when using discounted-reward MDPs to approximate average-reward solutions. This underscores the necessity of methods specifically designed for the average-reward setting.

Comment 4. Adaptation to stochastic policy gradients could be considered.

Thank you for this valuable question. Our analysis assumes a model-based setting in which the MDP structure is known except for the transition kernel, so the policy gradient can be computed exactly. When only stochastic policy gradients are available, additional challenges arise in modeling and in ensuring robustness under broader uncertainty. Extending our results to this case is a promising direction for future work, for example by incorporating robust temporal-difference methods such as [9].

[1] Wiesemann W., Kuhn D., & Rustem B. 2013. Robust Markov decision processes.

[2] Ding D., Wei C. Y., Zhang K., & Jovanovic M. 2022. Independent policy gradient for large-scale markov potential games: Sharper rates, function approximation, and game-agnostic convergence.

[3] Jin C., Netrapalli P., & Jordan M. 2020. What is local optimality in nonconvex-nonconcave minimax optimization?

[4] Davis D., & Drusvyatskiy D. 2019. Stochastic model-based minimization of weakly convex functions.

[5] Li M., Kuhn, D., & Sutter T. 2023. Policy gradient algorithms for robust mdps with non-rectangular uncertainty sets.

[6] Beck A. 2017. First-order methods in optimization.

[7] Ghalme G., Nair V., Patil V., & Zhou Y. 2021. Long-term resource allocation fairness in average markov decision process (amdp) environment.

[8] Patrick J., & Begen M. A. 2011. Markov decision processes and its applications in healthcare.

[9] Wang Y., & Zou S. 2022. Policy gradient method for robust reinforcement learning.