Policy Mirror Descent for Regularized Reinforcement Learning:
A Generalized Framework with Linear Convergence

Recent years have witnessed a surge of activities towards understanding the global convergence properties of policy gradient methods and their variants for both continuous and discrete RL problems, examples including Fazel et al., (2018); Bhandari and Russo, (2019); Agarwal et al., (2020); Zhang et al., 2021b ; Wang et al., (2019); Mei et al., 2020a ; Bhandari and Russo, (2020); Khodadadian et al., (2021); Liu et al., (2020); Mei et al., 2020a ; Agazzi and Lu, (2020); Xu et al., (2019); Wang et al., (2019); Cen et al., 2022b ; Mei et al., (2021); Liu et al., (2019); Wang et al., (2021); Zhang et al., 2020a ; Zhang et al., 2021a ; Zhang et al., 2020b ; Shani et al., (2019), among other things. Neu et al., (2017) provided the first interpretation of NPG methods as mirror descent (Nemirovsky and Yudin,, 1983), thereby enabling the adaptation of techniques for analyzing mirror descent to the studies of NPG-type algorithms such as TRPO (Shani et al.,, 2019; Tomar et al.,, 2020). It has been shown that the NPG method converges sub-linearly for unregularized MDPs with a fixed learning rate (Agarwal et al.,, 2020), and converges linearly if the learning rate is set adaptively (Khodadadian et al.,, 2021), via exact line search (Bhandari and Russo,, 2020), or following a geometrically increasing schedule (Xiao,, 2022). The global linear convergence of NPG holds more generally for an arbitrary fixed learning rate when entropy regularization is enforced (Cen et al., 2022b, ). Noteworthily, Li et al., (2023) established a lower bound indicating that softmax PG methods can take an exponential time—in the size of the state space—to converge, while the convergence rates of NPG-type methods are almost independent of the problem dimension. In addition, another line of recent works (Abbasi-Yadkori et al.,, 2019; Hao et al.,, 2021; Lazic et al.,, 2021) established regret bounds for approximate NPG methods—termed as KL-regularized approximate policy iteration therein—for infinite-horizen undiscounted MDPs, which are beyond the scope of the current paper.

Regularization has been suggested to the RL literature either through the lens of optimization (Dai et al.,, 2018; Agarwal et al.,, 2020), or through the lens of dynamic programming (Geist et al.,, 2019; Vieillard et al.,, 2020). Our work is clearly an instance of the former type. Several recent results in the literature merit particular attention: Agarwal et al., (2020) demonstrated sublinear convergence guarantees for PG methods in the presence of relative entropy regularization, Mei et al., 2020b established linear convergence of entropy-regularized PG methods, whereas Cen et al., 2022b derived an almost dimension-free linear convergence theory for NPG methods with entropy regularization. Most of the existing literature focused on the entropy regularization or KL-type regularization, and the studies of general regularizers had been quite limited until the recent work Lan, (2022). The regularized MDP problems are also closely related to the studies of constrained MDPs, as both types of problems can be employed to model/promote constraint satisfaction in RL, as recently investigated in, e.g., Chow et al., 2018a ; Efroni et al., (2020); Ding et al., (2021); Yu et al., (2019); Xu et al., (2020). Note, however, that it is difficult to directly compare our algorithm with these methods, due to drastically different formulations and settings.

The focus of this paper is a discounted infinite-horizon Markov decision process, as represented by $\mathcal{M}=(\mathcal{S},\mathcal{A},P,r,\gamma)$ (Bertsekas,, 2017). Here, $\mathcal{S}\coloneqq\{1,\cdots,|\mathcal{S}|\}$ is the state space, $\mathcal{A}\coloneqq\{1,\cdots,|\mathcal{A}|\}$ is the action space, $\gamma\in[0,1)$ is the discount factor, $P:\mathcal{S}\times\mathcal{A}\to\Delta(\mathcal{S})$ is the probability transition matrix (so that $P(\cdot\,|\,s,a)$ is the transition probability from state $s$ upon execution of action $a$), whereas $r:\mathcal{S}\times\mathcal{A}\to[0,1]$ is the reward function (so that $r(s,a)$ indicates the immediate reward received in state $s$ after action $a$ is executed). Here, we focus on finite-state and finite-action scenarios, meaning that both $|\mathcal{S}|$ and $|\mathcal{A}|$ are assumed to be finite. A policy $\pi:\mathcal{S}\to\Delta(\mathcal{A})$ specifies a possibly randomized action selection rule, namely, $\pi(\cdot\,|\,s)$ represents the action selection probability in state $s$.

For any policy $\pi$, we define the associated value function $V^{\pi}:\mathcal{S}\to\mathbb{R}$ as follows

