Policy Optimization for Robust Average Reward MDPs

We thank the Reviewer for the helpful and insightful feedback. Below we provide point-to-point responses to the weaknesses and questions.

W1. Present assumption 2.1 earlier. In the revision, we will introduce the unichain assumption earlier in abstract and emphasize that our contributions are based on the unichain assumption.

W2. Polynomial time complexity of the algorithm. Since the complexity of our algorithm is in the order of $\mathcal{O}(1/\epsilon)$ not $\mathcal{O}(\log(1/\epsilon))$, the polynomial-time complexity claim is not accurate. We will reformulate our claims in the revision. Though our algorithm doesn't have polynomial-time complexity, our algorithm is still the first algorithm that has finite time complexity for robust average reward MDPs.

Since $M$ and $C_{PL}$ depend on the structure of the underlying MDPs, characterizing the upper bound of $M$ and $C_{PL}$ is non-trivial. We would like to emphasize that similar constants also appear in prior works on robust discounted reward MDPs and average reward MDPs ([3], [4], [5], [6]). For the constant $L_\pi$, when the underlying MDP is ergodic, the upper bound of $L_\pi$ is characterized in [6], which is also depends on the structure of the MDP. From the Lemma 14 in [6], we have that $L_\pi$ is in the order of $\mathcal{O}\Big(\frac{C_e^2S^2}{(1-\lambda)^2}\sqrt{A}\Big)$, where $C_e$ and $\lambda$ are constants that characterize the geometric ergodicity of the MDPs.

W3. Clarify definitions of average reward. In [2], the Cesaro limit sense relative value function was defined for periodic chain. This is because the Markov chain is not assumed to be unichain in [2]. When the Markov chain is unichain, the Markov chain only contains a single recurrent class and some transient states, and the average reward is independent of the initial state. Since there is only one recurrent class, the relative value function does not need to involve the Cesaro average.

Q1. Compare with an existing work with $\mathcal{O}(1/\epsilon)$ convergence rate. In [1], the $\mathcal{O}(1/\epsilon)$ convergence rate was achieved for robust average reward MDPs when the uncertainty set is a bounded parameter MDP. Specifically, for each state action pair $(s, a)$, the transition probability $p_s^a(s^\prime)$ is bounded between a lower bound $l_s^a(s^\prime)$ and an upper bound $u_s^a(s^\prime)$. In our paper, we consider more general uncertainty sets. Moreover, [1] studies the convergence rate by building the connection between discounted reward MDPs and average reward MDPs while our paper focuses on policy gradient-based method. We discussed [1] in the related works section in our paper. We will also compare the $\mathcal{O}(1/\epsilon)$ convergence rate in the revision.

Q2. Compare with policy gradient method for discounted MDPs with large discount factor. There are recently works on policy gradient methods for robust discounted MDPs. In [3], the robust MDPs are considered under the $R$-contamination uncertainty set. The robust policy gradient converges to the global optimum with iteration complexity $\mathcal{O}(1/\epsilon^3)$. In [4], the double-loop robust policy gradient was proposed and was further proved to converge to the global optimum with complexity $\mathcal{O}(1/\epsilon^4)$. In [6], with sufficiently large step size, the robust policy mirror descent converges to the global optimum with complexity $\mathcal{O}(1/\epsilon)$. To the best of the author's knowledge, there is no works building the connection between the robust discounted reward and robust average reward. Therefore, it's non-trivial to compare the complexity of our algorithm with the policy gradient methods for discounted MDPs combined with the large discount factor.

The authors of [7] shows that there exists a discount factor $\gamma_{bw}\in(0, 1)$ such that any policy that is $\gamma$-discount optimal for $\gamma\in(\gamma_{bw}, 1)$ is also Blackwell optimal when the uncertainty set is definable. However, how to choose $\gamma_{bw}$ is still an open problem. In [7], only value iteration based algorithms are proposed and no stopping criterion is introduced for their algorithms.

