Regret Analysis of Policy Gradient Algorithm for Infinite Horizon Average Reward Markov Decision Processes

Reinforcement Learning (RL) describes a class of problems where a learner repeatedly interacts with an unknown environment with the intention of maximizing the cumulative sum of rewards. This model has found its application in a wide array of areas, ranging from networking to transportation to epidemic control (Geng et al. 2020; Al-Abbasi, Ghosh, and Aggarwal 2019; Ling, Mondal, and Ukkusuri 2023). RL problems are typically analysed via three distinct setups $-$ episodic, infinite horizon discounted reward, and infinite horizon average reward. Among these, the infinite horizon average reward setup holds particular significance in real-world applications (including those mentioned above) due to its alignment with many practical scenarios and its ability to capture essential long-term behaviors. However, scalable algorithms in this setup have not been widely studied. This paper provides an algorithm in the infinite horizon average reward setup with general parametrized policies which yields sub-linear regret guarantees. We would like to mention that this result is the first of its kind in the average reward setting.

There are two major approaches to solving an RL problem. The first one, known as the model-based approach, involves constructing an estimate of the transition probabilities of the underlying Markov Decision Process (MDP). This estimate is subsequently leveraged to derive policies (Auer, Jaksch, and Ortner 2008; Agrawal and Jia 2017; Ouyang et al. 2017; Fruit et al. 2018). It is worth noting that model-based techniques encounter a significant challenge – these algorithms demand a substantial memory to house the model parameters. Consequently, their practical application is hindered when dealing with large state spaces. An alternative strategy is referred to as model-free algorithms. These methods either directly estimate the policy function or maintain an estimate of the $Q$ function, which are subsequently employed for policy generation (Mnih et al. 2015; Schulman et al. 2015; Mnih et al. 2016). The advantage of these algorithms lies in their adaptability to handle large state spaces.

In the average reward MDP, which is the setting considered in our paper, one of the key performance indicators of an algorithm is the expected regret. It has been theoretically demonstrated in (Auer, Jaksch, and Ortner 2008) that the expected regret of any algorithm for a broad class of MDPs is lower bounded by $\Omega(\sqrt{T})$ where $T$ denotes the length of the time horizon. Many model-based algorithms, such as, (Auer, Jaksch, and Ortner 2008; Agrawal and Jia 2017) achieve this bound. Unfortunately, the above algorithms are designed to be applicable solely in the tabular setup. Recently, (Wei et al. 2021) proposed a model-based algorithm for the linear MDP setup that is shown to achieve the optimal regret bound. On the other hand, (Wei et al. 2020) proposed a model-free $Q$-estimation-based algorithm that achieves the optimal regret in the tabular setup.

One way to extend algorithms beyond the tabular setting is via policy parameterization. Here, the policies are indexed by parameters (via, for example, neural networks), and the learning process is manifested by updating these parameters using some update rule (such as gradient descent). Such algorithms are referred to as policy gradient (PG) algorithms. Interestingly, the analysis of PG algorithms is typically restricted within the discounted reward setup. For example, (Agarwal et al. 2021) characterized the sample complexity of PG and Natural PG (NPG) with softmax and tabular parameterization. Sample complexity results for general parameterization are given by (Liu et al. 2020; Ding et al. 2020). However, the sub-linear regret analysis of a PG-based algorithm with general parameterization in the average reward setup, to the best of our knowledge, has not been studied in the literature. This paper aims to bridge this gap by addressing this crucial problem.

As discussed in the introduction, the reinforcement learning problem has been widely studied recently for infinite horizon discounted reward cases or the episodic setting. For example, (Jin et al. 2018) proposed the model-free UCB-Q learning and showed a $\mathcal{O}(\sqrt{T})$ regret in the episodic setting. In the discounted reward setting, (Ding et al. 2020) achieved $\mathcal{O}(\epsilon^{-2})$ sample complexity for the softmax parametrization using the Natural Policy Gradient algorithm whereas (Mondal and Aggarwal 2023; Fatkhullin et al. 2023) exhibited the same complexity for the general parameterization. However, the regret analysis or the global convergence of the average reward infinite horizon case is much less investigated.