which can be viewed as the utility function we wish to maximize. Here, the expectation is taken over the randomness of the MDP trajectory $\{(s_{t},a_{t})\}_{t\geq 0}$ induced by policy $\pi$. Similarly, when the initial action $a$ is fixed, we can define the action-value function (or Q-function) as follows

As a well-known fact, the policy gradient of $V^{\pi}$ (w.r.t. the policy $\pi$) admits the following closed-form expression (Sutton et al., (2000))

Here, $d_{s_{0}}^{\pi}\in\Delta(\mathcal{S})$ is the so-called discounted state visitation distribution defined as follows

where $\mathbb{P}^{\pi}(s_{t}=s\,|\,s_{0})$ denotes the probability of $s_{t}=s$ when the MDP trajectory $\{s_{t}\}_{t\geq 0}$ is generated under policy $\pi$ given the initial state $s_{0}$.

Furthermore, the optimal value function and the optimal Q-function are defined and denoted by

It is well known that there exists at least one optimal policy, denoted by $\pi^{\star}$, that simultaneously maximizes the value function and the Q-function for all state-action pairs (Agarwal et al.,, 2019).

In practice, the agent is often asked to design policies that possess certain structural properties in order to be cognizant of system constraints such as safety and operational constraints, as well as encourage exploration during the optimization/learning stage. A natural strategy to achieve these is to resort to the following regularized value function w.r.t. a given policy $\pi$ (Neu et al.,, 2017; Mei et al., 2020b, ; Cen et al., 2022b, ; Lan,, 2022):

where $h_{s}:\Delta_{\zeta}(\mathcal{A})\to\mathbb{R}$ stands for a convex and possibly nonsmooth regularizer for state $s$, $\tau>0$ denotes the regularization parameter, and $d_{s}^{\pi}(\cdot)$ is defined in (5). Here, for technical convenience, we assume throughout that $h_{s}(\cdot)$ ($s\in\mathcal{S}$) is well-defined over an “$\zeta$-neighborhood” of the probability simplex $\Delta(\mathcal{A})$ defined as follows

where $\zeta>0$ can be an arbitrary constant. For instance, entropy regularization adopts the choice $h_{s}(p)=\sum_{i\in\mathcal{A}}p_{i}\log p_{i}$ for all $s\in\mathcal{S}$ and $p\in\Delta(\mathcal{A})$, which coincides with the negative Shannon entropy of a probability distribution. Similar, a KL regularization adopts the choice $h_{s}(p)=\mathsf{KL}(p\,\|\,p_{\mathsf{ref}})$, which penalizes the distribution $p$ that deviates from the reference $p_{\mathsf{ref}}$. As another example, a weighted $\ell_{1}$ regularization adopts the choice $h_{s}(p)=\sum_{i\in\mathcal{A}}w_{s,i}p_{i}$ for all $s\in\mathcal{S}$ and $p\in\Delta(\mathcal{A})$, where $w_{s,i}\geq 0$ is the cost of taking action $i$ at state $s$, and the regularizer $h_{s}(\pi(\cdot|s))$ captures the expected cost of the policy $\pi$ in state $s$. Throughout this paper, we impose the following assumption.

Following the convention in prior literature (e.g., Mei et al., 2020b ), we also define the corresponding regularized Q-function as follows:

The optimal regularized value function $V^{\star}_{\tau}$ and the corresponding optimal policy $\pi_{\tau}^{\star}$ are defined respectively as follows:

It is worth noting that Puterman, (2014) asserts the existence of an optimal policy $\pi_{\tau}^{\star}$ that achieves (10) simultaneously for all $s\in\mathcal{S}$. Correspondingly, we shall also define the resulting optimal regularized Q-function as

To better elucidate our algorithmic idea, let us first briefly review the design of classical mirror descent—originally proposed by Nemirovsky and Yudin, (1983)—in the optimization literature. Consider the following composite model:

where the objective function consists of two components. The first component is assumed to be differentiable, while the second component $h(\cdot)$ can be more general and is commonly employed to model some sort of regularizers. To solve this composite problem, one variant of mirror descent adopts the following update rule (see also Beck, (2017); Duchi et al., (2010)):

where $\eta>0$ is the learning rate or step size, and $D_{h}(\cdot,\cdot)$ is the Bregman divergence defined in (1). Note that the first term within the curly brackets of (12) can be safely discarded as it is a constant given $x^{(k)}$. In words, the above update rule approximates $f(x)$ via its first-order Taylor expansion $f\big{(}x^{(k)}\big{)}+\big{\langle}\nabla f(x^{(k)}),x\big{\rangle}$ at the point $x^{(k)}$, employs the Bregman divergence $D_{h}$ to monitor the difference between the new iterate and the current iterate $x^{(k)}$, and attempts to optimize such (properly monitored) approximation instead. While one can further generalize the Bregman divergence to $D_{\omega}$ for a different generator $\omega$, we shall restrict attention to the case with $h=\omega$ in the current paper.