Q3. Bypass non-smoothness of robust value function in discounted MDPs. For discounted reward MDPs, different methods are proposed to bypass the non-smoothness of the robust value function. In [3], the total-variation uncertainty set is studied and the non-smoothed robust value function is smoothed by approximating the max operator using the LogSumExp (LSE) operator. In [4], the non-smoothness of the robust value function is bypassed by introducing the Moreau envelope function, which can also be viewed as a smoothed version of the robust value function. In [5], the authors develop a similar bound as the Lemma 4.1 in our paper by applying the Bellman equation of robust value function, and further apply it to show the convergence of the mirror descent algorithm. However, it is not applicable in our case as the average reward doesn't satisfy the Bellman equation.

A sufficient condition for the robust value function to be smooth is that there exists a constant $C$ such that $\|d\_s\^{\pi, \mathsf{P}\_{\pi}} - d\_s\^{\pi\^\prime, \mathsf{P}\_{\pi\^\prime}}\|\leq C\\|\pi-\pi\^\prime\\|$, which depends on the geometry of the uncertainty set. Since it involves the worst-case transition kernel and the closed-form expression doesn't exist, it unclear for what geometry of the uncertainty set is the robust value function smooth.

Q4. Clarify the definition of relative value function. Please refer to the response of W3

Q5. Average cost. In the revision, we will replace rewards with costs.

Q6. Introduce the equations for $d^\pi$. We provide the equations for $d^\pi_\mathsf{P}$ in the revision: $d^\pi_\mathsf{P}$ is the stationary measure of $s$ under the transition kernel $\mathsf{P}$ satisfies that $d^\pi_\mathsf{P} \mathsf{P} = d^\pi_\mathsf{P}$.

Q7. Typo. Fixed.

W1. Listed in the question section. Please refer to the response of Q1-Q6.

Q1. Unichain assumption should be placed earlier. We thank the reviewer for this comment. In the revision, we will introduce the unichain assumption earlier in the paper and emphasize that our contributions are based on the unichain assumption.

Q2. Uniqueness of $d_\mathcal{P}^\pi$. The worst-case transition kernel may not be unique for general uncertainty sets, and thus $d_\mathcal{P}^\pi$ may not be unique.

Q3. Uniqueness of $d_\mathcal{P}^\pi$ and $Q_\mathcal{P}^\pi$. Though the worst-case transition kernel may not be unique, $Q\_\mathcal{P}\^\pi$ is unique. From the definition of the worst-case transition kernel, for a worst-case transition kernel $\mathsf{P}\_\pi$ of a policy $\pi$, we have that $Q\_{\mathsf{P}\_\pi}\^\pi = \max\_{\mathsf{P}\in\mathcal{P}} Q\_{\mathsf{P}}\^\pi$. Therefore, the robust action value function is unique. For the sub-gradient, the worst-case transition kernel of $d\_\mathcal{P}\^\pi$ and $Q\_\mathcal{P}\^\pi$ can be different, but this does not affect our results.

Q4. Boundedness of $C_{PL}$. We agree with the reviewer that in the unichain setting, $C_{PL}$ might not be finite if transient state exists. The constant $C_{PL}$ depends on the structure of the underlying MDPs and also appears in prior works on discounted reward MDPs and non-robust average reward MDPs ([1], [2], [3], [4]). If the Markov chain is ergodic, i.e., there is only a single recurrent class, then $C_{PL}$ is guaranteed to be finite.

Q5. Boundedness of $M$. Similar as $C_{PL}$, the constant $C_{PL}$ depends on the structure of the underlying MDPs. If the Markov chain is ergodic, then $M$ is guaranteed to be finite. The definition of $M$ is our paper is accurate. For the denominator, the subscript $\mathsf{P}$ denotes the transition kernel that we maximize. In the numerator, the subscript $\mathsf{P}_\pi$ denotes the worst-case transition kernel of $\pi$.

