Policy Mirror Descent for Regularized Reinforcement Learning: A Generalized Framework with Linear Convergence
Wenhao Zhan, Shicong Cen, Baihe Huang, Yuxin Chen, Jason D. Lee, Yuejie Chi
Introduction
Policy optimization lies at the heart of recent successes of reinforcement learning (RL) (Mnih et al., 2015). In its basic form, the optimal policy of interest, or a suitably parameterized version, is learned by attempting to maximize the value function in a Markov decision processes (MDP). For the most part, the maximization step is carried out by means of first-order optimization algorithms amenable to large-scale applications, whose foundations were set forth in the early works of Williams, 1992; Sutton et al., 2000. A partial list of widely adopted variants in modern practice includes policy gradient (PG) methods (Sutton et al., 2000), natural policy gradient (NPG) methods (Kakade, 2002), TRPO (Schulman et al., 2015), PPO (Schulman et al., 2017), soft actor-critic methods (Haarnoja et al., 2018), to name just a few. In comparison with model-based and value-based approaches, this family of policy-based algorithms offers a remarkably flexible framework that accommodates both continuous and discrete action spaces, and lends itself well to the incorporation of powerful function approximation schemes like neural networks. In stark contrast to its practical success, however, theoretical understanding of policy optimization remains severely limited even for the tabular case, largely owing to the ubiquitous nonconvexity issue underlying the objective function.
In practice, there are often competing objectives and additional constraints that the agent has to deal with in conjunction with maximizing values, which motivate the studies of regularization techniques in RL. In what follows, we isolate a few representative examples.
Promoting exploration. In the face of large problem dimensions and complex dynamics, it is often desirable to maintain a suitable degree of randomness in the policy iterates, in order to encourage exploration and discourage premature convergence to sub-optimal policies. A popular strategy of this kind is to enforce entropy regularization (Williams and Peng, 1991), which penalizes policies that are not sufficiently stochastic. Along similar lines, the Tsallis entropy regularization (Chow et al., 2018b; Lee et al., 2018) further promotes sparsity of the learned policy while encouraging exploration, ensuring that the resulting policy does not assign non-negligible probabilities to too many sub-optimal actions.
Safe RL. In a variety of application scenarios such as industrial robot arms and self-driving vehicles, the agents are required to operate safely both to themselves and the surroundings (Amodei et al., 2016; Moldovan and Abbeel, 2012); for example, certain actions might be strictly forbidden in some states. One way to incorporate such prescribed operational constraints is through adding a regularizer (e.g., a properly chosen log barrier or indicator function tailored to the constraints) to explicitly account for the constraints.
Cost-sensitive RL. In reality, different actions of an agent might incur drastically different costs even for the same state. This motivates the design of new objective functions that properly trade off the cumulative rewards against the accumulated cost, which often take the form of certain regularized value functions.
Viewed in this light, it is of imminent value to develop a unified framework towards understanding the capability and limitations of regularized policy optimization. While a recent line of works (Agarwal et al., 2020; Mei et al., 2020b; Cen et al., 2022b) have looked into specific types of regularization techniques such as entropy regularization, existing convergence theory remains highly inadequate when it comes to a more general family of regularizers.
2 Main contributions
The current paper focuses on policy optimization for regularized RL in a -discounted infinite horizon Markov decision process (MDP) with state space , action space , and reward function . The goal is to find an optimal policy that maximizes a regularized value function. Informally speaking, the regularized value function associated with a given policy takes the following form:
where denotes the original (unregularized) value function, is the regularization parameter, denotes a convex regularizer employed to regularize the policy in state , and the expectation is taken over certain marginal state distribution w.r.t. the MDP (to be made precise in Section 2.1). It is noteworthy that this paper does not require the regularizer to be either strongly convex or smooth.
In order to maximize the regularized value function (9b), Lan, 2022 exhibited a seminal algorithm called Policy Mirror Descent (PMD), which can be viewed as an adaptation of the mirror descent algorithm (Nemirovsky and Yudin, 1983; Beck and Teboulle, 2003) to the realm of policy optimization. In particular, PMD subsumes the natural policy gradient (NPG) method (Kakade, 2002) as a special case. To further generalize PMD (Lan, 2022), we propose an algorithm called Generalized Policy Mirror Descent (GPMD). In each iteration, the policy is updated for each state in parallel via a mirror-descent style update rule. In sharp contrast to Lan, 2022 that considered a generic Bregman divergence, our algorithm selects the Bregman divergence adaptively in cognizant of the regularizer, which leads to complementary perspectives and insights. Several important features and theoretical appeal of GPMD are summarized as follows.
GPMD substantially broadens the range of (provably effective) algorithmic choices for regularized RL, and subsumes several well-known algorithms as special cases. For example, it reduces to regularized policy iteration (Geist et al., 2019) when the learning rate tends to infinity, and subsumes entropy-regularized NPG methods as special cases if we take the Bregman divergence to be the Kullback-Leibler (KL) divergence (Cen et al., 2022b).
Assuming exact policy evaluation and perfect policy update in each iteration, GPMD converges linearly—in a dimension-free fashion— over the entire range of the learning rate . More precisely, it converges to an -optimal regularized Q-function in no more than an order of
iterations (up to some logarithmic factor). Encouragingly, this appealing feature is valid for a broad family of convex and possibly nonsmooth regularizers.
The intriguing convergence guarantees are robust in the face of inexact policy evaluation and imperfect policy updates, namely, the algorithm is guaranteed to converge linearly at the same rate until an error floor is hit. See Section 3.2 for details.
Numerical experiments are provided in Section 5 to demonstrate the practical applicability and appealing performance of the proposed GPMD algorithm.
Finally, we find it helpful to briefly compare the above findings with prior works. As soon as the learning rate exceeds , the iteration complexity of our algorithm is at most on the order of , thus matching that of regularized policy iteration (Geist et al., 2019). In comparison to Lan, 2022, our work sets forth a different framework to analyze mirror-descent type algorithms for regularized policy optimization, generalizing and refining the approach in Cen et al., 2022b far beyond entropy regularization. When constant learning rates are employed, the linear convergence of PMD (Lan, 2022) critically requires the regularizer to be strongly convex, with only sublinear convergence theory established for convex regularizers. In contrast, we establish the linear convergence of GPMD under constant learning rates even in the absence of strong convexity. Furthermore, for the special case of entropy regularization, the stability analysis of GPMD also significantly improves over the prior art in Cen et al., 2022b, preventing the error floor from blowing up when the learning rate approaches zero, as well as incorporating the impact of optimization error that was previously uncaptured. More detailed comparisons with Lan, 2022 and Cen et al., 2022b can be found in Section 3.
3 Related works
Before embarking on our algorithmic and theoretic developments, we briefly review a small sample of other related works.
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.
4 Notation
Let us introduce several notation that will be adopted throughout. For any set , we denote by the cardinality of a set , and let indicate the probability simplex over the set . For any convex and differentiable function , the Bregman divergence generated by is defined as
For any convex (but not necessarily differentiable) function , we denote by the subdifferential of . Given two probability distributions and over , the KL divergence from to is defined as . For any vectors and , the notation (resp. ) means that () for every . We shall also use (resp. ) to denote the all-one (resp. all-zero) vector whenever it is clear from the context.
Model and algorithms
The focus of this paper is a discounted infinite-horizon Markov decision process, as represented by (Bertsekas, 2017). Here, is the state space, is the action space, is the discount factor, is the probability transition matrix (so that is the transition probability from state upon execution of action ), whereas is the reward function (so that indicates the immediate reward received in state after action is executed). Here, we focus on finite-state and finite-action scenarios, meaning that both and are assumed to be finite. A policy specifies a possibly randomized action selection rule, namely, represents the action selection probability in state .
which can be viewed as the utility function we wish to maximize. Here, the expectation is taken over the randomness of the MDP trajectory induced by policy . Similarly, when the initial action is fixed, we can define the action-value function (or Q-function) as follows
As a well-known fact, the policy gradient of (w.r.t. the policy ) admits the following closed-form expression (Sutton et al., 2000)
Here, is the so-called discounted state visitation distribution defined as follows
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 , 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 (Neu et al., 2017; Mei et al., 2020b; Cen et al., 2022b; Lan, 2022):
Consider an arbitrarily small constant . For for any , suppose that is convex and
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 and the corresponding optimal policy are defined respectively as follows:
It is worth noting that Puterman, 2014 asserts the existence of an optimal policy that achieves (10) simultaneously for all . Correspondingly, we shall also define the resulting optimal regularized Q-function as
2 Algorithm: generalized policy mirror descent
Motivated by PMD (Lan, 2022), we put forward a generalization of PMD that selects the Bregman divergence in cognizant of the regularizer in use. A thorough comparison with Lan, 2022 will be provided after introducing our generalized PMD algorithm.
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 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 is the learning rate or step size, and 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 . In words, the above update rule approximates via its first-order Taylor expansion at the point , employs the Bregman divergence to monitor the difference between the new iterate and the current iterate , and attempts to optimize such (properly monitored) approximation instead. While one can further generalize the Bregman divergence to for a different generator , we shall restrict attention to the case with 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 denotes our policy estimate in the -th iteration.
To begin with, suppose for simplicity that is differentiable everywhere. In the -th iteration, a natural MD scheme that comes into mind for solving (7)—namely, for a given initial state —is the following update rule:
for every state , which is a direct application of (12) to our setting. Here, we start with a learning rate , and obtain simplification by replacing with . Notably, the update strategy (14) is invariant to the initial state , akin to natural policy gradient methods (Agarwal et al., 2020).
This update rule is well-defined for, say, the case when is the negative entropy, since the algorithm guarantees all the time and hence is always differentiable w.r.t. the -th iterate (see Cen et al., 2022b). In general, however, it is possible to encounter situations when the gradient of does not exist on the boundary (e.g., when 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 by the following metric
where the last line is valid since . As a result, everything boils down to identifying a vector that falls within upon global shift.
Towards this, we propose the following iterative rule for designing such a sequence of vectors as surrogates for the subgradient of :
where is updated as a convex combination of the previous and , where more emphasis is put on when the learning rate is large. As asserted by the following lemma, the above vectors we construct satisfy the desired property, i.e., lying within the subdifferential of under suitable global shifts. It is worth mentioning that these global shifts 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 is taken as the KL divergence, GPMD reduces to the well-renowned NPG algorithm (Kakade, 2002) when (no regularization), and to the NPG algorithm with entropy regularization analyzed in (Cen et al., 2022b) when is taken as the negative Shannon entropy.
When (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 is also .
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 is, however, unspecified. Furthermore, GPMD adaptively sets this term to be the Bregman divergence generated by the regularizer in use, together with a carefully designed recursive update rule (cf. (17)) to compute surrogates for the subgradient of 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 ; details to follow momentarily.
Main results
This section presents our convergence guarantees for the GPMD method presented in Algorithm 1. We shall start with the idealized case assuming that the update rule can be precisely implemented, and then discuss how to generalize it to the scenario with imperfect policy evaluation.
To start with, let us pin down the convergence behavior of GPMD, assuming that accurate evaluation of the policy is available and the subproblem (20a) can be solved perfectly. Here and below, we shall refer to the algorithm in this case as exact GPMD. Encouragingly, exact GPMD provably achieves global linear convergence from an arbitrary initialization, as asserted by the following theorem.
Suppose that Assumption 1 holds. Consider any learning rate , and set . Then the iterates of Algorithm 1 satisfy
for all , where .
Our theorem confirms the fast global convergence of the GPMD algorithm, in terms of both the resulting regularized Q-value (if is convex) and the policy estimate (if is strongly convex). In summary, it takes GPMD no more than
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 and —even when the regularizer is only convex but not strongly convex. In contrast, Lan, 2022 provided only sublinear convergence guarantees (with an iteration complexity proportional to ) for the case with convex regularizers, provided that constant learning rates are adopted. In fact, Lan, 2022 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 . In comparison, linear convergence was established in Lan, 2022 only when the learning rates are sufficiently large and when is -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 when the regularizer is strongly convex and the step size satisfies , while the contraction rate of GPMD is under the full range of the step size, which is slower but approaches the contraction rate of PMD as goes to infinity. Therefore, in the limit , both GPMD and PMD achieve the contraction rate . As soon as , their iteration complexities are on the same order.
While our primary focus is to solve the regularized RL problem, one might be tempted to apply GPMD as a means to solve unregularized RL; for instance, one might run GPMD with the regularization parameter diminishing gradually in order to approach a policy with the desired accuracy. We leave the details to Appendix C.
2 Convergence of approximate GPMD
In reality, however, it is often the case that GPMD cannot be implemented in an exact manner, either because perfect policy evaluation is unavailable or because the subproblem (20a) cannot be solved exactly. To accommodate these practical considerations, this subsection generalizes our previous result by permitting inexact policy evaluation and non-zero optimization error in solving (20a). The following assumptions make precise this imperfect scenario.
Suppose for any , we have access to an estimate obeying
where is defined in (15). Suppose there exists an oracle , which is capable of returning such that
Note that the oracle in Assumption 3 can be implemented efficiently in practice via various first-order methods (Beck, 2017). Under Assumptions 2 and 3, we can modify Algorithm 1 by replacing with the estimate , and invoking the oracle to solve the subproblem (20a) approximately. The whole procedure, which we shall refer to as approximate GPMD, is summarized in Algorithm 2.
The following theorem uncovers that approximate GPMD converges linearly—at the same rate as exact GPMD—before an error floor is hit.
Suppose that Assumptions 1, 2 and 3 hold. Consider any learning rate . Then the iterates of Algorithm 2 satisfy
where , is defined in Theorem 1, and
In the special case where and , Algorithm 2 reduces to regularized policy iteration, and the convergence result can be simplified as follows
In particular, when is taken as the negative entropy, our result strengthens the prior result established in Cen et al., 2022b for approximate entropy-regularized NPG method with over a wide range of learning rates. Specifically, the error bound in Cen et al., 2022b reads , where the second term in the bracket scales inversely with respect to and therefore grows unboundedly as approaches . In contrast, (29) and (30) suggest a bound , which is independent of the learning rate in use and thus prevents the error bound from blowing up when the learning rate approaches . Indeed, our result improves over the prior art Cen et al., 2022b whenever .
One might naturally ask how many samples are sufficient to learn an -optimal regularized Q-function, by leveraging sample-based policy evaluation algorithms in GPMD. Notice that it is straightforward to consider an expected version of Assumption 2 as following:
Roughly speaking, approximate GPMD is guaranteed to converge linearly to an error bound that scales linearly in both the policy evaluation error and the optimization error , thus confirming the stability of our algorithm vis-à-vis imperfect implementation of the algorithm. As before, our theory improves upon prior works by demonstrating linear convergence for a full range of learning rates even in the absence of strong convexity and smoothness.
Analysis for exact GPMD (Theorem 1)
In this section, we present the analysis for our main result in Theorem 1, which follows a different framework from Lan, 2022. Here and throughout, we shall often employ the following shorthand notation when it is clear from the context:
in addition to those already defined in (13).
In this subsection, we single out a few basic results that underlie the proof of our main theorems.
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).
For any and any , Algorithm 1 achieves
Interestingly, the above monotonicity holds simultaneously for all state-action pairs, and hence can be understood as a kind of pointwise monotonicity.
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 , while ours is defined w.r.t. .
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.
For any , the operator defined in (34) satisfies the following properties:
The optimal regularized -function is a fixed point of , that is,
2 Proof of Theorem 1
With the assistance of the above preparations, we are ready to elucidate how to characterize the convergence behavior of . Recalling the update rule of (cf. (20c)), we can deduce that
with , thus indicating that
Interestingly, there exists an intimate connection between and that allows us to bound the former term by the latter. This is stated in the following lemma, with the proof postponed to Appendix A.4.
Set . The iterates of Algorithm 1 satisfy
The above inequalities (37) and (38) can be succinctly described via a useful linear system with two variables and , that is,
This forms the basis for proving Theorem 1.
Before proceeding, we note that a linear system similar to (39) has been analyzed in Cen et al., 2022b. We intend to apply the following properties that have been derived therein:
Substituting (41c) and (41b) into (41a) and rearranging terms, we reach
which taken together with the definition of gives
where (i) results from . Plugging (43) into (44) completes the proof for (21b).
where the second line is valid since . This taken together with (43) gives rise to the advertised bound
Numerical experiments
In this section, we provide some simple numerical experiments to corroborate the effectiveness of the GPMD algorithm.
While Shannon entropy is a popular choice of regularization, the discrepancy between the value function of the regularized MDP and the unregularized counterpart scales as . In addition, the optimal policy under Shannon entropy regularization assigns positive mass to all actions and is hence non-sparse. To promote sparsity and obtain better control of the bias induced by regularization, Lee et al., 2018; Lee et al., 2019 proposed to employ the Tsallis entropy (Tsallis, 1988) as an alternative. To be precise, for any vector , the associated Tsallis entropy is defined as
where is often referred to as the entropic-index. When , the Tsallis entropy reduces to the Shannon entropy.
We now evaluate numerically the performance of PMD and GPMD when applied to a randomly generated MDP with and . Here, the transition probability kernel and the reward function are generated as follows. For each state-action pair , we randomly select states to form a set , and set if , and otherwise. The reward function is generated by , where and are independent uniform random variables over $h_{s}(p)=-\mathsf{Tsallis}_{2}(p)s\in\mathcal{S}\tau=0.001$. As can be seen from the numerical results displayed in Figure 1(a), GPMD enjoys a faster convergence rate compared to PMD.
2 Constrained RL
In reality, an agent with the sole aim of maximizing cumulative rewards might sometimes end up with unintended or even harmful behavior, due to, say, improper design of the reward function, or non-perfect simulation of physical laws. Therefore, it is sometimes necessary to enforce proper constraints on the policy in order to prevent it from taking certain actions too frequently.
To simulate this problem, we first solve a MDP with and , generated in the same way as in the previous subsection. We then pick 10 state-action pairs from the support of the optimal policy at random to form a set . We can ensure that for all by adding the following log-barrier regularization with :
Numerical comparisons of PMD and GPMD when applied this problem are plotted in Figure 1(b). It is observed that PMD methods stall after reaching an error floor on the order of , while GPMD methods are able to converge to the optimal policy efficiently.
Discussion
The present paper has introduced a generalized framework of policy optimization tailored to regularized RL problems. We have proposed a generalized policy mirror descent (GPMD) algorithm that achieves dimension-free linear convergence, which covers an entire range of learning rates and accommodates convex and possibly nonsmooth regularizers. Numerical experiments have been conducted to demonstrate the utility of the proposed GPMD algorithm. Our approach opens up a couple of future directions that are worthy of further exploration. For example, the current work restricts its attention to convex regularizers and tabular MDPs; it is of paramount interest to develop policy optimization algorithms when the regularizers are nonconvex and when sophisticated policy parameterization—including function approximation—is adopted. Understanding the sample complexities of the proposed algorithm—when the policies are evaluated using samples collected over an online trajectory—is crucial in sample-constrained scenarios and is left for future investigation. Furthermore, it might be worthwhile to extend the proposed algorithm to accommodate multi-agent RL, with a representative example being regularized multi-agent Markov games (Cen et al., 2021; Zhao et al., 2022; Cen et al., 2022a; Cen et al., 2022c).
Acknowledgements
S. Cen and Y. Chi are supported in part by the grants ONR N00014-19-1-2404, NSF CCF-2106778, DMS-2134080, CCF-1901199, CCF-2007911, and CNS-2148212. S. Cen is also gratefully supported by Wei Shen and Xuehong Zhang Presidential Fellowship, and Nicholas Minnici Dean’s Graduate Fellowship in Electrical and Computer Engineering at Carnegie Mellon University. W. Zhan and Y. Chen are supported in part by the Google Research Scholar Award, the Alfred P. Sloan Research Fellowship, and the grants AFOSR FA9550-22-1-0198, ONR N00014-22-1-2354, NSF CCF-2221009, CCF-1907661, IIS-2218713, and IIS-2218773. W. Zhan and J. Lee are supported in part by the ARO under MURI Award W911NF-11-1-0304, the Sloan Research Fellowship, NSF CCF 2002272, NSF IIS 2107304, and an ONR Young Investigator Award.
Appendix A Proof of key lemmas
for any policy and and any , whenever it is clear from the context.
We start by relaxing the probability simplex constraint (i.e., ) in (20a) with a simpler linear constraint as follows
To justify the validity of dropping the non-negative constraint, we note that for any obeying for some , our assumption on (see Assumption 1) leads to , which cannot possibly be the optimal solution. This confirms the equivalence between (20a) and (47).
Observe that the Lagrangian w.r.t. (47) is given by
Rearranging terms and making use of the construction (17), we are left with
thus concluding the proof of the first claim (18).
We now turn to the second claim (19). In view of the property (36), we have
This optimization problem is equivalent to
which can be verified by repeating a similar argument for (47). The Lagrangian associated with (48) is
A.2 Proof of Lemma 2
We start by introducing the performance difference lemma that has previously been derived in Lan, 2022. For the sake of self-containedness, we include a proof of this lemma in Appendix A.2.1.
For any two policies and , we have
where has been defined in (5).
Armed with Lemma 5, one can readily rewrite the difference between two consecutive iterates as follows
It then comes down to studying the right-hand side of the relation (A.2), which can be accomplished via the following “three-point” lemma. The proof of this lemma can be found in Appendix A.2.2.
For any and any vector , we have
Taking in Lemma 6 and combining it with (A.2), we arrive at
for any , thus establishing the advertised pointwise monotonicity w.r.t. the regularized value function.
When it comes to the regularized Q-function, it is readily seen from the definition (9a) that
for any , where the last line is valid since . This concludes the proof.
For any two policies and , it follows from the definition (7) of that
where the penultimate line comes from the definition (9a). To see why the last line of (51) is valid, we make note of the following identity
where the first identity results from the relation (9b), and the last relation holds since . The last line of (51) then follows immediately from the relation (52) and the definition (5) of .
A.2.2 Proof of Lemma 6
For any state , we make the observation that
where the first and the fourth steps invoke the definition (15) of the generalized Bregman divergence and the last line results from the update rule (20c). Rearranging terms, we are left with
Adding the term to both sides of this identity leads to
as claimed, where the last line makes use of the definition (15).
A.3 Proof of Lemma 3
In the sequel, we shall prove each claim in Lemma 3 separately.
where (a) arises from the elementary fact .
where the first identity results from (9), and the second line arises from the maximizing property of (see (53)).
Note that the right-hand side of (54) involves the term , which can be further upper bounded via the same argument for (54). Successively repeating this upper bound argument (and the expansion) eventually allows one to obtain
However, the fact that is the optimal policy necessarily implies the following reverse inequality:
To finish up, it suffices to show that . To this end, it is observed that
where (b) utilizes the fact (55), (c) follows from the definition (53) of , and the last identity is a consequence of the definition (34) of . The above results taken collectively demonstrate that as claimed.
A.4 Proof of Lemma 4
Recall that . In view of the relation (9), one obtains
This combined with the fixed-point condition (36) allows us to derive
In what follows, we control each term on the right-hand side of (A.4) separately.
The condition (57) can then be interpreted as the optimality condition w.r.t. the program (58) and , meaning that
In addition, for any vector that does not obey , Assumption 1 implies that , and hence cannot possibly be the optimal solution to . This together with (59) essentially implies that
where the last step results from the contraction property (35) in Lemma 3.
Recall that . Invoking the monotonicity property in Lemma 2 and the update rule (20c), we obtain
Repeating this lower bound argument then yields
Substituting (61) and (62) into (A.4) gives
for all , thus concluding the proof.
Appendix B Analysis for approximate GPMD (Theorem 2)
The proof consists of three steps: (i) evaluating the performance difference between and , (ii) establishing a linear system to characterize the error dynamic, and (iii) analyzing this linear system to derive global convergence guarantees. We shall describe the details of each step in the sequel. As before, we adopt the notational convention (46) whenever it is clear from the context.
When only approximate policy evaluation is available, we are no longer guaranteed to have pointwise monotonicity as in the case of Lemma 2. Fortunately, we are still able to establish an approximate versioin of Lemma 2, as stated below.
For all and all , we have
In words, while monotonicity is not guaranteed, this lemma precludes the possibility of being be much smaller than , as long as both and are reasonably small.
Let be the exact solution of the following problem
With this auxiliary policy iterate in mind, we start by decomposing into the following three parts:
where the first identity arises from the performance difference lemma (cf. Lemma 5). To continue, we seek to control each part of (65) separately.
Regarding the first term of (65), replacing (resp. ) by (resp. ) in Lemma 6 indicates that
As for the second term of (65), the definition of the oracle (see Assumption 3) guarantees that
for any . Rearranging terms, we are left with
appears in both (66) and (68), which can be canceled out when summing these two equalities. Specifically, adding (66) and (68) gives
Substituting this into (65) and invoking the elementary inequality thus lead to
where the last line makes use of Assumption 2 and the fact .
Following the discussion in Lemma 1, we can see that with some constant for all . This together with the convexity of (see (15)) guarantees that
for any , thus implying that the first term of (69) is non-negative. It remains to control the second term in (69). Towards this, a little algebra gives
Here, the first and the third lines follow from the definition (15), the second inequality comes from the construction (27), whereas the last step invokes the definition of the oracle (25). Substitution of (70) and (71) into (69) gives
Additionally, the strong convexity assumption also implies that
where the third line results from Young’s inequality, and the final step follows from (73). We can develop a similar lower bound on as well. Taken together, these lower bounds give
Combining the above two inequalities with (72), we arrive at the advertised bound
B.2 Step 2: connecting the algorithm dynamic with a linear system
Now we are ready to discuss how to control . In short, we intend to establish the connection among several intertwined quantities, and identify a simple linear system that captures the algorithm dynamic.
From the definition of in (27), we have
where the last inequality is a consequence of Assumption 2.
Applying the definition in (27) once again, we obtain
Here, the last step of (75) follows from Assumption 2 as well as the following relation:
where we have made use of Lemma 7. Taking the maximum over on both sides of (75) yields
To begin with, let us decompose into several parts. Invoking the relation (36) in Lemma 3 as well as the property (9), we reach
In the sequel, we control the three terms in (78) separately.
To begin with, we repeat a similar argument as for (61) to show that
The second term of (78) can be bounded by applying (71) with replaced by :
As for the third term of (78), taking the maximum over all gives
Taken together, the above bounds and the decomposition (78) lead to
Combining (B.2),(77) and (79), we reach the following linear system
This linear system of three variables captures how the estimation error progresses as the iteration count increases.
B.3 Step 3: linear system analysis
In this step, we analyze the behavior of the linear system (80) derived above. Observe that the eigenvalues and respective eigenvectors of the matrix are given by
Armed with these, we can decompose in terms of the eigenvectors of as follows
Using the decomposition in (B.3) and (85) and applying the system relation (80) recursively, we can derive
Recognizing that the first two entries of are non-positive, we can discard the term involving and obtain
Making use of the fact that , we can conclude
Turning to , by a similar argument as (44), we have
where the third step results from the standard three-point lemma. To control the second term, we rearrange terms in (67) and reach at
For the remaining terms, recall that with some constant . So we have
Appendix C Adaptive GPMD
In this section, we present adaptive GPMD, an adaptive variant of GPMD that computes optimal policies of the original MDP without the need of specifying the regularization parameter in advance. In a nutshell, the proposed adaptive GPMD algorithm is a stage-based algorithm. In the -th stage, we execute GPMD (i.e., Algorithm 1) with regularization parameter for iterations. In what follows, we shall denote by and the -th iterates in the -th stage. At the end of each stage, Adaptive GPMD will halve the regularization parameter , and in the meantime, double (i.e., the auxiliary vector corresponding to some subgradient up to global shift) and use it to as the initial vector for the next stage. To ensure that still lies within the set of subgradients up to global shift, we solve the sub-optimization problem (86) to obtain as the initial policy iterate for the next stage. The whole procedure is summarized in Algorithm 3.
To help characterize the discrepancy of the value functions due to regularization, we assume boundedness of the regularizers as follows.
Suppose that there exists some quantity such that holds for all and all .
The following theorem demonstrates that Algorithm 3 is capable of finding an -optimal policy for the unregularized MDP within an order of stages. To simplify notation, we abbreviate , and as , and , respectively, as long as it is clear from the context.
Suppose that Assumptions 1 and 4 hold. For any learning rate and any stage , the iterates of Algorithm 3 satisfy
As a direct implication of Theorem 3, it suffices to run Algorithm 3 with stages, resulting in a total iteration complexity of at most
In comparison, we recall from Theorem 1 that: directly running GPMD with regularization parameter leads to an iteration complexity of
When focusing on the term , (87) improves upon (88) by a factor of .
To begin with, we make note of the fact that, for any ,
Next, we demonstrate how to control . The definition of implies the existence of some constant such that
By invoking the convergence results of GPMD (cf. (43)), we obtain: for all ,
where . To proceed, we follow similar arguments in (44) and show that
where the first step invokes the regularized performance difference lemma (Lemma 5). It then follows that
Next, we aim to prove by induction that . Clearly, this claim holds trivially for the base case with . Next, supposing that the claim holds for some , we would like to prove it for as well. Towards this end, observe that
When , we arrive at
which verifies the claim for . Substitution back into (94) leads to
Combining (95) with (90) concludes the proof. ∎