SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Emmeran Johnson
emmeran.johnson17@imperial.ac.uk
Ciara Pike-Burke
c.pike-burke@imperial.ac.uk
Policy Mirror Descent (PMD) is a general family of algorithms that covers a wide range of novel and fundamental methods in reinforcement learning. Motivated by the instability of policy iteration (PI) with inexact policy evaluation, PMD algorithmically regularises the policy improvement step of PI. With exact policy evaluation, PI is known to converge linearly with a rate given by the discount factor  $\gamma$  of a Markov Decision Process. In this work, we bridge the gap between PI and PMD with exact policy evaluation and show that the dimension-free  $\gamma$ -rate of PI can be achieved by the general family of unregularised PMD algorithms under an adaptive step-size. We show that both the rate and step-size are unimprovable for PMD: we provide matching lower bounds that demonstrate that the  $\gamma$ -rate is optimal for PMD methods as well as PI, and that the adaptive step-size is necessary for PMD to achieve it. Our work is the first to relate PMD to rate-optimality and step-size necessity. Our study of the convergence of PMD avoids the use of the performance difference lemma, which leads to a direct analysis of independent interest. We also extend the analysis to the inexact setting and establish the first dimension-optimal sample complexity for unregularised PMD under a generative model, improving upon the best-known result.
The problem of finding an optimal policy in tabular discounted Markov Decision Processes (MDPs) was classically solved using dynamic programming approaches such as policy iteration (PI) and value iteration (VI) [29, 35]. These methods are well understood theoretically and are guaranteed to converge linearly to the optimal policy in the tabular setting with a rate equal to the discount factor  $\gamma$  of the MDP [9]. Recently, increased interest has been devoted to the study of policy-gradient (PG) approaches based on optimising a parameterised policy with respect to an objective [36, 23, 20].
Given their popularity, it is of interest to better understand PG methods and determine if their guarantees match those of classical algorithms in tabular MDPs. Among the recent works focused on understanding these methods in the tabular setting, [38] study a general family of algorithms known as Policy Mirror Descent (PMD). PMD algorithmically regularises the policy improvement step of PI and as such can be seen as a version of regularised PI, without actually regularising the objective of interest. It is also viewed as a policy-gradient method through its connection to mirror descent [8]. Linear convergence of PMD was established by [38], though their rate depends on an instance-
dependent factor that can scale with the dimension of the problem, such as the size of the state space. For a specific instance of PMD known as Natural Policy Gradient (NPG), [22] showed that an instance-independent  $\gamma$ -rate is achievable, although their results do not cover general PMD. In MDPs where the objective is regularised, the  $\gamma$ -rate has been established for PMD [13, 24, 41]. The classical approaches (PI and VI) achieve the  $\gamma$ -rate without regularisation, revealing that regularisation is, in general, not necessary for algorithms to reach the  $\gamma$ -rate. This motivates the following questions:
Can the classical linear  $\gamma$ -rate be matched by unregularised policy-gradient algorithms? And what is the best rate that unregularised policy-gradient methods can achieve?
For PMD, our work answers the first question positively and answers the second by establishing that the  $\gamma$ -rate is in fact the best rate achievable for PMD as well as for a more general family of algorithms (see Section 4.1). PMD allows for a choice of a mirror map that specifies different algorithms. Among these, NPG and PI are two ubiquitous instances of PMD each corresponding to their own mirror map. However, PMD is much more general and other mirror maps will lead to alternative algorithms endowed with the guarantees of PMD that we establish in this paper. In particular, the correspondence of mirror maps with exponential families [7] allows us to specify a wealth of valid mirror maps. This illustrates that PMD is a general framework that encompasses a wide range of novel but also fundamental algorithms, and motivates the study of its convergence guarantees. In this work, we make the following contributions and summarise them in Table 1,
- We recover the  $\gamma$ -rate for the general family of PMD algorithms with exact policy evaluations under an adaptive size (see the third bullet point below). In particular, Theorem 4.1 establishes the following bound in  $\ell_{\infty}$ -norm for the value  $V^{\pi^k}$  of the policy  $\pi^k$  after  $k$  iterations of PMD compared to the value  $V^{\pi^*}$  of an optimal policy  $\pi^*$ ,
$$
\| V ^ {\pi^ {\star}} - V ^ {\pi^ {k}} \| _ {\infty} \leq \frac {2}{1 - \gamma} \gamma^ {k},
$$
providing guarantees for any starting-state distribution. This matches the rate of VI and PI as well as the best known rates for PMD on regularised MDPs. This is also the first fully dimension-independent linear convergence result for unregularised PMD, by which we mean that there is no dependence on the size of the state space or the action space.
- We provide a matching non-asymptotic lower-bound in Theorem 4.2, establishing the  $\gamma$ -rate as the optimal rate for PMD methods in early iterations. Our results show that a particular choice of learning rate allows PMD to reach this lower-bound exactly. Note that an asymptotic lower-bound is not possible due to the exact convergence of PI in finite iterations (see Section 4.1 for a discussion of the significance of this non-asymptotic lower-bound).
- The  $\gamma$ -rate for PMD in Theorem 4.1 relies on an adaptive step-size, where the adaptivity comes from the fact that the step-size depends on the policy at the current iteration (see Section 4). In Theorem 4.3 we show that this adaptivity is necessary for PMD to achieve the  $\gamma$ -rate, establishing our step-size as both sufficient and necessary.
- We establish a novel theoretical analysis that avoids the use of the performance difference lemma [19]. This leads to a simple analysis and avoids needing to deal with visitation distribution mismatches that lead to dimension dependence in prior work.
- By extending our analysis to the inexact setting, with an approach similar to that of in [38], we establish an instance-independent sample complexity of  $\tilde{O}(|S||\mathcal{A}|(1 - \gamma)^{-8}\varepsilon^{-2})$  under a generative model, where the notation  $\tilde{O}()$  hides poly-logarithmic factors,  $\mathcal{S}$  is the state space of the MDP,  $\mathcal{A}$  is the action space and  $\varepsilon$  is the required accuracy. This improves on the previous best known sample complexity for PMD by removing the dependence on a distribution mismatch coefficient that can scale with problem-dependent quantities such as the size of the state space. More generally, we highlight that the analysis we establish in the exact setting can easily be combined with any other scheme for estimating the Q functions (see Section 5), paving the way for further improvements in instance-independent sample complexity results should more efficient estimation procedures be developed.
Our contributions are primarily on establishing the optimal rate for general (not just NPG) exact PMD where we assume access to the true action-values of policies, and the simplicity of the analysis. The sample complexity result in the inexact setting illustrates how our analysis can be easily extended to
Table 1: Comparison of contributions with prior work that study PMD. [22] focus on NPG, an instance of PMD for a specific mirror map (see Section 3). Their analysis is fundamentally different to ours as it exploits the closed-form update of NPG. Their step-size is similar to ours, although it has a dependence on a sub-optimality gap (see Section 4). Linear  $\gamma$ -rate holds if the linear convergence is with the  $\gamma$ -rate, not only linear convergence. The  $\ell_{\infty}$ -bound is satisfied if it holds for  $\|V^{\pi^{\star}} - V^{\pi^{k}}\|_{\infty}$ . Dimension independence is satisfied when there is no instance for which the bound can scale with the size of the state or action space. We compare these works in more detail in Section 4.
Linear γ-Rate	General Mirror Map	l∞ Bound	Dimension Independent	Matching Lower-Bound	Step-Size Necessity
[22]	✓	×	✓	✓	×	×
[38]	×	✓	×	×	×	×
This work	✓	✓	✓	✓	✓	✓
obtain improved results for inexact PMD. These contributions advance the theoretical understanding of PMD and rule out searching for instances of PMD that may improve on PI or NPG in the exact setting. The algorithmic differences between instances of PMD do not affect the performance of the algorithm beyond the step-size condition.
2.1 Convergence rates for exact policy mirror descent
We first consider the setting where exact policy evaluation is assumed. In this setting, several earlier works have sub-linear convergence results for PMD [17, 33] and NPG specifically [3, 1], though these have since been improved to linear convergence results as discussed below. Note that the work of [1] refers to their algorithm as "Poltex", which shares the same update as NPG.
A line of work has considered PG methods applied to regularised MDPs. In this setting, linear convergence has been established for NPG with entropy regularisation [13], PMD with strongly-convex regularisers [24] and PMD with convex non-smooth regularisers [41]. The rates of convergence are either exactly  $\gamma$  or can be made arbitrarily close to  $\gamma$  by letting the step-size go to infinity.
In the setting of unregularised MDPs, which is the focus of this paper, linear convergence of the special case of NPG was established [11, 22] under an adaptive step-size similar to ours that depends on the current policy at each step. The bound of [11] has an additive asymptotic error-floor that can be made arbitrarily small by making the step-size larger, while for a similar step-size [22] does not have this term so we focus on this work. Their analysis relies on a link between NPG and PI and consists of bounding the difference in value between iterates of both methods. [11] also establish linear convergence for a number of algorithms including PMD, although it is in the idealised setting of choosing the step size at each iteration that leads to the largest increase in value. This step-size choice will make PMD at least as good as PI since arbitrarily large step-sizes can be chosen and PMD with an infinite step-size converges to a PI update. Since PI converges linearly, so will PMD. This does not establish linear convergence of PMD for step-sizes with a closed-form expression. However, linear convergence for unregularised general PMD was recently established by [38] under a geometrically increasing step-size. In general, their rate is instance-dependent and may scale with problem dependent quantities such as the size of the state space. For general starting-state distributions, this same instance-dependent rate was established by [27] for a variant of PMD which augments the update with an added regularisation term. We focus our comparison on the work of [38] rather than this work as the guarantees are equivalent in both but [27] do not directly study PMD and have a more complicated algorithm. A summary of our results compared to those of [22] and [38] is presented in Table 1 and discussed in more detail in Section 4. In terms of optimality, [22] provide a lower-bound for constant step-size NPG, though it only applies to MDPs with a single-state, which can be solved in a single iteration with exact policy evaluation as the step-size goes to infinity (for which the lower-bound goes to 0). We provide a lower-bound in Theorem 4.2 that applies to PMD with arbitrary step-size on an MDP with any finite state space. To the best of our knowledge, prior to this work no lower-bound has been established in this general setting.
2.2 Sample complexity of inexact policy mirror descent
Sample complexity in the inexact policy evaluation setting refers to the number of samples needed to guarantee an  $\varepsilon$ -optimal policy is returned when we no longer have access to the Q-values exactly. We here give an outline of results, typically established in high-probability, under a generative model that we formally present in Section 5. The lower bound on the sample complexity in this setting was shown to be of  $\tilde{\Omega}\left(\frac{|S||A|}{(1 - \gamma)^3\varepsilon^2}\right)$  by [6]. This lower-bound can be reached by model-based approaches [2, 26] and model-free approaches [34, 37].
The sample-complexity for PG methods has been recently studied in [39]. Under a generative model, some works have considered PMD or NPG under various types of regularisation [13, 24, 41]. We focus on unregularised methods, for which results for PMD or its instances on tabular MDPs under a generative model are limited. There are works that obtain sample complexity results for NPG [3, 28] and for PMD [33] though they do not attain the optimal  $\varepsilon$ -dependence of  $\mathcal{O}(\varepsilon^{-2})$ . [25] show that a variant of PI, a special case of PMD, achieves the optimal  $\varepsilon$ -dependence of  $\mathcal{O}(\varepsilon^{-2})$ . More recently, [38] show that the general family of PMD methods match the  $\mathcal{O}(\varepsilon^{-2})$  sample complexity with a factor of  $(1 - \gamma)^{-8}$ . Our result for the inexact setting shares the same dependence on  $\varepsilon$  and  $1 - \gamma$  as [38] but removes an instance-dependent quantity which can depend on the size of the state space. Further comparison to the result in [38] is given in Section 5. Beyond tabular MDPs and generative models, [15] study NPG under linear function approximation and off-policy sampling, though their results imply worse sample complexities when restricted to tabular MDPs under a generative model. PMD under linear function approximation [40, 4] and general function approximation [5] have also been studied and results similar to [38] were obtained in those settings.
A Markov Decision Process (MDP) is a discrete-time stochastic process, comprised of a set of states  $S$ , a set of actions  $\mathcal{A}$ , a discount factor  $\gamma \in [0,1)$  and for each state-action pair  $(s,a)$  a next-state transition function  $p(\cdot |s,a) \in \Delta(S)$  and a (assumed here deterministic) reward function  $r(s,a) \in [0,1]$ .  $\Delta(\mathcal{X})$  denotes the probability simplex over a set  $\mathcal{X}$ . We consider both  $S$  and  $\mathcal{A}$  to be finite, which is known as the tabular setting. In a state  $s$ , an agent chooses an action  $a$ , which gives them a reward  $r(s,a)$  and transitions them to a new state according to the transition function  $p(\cdot |s,a)$ . Once they are in a new state, the process continues. The actions chosen by an agent are formalised through policies. A policy  $\pi : S \to \Delta(\mathcal{A})$  is a mapping from a state to a distribution over actions. We will often write it as an element in  $\Pi = \Delta(\mathcal{A})^{|S|}$ . In each state  $s \in S$ , an agent following policy  $\pi$  chooses an action  $a \in \mathcal{A}$  according to  $\pi_s = \pi(\cdot |s) \in \Delta(\mathcal{A})$ .
In this work, the goal is to learn how to behave in a  $\gamma$ -discounted infinite-horizon MDP. We measure the performance of a policy with respect to the value function  $V^{\pi}:\mathcal{S}\to \mathbb{R}$
$$
V ^ {\pi} (s) = \mathbb {E} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r \left(s _ {t}, a _ {t}\right) | \pi , s _ {0} = s \right],
$$
where  $s_t, a_t$  are the state and action in time-step  $t$  and the expectation is with respect to both the randomness in the transitions and the choice of actions under policy  $\pi$ . This is a notion of long-term reward that describes the discounted rewards accumulated over future time-steps when following policy  $\pi$  and starting in state  $s$ . For a distribution over starting states  $\rho \in \Delta(S)$ , we write  $V^{\pi}(\rho) = \sum_{s \in S} \rho(s) V^{\pi}(s)$  for the expected value when starting in a state distributed according to  $\rho$ . It is also useful to work with the state-action value  $Q^{\pi}: S \times \mathcal{A} \to \mathbb{R}$ :
$$
Q ^ {\pi} (s, a) = \mathbb {E} \Big [ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r (s _ {t}, a _ {t}) | \pi , s _ {0} = s, a _ {0} = a \Big ],
$$
which is similar to  $V^{\pi}$ , with the additional constraint of taking action  $a$  in the first time-step. We will often write  $V^{\pi} \in \mathbb{R}^{|\mathcal{S}|}$  (resp.  $Q^{\pi} \in \mathbb{R}^{|\mathcal{S}| \times |\mathcal{A}|}$ ) to refer to the vector form, where each entry represents the value (resp. action-value) in that state (resp. state-action pair). Similarly, we write  $Q_{s}^{\pi} \in \mathbb{R}^{|\mathcal{A}|}$  for the vector of action-values in state  $s$ . The following expressions, which relate  $Q^{\pi}$  and  $V^{\pi}$  in terms of each other and when combined give the Bellman equations [9], follow from their definitions above,
$$
V ^ {\pi} (s) = \langle Q _ {s} ^ {\pi}, \pi_ {s} \rangle , \quad Q ^ {\pi} (s, a) = r (s, a) + \gamma \sum_ {s ^ {\prime} \in \mathcal {S}} p (s ^ {\prime} | s, a) V ^ {\pi} (s ^ {\prime}).
$$
We now define the discounted visitation-distribution for starting state  $s'$  and policy  $\pi$ ,
$$
d _ {s ^ {\prime}} ^ {\pi} (s) = (1 - \gamma) \sum_ {t = 0} ^ {\infty} \gamma^ {t} \mathbb {P} ^ {\pi} \left(s _ {t} = s \mid s _ {0} = s ^ {\prime}\right), \tag {1}
$$
which plays an important part in the study of PG methods. Note that  $\mathbb{P}^{\pi}(s_t = s|s_0 = s')$  is the probability of being in state  $s$  at time  $t$  when starting in state  $s'$  and following policy  $\pi$ . We also write  $d_{\rho}^{\pi}(s) = \sum_{s' \in S} \rho(s') d_{s'}^{\pi}(s)$ .
One of the main aims of reinforcement learning is to find a policy  $\pi$  that maximises  $V^{\pi}$ . It is known that there exists a deterministic policy that simultaneously maximises  $V^{\pi}$  and  $Q^{\pi}$  for all states and actions [9]. We call such a policy an optimal policy and denote it by  $\pi^{\star}$ . We are interested in finding an  $\varepsilon$ -optimal policy, i.e a policy  $\pi$  such that  $\| V^{\pi^{\star}} - V^{\pi} \|_{\infty} < \varepsilon$ .
3.1 Exact policy mirror descent
We are interested in PG methods that are based on optimising a parameterised policy  $\pi_{\theta}$  with respect to  $V^{\pi_{\theta}}(\rho)$  for some  $\rho \in \Delta(S)$ . In the tabular setting, we can use the direct parameterisation of a policy  $\pi_{\theta}$ , which associates a parameter to each state-action pair, i.e. we have  $\pi_{\theta}(a|s) = \theta_{s,a}$ . We will drop the subscript  $\theta$  for notational convenience. The gradient of the value function with respect to this parameterisation [36] is given for each state-action pair (s,a) by
$$
\frac {\partial}{\partial \pi (a | s)} V ^ {\pi} (\rho) = \frac {1}{1 - \gamma} d _ {\rho} ^ {\pi} (s) Q ^ {\pi} (s, a). \tag {2}
$$
Mirror Descent (MD, [8]) carries out gradient descent in a geometry that is non-Euclidean. Using  $-V^{\pi}(\rho)$  as the minimising objective, the proximal perspective of MD gives an update of the form
$$
\pi^ {k + 1} = \operatorname {a r g m i n} _ {p \in \Pi} \left\{- \eta_ {k} \langle \nabla V ^ {\pi^ {k}} (\rho), p \rangle + D _ {h} (p, \pi^ {k}) \right\} \tag {3}
$$
where  $h: \operatorname{dom} h \to \mathbb{R}$  is the mirror map (with  $\Pi \subset \operatorname{dom} h$ ) and  $D_h$  is the Bregman divergence generated by  $h$ . We require  $h$  to be of Legendre type [30], i.e. strictly convex and essentially smooth (differentiable and  $\| \nabla h(x_k) \| \to \infty$  for any sequence  $x_k$  converging to a point on the boundary of  $\operatorname{dom} h$ ) on the relative interior (rint) of  $\operatorname{dom} h$ . The Bregman Divergence is defined as
$$
D _ {h} \left(\pi , \pi^ {\prime}\right) = h (\pi) - h \left(\pi^ {\prime}\right) - \langle \nabla h \left(\pi^ {\prime}\right), \pi - \pi^ {\prime} \rangle \quad \text {f o r} \pi , \pi^ {\prime} \in \operatorname {d o m} h.
$$
As the objective  $V^{\pi}(\rho)$  is non-convex in general [3], usual techniques from convex theory [12] are not applicable.
The presence of the visitation-distribution term  $d_{\rho}^{\pi}(s)$  in the gradient of the objective in (2) can slow down learning because it can lead to vanishingly small gradients when states are infrequently visited under the current policy  $\pi$  [3]. To circumvent this issue, Policy Mirror Descent (PMD) [24, 33, 38] applies a variant of update (3) with a weighted Bregman divergence  $D_h^{\mathrm{PMD}}$  that matches the visitation distribution factors of the gradient  $D_h^{\mathrm{PMD}}(p, \pi^k) = \sum_s d_{\rho}^{\pi^k}(s) D_h(p_s, \pi_s^k)$  where the mirror map  $h$  is now defined on a subset of  $\mathbb{R}^{|\mathcal{A}|}$ . The resulting update has for all states a factor of  $d_{\rho}^{\pi^k}(s)$  in both terms. The minimisation can then be applied for each state individually to get the PMD update
$$
\pi_ {s} ^ {k + 1} = \operatorname {a r g m i n} _ {p \in \Delta (\mathcal {A})} \left\{- \eta_ {k} \langle Q _ {s} ^ {\pi^ {k}}, p \rangle + D _ {h} (p, \pi_ {s} ^ {k}) \right\} \tag {4}
$$
for all states  $s$ . We will often add a superscript  $k$  to any quantity that is associated to  $\pi^k$ . For example,  $V^k(s) = V^{\pi^k}(s)$ . Similarly for  $\pi^\star$  and the superscript  $\star$ . Exact PMD iteratively applies update (4) for some sequence of step-sizes  $\eta_k > 0$  and initial policy  $\pi^0 \in \mathrm{rint} \, \Pi$ . We call this algorithm exact because we assume access to the true state-action values  $Q^k$ .
The update (4) of PMD considered in this work uses the true action-value  $Q^{\pi}$ . In prior work, PMD has sometimes been applied to regularised MDPs [24] where the action-value is augmented with regularisation and is no longer the true action-value. This is a different algorithm that converges to a policy that is optimal in the regularised MDP but not in the original unregularised MDP.
PMD is a general family that covers many algorithms, specified by the choice of mirror map  $h$ . These will inherit the guarantees of PMD, which motivates the study of the convergence guarantees of
PMD beyond specific instances. Taking  $h$  to be the negative entropy yields NPG, whose theoretical properties have attracted a lot of interest [3, 13, 22]. With a null Bregman Divergence, PMD recovers PI. PI is generated by a constant mirror map, which is not of Legendre type but the analysis still applies so all results on PMD remain valid for PI. In fact, PMD can be viewed as a form of regularised PI since the update (4) converges to a PI update as  $\eta_k \to \infty$ , regardless of the mirror map. Beyond these, providing mirror maps that generate other Bregman Divergences will lead to different algorithms. In particular, every exponential family has a corresponding mirror map generating a unique Bregman Divergence [7], highlighting the generality of PMD.
4 Main results for exact policy mirror descent
In this section, we present our main results on the convergence of exact PMD. We first introduce some relevant notation. Fix a state  $s \in S$  and an integer  $k \geq 0$ . Let  $\mathcal{A}_s^k = \{a \in \mathcal{A} : Q^k(s, a) = \max_{a' \in \mathcal{A}} Q^k(s, a')\}$  denote the set of optimal actions in state  $s$  under policy  $\pi^k$ . Denote by  $\widetilde{\Pi}_s^{k+1}$  the set of greedy policies w.r.t  $Q_s^k$  in state  $s$ , i.e.  $\widetilde{\Pi}_s^{k+1} = \left\{ p \in \Delta(\mathcal{A}) : \sum_{a \in \mathcal{A}_s^k} p(a) = 1 \right\}$ . We are now ready to state our main result in the setting of exact PMD, which is proved in Section 6.
Theorem 4.1. Let  $\{c_k\}_{k\in \mathbb{Z}_{\geq 0}}$  be a sequence of positive reals. Consider applying iterative updates of (4) with  $\pi^0\in \mathrm{rint}\Pi$  and step-sizes satisfying for all  $k\geq 0$