For infinite horizon average reward MDPs, (Auer, Jaksch, and Ortner 2008) proposed a model-based Upper confidence Reinforcement learning (UCRL2) algorithm and established that it obeys a $\tilde{\mathcal{O}}(\sqrt{T})$ regret bound. (Agrawal and Jia 2017) proposed posterior sampling-based approaches for average reward MDPs. (Wei et al. 2020) proposed the optimistic-Q learning algorithm which connects the discounted reward and average reward setting together to show $\mathcal{O}(T^{3/4})$ regret in weakly communicating average reward case and another online mirror descent algorithm which achieves $\mathcal{O}(\sqrt{T})$ regret in the ergodic setting. For the linear MDP setting, (Wei et al. 2021) proposed three algorithms, including the MDP-EXP2 algorithm which achieves $O(\sqrt{T})$ regret under the ergodicity assumption. These works have been summarized in Table 1. We note that the assumption of weakly communicating MDP is the minimum assumption needed to have sub-linear regret results. However, it is much more challenging to work with this assumption in the general parametrized setting because of the following reasons. Firstly, there is no guarantee that the state distribution will converge to the steady distribution exponentially fast which is the required property to show that the value functions are bounded by the mixing time. Secondly, it is unclear how to obtain an asymptotically unbiased estimate of the policy gradient. Thus, we assume ergodic MDP in this work following other works in the literature (Pesquerel and Maillard 2022; Gong and Wang 2020). MDPs with constraints have also been recently studied for model-based (Agarwal, Bai, and Aggarwal 2022b, a), model-free tabular (Wei, Liu, and Ying 2022; Chen, Jain, and Luo 2022), and linear MDP setup (Ghosh, Zhou, and Shroff 2022).

However, all of the above algorithms are designed for the tabular setting or the linear MDP assumption, and none of them uses a PG algorithm with the general parametrization setting. In this paper, we propose a PG algorithm for ergodic MDPs with general parametrization and analyze its regret. Our algorithm can, therefore, be applied beyond the tabular or linear MDP setting.

In this paper, we consider an infinite horizon reinforcement learning problem with an average reward criterion, which is modeled by the Markov Decision Process (MDP) written as a tuple $\mathcal{M}=(\mathcal{S},\mathcal{A},r,P,\rho)$ where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space of size $A$, $r:\mathcal{S}\times\mathcal{A}\rightarrow[0,1]$ is the reward function, $P:\mathcal{S}\times\mathcal{A}\rightarrow\Delta^{|\mathcal{S}|}$ is the state transition function where $\Delta^{|\mathcal{S}|}$ denotes the probability simplex with dimension $|\mathcal{S}|$, and $\rho:\mathcal{S}\rightarrow[0,1]$ is the initial distribution of states. A policy $\pi:\mathcal{S}\rightarrow\Delta^{|\mathcal{A}|}$ decides the distribution of the action to be taken given the current state. For a given policy, $\pi$ we define the long-term average reward as follows.

where the expectation is taken over all state-action trajectories that are generated by following the action execution process, $a_{t}\sim\pi(\cdot|s_{t})$ and the state transition rule, $s_{t+1}\sim P(\cdot|s_{t},a_{t})$, $\forall t\in\{0,1,\cdots\}$. To simplify notations, we shall drop the dependence on $\rho$ whenever there is no confusion. We consider a parametrized class of policies, $\Pi$ whose each element is indexed by a $\mathrm{d}$-dimensional parameter, $\theta\in\Theta$ where $\Theta\subset\mathbb{R}^{\mathrm{d}}$. Our goal is to solve the following optimization problem.