We are now ready to present the algorithm we come up with, which is an extension of the PMD algorithm (Lan,, 2022). For notational simplicity, we shall write

throughout the paper, where $\pi^{(k)}$ denotes our policy estimate in the $k$-th iteration.

To begin with, suppose for simplicity that $h_{s}(\cdot)$ is differentiable everywhere. In the $k$-th iteration, a natural MD scheme that comes into mind for solving (7)—namely, $\text{maximize}_{\pi}V_{\tau}^{\pi}(s_{0})$ for a given initial state $s_{0}$—is the following update rule:

for every state $s\in\mathcal{S}$, which is a direct application of (12) to our setting. Here, we start with a learning rate $\eta^{\prime}$, and obtain simplification by replacing $\eta^{\prime}$ with ${\eta(1-\gamma)}/{d_{s_{0}}^{(k)}(s)}$. Notably, the update strategy (14) is invariant to the initial state $s_{0}$, akin to natural policy gradient methods (Agarwal et al.,, 2020).

This update rule is well-defined for, say, the case when $h_{s}$ is the negative entropy, since the algorithm guarantees $\pi^{(k)}>0$ all the time and hence $h_{s}$ is always differentiable w.r.t. the $k$-th iterate (see Cen et al., 2022b ). In general, however, it is possible to encounter situations when the gradient of $h_{s}$ does not exist on the boundary (e.g., when $h_{s}$ represents a certain indicator function). To cope with such cases, we resort to a generalized version of Bregman divergence (e.g., Kiwiel, (1997); Lan et al., (2011); Lan and Zhou, (2018)). To be specific, we attempt to replace the usual Bregman divergence $D_{h_{s}}(p,q)$ by the following metric

where $g_{s}$ can be any vector falling within the subdifferential $\partial h_{s}(q)$. Here, the non-negativity condition in (15) follows directly from the definition of the subgradient for any convex function. The constraint on $g_{s}$ can be further relaxed by exploiting the requirement $p,q\in\Delta(\mathcal{A})$. In fact, for any vector $\xi_{s}=g_{s}-c_{s}1$ (with $c_{s}\in\mathbb{R}$ some constant and $1$ the all-one vector), one can readily see that

where the last line is valid since $1^{\top}p=1^{\top}q=1$. As a result, everything boils down to identifying a vector $\xi_{s}$ that falls within $\partial h_{s}(q)$ upon global shift.

Towards this, we propose the following iterative rule for designing such a sequence of vectors as surrogates for the subgradient of $h_{s}$:

where $\xi^{(k+1)}(s,\cdot)$ is updated as a convex combination of the previous $\xi^{(k)}(s,\cdot)$ and ${Q}^{(k)}_{\tau}(s,\cdot)$, where more emphasis is put on ${Q}^{(k)}_{\tau}(s,\cdot)$ when the learning rate $\eta$ is large. As asserted by the following lemma, the above vectors $\xi^{(k)}(s,\cdot)$ we construct satisfy the desired property, i.e., lying within the subdifferential of $h_{s}$ under suitable global shifts. It is worth mentioning that these global shifts $\{c_{s}^{(k)}\}$ only serve as an aid to better understand the construction, but are not required during the algorithm updates.

Thus far, we have presented all crucial ingredients of our algorithm. The whole procedure is summarized in Algorithm 1, and will be referred to as Generalized Policy Mirror Descent (GPMD) throughout the paper. Interestingly, several well-known algorithms can be recovered as special cases of GPMD:

• •

  When the Bregman divergence $D_{h_{s}}(\cdot,\cdot)$ is taken as the KL divergence, GPMD reduces to the well-renowned NPG algorithm (Kakade,, 2002) when $\tau=0$ (no regularization), and to the NPG algorithm with entropy regularization analyzed in (Cen et al., 2022b, ) when $h_{s}(\cdot)$ is taken as the negative Shannon entropy.

• •

  When $\eta=\infty$ (no divergence), GPMD reduces to regularized policy iteration in Geist et al., (2019); in particular, GPMD reduces to the standard policy iteration algorithm if in addition $\tau$ is also $0$.