Q6. Proof of Lemma 4.4. In Lemma 14 of [4], the $L_\pi$-Lipschitz of the relative value function is proved. In their paper, $L_\pi = 2C_m^2C_p\kappa_r + 2C_mC_r$. For $C_p$ and $\kappa_r$, the proof of Lemma 18 in [4] shows that the ergodicity of the Markov chain is not required and the results can be extended to the unichain setting. We only need to show that $C_m$ is finite for unchain. In [4], $C_m$ is defined as the maximum of the operator norm of the matrix $(I-\Phi P_\pi)^{-1}$ across all policies $\pi\in\Pi$. For any $\pi\in\Pi$, following the proofs in Section A.2 of [4], we have that the eigenvalues of $(I-\Phi P_\pi)$ is non-zero for unchain since unchain only has a single recurrent class and the stationary distribution exists. Therefore, for any $\pi\in\Pi$, $\|(I-\Phi P_\pi)^{-1}\| < \infty$. Moreover, since the collection of all policy $\Pi$ is compact, we have that the maximum of the operator norm of the matrix $(I-\Phi P_\pi)^{-1}$ across all policies $\pi\in\Pi$ is finite. Therefore, $L_\pi$ exists for unchains. We will add the proof sketch of Lemma 4.4 in the revision.

Reference. [1] Y. Wang and S. Zou. Policy gradient method for robust reinforcement learning. In Proc. International Conference on Machine Learning (ICML), volume 162, pages 23484–23526. PMLR, 2022.

[2] Q. Wang, C. P. Ho, and M. Petrik. Policy gradient in robust mdps with global convergence415 guarantee, 2023.

[3] Y. Li, G. Lan, and T. Zhao. First-order policy optimization for robust markov decision process. arXiv preprint arXiv:2209.10579, 2022.

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

We thank the Reviewer for the helpful and insightful feedback. Below we provide point-to-point responses to the weaknesses and questions.

W1. Didn't mention the assumptions for the lemma/theorem statements. We thank the reviewer for this comment. In the revision, we will rephrase the statements of lemmas and theorems to include the assumptions we use. Specifically, Theorem 4.6: Under Assumption 2.1, set the step size $\eta_k \geq \eta_{k-1}\Big(1-\frac{1}{M}\Big)^{-1}M$, the robust policy mirror descent satisfies $g\_\mathcal{P}\^{\pi\_k} - g\_\mathcal{P}\^\* \leq (1-\frac{1}{M})\^k(g\_\mathcal{P}^{\pi\_0} - g\_\mathcal{P}\^\*) + (1-\frac{1}{M})\^{k-1}\frac{1}{M\eta\_0}\mathbb{E}\_{s\sim d\_{\mathsf{P}\_{\pi\_0}}\^{\pi\^\*}}[\|\pi\^\*(\cdot|s) - \pi\_0(\cdot|s)\|\^2].$

For Theorem 4.7: Under Assumption 2.1, let step size $\eta = \frac{1}{L_\pi}$ for all $k\geq 1$, we have that the each iteration of the robust policy mirror descent satisfies g\_\mathcal{P}\^{\pi\_k} - g\_\mathcal{P}\^\* \leq \max\Big\\{\frac{4L_\pi}{\omega k},(\frac{\sqrt{2}}{2})^k\big(g_\mathcal{P}^{\pi_0} - g_\mathcal{P}^*\big)\Big\\}, where $\omega = (2\sqrt{2S}C_{PL})^{-2}$.

W2. Explain the tools used from prior works at one place. We thank the reviewer for this comment. We summarize the results that are used from prior works in Section 1.2. We will add the following paragraph in the revision:

In [1], the Fr'echet sub-gradient has been derived for robust discounted reward MDPs. However, the Fr'echet sub-gradient for robust discounted reward MDPs in [1] can not be extended to the average setting since in [1], the discounted factor $\gamma$ is required to be strictly less than 1. In [2], the performance difference lemme was derived for average reward MDPs. In our paper, we provide a lower bound on the difference of robust average reward between two policies. Such a bound was also derived in [1] for robust discounted value function by applying the Bellman equation of robust value function, which is not applicable in our case as the average reward itself doesn’t satisfy the Bellman equation. We also extend the Lipschitz property of the non-robust relative value function in [3] to the unichain setting. In [4] and [5], the convergence rate of the policy gradient descent for non-robust discounted reward MDPs was derived. In their analyses, the smoothness of the value function is required. However, the robust average reward might not be smooth. Therefore, the approaches applied in [4] and [5] can not be directly extended to robust average reward MDPs.