A policy $\pi_{\theta}$ induces a transition function $P^{\pi_{\theta}}:\mathcal{S}\rightarrow\Delta^{|\mathcal{S}|}$ as $P^{\pi_{\theta}}(s,s^{\prime})=\sum_{a\in\mathcal{A}}P(s^{\prime}|s,a)\pi_{\theta}(a|s)$, $\forall s,s^{\prime}\in\mathcal{S}$. If $\mathcal{M}$ is such that for every policy $\pi$, the induced function, $P^{\pi}$ is irreducible, and aperiodic, then $\mathcal{M}$ is called ergodic.

Ergodicity is commonly applied in the analysis of MDPs (Pesquerel and Maillard 2022; Gong and Wang 2020). It is well known that if $\mathcal{M}$ is ergodic, then $\forall\theta\in\Theta$, there exists a unique stationary distribution, $d^{\pi_{\theta}}\in\Delta^{|\mathcal{S}|}$ defined as,

Note that under the assumption of ergodicity, $d^{\pi_{\theta}}$ is independent of the initial distribution, $\rho$, and satisfies $P^{\pi_{\theta}}d^{\pi_{\theta}}=d^{\pi_{\theta}}$. In this case, we can write the average reward as follows.

Hence, the average reward $J(\theta)$ is also independent of the initial distribution, $\rho$. Furthermore, $\forall\theta\in\Theta$, there exist a function $Q^{\pi_{\theta}}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ such that the following Bellman equation is satisfied $\forall(s,a)\in\mathcal{S}\times\mathcal{A}$.

where the state value function, $V^{\pi_{\theta}}:\mathcal{S}\rightarrow\mathbb{R}$ is defined as,

Note that if $(\ref{eq_bellman})$ is satisfied by $Q^{\pi_{\theta}}$, then it is also satisfied by $Q^{\pi{\theta}}+c$ for any arbitrary constant, $c$. To define these functions uniquely, we assume that $\sum_{s\in\mathcal{S}}d^{\pi_{\theta}}(s)V^{\pi_{\theta}}(s)=0$. In this case, $V^{\pi_{\theta}}(s)$ can be written as follows $\forall s\in\mathcal{S}$.

where $\mathbf{E}_{\theta}[\cdot]$ denotes expectation over all trajectories induced by the policy $\pi_{\theta}$. Similarly, $\forall(s,a)\in\mathcal{S}\times\mathcal{A}$, $Q^{\pi_{\theta}}(s,a)$ can be uniquely written as,

Additionally, we define the advantage function $A^{\pi_{\theta}}:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ such that $\forall(s,a)\in\mathcal{S}\times\mathcal{A}$,

Ergodicity also implies the existence of a finite mixing time. In particular, if $\mathcal{M}$ is ergodic, then the mixing time is defined as follows.

Mixing time is a measure of how fast the MDP reaches close to its stationary distribution if the same policy is kept on being executed repeatedly. We also define the hitting time as follows.

Let, $J^{*}\triangleq\sup_{\boldsymbol{\theta}\in\Theta}J(\theta)$. For a given MDP $\mathcal{M}$ and a time horizon $T$, the regret of an algorithm $\mathbb{A}$ is defined as follows.

where the action, $a_{t}$, $t\in\{0,1,\cdots\}$ is chosen by following the algorithm, $\mathbb{A}$ based on the trajectory up to time, $t$, and the state, $s_{t+1}$ is obtained by following the state transition function, $P$. Wherever there is no confusion, we shall simplify the notation of regret to $\mathrm{Reg}_{T}$. The goal of maximizing $J(\cdot)$ can be accomplished by designing an algorithm that minimizes the regret.

In this section, we discuss a policy-gradient-based algorithm in the average reward RL settings. For simplicity, we assume that the set of all policy parameters is $\Theta=\mathbb{R}^{\mathrm{d}}$. The standard policy gradient algorithm iterates the policy parameter $\theta$ as follows $\forall k\in\{1,2,\cdots\}$ starting with an initial guess $\theta_{1}$.

where $\alpha$ is the parameter learning rate. The following result is well-known in the literature (Sutton et al. 1999).