Before continuing, let us take a moment to point out the key differences between our algorithm GPMD and the PMD algorithm proposed in Lan, (2022) in terms of algorithm designs. Although the primary exposition of PMD in Lan, (2022) fixes the Bregman divergence as the KL divergence, the algorithm also works in the presence of a generic Bregman divergence, whose relationship with the regularizer $h_{s}$ is, however, unspecified. Furthermore, GPMD adaptively sets this term to be the Bregman divergence generated by the regularizer $h_{s}$ in use, together with a carefully designed recursive update rule (cf. (17)) to compute surrogates for the subgradient of $h_{s}$ to facilitate implementation. Encouragingly, this specific choice leads to a tailored performance analysis of GPMD, which was not present in and instead complementary with that of PMD (Lan,, 2022). In truth, our theory offers linear convergence guarantees for more general scenarios by adapting to the geometry of the regularizer $h_{s}$; details to follow momentarily.

To make clear our contributions, it is helpful to compare Theorem 1 with the theory for the state-of-the-art algorithm PMD in Lan, (2022).

• •

  Linear convergence for convex regularizers under constant learning rates. Suppose that constant learning rates are adopted for both GPMD and PMD. Our finding reveals that GPMD enjoys global linear convergence—in terms of both $\|{Q}^{\star}_{\tau}-{Q}^{(k+1)}_{\tau}\|_{\infty}$ and $\|{V}^{\star}_{\tau}-{V}^{(k+1)}_{\tau}\|_{\infty}$—even when the regularizer $h_{s}(\cdot)$ is only convex but not strongly convex. In contrast, Lan, (2022, Theorem 2) provided only sublinear convergence guarantees (with an iteration complexity proportional to $1/\varepsilon$) for the case with convex regularizers, provided that constant learning rates are adopted.²²2In fact, Lan, (2022, Theorem 3) suggests using a vanishing strongly convex regularization, as well as a corresponding increasing sequence of learning rates, in order to enable linear convergence for non-strongly-convex regularizers.

• •

  A full range of learning rates. Theorem 1 reveals linear convergence of GPMD for a full range of learning rates, namely, our result is applicable to any $\eta>0$. In comparison, linear convergence was established in Lan, (2022) only when the learning rates are sufficiently large and when $h_{s}$ is $1$-strongly convex w.r.t. the KL divergence. Consequently, the linear convergence results in Lan, (2022) do not extend to several widely used regularizers such as negative Tsallis entropy and log-barrier functions (even after scaling), which are, in contrast, covered by our theory. It is worth noting that the case with small-to-medium learning rates is often more challenging to cope with in theory, given that its dynamics could differ drastically from that of regularized policy iteration.

• •

  Further comparison of rates under large learning rates. (Lan,, 2022, Theorem 1) achieves a contraction rate of $\gamma$ when the regularizer is strongly convex and the step size satisfies $\eta\geq\frac{1-\gamma}{\gamma\tau}$, while the contraction rate of GPMD is $1-\frac{\eta\tau}{1+\eta\tau}(1-\gamma)$ under the full range of the step size, which is slower but approaches the contraction rate $\gamma$ of PMD as $\eta$ goes to infinity. Therefore, in the limit $\eta\to\infty$, both GPMD and PMD achieve the contraction rate $\gamma$. As soon as $\eta\geq 1/\tau$, their iteration complexities are on the same order.

To begin with, we demonstrate that GPMD enjoys a sort of monotonic improvements concerning the updates of both the value function and the Q-function, as stated in the following lemma. This lemma can be viewed as a generalization of the well-established policy improvement lemma in the analysis of NPG (Agarwal et al.,, 2020; Cen et al., 2022b, ) as well as PMD (Lan,, 2022).

Interestingly, the above monotonicity holds simultaneously for all state-action pairs, and hence can be understood as a kind of pointwise monotonicity.

Another key ingredient of our proof lies in the use of a generalized Bellman operator $\mathcal{T}_{\tau,h}:\mathbb{R}^{|\mathcal{S}||\mathcal{A}|}\rightarrow\mathbb{R}^{|\mathcal{S}||\mathcal{A}|}$ associated with the regularizer $h=\{h_{s}\}_{s\in\mathcal{S}}$. Specifically, for any state-action pair $(s,a)$ and any vector $Q\in\mathbb{R}^{|\mathcal{S}||\mathcal{A}|}$, we define

It is worth noting that this definition shares similarity with the regularized Bellman operator proposed in Geist et al., (2019), where the operator defined there is targeted at $V_{\tau}$, while ours is defined w.r.t. $Q_{\tau}$.

The importance of this generalized Bellman operator is two-fold: it enjoys a desired contraction property, and its fixed point corresponds to the optimal regularized Q-function. These are generalizations of the properties for the classical Bellman operator, and are formally stated in the following lemma. The proof is deferred to Appendix A.3.