Q1. Details on related works. [6] and [10] build the connection between the discounted reward MDPs and average reward MDPs and prove the existence of Blackwell optimal policies. The algorithms based on value iteration are further proposed. [7] is an extension of [6] with parts of the state space having arbitrary transitions and other parts are purely stochastic. In [8] and [9], the model-based and model-free robust average reward MDPs are studied and the robust relative value iteration algorithms are proposed. We will add this discussion in the revision

Q2. Boundedness assumption on V and Q. We thank the reviewer for pointing this out. Since $V$ and $Q$ are unique up to an additive constant, similar as [3], we consider the projection of the value function onto the subspace orthogonal to the $**1**$ vector. In this case, our theoretical results still hold.

Q3. Sufficient condition and motivation for assumption 2.1. The sufficient condition for Assumption 2.1 is that for any $\pi\in\Pi$ and $\mathsf{P}\in \mathcal{P}$, the Markov chain only contains a single recurrent class and some transient states. The unichian assumption is important and widely used since under the unichain assumption, the robust average reward is identical for every starting state. To derive Lemma 4.1, we leverage the fact that the robust average reward is independent of the initial state. If the unichain assumption doesn't hold, the inequality may not hold, and this is a key step to derive the global optimality of our algorithm.

In practice, it is often the case that only unichains are of interest to the decision-making problem. For example, the nominal transition kernel is obtained from samples, and is a unichain. Then, the true transition kernel must be a unichain, and therefore to obtain a robust policy that works well on the true transition kernel, it is sufficient to only consider unichain. Even in standard MDP, extending results from unichain to e.g., multichain problems is very challenging. Extending our results under a relaxed assumption beyond unichain is even more challenging, and it is of future interest.

Q4. Relation between $C_{PL}$ and $M$. The relation between $C_{PL}$ and $M$ is unclear. We have that M = \sup\_{\pi, \mathsf{P}\in\mathcal{P}}\Big\\|\frac{d\_{\mathsf{P}\_\pi}\^{\pi\^\*}}{d\_\mathsf{P}\^\pi}\Big\\|\_\infty = \sup\_{\pi, s, \mathsf{P}\in\mathcal{P}} \frac{d\_{\mathsf{P}\_\pi}\^{\pi^*}(s)}{d\_\mathsf{P}\^\pi(s)} and $C\_{PL} = \max\_{\pi, s} \frac{d\_{\mathcal{P}}\^{\pi\^\*}(s)}{d\_\mathcal{P}\^\pi(s)} = \max\_{\pi, s} \frac{d\_{\mathsf{P}\_{\pi\^\*}}\^{\pi\^\*}(s)}{d\_{\mathsf{P}\_{\pi}}\^\pi(s)}$. The sup of $M$ and $C_{PL}$ can be achieved at different $\pi$. In the numerator, it's unclear if $d\_{\mathsf{P}\_\pi}\^{\pi\^\*}(s)$ is larger than $d\_{\mathsf{P}\_{\pi\^\*}}\^{\pi\^\*}(s)$ or not.

Q5. Objective and transition kernel of baseline algorithm. For the baseline algorithm, it's trained under the known nominal transition kernel and the objective is to maximize the non-robust average reward.

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.

We thank the Reviewer for the helpful and insightful feedback. Below we provide point-to-point responses to the weaknesses and questions.

W1. Related works on robust average cost MDPs. We thank the reviewer for pointing this out. In the revision, we modify our statement as 'We characterized the first policy based global convergence bounds with general uncertainty sets for robust average cost MDPs.' We also compare our work with the related works [Borkar 2010], [BS02], [MMS24], [MMS23]. We add the following paragraph in Section 1.2 in the revision:

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 [Borkar 2010], the Q-learning and the actor-critic method are described and the asymptotic performance are characterized for risk sensitive robust MDPs. In [BS02], the value iteration and policy iteration algorithms are also analyzed for risk sensitive MDPs. Recently, in [MMS23], 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 [MMS24]. In our paper, we study the robust average reward MDPs with general uncertainty sets and characterize the global convergence of our algorithm.