Typically we have access neither to $P$, the state transition function to compute the required expectation nor to the functions $V^{\pi_{\theta}}$, $Q^{\pi_{\theta}}$. In the absence of this knowledge, computation of gradient, therefore, becomes a difficult job. In the subsequent discussion, we shall demonstrate how the gradient can be estimated using sampled trajectories. Our policy gradient-based algorithm is described in Algorithm 1.

The algorithm proceeds in multiple epochs with the length of each epoch being $H=16t_{\mathrm{hit}}t_{\mathrm{mix}}\sqrt{T}(\log T)^{2}$. Observe that the algorithm is assumed to be aware of $T$. This assumption can be easily relaxed invoking the well-known doubling trick (Lattimore and Szepesvári 2020). We also assume that the values of $t_{\mathrm{mix}}$, and $t_{\mathrm{hit}}$ are known to the algorithm. Similar presumptions have been used in the previous literature (Wei et al. 2020). In the $k$th epoch, the algorithm generates a trajectory of length $H$, denoted as $\mathcal{T}_{k}=\{(s_{t},a_{t})\}_{t=(k-1)H}^{kH-1}$, by following the policy $\pi_{\theta_{k}}$. We utilise the policy parameter $\theta_{k}$ and the trajectory $\mathcal{T}_{k}$ in Algorithm 2 to compute the estimates $\hat{V}^{\pi_{\theta_{k}}}(s)$, and $\hat{Q}^{\pi_{\theta_{k}}}(s,a)$ for a given state-action pair $(s,a)$. The algorithm searches the trajectory $\mathcal{T}_{k}$ to locate disjoint sub-trajectories of length $N=4t_{\mathrm{mix}}(\log T)$ that start with the given state $s$ and are at least $N$ distance apart. Let $i$ be the number of such sub-trajectories and the sum of rewards in the $j$th such sub-trajectory be $y_{j}$. Then $\hat{V}^{\pi_{\theta_{k}}}(s)$ is computed as,

The sub-trajectories are kept at least $N$ distance apart to ensure that the samples $\{y_{j}\}_{j=1}^{i}$ are fairly independent. The estimate $\hat{Q}^{\pi_{\theta}}(s,a)$, on the other hand, is given as,

where $\tau_{j}$ is the starting time of the $j$th chosen sub-trajectory. Finally, the advantage value is estimated as,

This allows us to compute an estimate of the policy gradient as follows.

where $t_{k}=(k-1)H$ is the starting time of the $k$th epoch. The policy parameters are updated via (20). In the following lemma, we show that $\hat{A}^{\pi_{\theta_{k}}}(s,a)$ is a good estimator of $A^{\pi_{\theta_{k}}}(s,a)$.

Lemma 2 establishes that the $L_{2}$ error of our proposed estimator can be bounded above as $\tilde{\mathcal{O}}(1/\sqrt{T})$. As we shall see later, this result can be used to bound the estimation error of the gradient. It is worthwhile to point out that $\hat{V}^{\pi_{\theta_{k}}}(s)$ and $\hat{Q}^{\pi_{\theta_{k}}}(s,a)$ defined in $(\ref{eq_V_estimate})$, $(\ref{eq_reward_tracking})$ respectively, may not themselves be good estimators of their target quantities although their difference is one. We would also like to mention that our Algorithm 2 is inspired by Algorithm 2 of (Wei et al. 2020). The main difference is that we take the episode length to be $H=\tilde{\mathcal{O}}(\sqrt{T})$ while in (Wei et al. 2020), it was chosen to be $\tilde{\mathcal{O}}(1)$. This extra $\sqrt{T}$ factor makes the estimation error a decreasing function of $T$.

In this section, we show that our proposed Algorithm 1 converges globally. This essentially means that the parameters $\{\theta_{k}\}_{k=1}^{\infty}$ are such that the sequence $\{J(\theta_{k})\}_{k=1}^{\infty}$, in certain sense, approaches the optimal average reward, $J^{*}$. Such convergence will be later useful in bounding the regret of our algorithm. Before delving into the analysis, we would like to first point out a few assumptions that are needed to establish the results.