W2. Evaluate the robust relative value function. The robust relative value function can be evaluated by generalizing the relative value iteration to the robust setting. The algorithm is presented as follows:

Algorithm parameters: $V_0$, $\epsilon$, and reference state $s^* \in \mathcal{S}$
Initialize $w_0 \leftarrow V_0 - V_0(s^*)**1**$

Loop for each step: $sp(w_t - w_{t-1})\geq \epsilon$
$\quad$Loop for each $s \in \mathcal{S}$:
$\qquad V\_{t+1} \leftarrow \mathbb{E}\_{a\sim\pi}[r(s, a) + \sigma\_{\mathcal{P}\_s\^a}(w\_t)]$
$\qquad w\_{t+1}(s) \leftarrow V\_{t+1}(s) - V\_{t+1}(s^*)$
return $V_t$

In the algorithm, $\\sigma\_{\\mathcal{P}\_s\^a}(w\_t) = \\max_{p\\in\\mathcal{P}\_s\^a}p\\top w\_t$, and can be solved in the dual formulation for various uncertainty sets [Iyengar05]. The returned $V\_t$ is the robust relative value function and $V\_t(s\^*)$ is the robust average reward.

The assumption that an oracle that outputs the robust value function exists is also presented in robust discounted reward MDPs, for example [Li22].

W3. Unichain assumption may lead to infinite $C_{PL}$. We agree with the reviewer that the unichain assumption may not guarantee that $C_{PL}$ is finite, and moreover, ergodicity is a sufficient condition to guarantee $C_{PL}$ is finite.

W4. Determine the increasing step size $M$. Since the definition of $M$ involves taking the sup over all transition kernels in the uncertainty set, theoretically determining $M$ is challenging. In practice, we only require an upper bound for $M$ so that the convergence result holds in Theorem 4.6. Therefore, in practice, we can run simulation to find an upper bound of $M$ and fine tune it. For example, we can numerically solve $\min_{\pi, P\in\mathcal{P}}\\|d_P^\pi\\|_\infty$ when the ergodicity holds, which is an upper bound of $M$.

Q1. Closed from expression for the value function. For the widely used metrics like the total variation distance, $\chi$ -square distance and the KL divergence, the closed-form expression exists for $\max\_{\mathsf{P}\in\mathcal{P}}\sum\_{s\^\prime \in \mathcal{S}}\mathsf{P}\_{s, s\^\prime}\^aV\_{\mathcal{P}}\^\pi(s\^\prime)$ [Iyengar05]. The dual formulation depends on the structure of the uncertainty set. However, if the value function has closed-form solution as the exponential cost MDPs is unclear. Moreover, our algorithm and theoretical results do not depend on the specific structure of the uncertainty set and the closed-form solution of the value function.

Q2. Policy gradient or mirror descent. We don't refer to our algorithm as policy gradient is due to the fact that in our algorithm, we replace the policy (sub)-gradient $\nabla g_\mathcal{P}^{\pi_k}$ by $Q_\mathcal{P}^{\pi_k}$. Since our algorithm is the mirror descent with a specific Bregman divergence and $Q_\mathcal{P}^{\pi_k}$ is not the (sub)-gradient, we refer to our algorithm as policy mirror descent.

Q3. Average cost MDPs. We thank the reviewer for this comment. In our paper, we adopt a minimization formulation to align with conventions in the optimization literature. In the revision, we will replace rewards with costs.

Q4. Typos. We thank the reviewer for pointing this out. Fixed.

Reference. [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.

[Iyengar05] Garud N. Iyengar (2005). Robust Dynamic Programming. Mathematics of Operations Research.

[Li22] Yan Li, Guanghui Lan, Tuo Zhao. (2022). First-order Policy Optimization for Robust Markov Decision Process.