One can immediately see that by combining Assumption 2 with Lemma 2 and using the definition of the gradient estimator as given in $(\ref{eq_grad_estimate})$, we arrive at the following important result.

Lemma 3 claims that the error in estimating the gradient can be bounded above as $\tilde{\mathcal{O}}(1/\sqrt{T})$. This result will be used in proving the global convergence of our algorithm.

Assumption 4 is also commonly used in the policy gradient analysis (Liu et al. 2020). This is satisfied by the Gaussian policy with a linearly parameterized mean.

In the discounted reward setup, one key result is the performance difference lemma. In the averaged reward setting, this is derived as stated below.

Using Lemma 4, we present a general framework for convergence analysis of the policy gradient algorithm in the averaged reward case as dictated below. This is inspired by the convergence analysis of (Liu et al. 2020) for the discounted reward MDPs.

Lemma 5 bounds the optimality error of any gradient ascent algorithm as a function of intermediate gradient norms. Note the presence of $\epsilon_{\mathrm{bias}}$ in the upper bound. Clearly, for a severely restricted policy class where $\epsilon_{\mathrm{bias}}$ is significant, the optimality bound becomes poor. Consider the expectation of the second term in (28). Note that,

where $(a)$ uses Assumption 4. The expectation of the third term in (28) can be bounded as follows.

In both (29), (30), the terms related to $\left\lVert\omega_{k}-\nabla_{\theta}J(\theta_{k})\right\rVert^{2}$ are bounded by Lemma 3. To bound the term $\left\lVert\nabla_{\theta}J(\theta_{k})\right\rVert^{2}$, we use the following lemma.

Using Lemma 3, we obtain the following inequality.

Applying Lemma 3 and (31) in (30), we finally get,

Similarly, using $(\ref{eq_second_term_bound})$, we deduce the following.

Inequalities (32) and (33) lead to the following result.

Theorem 1 dictates that the sequence $\{J(\theta_{k})\}_{k=1}^{K}$ generated by Algorithm 1 converges to $J^{*}$ with a convergence rate of $\mathcal{O}(T^{-\frac{1}{4}}+\sqrt{\epsilon_{\mathrm{bias}}})$. Alternatively, one can say that in order to achieve an optimality error of $\epsilon+\sqrt{\epsilon_{\mathrm{bias}}}$, it is sufficient to choose $T=\mathcal{O}\left(\epsilon^{-4}\right)$. It matches the state-of-the-art sample complexity bound of the policy gradient algorithm with general parameterization in the discounted reward setup (Liu et al. 2020).

In this section, we demonstrate how the convergence analysis in the previous section can be used to bind the expected regret of our proposed algorithm. Note that the regret can be decomposed as follows.

where $\mathcal{I}_{k}\triangleq\{(k-1)H,\cdots,kH-1\}$. Note that the expectation of the first term in $(\ref{reg_decompose})$ can be bounded using Theorem 34. The expectation of the second term can be expressed as follows,

where $(a)$ follows from Bellman equation and $(b)$ utilises the following facts: $\mathbf{E}[V^{\pi_{\theta_{k}}}(s_{t+1})]=\mathbf{E}_{s^{\prime}\sim P(\cdot|s_{t},a_{t})}[V^{\pi_{\theta_{k}}}(s^{\prime})]$ and $\mathbf{E}[V^{\pi_{\theta_{k}}}(s_{t})]=\mathbf{E}[Q^{\pi_{\theta_{k}}}(s_{t},a_{t})]$. The term, $P$ in (36) can be bounded using Lemma 7 (stated below). Moreover, the term $Q$ can be upper bounded as $\mathcal{O}(t_{\mathrm{mix}})$ as clarified in the appendix.

Lemma 7 can be interpreted as the stability results of our algorithm. It essentially states that the policy parameters are updated such that the average difference between consecutive average reward and value functions decreases with the horizon, $T$. Using the above result, we now prove our regret guarantee.