Title: A Novel Framework for Policy Mirror Descent with General Parameterization and Linear Convergence
Venue: NeurIPS 2023
Accepted: yes
Keywords: Theory for Reinforcement Learning

Policy Optimization

Policy Gradient

Mirror Descent.
Abstract: Modern policy optimization methods in reinforcement learning, such as TRPO and PPO, owe their success to the use of parameterized policies. However, while theoretical guarantees have been established for this class of algorithms, especially in the tabular setting, the use of general parameterization schemes remains mostly unjustified. In this work, we introduce a novel framework for policy optimization based on mirror descent that naturally accommodates general parameterizations. The policy class induced by our scheme recovers known classes, e.g., softmax, and generates new ones depending on the choice of mirror map. Using our framework, we obtain the first result that guarantees linear convergence for a policy-gradient-based method involving general parameterization. To demonstrate the ability of our framework to accommodate general parameterization schemes, we provide its sample complexity when using shallow neural networks, show that it represents an improvement upon the previous best results, and empirically validate the effectiveness of our theoretical claims on classic control tasks.

A Novel Framework for Policy Mirror Descent with General Parameterization and Linear Convergence

Abstract

Modern policy optimization methods in reinforcement learning, such as TRPO and PPO, owe their success to the use of parameterized policies. However, while theoretical guarantees have been established for this class of algorithms, especially in the tabular setting, the use of general parameterization schemes remains mostly unjustified. In this work, we introduce a framework for policy optimization based on mirror descent that naturally accommodates general parameterizations. The policy class induced by our scheme recovers known classes, e.g., softmax, and generates new ones depending on the choice of mirror map. Using our framework, we obtain the first result that guarantees linear convergence for a policy-gradient-based method involving general parameterization. To demonstrate the ability of our framework to accommodate general parameterization schemes, we provide its sample complexity when using shallow neural networks, show that it represents an improvement upon the previous best results, and empirically validate the effectiveness of our theoretical claims on classic control tasks.

1 Introduction

Policy optimization is one of the most widely-used classes of algorithms for reinforcement learning (RL). Among policy optimization techniques, policy gradient (PG) methods [e.g., 92, 87, 49, 6] are gradient-based algorithms that optimize the policy over a parameterized policy class and have emerged as a popular class of algorithms for RL [e.g., 42, 75, 12, 69, 78, 80, 54].

The design of gradient-based policy updates has been key to achieving empirical success in many settings, such as games [9] and autonomous driving [81]. In particular, a class of PG algorithms that has proven successful in practice consists of building updates that include a hard constraint (e.g., a trust region constraint) or a penalty term ensuring that the updated policy does not move too far from the previous one. Two examples of algorithms belonging to this category are trust region policy optimization (TRPO) [78], which imposes a Kullback-Leibler (KL) divergence [53] constraint on its updates, and policy mirror descent (PMD) [e.g. 88, 54, 93, 51, 89], which applies mirror descent (MD) [70] to RL. Shani et al. [82] propose a variant of TRPO that is actually a special case of PMD, thus linking TRPO and PMD.

From a theoretical perspective, motivated by the empirical success of PMD, there is now a concerted effort to develop convergence theories for PMD methods. For instance, it has been established that PMD converges linearly to the global optimum in the tabular setting by using a geometrically increasing step-size [54, 93], by adding entropy regularization [17], and more generally by adding convex regularization [100]. Linear convergence of PMD has also been established for the negative

Department of Statistics University of Oxford ♣❛/a116/a114✐❝❦✳/a114❡❜❡/a115❝❤✐♥✐❅/a115/a116❛/a116/a115✳♦①✳❛❝✳✉❦

entropy mirror map in the linear function approximation regime, i.e., for log-linear policies, either by adding entropy regularization [15], or by using a geometrically increasing step-size [20, 2, 98]. The proofs of these results are based on specific policy parameterizations, i.e., tabular and log-linear, while PMD remains mostly unjustified for general policy parameterizations and mirror maps, leaving out important practical cases such as neural networks. In particular, it remains to be seen whether the theoretical results obtained for tabular policy classes transfer to this more general setting.

In this work, we introduce Approximate Mirror Policy Optimization (AMPO), a novel framework designed to incorporate general parameterization into PMD in a theoretically sound manner. In summary, AMPO is an MD-based method that recovers PMD in different settings, such as tabular MDPs, is capable of generating new algorithms by varying the mirror map, and is amenable to theoretical analysis for any parameterization class. Since the MD update can be viewed as a two-step procedure, i.e., a gradient update step on the dual space and a mapping step onto the probability simplex, our starting point is to define the policy class based on this second MD step (Definition 3.1). This policy class recovers the softmax policy class as a special case (Example 3.2) and accommodates any parameterization class, such as tabular, linear, or neural network parameterizations. We then develop an update procedure for this policy class based on MD and PG.

We provide an analysis of AMPO and establish theoretical guarantees that hold for any parameterization class and any mirror map. More specifically, we show that our algorithm enjoys quasi-monotonic improvements (Proposition 4.2), sublinear convergence when the step-size is non-decreasing, and linear convergence when the step-size is geometrically increasing (Theorem 4.3). To the best of our knowledge, AMPO is the first policy-gradient-based method with linear convergence that can accommodate any parameterization class. Furthermore, the convergence rates hold for any choice of mirror map. The generality of our convergence results allows us not only to unify several current best-known results with specific policy parameterizations, i.e., tabular and log-linear, but also to achieve new state-of-the-art convergence rates with neural policies. Tables 1 and 2 in Appendix A.2 provide an overview of our results. We also refer to Appendix A.2 for a thorough literature review.

The key point of our analysis is Lemma 4.1, which allows us to keep track of the errors incurred by the algorithm (Proposition 4.2). It is an application of the three-point descent lemma by [19, Lemma 3.2] (see also Lemma F.2), which is possible thanks to our formulations of the policy class (Definition 3.1) and the policy update (Line 2 of Algorithm 1). The convergence rates of AMPO are obtained by building on Lemma 4.1 and leveraging modern PMD proof techniques [93].

In addition, we show that for a large class of mirror maps, i.e., the ω -potential mirror maps in Definition 3.5, AMPO can be implemented in ˜ O ( |A| ) computations. We give two examples of mirror maps belonging to this class, Examples 3.6 and 3.7, that illustrate the versatility of our framework. Lastly, we examine the important case of shallow neural network parameterization both theoretically and empirically. In this setting, we provide the sample complexity of AMPO, i.e., ˜ O ( ε -4 ) (Corollary 4.5), and show how it improves upon previous results.

2 Preliminaries

Let M = ( S , A , P, r, γ, µ ) be a discounted Markov Decision Process (MDP), where S is a possibly infinite state space, A is a finite action space, P ( s ′ | s, a ) is the transition probability from state s to s ′ under action a , r ( s, a ) ∈ [0 , 1] is a reward function, γ is a discount factor, and µ is a target state distribution. The behavior of an agent on an MDP is then modeled by a policy π ∈ (∆( A )) S , where a ∼ π ( · | s ) is the density of the distribution over actions at state s ∈ S , and ∆( A ) is the probability simplex over A .

Given a policy π , let V π : S → R denote the associated value function . Letting s t and a t be the current state and action at time t , the value function V π is defined as the expected discounted cumulative reward with the initial state s 0 = s , namely,

Now letting V π ( µ ) := E s ∼ µ [ V π ( s )] , our objective is for the agent to find an optimal policy

As with the value function, for each pair ( s, a ) ∈ S × A , the state-action value function, or Q-function , associated with a policy π is defined as

We also define the discounted state visitation distribution by

where P ( s t = s | π, s 0 ) represents the probability of the agent being in state s at time t when following policy π and starting from s 0 . The probability d π µ ( s ) represents the time spent on state s when following policy π .

The gradient of the value function V π ( µ ) with respect to the policy is given by the policy gradient theorem (PGT) [87]:

2.1 Mirror descent

The first tools we recall from the mirror descent (MD) framework are mirror maps and Bregman divergences [14, Chapter 4]. Let Y ⊆ R |A| be a convex set. A mirror map h : Y → R is a strictly convex, continuously differentiable and essentially smooth function 1 such that ∇ h ( Y ) = R |A| . The convex conjugate of h , denoted by h ∗ , is given by

The gradient of the mirror map ∇ h : Y → R |A| allows to map objects from the primal space Y to its dual space R |A| , x ↦→∇ h ( x ) , and viceversa for ∇ h ∗ , i.e., x ∗ ↦→∇ h ∗ ( x ∗ ) . In particular, from ∇ h ( Y ) = R |A| , we have: for all ( x, x ∗ ) ∈ Y × R |A| ,

Furthermore, the mirror map h induces a Bregman divergence [13, 18] , defined as

where D h ( x, y ) ≥ 0 for all x, y ∈ Y . We can now present the standard MD algorithm [70, 14]. Let X ⊆ Y be a convex set and V : X → R be a differentiable function. The MD algorithm can be formalized 2 as the following iterative procedure in order to solve the minimization problem min x ∈X V ( x ) : for all t ≥ 0 ,

where η t is set according to a step-size schedule ( η t ) t ≥ 0 and Proj h X ( · ) is the Bregman projection

Precisely, at time t , x t ∈ X is mapped to the dual space through ∇ h ( · ) , where a gradient step is performed as in (5) to obtain y t +1 . The next step is to map y t +1 back in the primal space using ∇ h ∗ ( · ) . In case ∇ h ∗ ( y t +1 ) does not belong to X , it is projected as in (6).

3 Approximate Mirror Policy Optimization

The starting point of our novel framework is the introduction of a novel parameterized policy class based on the Bregman projection expression recalled in (7).

Algorithm 1: Approximate Mirror Policy Optimization

Input: Initial policy π 0 , mirror map h , parameterization class F Θ , iteration number T , step-size schedule ( η t ) t ≥ 0 , state-action distribution sequence ( v t ) t ≥ 0 .

for t = 0 to T -1 do

Definition 3.1. Given a parameterized function class F Θ = { f θ : S × A → R , θ ∈ Θ } , a mirror map h : Y → R , where Y ⊆ R |A| is a convex set with ∆( A ) ⊆ Y , and η > 0 , the Bregman projected policy class associated with F Θ and h consists of all the policies of the form:

where for all s ∈ S , π θ s , f θ s ∈ R |A| denote vectors [ π θ ( a | s )] a ∈A and [ f θ ( s, a )] a ∈A , respectively.

In this definition, the policy is induced by a mirror map h and a parameterized function f θ , and is obtained by mapping f θ to Y with the operator ∇ h ∗ ( · ) , which may not result in a well-defined probability distribution, and is thus projected on the probability simplex ∆( A ) . Note that the choice of h is key in deriving convenient expressions for π θ . The Bregman projected policy class contains large families of policy classes. Below is an example of h that recovers a widely used policy class [7, Example 9.10].

Example 3.2 (Negative entropy) . If Y = R |A| + and h is the negative entropy mirror map, i.e., h ( π ( ·| s )) = ∑ a ∈A π ( a | s ) log( π ( a | s )) , then the associated Bregman projected policy class becomes

where the exponential and the fraction are element-wise and ∥·∥ 1 is the ℓ 1 norm. In particular, when f θ ( s, a ) = θ s,a , the policy class (8) becomes the tabular softmax policy class; when f θ is a linear function, (8) becomes the log-linear policy class; and when f θ is a neural network, (8) becomes the neural policy class defined by Agarwal et al. [1]. We refer to Appendix C.1 for its derivation.

We now construct an MD-based algorithm to optimize V π θ over the Bregman projected policy class associated with a mirror map h and a parameterization class F Θ by adapting Section 2.1 to our setting. We define the following shorthand: at each time t , let π t := π θ t , f t := f θ t , V t := V π t , Q t := Q π t , and d t µ := d π t µ . Further, for any function y : S × A → R and distribution v over S × A , let y s := y ( s, · ) ∈ R |A| and ∥ y ∥ 2 L 2 ( v ) = E ( s,a ) ∼ v [( y ( s, a )) 2 ] . Ideally, we would like to execute the exact MD algorithm: for all t ≥ 0 and for all s ∈ S ,

Here, (10) reflects our Bregman projected policy class 3.1. However, we usually cannot perform the update (9) exactly. In general, if f θ belongs to a parameterized class F Θ , there may not be any θ t +1 ∈ Θ such that (9) is satisfied for all s ∈ S .

To remedy this issue, we propose Approximate Mirror Policy Optimization (AMPO), described in Algorithm 1. At each iteration, AMPO consists in minimizing a surrogate loss and projecting the result onto the simplex to obtain the updated policy. In particular, the surrogate loss in Line 1 of Algorithm 1 is a standard regression problem where we try to approximate Q t + η -1 t ∇ h ( π t ) with f t +1 , and has been studied extensively when f θ is a neural network [3, 35]. We can then readily use (10) to update π t +1 within the Bregman projected policy class defined in 3.1, which gives Line 2 of Algorithm 1.

To better illustrate the novelty of our framework, we provide below two remarks on how the two steps of AMPO relate and improve upon previous works.

Remark 3.3 . Line 1 associates AMPO with the compatible function approximation framework developed by [87, 42, 1], as both frameworks define the updated parameters θ t +1 as the solution to a regression problem aimed at approximating the current Q -function Q t . A crucial difference is that, Agarwal et al. [1] approximate Q t linearly with respect to ∇ θ log π t (see (61)), while in Line 1 we approximate Q t and the gradient of the mirror map of the previous policy with any function f θ . This generality allows our algorithm to achieve approximation guarantees for a wider range of assumptions on the structure of Q t . Furthermore, the regression problem proposed by Agarwal et al. [1] depends on the distribution d t µ , while ours has no such constraint and allows off-policy updates involving an arbitrary distribution v t . See Appendix E for more details.

Remark 3.4 . Line 2 associates AMPO to previous approximations of PMD [89, 88]. For instance, Vaswani et al. [89] aim to maximize an expression equivalent to

where Π(Θ) is a given parameterized policy class, while the Bregman projection step of AMPO can be rewritten as

We formulate this result as Lemma F.1 in Appendix F with a proof. When the policy class Π(Θ) is the entire policy space ∆( A ) S , (11) is equivalent to the two-step procedure (9)-(10) thanks to the PGT (3). A derivation of this observation is given in Appendix B for completeness. The issue with the update in (11), which is overcome by (12), is that Π(Θ) in (11) is often a non-convex set, thus the three-point-descent lemma [19] cannot be applied. The policy update in (12) circumvents this problem by defining the policy implicitly through the Bregman projection, which is a convex problem and thus allows the application of the three-point-descent lemma [19]. We refer to Appendix F for the conditions of the three-point-descent lemma in details.

3.1 ω -potential mirror maps

In this section, we provide a class of mirror maps that allows to compute the Bregman projection in Line 2 with ˜ O ( |A| ) operations and simplifies the minimization problem in Line 1.

Definition 3.5 ( ω -potential mirror map [50]) . For u ∈ ( -∞ , + ∞ ] , ω ≤ 0 , let an ω -potential be an increasing C 1 -diffeomorphism ϕ : ( -∞ , u ) → ( ω, + ∞ ) such that

For any ω -potential ϕ , the associated mirror map h ϕ , called ω -potential mirror map , is defined as

Thanks to Krichene et al. [50, Proposition 2] (see Proposition C.1 as well), the policy π t +1 in Line 2 induced by the ω -potential mirror map can be obtained with ˜ O ( |A| ) computations and can be written as

where λ t +1 s ∈ R is a normalization factor to ensure ∑ a ∈A π t +1 ( a | s ) = 1 for all s ∈ S , and σ ( z ) = max( z, 0) for z ∈ R . We call this policy class the ω -potential policy class. By using (13) and the definition of the ω -potential mirror map h ϕ , the minimization problem in Line 1 is simplified to be

where max( · , · ) is applied element-wisely. The ω -potential policy class allows AMPO to generate a wide range of applications by simply choosing an ω -potential ϕ . In fact, it recovers existing approaches to policy optimization, as we show in the next two examples.

Example 3.6 (Squared ℓ 2 -norm) . If Y = R |A| and ϕ is the identity function, then h ϕ is the squared ℓ 2 -norm, that is h ϕ ( π s ) = ∥ π s ∥ 2 2 / 2 , and ∇ h ϕ is the identity function. So, Line 1 in Algorithm 1 becomes

The ∇ h ∗ ϕ also becomes the identity function, and the policy update is given for all s ∈ S by

which is the Euclidean projection on the probability simplex. In the tabular setting, where S and A are finite and f θ ( s, a ) = θ s,a , (15) can be solved exactly with the minimum equal to zero, and Equations (15) and (16) recover the projected Q-descent algorithm [93]. As a by-product, we generalize the projected Q-descent algorithm from the tabular setting to a general parameterization class F Θ , which is a novel algorithm in the RL literature.

Example 3.7 (Negative entropy) . If Y = R |A| + and ϕ ( x ) = exp( x -1) , then h ϕ is the negative entropy mirror map from Example 3.2 and Line 1 in Algorithm 1 becomes

Consequently, based on Example 3.2, we have π t +1 s ∝ exp( η t f t +1 s ) for all s ∈ S . In this example, AMPO recovers tabular NPG [82] when f θ ( s, a ) = θ s,a , and Q-NPG with log-linear policies [98] when f θ and Q t are linear functions for all t ≥ 0 .

We refer to Appendix C.2 for detailed derivations of the examples in this section and an efficient implementation of the Bregman projection step. In addition to the ℓ 2 -norm and the negative entropy, several other mirror maps that have been studied in the optimization literature fall into the class of ω -potential mirror maps, such as the Tsallis entropy [74, 57] and the hyperbolic entropy [34], as well as a generalization of the negative entropy [50]. These examples illustrate how the class of ω -potential mirror maps recovers known methods and can be used to explore new algorithms in policy optimization. We leave the study of the application of these mirror maps in RL as future work, both from an empirical and theoretical point of view, and provide additional discussion and details in Appendix C.2.

4 Theoretical analysis

In our upcoming theoretical analysis of AMPO, we rely on the following key lemma.

Lemma 4.1. For any policies π and ¯ π , for any function f θ ∈ F Θ and for η > 0 , we have

where ˜ π is the Bregman projected policy induced by f θ and h according to Definition 3.1, that is ˜ π s = argmin p ∈ ∆( A ) D h ( p, ∇ h ∗ ( ηf θ s )) for all s ∈ S .

The proof of Lemma 4.1 is given in Appendix D.1. Lemma 4.1 describes a relation between any two policies and a policy belonging to the Bregman projected policy class associated with F Θ and h . As mentioned in Remark 3.4, Lemma 4.1 is the direct consequence of (12) and can be interpreted as an application of the three-point descent lemma [19], while it cannot be applied to algorithms based on the update in (11) [88, 89] due to the non-convexity of the optimization problem (see also Appendix F). Notice that Lemma 4.1 accommodates naturally with general parameterization also thanks to (12). In contrast, similar results have been obtained and exploited for specific policy and mirror map classes [93, 60, 36, 98], while our result allows any parameterization class F Θ and any choice of mirror map, thus greatly expanding the scope of applications of the lemma. A similar result for general parameterization has been obtained by Lan [55, Proposition 3.5] in the setting of strongly convex mirror maps.

Lemma 4.1 becomes useful when we set ¯ π = π t , f θ = f t +1 , η = η t and π = π t or π = π ⋆ . In particular, when η t f t +1 s -∇ h ( π t s ) ≈ η t Q π s , Lemma 4.1 allows us to obtain telescopic sums and recursive relations, and to handle error terms efficiently, as we show in Appendix D.

4.1 Convergence for general policy parameterization

In this section, we consider the parameterization class F Θ and the fixed but arbitrary mirror map h . We show that AMPO enjoys quasi-monotonic improvement and sublinear or linear convergence, depending on the step-size schedule. The first step is to control the approximation error of AMPO.

(A1) ( Approximation error ). There exists ε approx ≥ 0 such that, for all times t ≥ 0 ,

where ( v t ) t ≥ 0 is a sequence of distributions over states and actions and the expectation is taken over the randomness of the algorithm that obtains f t +1 .

Assumption (A1) is common in the conventional compatible function approximation approach 5 [1]. It characterizes the loss incurred by Algorithm 1 in solving the regression problem in Line 1. When the step-size η t is sufficiently large, Assumption (A1) measures how well f t +1 approximates the current Q-function Q t . Hence, ε approx depends on both the accuracy of the policy evaluation method used to obtain an estimate of Q t [86, 79, 28] and the error incurred by the function f θ ∈ F Θ that best approximates Q t , that is the representation power of F Θ . Later in Section 4.2, we show how to solve the minimization problem in Line 1 when F Θ is a class of shallow neural networks so that Assumption (A1) holds. We highlight that Assumption (A1) is weaker than the conventional assumptions [1, 98], since we do not constrain the minimization problem to be linear in the parameters (see (61)). We refer to Appendix A.2 for a discussion on its technical novelty and Appendix G for a relaxed version of the assumption.

As mentioned in Remark 3.3, the distribution v t does not depend on the current policy π t for all times t ≥ 0 . Thus, Assumption (A1) allows for off-policy settings and the use of replay buffers [68]. We refer to Appendix A.3 for details. To quantify how the choice of these distributions affects the error terms in the convergence rates, we introduce the following coefficient.

(A2) ( Concentrability coefficient ). There exists C v ≥ 0 such that, for all times t ,

The concentrability coefficient C v quantifies how much the distribution v t overlaps with the distributions ( d ⋆ µ , π ⋆ ) , ( d t +1 µ , π t +1 ) , ( d ⋆ µ , π t ) and ( d t +1 µ , π t ) . It highlights that the distribution v t should have full support over the environment, in order to avoid large values of C v . Assumption (A2) is weaker than the previous best-known concentrability coefficient [98, Assumption 9], in the sense that we have the full control over v t . We refer to Appendix H for a more detailed discussion. We can now present our first result on the performance of Algorithm 1.

Proposition 4.2 (Quasi-monotonic updates) . Let (A1) , (A2) be true. We have, for all t ≥ 0 ,

where the expectation is taken over the randomness of AMPO.

We refer to Appendix D.3 for the proof. Proposition 4.2 ensures that an update of Algorithm 1 cannot lead to a performance degradation, up to an error term. The next assumption concerns the coverage of the state space for the agent at each time t .

(A3) ( Distribution mismatch coefficient ). Let d ⋆ µ := d π ⋆ µ . There exists ν µ ≥ 0 such that

Since d t µ ( s ) ≥ (1 -γ ) µ ( s ) for all s ∈ S , obtained from the definition of d µ in (2), we have that

where assuming boundedness for the term on the right-hand side is standard in the literature on both the PG [e.g., 101, 90] and NPG convergence analysis [e.g., 1, 15, 93]. We refer to Appendix I for details. It is worth mentioning that the quasi-monotonic improvement in Proposition 4.2 holds without (A3) .

We define the weighted Bregman divergence between the optimal policy π ⋆ and the initial policy π 0 as D ⋆ 0 = E s ∼ d ⋆ µ [ D h ( π ⋆ s , π 0 s )] . We then have our main results below.

Theorem 4.3. Let (A1) , (A2) and (A3) be true. If the step-size schedule is non-decreasing, i.e., η t ≤ η t +1 for all t ≥ 0 , the iterates of Algorithm 1 satisfy: for every T ≥ 0 ,

Furthermore, if the step-size schedule is geometrically increasing, i.e., satisfies

we have: for every T ≥ 0 ,

Theorem 4.3 is, to the best of our knowledge, the first result that establishes linear convergence for a PG-based method involving general policy parameterization. For the same setting, it also matches the previous best known O (1 /T ) convergence [55], without requiring regularization. Lastly, Theorem 4.3 provides a convergence rate for a PMD-based algorithm that allows for arbitrary mirror maps and policy parameterization without requiring the assumption on the approximation error to hold in ℓ ∞ -norm, in contrast to Lan [55]. We give here a brief discussion of Theorem 4.3 w.r.t. previous results and refer to Tables 1 and 2 in Appendix A.2 for a detailed comparison.

In terms of iteration complexity, Theorem 4.3 recovers the best-known convergence rates in the tabular setting [93], for both non-decreasing and exponentially increasing step-size schedules. While considering a more general setting, Theorem 4.3 matches or improves upon the convergence rate of previous work on policy gradient methods for non-tabular policy parameterizations that consider constant step-size schedules [60, 82, 61, 90, 1, 89, 16, 55], and matches the convergence speed of previous work that employ NPG, log-linear policies, and geometrically increasing step-size schedules [2, 98].

In terms of generality, the results in Theorem 4.3 hold without the need to implement regularization [17, 100, 15, 16, 54], to impose bounded updates or smoothness of the policy [1, 61], to restrict the analysis to the case where the mirror map h is the negative entropy [60, 36], or to make ℓ ∞ -norm assumptions on the approximation error [55]. We improve upon the latest results for PMD with general policy parameterization by Vaswani et al. [89], which only allow bounded step-sizes, where the bound can be particularly small, e.g., (1 -γ ) 3 / (2 γ |A| ) , and can slow down the learning process.

When S is a finite state space, a sufficient condition for ν µ in (A3) to be bounded is requiring µ to have full support on S . If µ does not have full support, one can still obtain linear convergence for V ⋆ ( µ ′ ) -V T ( µ ′ ) , for an arbitrary state distribution µ ′ with full support, and relate this quantity to V ⋆ ( µ ) -V T ( µ ) . We refer to Appendix I for a detailed discussion on the distribution mismatch coefficient.

Intuition. An interpretation of our theory can be provided by connecting AMPO to the Policy Iteration algorithm (PI), which also enjoys linear convergence. To see this, first recall (12)

Secondly, solving Line 1 of Algorithm 1 leads to f t +1 s -η -1 t ∇ h ( π t s ) ≈ Q t s . When the step-size η t →∞ , that is η -1 t → 0 , the above viewpoint of the AMPO policy update becomes

which is the PI algorithm. Here we ignore the Bregman divergence term D h ( π, π t s ) , as it is multiplied by 1 /η t , which goes to 0. So AMPO behaves more and more like PI with a large enough step-size and thus is able to converge linearly like PI.

Proof idea. We provide a sketch of the proof here; the full proof is given in Appendix D. In a nutshell, the convergence rates of AMPO are obtained by building on Lemma 4.1 and leveraging modern PMD proof techniques [93]. Following the conventional compatible function approximation approach [1], the idea is to write the global optimum convergence results in an additive form, that is sub-optimality gap ≤ optimization error + approximation error .

The separation between the two errors is allowed by Lemma 4.1, while the optimization error is bounded through the PMD proof techniques from Xiao [93] and the approximation error is characterized by Assumption (A1) . Overall, the proof consists of three main steps.

which characterizes the improvement of the updated policy.

Step 2. Assumption (A1) , Step 1, the performance difference lemma (Lemma D.4), and Lemma 4.1 with ¯ π = π t , f θ = f t +1 , η = η t , ˜ π = π t +1 , and π s = π ⋆ s permit us to obtain the following.

Proposition 4.4. Let ∆ t := V ⋆ ( µ ) -V t ( µ ) . For all t ≥ 0 , we have

Step 3. Proposition 4.4 leads to sublinear convergence using a telescoping sum argument, and to linear convergence by properly defining step-sizes and by rearranging terms into the following contraction,

4.2 Sample complexity for neural network parameterization

Neural networks are widely used in RL due to their empirical success in applications [67, 68, 84]. However, few theoretical guarantees exist for using this parameterization class in policy optimization [60, 90, 16]. Here, we show how we can use our framework and Theorem 4.3 to fill this gap by deriving a sample complexity result for AMPO when using neural network parameterization. We consider the case where the parameterization class F Θ from Definition 3.1 belongs to the family of shallow ReLU networks, which have been shown to be universal approximators [38, 4, 27, 39]. That is, for ( s, a ) ∈ ( S × A ) ⊆ R d , define f θ ( s, a ) = c ⊤ σ ( W ( s, a ) + b ) with θ = ( c, W, b ) , where σ ( y ) = max( y, 0) for all y ∈ R is the ReLU activation function and is applied element-wisely, c ∈ R m , W ∈ R m × d and b ∈ R m .

At each iteration t of AMPO, we set v t = d t µ and solve the regression problem in Line 1 of Algorithm 1 through stochastic gradient descent (SGD). In particular, we initialize entry-wise W 0 and b as i.i.d. random Gaussian variables from N (0 , 1 /m ) , and c as i.i.d. random Gaussian variables from N (0 , ϵ A ) with ϵ A ∈ (0 , 1] . Assuming access to a simulator for the distribution v t , we run SGD for K steps on the matrix W , that is, for k = 0 , . . . , K -1 ,

where f ( k ) ( s, a ) = c ⊤ σ (( W 0 + W k )( s, a ) + b ) , ( s, a ) ∼ v t and ̂ Q t ( s, a ) is an unbiased estimate of Q t ( s, a ) obtained through Algorithm 4. We can then present our result on the sample complexity of AMPO for neural network parameterization, which is based on our convergence Theorem 4.3 and an analysis of neural networks by Allen-Zhu et al. [3, Theorem 1].

Corollary 4.5. In the setting of Theorem 4.3, let the parameterization class F Θ consist of sufficiently wide shallow ReLU neural networks. Using an exponentially increasing step-size and solving the minimization problem in Line 1 with SGD as in (19) , the number of samples required by AMPO to find an ε -optimal policy with high probability is ˜ O ( C 2 v ν 5 µ /ε 4 (1 -γ ) 6 ) , where ε has to be larger than a fixed and non-vanishing error floor.

We provide a proof of Corollary 4.5 and an explicit expression for the error floor in Appendix J. Note that the sample complexity in Corollary 4.5 might be impacted by an additional poly ( ε -1 ) term. We refer to Appendix J for more details and an alternative result (Corollary J.4) which does not include an additional poly ( ε -1 ) term, enabling comparison with prior works.

5 Numerical experiments

We provide an empirical evaluation of AMPO in order to validate our theoretical findings. We note that the scope of this work is mainly theoretical and that we do not aim at establishing state-of-theart results in the setting of deep RL. Our implementation is based upon the PPO implementation

from PureJaxRL [63], which obtains the estimates of the Q -function through generalized advantage estimation (GAE) [79] and performs the policy update using ADAM optimizer [48] and mini-batches. To implement AMPO, we (i) replaced the PPO loss with the expression to minimize in Equation (14), (ii) replaced the softmax projection with the Bregman projection, (iii) saved the constants λ along the sampled trajectories in order to compute Equation (14). The code is available here.

0

100

200

300

400

500

Iteration

t

100

200

300

400

500

V t

AMPO performance on CartPole-v1

AMPO with

φ

(

x

) =

e x

AMPO with

φ

(

x

) = sinh(

x

)

AMPO with

φ

(

x

) =

x

AMPO with Tsallis entropy

(

q

=0

.

5)

AMPO with Tsallis entropy

(

q

=1

.

5)

PPO

0

100

200

300

400

500

Iteration

t

500

400

300

200

100

V t

AMPO performance on Acrobot-v1

AMPO with

φ

(

x

) =

e x

AMPO with

φ

(

x

) = sinh(

x

)

AMPO with

φ

(

x

) =

x

AMPO with Tsallis entropy

(

q

=0

.

5)

AMPO with Tsallis entropy

(

q

=1

.

5)

PPO

In Figure 1, we show the averaged performance over 100 runs of AMPO in two classic control environments, i.e. CartPole and Acrobot, in the setting of ω -potential mirror maps. In particular, we choose: ϕ ( x ) = e x , which corresponds to the negative entropy (Example 3.7); ϕ ( x ) = sinh( x ) , which corresponds to the hyperbolic entropy [34, see also (41)]; ϕ ( x ) = x , which corresponds to the Euclidean norm (Example 3.6); and the Tsallis entropy for two values of the entropic index q [74, 57, see also (40)]. We refer to Appendix C.2 for a detailed discussion on these mirror maps. We set the step-size to be constant and of value 1. For a comparison, we also plot the averaged performance over 100 runs of PPO.

The plots in Figure 1 confirm our results on the quasi-monotonicity of the updates of AMPO and on its convergence to the optimal policy. We observe that instances of AMPO with different mirror maps are very competitive as compared to PPO. We also note that, despite the convergence rates in Theorem 4.3 depend on the mirror map only in terms of a D ⋆ 0 term, different mirror maps may result in different convergence speeds and error floors in practice. In particular, our experiments suggest that the negative entropy mirror map may not be the best choice for AMPO, and that exploring different mirror maps is a promising direction of research.

6 Conclusion

We have introduced a novel framework for RL which, given a mirror map and any parameterization class, induces a policy class and an update rule. We have proven that this framework enjoys sublinear and linear convergence for non-decreasing and geometrically increasing step-size schedules, respectively. Future venues of investigation include studying the sample complexity of AMPO in on-policy and off-policy settings other than neural network parameterization, exploiting the properties of specific mirror maps to take advantage of the structure of the MDP and efficiently including representation learning in the algorithm. We refer to Appendix A.3 for a thorough discussion of future work. We believe that the main contribution of AMPO is to provide a general framework with theoretical guarantees that can help the analysis of specific algorithms and MDP structures. AMPO recovers and improves several convergence rate guarantees in the literature, but it is important to keep in consideration how previous works have exploited particular settings, while AMPO tackles the most general case. It will be interesting to see whether these previous works combined with our fast linear convergence result can derive new efficient sample complexity results.

Acknowledgments and Disclosure of Funding

We thank the anonymous reviewers for their helpful comments. Carlo Alfano was supported by the Engineering and Physical Sciences Research Council and thanks G-Research for partly funding attendance to NeurIPS 2023.

References

Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research , 2021. (Cited on pages 4, 5, 7, 8, 20, 22, 24, 28, 37, 38, 39, and 41.)

Carlo Alfano and Patrick Rebeschini. Linear convergence for natural policy gradient with log-linear policy parametrization. arXiv preprint arXiv:2209.15382 , 2022. (Cited on pages 2, 8, 20, 21, 22, 25, 28, 38, and 39.)

Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparameterized neural networks, going beyond two layers. Advances in Neural Information Processing Systems , 2019. (Cited on pages 4, 9, 40, and 41.)

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning , 2019. (Cited on pages 9 and 43.)

Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation , 1998. (Cited on page 21.)

Jonathan Baxter and Peter L. Bartlett. Infinite-horizon policy-gradient estimation. Journal of Artificial Intelligence Research , 2001. (Cited on page 1.)

Amir Beck. First-Order Methods in Optimization . SIAM-Society for Industrial and Applied Mathematics, 2017. (Cited on pages 4 and 23.)

Amir Beck and Marc Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters , 2003. (Cited on pages 21 and 26.)

Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Chris Hesse, et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint arXiv:1912.06680 , 2019. (Cited on page 1.)

Jalaj Bhandari and Daniel Russo. Global optimality guarantees for policy gradient methods. arXiv preprint arXiv:1906.01786 , 2019. (Cited on page 23.)

Jalaj Bhandari and Daniel Russo. On the linear convergence of policy gradient methods for finite MDPs. In International Conference on Artificial Intelligence and Statistics , 2021. (Cited on pages 22, 23, and 25.)

Shalabh Bhatnagar, Richard S. Sutton, Mohammad Ghavamzadeh, and Mark Lee. Natural actor-critic algorithms. Automatica , 2009. (Cited on pages 1 and 22.)

Lev M. Bregman. The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming. USSR Computational Mathematics and Mathematical Physics , 1967. (Cited on page 3.)

Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning , 2015. (Cited on pages 3, 25, 31, and 32.)

Semih Cayci, Niao He, and R Srikant. Linear convergence of entropy-regularized natural policy gradient with linear function approximation. arXiv preprint arXiv:2106.04096 , 2021. (Cited on pages 2, 7, 8, 22, 25, 38, and 39.)

Semih Cayci, Niao He, and R Srikant. Finite-time analysis of entropy-regularized neural natural actor-critic algorithm. arXiv preprint arXiv:2206.00833 , 2022. (Cited on pages 8, 9, 20, 22, 24, 39, 43, and 44.)

Shicong Cen, Chen Cheng, Yuxin Chen, Yuting Wei, and Yuejie Chi. Fast global convergence of natural policy gradient methods with entropy regularization. Operations Research , 2021. (Cited on pages 1, 8, 22, 25, and 39.)

Yair Censor and Stavros A. Zenios. Parallel Optimization: Theory, Algorithms, and Applications . Oxford University Press, USA, 1997. (Cited on page 3.)

Gong Chen and Marc Teboulle. Convergence analysis of a proximal-like minimization algorithm using bregman functions. SIAM Journal on Optimization , 1993. (Cited on pages 2, 5, 6, 21, 31, and 37.)

Zaiwei Chen and Siva Theja Maguluri. Sample complexity of policy-based methods under offpolicy sampling and linear function approximation. In International Conference on Artificial Intelligence and Statistics , 2022. (Cited on pages 2, 22, 25, and 39.)

Zaiwei Chen, Sajad Khodadadian, and Siva Theja Maguluri. Finite-sample analysis of offpolicy natural actor-critic with linear function approximation. IEEE Control Systems Letters , 2022. (Cited on page 38.)

Ashok Cutkosky and Francesco Orabona. Momentum-based variance reduction in non-convex sgd. In Advances in Neural Information Processing Systems , 2019. (Cited on pages 22 and 23.)

Dongsheng Ding, Kaiqing Zhang, Tamer Basar, and Mihailo Jovanovic. Natural policy gradient primal-dual method for constrained markov decision processes. In Advances in Neural Information Processing Systems , 2020. (Cited on page 23.)

Yuhao Ding, Junzi Zhang, and Javad Lavaei. On the global optimum convergence of momentum-based policy gradient. In International Conference on Artificial Intelligence and Statistics , 2022. (Cited on page 23.)

Simon Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudik, and John Langford. Provably efficient rl with rich observations via latent state decoding. In International Conference on Machine Learning , 2019. (Cited on page 23.)

Simon Du, Sham Kakade, Jason Lee, Shachar Lovett, Gaurav Mahajan, Wen Sun, and Ruosong Wang. Bilinear classes: A structural framework for provable generalization in RL. In International Conference on Machine Learning , 2021. (Cited on page 23.)

Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over-parameterized neural networks. In International Conference on Learning Representations , 2019. (Cited on page 9.)

Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, Iain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. In International conference on machine learning , 2018. (Cited on page 7.)

Ilyas Fatkhullin, Jalal Etesami, Niao He, and Negar Kiyavash. Sharp analysis of stochastic optimization under global Kurdyka-łojasiewicz inequality. In Advances in Neural Information Processing Systems , 2022. (Cited on page 23.)

Ilyas Fatkhullin, Anas Barakat, Anastasia Kireeva, and Niao He. Stochastic policy gradient methods: Improved sample complexity for fisher-non-degenerate policies. arXiv preprint arXiv:2302.01734 , 2023. (Cited on page 23.)

Maryam Fazel, Rong Ge, Sham Kakade, and Mehran Mesbahi. Global convergence of policy gradient methods for the linear quadratic regulator. In International Conference on Machine Learning , 2018. (Cited on page 22.)

Boyan Gao, Henry Gouk, Hae Beom Lee, and Timothy M Hospedales. Meta mirror descent: Optimiser learning for fast convergence. arXiv preprint arXiv:2203.02711 , 2022. (Cited on page 31.)

Matthieu Geist, Bruno Scherrer, and Olivier Pietquin. A theory of regularized Markov decision processes. In International Conference on Machine Learning , 2019. (Cited on pages 19, 20, 22, 23, and 24.)

Udaya Ghai, Elad Hazan, and Yoram Singer. Exponentiated gradient meets gradient descent. In International Conference on Algorithmic Learning Theory , 2020. (Cited on pages 6, 10, and 30.)

Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning . MIT press, 2016. (Cited on page 4.)

Yuzheng Hu, Ziwei Ji, and Matus Telgarsky. Actor-critic is implicitly biased towards high entropy optimal policies. In International Conference on Learning Representations , 2022. (Cited on pages 6, 8, 20, 22, and 24.)

Feihu Huang, Shangqian Gao, and Heng Huang. Bregman gradient policy optimization. In International Conference on Learning Representations , 2022. (Cited on pages 22 and 24.)

Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems , 2018. (Cited on page 9.)

Ziwei Ji, Matus Telgarsky, and Ruicheng Xian. Neural tangent kernels, transportation mappings, and universal approximation. In International Conference on Learning Representations , 2019. (Cited on pages 9, 38, 43, and 44.)

Chi Jin, Zhuoran Yang, Zhaoran Wang, and Michael I Jordan. Provably efficient reinforcement learning with linear function approximation. In Conference on Learning Theory , 2020. (Cited on pages 22 and 23.)

Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In International Conference on Machine Learning , 2002. (Cited on page 33.)

Sham M. Kakade. A natural policy gradient. Advances in Neural Information Processing Systems , 2002. (Cited on pages 1, 5, and 21.)

William Karush. Minima of functions of several variables with inequalities as side conditions. Master's thesis, Department of Mathematics, University of Chicago, Chicago, IL, USA, 1939. (Cited on page 26.)

Michael J. Kearns and Daphne Koller. Efficient reinforcement learning in factored mdps. In International Joint Conference on Artificial Intelligence , 1999. (Cited on page 23.)

Sajad Khodadadian, Zaiwei Chen, and Siva Theja Maguluri. Finite-sample analysis of offpolicy natural actor-critic algorithm. In International Conference on Machine Learning , 2021. (Cited on pages 22 and 24.)

Sajad Khodadadian, Prakirt Raj Jhunjhunwala, Sushil Mahavir Varma, and Siva Theja Maguluri. On the linear convergence of natural policy gradient algorithm. In IEEE Conference on Decision and Control , 2021. (Cited on pages 22 and 25.)

Sajad Khodadadian, Prakirt Raj Jhunjhunwala, Sushil Mahavir Varma, and Siva Theja Maguluri. On linear and super-linear convergence of natural policy gradient algorithm. Systems and Control Letters , 2022. (Cited on pages 22 and 25.)

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference on Learning Representations , 2015. (Cited on page 10.)

Vijay Konda and John Tsitsiklis. Actor-critic algorithms. In Advances in Neural Information Processing Systems , 2000. (Cited on page 1.)

Walid Krichene, Syrine Krichene, and Alexandre Bayen. Efficient bregman projections onto the simplex. In IEEE Conference on Decision and Control , 2015. (Cited on pages 5, 6, 27, 29, and 30.)

Jakub Grudzien Kuba, Christian A Schroeder De Witt, and Jakob Foerster. Mirror learning: A unifying framework of policy optimisation. In International Conference on Machine Learning , 2022. (Cited on pages 1, 20, and 23.)

Harold W. Kuhn and Albert W. Tucker. Nonlinear programming. In Proceedings of the Second Berkeley Symposium on Mathematical Statistics and Probability , 1951. (Cited on page 26.)

Solomon Kullback and Richard A. Leibler. On Information and Sufficiency. The Annals of Mathematical Statistics , 1951. (Cited on page 1.)

Guanghui Lan. Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and generalized problem classes. Mathematical programming , 2022. (Cited on pages 1, 8, 21, 22, 25, and 39.)

Guanghui Lan. Policy optimization over general state and action spaces. arXiv preprint arXiv:2211.16715 , 2022. (Cited on pages 6, 8, 21, and 24.)

Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020. (Cited on page 23.)

Yan Li and Guanghui Lan. Policy mirror descent inherently explores action space. arXiv preprint arXiv:2303.04386 , 2023. (Cited on pages 6, 10, and 30.)

Yan Li, Tuo Zhao, and Guanghui Lan. Homotopic policy mirror descent: Policy convergence, implicit regularization, and improved sample complexity. arXiv preprint arXiv:2201.09457 , 2022. (Cited on pages 22 and 25.)

Zhize Li, Hongyan Bao, Xiangliang Zhang, and Peter Richtarik. PAGE: A simple and optimal probabilistic gradient estimator for nonconvex optimization. In International Conference on Machine Learning , 2021. (Cited on page 23.)

Boyi Liu, Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural trust region/proximal policy optimization attains globally optimal policy. Advances in Neural Information Processing Systems , 2019. (Cited on pages 6, 8, 9, 22, and 24.)

Yanli Liu, Kaiqing Zhang, Tamer Basar, and Wotao Yin. An improved analysis of (variancereduced) policy gradient and natural policy gradient methods. Advances in Neural Information Processing Systems , 2020. (Cited on pages 8, 22, 23, 24, and 38.)

Stanislaw Lojasiewicz. Une propriété topologique des sous-ensembles analytiques réels. Les équations aux dérivées partielles , 1963. (Cited on page 22.)

Chris Lu, Jakub Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. Advances in Neural Information Processing Systems , 2022. (Cited on page 10.)

Saeed Masiha, Saber Salehkaleybar, Niao He, Negar Kiyavash, and Patrick Thiran. Stochastic second-order methods improve best-known sample complexity of SGD for gradient-dominated functions. In Advances in Neural Information Processing Systems , 2022. (Cited on page 23.)

Jincheng Mei, Chenjun Xiao, Csaba Szepesvari, and Dale Schuurmans. On the global convergence rates of softmax policy gradient methods. In International Conference on Machine Learning , 2020. (Cited on page 22.)

Jincheng Mei, Yue Gao, Bo Dai, Csaba Szepesvari, and Dale Schuurmans. Leveraging nonuniformity in first-order non-convex optimization. In International Conference on Machine Learning , 2021. (Cited on page 23.)

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602 , 2013. (Cited on page 9.)

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature , 2015. (Cited on pages 7 and 9.)

Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International Conference on Machine Learning , 2016. (Cited on page 1.)

Arkadi Nemirovski and David B. Yudin. Problem Complexity and Method Efficiency in Optimization . Wiley Interscience, 1983. (Cited on pages 1, 3, and 21.)

Yurii E. Nesterov and Boris T. Polyak. Cubic regularization of Newton method and its global performance. Mathematical Programming , 2006. (Cited on page 23.)

Gergely Neu, Anders Jonsson, and Vicenç Gómez. A unified view of entropy-regularized markov decision processes. arXiv preprint arXiv:1705.07798 , 2017. (Cited on page 22.)

Lam M. Nguyen, Jie Liu, Katya Scheinberg, and Martin Takáˇ c. SARAH: A novel method for machine learning problems using stochastic recursive gradient. In International Conference on Machine Learning , 2017. (Cited on pages 22 and 23.)

Francesco Orabona. A modern introduction to online learning, 2020. URL ❤/a116/a116♣/a115✿✴✴♦♣❡♥✳ ❜✉✳❡❞✉✴❤❛♥❞❧❡✴✷✶✹✹✴✹✵✾✵✵ . (Cited on pages 6, 10, 30, and 31.)

Jan Peters and Stefan Schaal. Natural actor-critic. Neurocomputing , 2008. (Cited on pages 1 and 22.)

Boris T. Polyak. Gradient methods for the minimisation of functionals. USSR Computational Mathematics and Mathematical Physics , 1963. (Cited on page 22.)

Bruno Scherrer. Approximate policy iteration schemes: A comparison. In International Conference on Machine Learning , 2014. (Cited on page 39.)

John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning , 2015. (Cited on pages 1, 19, and 22.)

John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. Highdimensional continuous control using generalized advantage estimation. In International Conference on Learning Representations , 2016. (Cited on pages 7 and 10.)

John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 , 2017. (Cited on pages 1, 19, and 22.)

Shai Shalev-Shwartz, Shaked Shammah, and Amnon Shashua. Safe, multi-agent, reinforcement learning for autonomous driving. arXiv preprint arXiv:1610.03295 , 2016. (Cited on page 1.)

Lior Shani, Yonathan Efroni, and Shie Mannor. Adaptive trust region policy optimization: Global convergence and faster rates for regularized MDPs. In AAAI Conference on Artificial Intelligence , 2020. (Cited on pages 1, 6, 8, 21, 22, and 24.)

David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International Conference on Machine Learning , 2014. (Cited on page 31.)

David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, Yutian Chen, Timothy Lillicrap, Fan Hui, Laurent Sifre, George van den Driessche, Thore Graepel, and Demis Hassabis. Mastering the game of Go without human knowledge. Nature , 2017. (Cited on page 9.)

Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Modelbased rl in contextual decision processes: Pac bounds and exponential improvements over model-free approaches. In Conference on Learning Theory , 2019. (Cited on page 23.)

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . MIT press, 2018. (Cited on pages 7 and 18.)

Richard S. Sutton, David A. McAllester, Satinder P. Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems , 1999. (Cited on pages 1, 3, and 5.)

Manan Tomar, Lior Shani, Yonathan Efroni, and Mohammad Ghavamzadeh. Mirror descent policy optimization. In International Conference on Learning Representations , 2022. (Cited on pages 1, 5, 6, and 19.)

Sharan Vaswani, Olivier Bachem, Simone Totaro, Robert Müller, Shivam Garg, Matthieu Geist, Marlos C Machado, Pablo Samuel Castro, and Nicolas Le Roux. A general class of surrogate functions for stable and efficient reinforcement learning. In International Conference on Artificial Intelligence and Statistics , 2022. (Cited on pages 1, 5, 6, 8, 19, and 23.)

Lingxiao Wang, Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural policy gradient methods: Global optimality and rates of convergence. In International Conference on Learning Representations , 2020. (Cited on pages 7, 8, 9, 22, and 24.)

Weiran Wang and Miguel A Carreira-Perpinán. Projection onto the probability simplex: An efficient algorithm with a simple proof, and an application. arXiv preprint arXiv:1309.1541 , 2013. (Cited on page 29.)

Ronald J. Williams and Jing Peng. Function optimization using connectionist reinforcement learning algorithms. Connection Science , 1991. (Cited on page 1.)

Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research , 2022. (Cited on pages 1, 2, 6, 7, 8, 9, 19, 20, 21, 22, 24, 25, 28, 29, 31, 33, 37, 38, 39, and 40.)

Tengyu Xu, Zhe Wang, and Yingbin Liang. Improving sample complexity bounds for (natural) actor-critic algorithms. In Advances in Neural Information Processing Systems , 2020. (Cited on pages 22 and 24.)

Long Yang, Yu Zhang, Gang Zheng, Qian Zheng, Pengfei Li, Jianhang Huang, and Gang Pan. Policy optimization with stochastic mirror descent. AAAI Conference on Artificial Intelligence , 2022. (Cited on pages 22 and 24.)

Batuhan Yardim, Semih Cayci, Matthieu Geist, and Niao He. Policy mirror ascent for efficient and independent learning in mean field games. In International Conference on Machine Learning , 2023. (Cited on page 23.)

Rui Yuan, Robert M. Gower, and Alessandro Lazaric. A general sample complexity analysis of vanilla policy gradient. In International Conference on Artificial Intelligence and Statistics , 2022. (Cited on pages 22 and 23.)

Rui Yuan, Simon Shaolei Du, Robert M. Gower, Alessandro Lazaric, and Lin Xiao. Linear convergence of natural policy gradient methods with log-linear policies. In International Conference on Learning Representations , 2023. (Cited on pages 2, 6, 7, 8, 20, 21, 22, 24, 25, 28, 31, 38, 39, and 40.)

Andrea Zanette, Ching-An Cheng, and Alekh Agarwal. Cautiously optimistic policy optimization and exploration with linear function approximation. In Conference on Learning Theory , 2021. (Cited on pages 20, 22, and 24.)

Wenhao Zhan, Shicong Cen, Baihe Huang, Yuxin Chen, Jason D. Lee, and Yuejie Chi. Policy mirror descent for regularized reinforcement learning: A generalized framework with linear convergence. SIAM Journal on Optimization , 33(2):1061-1091, 2023. (Cited on pages 1, 8, 22, 25, and 39.)

Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvari, and Mengdi Wang. Variational policy gradient method for reinforcement learning with general utilities. In Advances in Neural Information Processing Systems , 2020. (Cited on pages 7 and 23.)

Junyu Zhang, Chengzhuo Ni, Zheng Yu, Csaba Szepesvari, and Mengdi Wang. On the convergence and sample efficiency of variance-reduced policy gradient method. In Advances in Neural Information Processing Systems , 2021. (Cited on page 23.)

Appendix

Table of Contents

Here we provide the related work discussion, the deferred proofs from the main paper and some additional noteworthy observations.

A Related work

We provide an extended discussion for the context of our work, including a comparison of different PMDframeworks and a comparison of the convergence theories of PMD in the literature. Furthermore, we discuss future work, such as extending our analysis to the dual averaging updates and developing sample complexity analysis of AMPO.

A.1 Comparisons with other policy mirror descent frameworks

In this section, we give a comparison of AMPO with some of the most popular policy optimization algorithms in the literature. First, recall AMPO's update through (12), that is, for all s ∈ S ,

where η t f t +1 s -∇ h ( π t s ) ≈ η t Q t s following Line 1 of Algorithm 1. The proof of (20) can be found in Lemma F.1 in Appendix F.

Generalized Policy Iteration (GPI) [86]. The update consists of evaluating the Q-function of the policy and obtaining the new policy by acting greedily with respect to the estimated Q-function. That is, for all s ∈ S ,

AMPO behaves like GPI when we perfectly approximate f t +1 s to the value of Q t s (e.g. when we consider the tabular case) and η t → + ∞ (or η -1 t → 0 ) which is the case with the use of geometrically increasing step-size schedule.

Mirror Descent Modified Policy Iteration (MD-MPI) [33]. Consider the full policy space Π = ∆( A ) S . The MD-MPI's update is as follows:

In this case, the PMD framework of Xiao [93], which is a special case of AMPO, recovers MD-MPI with the fixed step-size η t = 1 . Consequently, Assumption (A1) holds with ε approx = 0 , and we obtain the sublinear convergence of MD-MPI through Theorem 4.3, which is

As explained later in Appendix H that the distribution mismatch coefficient ν µ is upper bounded by O ( 1 1 -γ ) , we obtain an average regret of MD-MPI as O ( 1 (1 -γ ) 2 T ) , which matches the convergence results in Geist et al. [33, Corollary 3].

Trust Region Policy Optimization (TRPO) [78]. The TRPO's update is as follows:

where A t s = Q t s -V t represents the advantage function, h is the negative entropy and δ > 0 . Like GPI, TRPO is equivalent to AMPO when at each time t , the admissible policy class is Π t = { π ∈ ∆( A ) S : E s ∼ d t µ D h ( π t s , π s ) ≤ δ } , and we perfectly approximate Q t s with η t → + ∞ .

Proximal Policy Optimization (PPO) [80]. The PPO's update consists of maximizing a surrogate function depending on the policy gradient with respect to the new policy. Namely,

with

where r π ( s, a ) = π ( s, a ) /π t ( s, a ) is the probability ratio between the current policy π t and the new one, and the function clip ( r π ( s, a ) , 1 ± ϵ ) clips the probability ratio r π ( s, a ) to be no more than 1+ ϵ and no less than 1 -ϵ . PPO has also a KL variation [80, Section 4], where the objective function L is defined as

where h is the negative entropy. In an exact setting and when Π = ∆( A ) S , the KL variation of PPO still differs from AMPO because it inverts the terms in the Bregman divergence penalty.

Mirror Descent Policy Optimization (MDPO.) [88]. The MDPO's update is as follows:

where Π is a parameterized policy class. While it is equivalent to AMPO in an exact setting and when Π = ∆( A ) S , as we show in Appendix B, the difference between the two algorithms lies in the approximation of the exact algorithm.

Functional Mirror Ascent Policy Gradient (FMA-PG) [89]. The FMA-PG's update is as follows:

The second line is obtained by the definition of V t and the policy gradient theorem (3). The discussion is the same as the previous algorithm.

Mirror Learning [51]. The on-policy version of the algorithm consists of the following update:

where Π( π t ) is a policy class that depends on the current policy π t and the drift functional D is defined as a map D : ∆( A ) × ∆( A ) → R such that D ( π s , ¯ π s ) ≥ 0 and ∇ π s D ( π s , ¯ π s ) ∣ ∣ π s =¯ π s = 0 . The drift functional D recovers the Bregman divergence as a particular case, in which case Mirror Learning is equivalent to AMPO in an exact setting and when Π = ∆( A ) S . Again, the main difference between the two algorithms lies in the approximation of the exact algorithm.

A.2 Discussion on related work

Our Contributions. Our work provides a framework for policy optimization - AMPO. For AMPO, we establish in Theorem 4.3 both O (1 /T ) convergence guarantee by using a non-decreasing step-size and linear convergence guarantee by using a geometrically increasing step-size. Our contributions to the prior literature on sublinear and linear convergence of policy optimization methods can be summarized as follows.

The generality of our framework allows Theorem 4.3 to unify previous results in the literature and generate new theoretically sound algorithms under one guise. Both the sublinear and the linear convergence analysis of natural policy gradient (NPG) with softmax tabular policies [93] or with log-linear policies [2, 98] are special cases of our general analysis. As mentioned in Appendix A.1, MD-MPI [33] in the tabular setting is also a special case of AMPO. Thus, Theorem 4.3 recovers the best-known convergence rates in both the tabular setting [33, 93] and the non-tabular setting [16, 2, 98]. AMPO also generates new algorithms by selecting mirror maps, such as the ϵ -negative entropy mirror map in Appendix C.2 associated with Algorithm 2, and generalizes the projected Q-descent algorithm [93] from the tabular setting to a general parameterization class F Θ .

As discussed in Section 4.1, the results of Theorem 4.3 hold for a general setting with fewer restrictions than in previous work. The generality of the assumptions of Theorem 4.3 allows the application of our theory to specific settings, where existing sample complexity analyses could be improved thanks to the linear convergence of AMPO. For instance, since Theorem 4.3 holds for any structural MDP, AMPO could be applied directly to the linear MDP setting to derive a sample complexity analysis of AMPO which could improve that of Zanette et al. [99] and Hu et al. [36]. As we discuss in Appendix A.3, this is a promising direction for future work.

From a technical point of view, our main contributions are: Definition 3.1 introduces a novel way of incorporating general parameterization into the policy; the update in Line 1 of Algorithm 1 simplifies the policy optimization step into a regression problem; and Lemma 4.1 establishes a key result for policies belonging to the class in Definition 3.1. Together, these innovations have allowed us to establish new state-of-the-art results in Theorem 4.3 by leveraging the modern PMD proof techniques of Xiao [93].

In particular, our technical novelty with respect to Xiao [93], Alfano and Rebeschini [2], and Yuan et al. [98] can be summarized as follows.

In terms of algorithm design, AMPO is an innovation. The PMD algorithm proposed by Xiao [93] is strictly limited to the tabular setting and, although it is well defined for any mirror map, it cannot include general parameterization. Alfano and Rebeschini [2] and Yuan et al. [98] propose a first generalization of the PMD algorithm in the function approximation regime thanks to the linear compatible function approximation framework [1], but are limited to considering the log-linear policy parameterization and the entropy mirror map. On the contrary, AMPOsolves the problem of incorporating general parameterizations in the policy thanks to Definition 3.1 and the extension of the compatible function approximation framework from linear to nonlinear, which corresponds to the parameter update in Line 1 of Algorithm 1. This innovation is key to the generality of the algorithm, as it allows AMPO to employ any mirror map and any parameterization class. Moreover, AMPO is computationally efficient for a large class of mirror maps (see Appendix C.2 and Algorithms 2 and 3). Our design is readily applied to deep RL, where the policy is usually parameterized by a neural network whose

last layer is a softmax transformation. Our policy definition can be implemented in this setting by replacing the softmax layer with a Bregman projection, as shown in Example 3.7.

Regarding the assumptions necessary for convergence guarantees, we have weaker assumptions. Xiao [93] requires an ℓ ∞ -norm on the approximation error of Q t , i.e., ∥ ̂ Q t -Q t ∥ ∞ ≤ ε approx , for all t ≤ T . Alfano and Rebeschini [2] and Yuan et al. [98] require an L 2 -norm bound on the error of the linear approximation of Q t , i.e., ∥ w ⊤ ϕ -Q t ∥ 2 L 2 ( v t ) ≤ ε approx for some feature mapping ϕ : S×A → R d and vector w ∈ R d , for all t ≤ T . Our approximation error ε approx in Assumption (A1) is an improvement since it does not require the bound to hold in ℓ ∞ -norm, and allows any regression model instead of linear function approximation, especially neural networks, which greatly increases the representation power of F Θ and expands the range of applications. We further relax Assumption (A1) in Appendix G and show that the approximation error bound can be larger for earlier iterations. In addition, we improve the concentrability coefficients of Yuan et al. [98] by defining the expectation under an arbitrary state-action distribution v t instead of the state-action visitation distribution with a fixed initial state-action pair (see Yuan et al. [98, Equation (4)]).

As for the analysis of the algorithm, while we borrow tools from Xiao [93], Alfano and Rebeschini [2], and Yuan et al. [98], our results are not simple extensions. In fact, without our work, it is not clear from Xiao [93], Alfano and Rebeschini [2], and Yuan et al. [98] whether PMD could have theoretical guarantees in a setting with general parameterization and an arbitrary mirror map. The two main problems on this front are the non-convexity of the policy class, which prevents the use of the three-point descent lemma by Chen and Teboulle [19, Lemma 3.2] (or by Xiao [93, Lemma 6]), and the fact that the three-point identity used by Alfano and Rebeschini [2, Equation 4] holds only for the negative entropy mirror map. Our Lemma 4.1 successfully addresses general policy parameterization and arbitrary mirror maps thanks to the design of AMPO. Additionally, we provide a sample complexity analysis of AMPO when employing shallow neural networks that improves upon previous state-of-the-art results in this setting. We further improve this sample complexity analysis in Appendix J, where we consider an approximation error assumption that is weaker than Assumption (A1) (see Appendix G).

We also include a comparison wih Lan [55]. Our diffences can be outlined in two points.

Lan [55] propose a PMD algorithm (Algorithm 2 in their paper) that can accommodate general parameterization and arbitrary mirror maps. As AMPO, it involves a two-step procedure where the first step is to find an approximation of Q t -η -1 t ∇ h ( π t ) and the second step is to find the policy through a Bregman projection. However, it is unclear how to implement their algorithm in practice, as they do not propose a specific method to perform either step. We provide an explicit implementation of AMPO and identify a class of mirror maps that is computationally efficient for AMPO (see Appendix C.2 and Algorithms 2 and 3).

In terms of theoretical analysis, they assume for their results that the approximation error is bounded in ℓ ∞ -norm over the action space. Let

where the expectation E [ f t +1 s ] is taken w.r.t. the stochasticity of the algorithm employed to obtain f t +1 . Lan [55] assume that both ε det and ε sto are bounded for all iterations t . In contrast, our assumptions are weaker as they are required to hold for the L 2 ( v ) -norm we define in Section 3. Additionally, Lan [55] establishes a O (1 / √ T ) convergence rate for their algorithm without regularization and a O (1 /T ) convergence rate in the regularized case, in both cases using bounded step-sizes. We improve upon these results by obtaining a O (1 /T ) convergence rate without regularization and a linear convergence rate.

Related literature. Recently, the impressive empirical success of policy gradient (PG)-based methods has catalyzed the development of theoretically sound algorithms for policy optimization. In particular, there has been a lot of attention around algorithms inspired by mirror descent (MD) [70, 8] and, more specifically, by natural gradient descent [5]. These two approaches led to policy mirror descent (PMD) methods [82, 54] and natural policy gradient (NPG) methods [42], which, as first shown by

Neu et al. [72], is a special case of PMD. For instance, PMD and NPG are the building blocks of the state-of-the-art policy optimization algorithms, TRPO [78] and PPO [80]. Leveraging various techniques from the MD literature, it has been established that PMD, NPG, and their variants converge to the global optimum in different settings. We refer to global optimum convergence as an analysis that guarantees that V ⋆ ( µ ) -E [ V T ( µ ) ] ≤ ϵ after T iterations with ϵ > 0 . As an important variant of NPG, we will also discuss the literature of the convergence analysis of natural actor-critic (NAC) [75, 12]. The comparison of AMPO with different methods will proceed from the tabular case to different function approximation regimes.

Sublinear convergence analyses of PMD, NPG and NAC. For softmax tabular policies, Shani et al. [82] establish a O (1 / √ T ) convergence rate for unregularized NPG and O (1 /T ) for regularized NPG. Agarwal et al. [1], Khodadadian et al. [45] and Xiao [93] improve the convergence rate for unregularized NPG and NAC to O (1 /T ) and Xiao [93] extends the same convergence rate to projected Q-descent. The same convergence rate is established by MD-MPI [33] through the PMD framework.

In the function approximation regime, Zanette et al. [99] and Hu et al. [36] achieve O (1 / √ T ) convergence rate by developing variants of PMD methods for the linear MDP [40] setting. The same O (1 / √ T ) convergence rate is obtained by Agarwal et al. [1] for both log-linear and smooth policies, while Yuan et al. [98] improve the convergence rate to O (1 /T ) for log-linear policies. For smooth policies, the convergence rate is later improved to O (1 /T ) either by adding an extra Fisher-non-degeneracy condition on the policies [61] or by analyzing NAC under Markovian sampling [94]. Yang et al. [95] and Huang et al. [37] consider Lipschitz and smooth policies [97], obtain O (1 / √ T ) convergence rates for PMD-type methods and faster O (1 /T ) convergence rates by applying the variance reduction techniques SARAH [73] and STORM [22], respectively. As for neural policy parameterization, Liu et al. [60] establish a O (1 / √ T ) convergence rate for two-layer neural PPO. The same O (1 / √ T ) convergence rate is established by Wang et al. [90] for two-layer neural NAC, which is later improved to O (1 /T ) by Cayci et al. [16], using entropy regularization.

We highlight that all of the above sublinear convergence analyses, for both softmax tabular policies and the function approximation regime, are obtained either by using a decaying step-size or a constant step-size. Under these step-size schemes, our AMPO's O (1 /T ) sublinear convergence rate is the state of the art: it recovers the best-known convergence rates in the tabular setting [33, 93] without regularization; it improves the O (1 / √ T ) convergence rate of Zanette et al. [99] and Hu et al. [36] to O (1 /T ) for the linear MDP setting; it recovers the best-known convergence rates for the log-linear policies [98]; it matches the O (1 /T ) sublinear convergence rate for smooth and Fisher-non-degenerate policies [61] and the same convergence rate of Yang et al. [95] and Huang et al. [37] for Lipschitz and smooth policies without introducing variance reduction techniques; it matches the previous best-known convergence result in the neural network settings [16] without regularization; lastly, it goes beyond all these results by allowing general parameterization. We refer to Table 1 for an overview of recent sublinear convergence analyses of NPG/PMD.

Linear convergence analysis of PMD, NPG, NAC and other PG methods. In the softmax tabular policy settings, the linear convergence guarantees of NPG and PMD are achieved by either adding regularization [17, 100, 54, 58] or by varying the step-sizes [11, 46, 47, 93].

In the function approximation regime, the linear convergence guarantees are achieved for NPG with log-linear policies, either by adding entropy regularization [15] or by choosing geometrically increasing step-sizes [2, 98]. It can also be achieved for NAC with log-linear policy by using adaptive increasing step-sizes [20].

Again, our AMPO's linear convergence rate is the state of the art: not only it recovers the bestknown convergence rates in both the tabular setting [93] and the log-linear policies [2, 98] without regularization [17, 100, 54, 58], nor adaptive step-sizes [11, 46, 47, 20], but also it achieves the new state-of-the-art linear convergence rate for PG-based methods with general parameterization, including the neural network parameterizations. We refer to Table 2 for an overview of recent linear convergence analyses of NPG/PMD.

Alternatively, by exploiting a Polyak-Lojasiewicz (PL) condition [76, 62], fast linear convergence results can be achieved for PG methods under different settings, such as linear quadratic control problems [31] and softmax tabular policies with entropy regularization [65, 97]. The PL condition

is extensively studied by Bhandari and Russo [10] to identify more general MDP settings. Like the cases of NPG and PMD, linear convergence of PG can also be obtained for the softmax tabular policy without regularization by choosing adaptive step sizes through exact line search [11] or by exploiting non-uniform smoothness [66]. When the PL condition is relaxed to other weaker conditions, PG methods combined with variance reduction methods such as SARAH [73] and PAGE [59] can also achieve linear convergence. This is shown by Fatkhullin et al. [29, 30] when the PL condition is replaced by the weak PL condition [97], which is satisfied by Fisher-non-degenerate policies [24]. It is also shown by Zhang et al. [102], where the MDP satisfies some hidden convexity property that contains a similar property to the weak PL condition studied by Zhang et al. [101]. Lastly, linear convergence is established for the cubic-regularized Newton method [71], a second-order method, applied to Fisher-non-degenerate policies combined with variance reduction [64].

Outside of the literature focusing on finite time convergence guarantees, Vaswani et al. [89] and Kuba et al. [51] provide a theoretical analysis for variations of PMD and show monotonic improvements for their frameworks. Additionally, Kuba et al. [51] give an infinite time convergence guarantee for their framework.

A.3 Future work

Our work opens several interesting research directions in both algorithmic and theoretical aspects.

From an algorithmic point of view, the updates in Lines 1 and 2 of AMPO are not explicit. This might be an issue in practice, especially for large scale RL problems. It would be interesting to design efficient regression solver for minimizing the approximation error in Line 1 of Algorithm 1. For instance, by using the dual averaging algorithm [7, Chapter 4], it could be possible to replace the term ∇ h ( π t s ) with f t s for all s ∈ S , to make the computation of the algorithm more efficient. That is, it could be interesting to consider the following variation of Line 1 in Algortihm 1:

Notice that (28) has the same update as (17), however, (28) is not restricted to using the negative entropy mirror map. To efficiently solve the regression problem in Line 1 of Algorithm 1, one may want to apply modern variance reduction techniques [73, 22, 59]. This has been done by Liu et al. [61] for NPG method.

From a theoretical point of view, it would be interesting to derive a sample complexity analysis for AMPO in specific settings, by leveraging its linear convergence. As mentioned for the linear MDP [40] in Appendix A.2, one can apply the linear convergence theory of AMPO to other structural MDPs, e.g., block MDP [25], factored MDP [44, 85], RKHS linear MDP and RKHS linear mixture MDP [26], to build new sample complexity results for these settings, since the assumptions of Theorem 4.3 do not impose any constraint on the MDP. On the other hand, it would be interesting to explore the interaction between the Bregman projected policy class and the expected Lipschitz and smooth policies [97] and the Fish-non-degenerate policies [61] to establish new improved sample complexity results in these settings, again thanks to the linear convergence theory of AMPO.

Additionally, it would be interesting to study the application of AMPO to the offline setting. In the main text, we have discussed how to extend Algorithm 1 and Theorem 4.3 to the offline setting, where v t can be set as the state-action distribution induced by an arbitrary behavior policy that generates the data. However, we believe that this direction requires further investigation. One of the major challenges of offline RL is dealing with the distribution shifts that stem from the mismatch between the trained policy π t and the behaviour policy. Several methods have been introduced to deal with this issue, such as constraining the current policy to be close to the behavior policy [56]. We leave introducing offline RL techniques in AMPO as future work.

Another direction for future work is extending the policy update of AMPO to mirror descent algorithm based on value iteration and Bellman operators, such as MD-MPI [33], in order to extend existing results to the general parameterization setting. Other interesting settings that have been addressed using the PMD framework are mean-field games [96] and constrained MDPs [23]. We hope to build on the existing literature for these settings and see whether our results can bring any improvements.

Finally, this work theoretically indicates that, perhaps the most important future work of PMD-type algorithms is to design efficient policy evaluation algorithms to make the estimation of the Q -function

as accurate as possible, such as using offline data for training, and to construct adaptive representation learning for F Θ to closely approximate Q -function, so that ϵ approx is guaranteed to be small. This matches one of the most important research questions for deep Q-learning type algorithms for general policy optimization problems.

B Equivalence of (9) -(10) and (11) in the tabular case

To demonstrate the equivalence between the two-step update (9)-(10) and the one-step update (11) for policy mirror descent in the tabular case, it is sufficient to validate the following lemma, which comes from the optimization literature. The proof of this lemma can be found in Bubeck [14, Chapter 4.2]. However, for the sake of completeness, we present the proof here.

Lemma B.1 (Right after Theorem 4.2 in Bubeck [14]) . Consider the mirror descent update in (5) -(6) for the minimization of a function V ( · ) , that is,

Then the mirror descent update can be rewritten as

Proof. From definition of the Bregman projection step, starting from (29) we have

where the second and the last lines are both obtained by the definition of the Bregman divergence.

The one-step update in (31) is often taken as the definition of mirror descent [8], which provides a proximal view point of mirror descent, i.e., a gradient step in the primal space with a regularization of Bregman divergence.

C AMPOfor specific mirror maps

In this section, we give the derivations for Example 3.2, which is based on the Karush-Kuhn-Tucker (KKT) conditions [43, 52], and then provide details about the ω -potential mirror map class from Section 3.1.

C.1 Derivation of Example 3.2

We give here the derivation of Example 3.2. Let h be the negative entropy mirror map, that is

For every state s ∈ S , we solve the minimization problem

through the KKT conditions. We formalize it as

where 1 denotes a vector in R |A| with coordinates equal to 1 element-wisely. The conditions then become

where log( π s ) is applied element-wisely, λ s ∈ R and c a s ∈ R are the dual variables, and c s devotes the vector [ c a s ] a ∈A . It is easy to verify that the solution

with λ s = log (∑ a ∈A exp( ηf θ ( s, a )) ) -1 and c a s = 0 for all a ∈ A , satisfies all the conditions.

When f θ ( s, a ) = θ s,a we obtain the tabular softmax policy π θ ( a | s ) ∝ exp( ηθ s,a ) . When f θ ( s, a ) = θ ⊤ ϕ ( s, a ) is a linear function, for θ ∈ R d and for a feature function ϕ : S × A → R d , we obtain the log-linear policy π θ ( a | s ) ∝ exp( ηθ ⊤ ϕ ( s, a )) . When f θ : S × A → R is a neural network, we obtain the softmax neural policy π θ ( a | s ) ∝ exp( ηf θ ( s, a )) .

C.2 More on ω -potential mirror maps

In this section, we provide details about the ω -potential mirror map class from Section 3.1, including the derivation of (14), several instantiations of ω -potential mirror map mentioned in Section 3.1 with their derivations, and an iterative algorithm to find approximately the Bregman projection induced by ω -potential mirror map when an exact solution is not available.

We give a different but equivalent formulation of Proposition 2 of Krichene et al. [50].

Proposition C.1. For u ∈ ( -∞ , + ∞ ] and ω ≤ 0 , an increasing C 1 -diffeomorphism ϕ : ( -∞ , u ) → ( ω, + ∞ ) is called an ω -potential if

Let the mirror map h ϕ be defined as

We have that π t s is a solution to the Bregman projection

if and only if there exist a normalization constant λ t s ∈ R such that

and ∑ a ∈A π t ( a | s ) = 1 , where for all s ∈ S and σ ( z ) = max( z, 0) for z ∈ R .

We can now use Proposition C.1 to derive (14).

Consider an ω -potential mirror map h ϕ associated with an ω -potential ϕ . By definition, we have

Plugging (32) into (33), we have

where the third line is obtained by using the increasing property of ϕ -1 , as ϕ is increasing. Finally, plugging the above expression of ∇ h ϕ ( π t s ) into Line 1, we obtain (14), which is

where the term λ t s is dropped, as it is constant over actions and does not affect the resulting policy.

Once (14) is obtained, we can instantiate AMPO for mirror maps belonging to this class. We highlight that due to the definition of the Bregman divergence, two mirror maps that only differ for a constant term are equivalent and generate the same algorithm. We start with the negative entropy, which leads to a closed solution for λ t s and therefore for the Bregman projection.

Negative entropy. Let ϕ ( x ) = exp( x -1) , which is an ω -potential with ω = 0 , u = + ∞ , and

The mirror map h ϕ becomes the negative entropy, as

and the associated Bregman divergence becomes the KL divergence, i.e., D h ϕ ( π s , ¯ π s ) = KL ( π s , ¯ π s ) . Equation (17) follows from Equation (14) by the fact that

which means that max( η t -1 f t , ϕ -1 (0) -λ t s ) = η t -1 f t element-wisely.

As we showed in Appendix C.1, the Bregman projection (Line 2 of Algorithm 1) for the negative entropy has a closed form which is π t +1 s ∝ exp( η t f t +1 s ) .

In NPG with softmax tabular policies [1, 93], at time t , the updates for the policies have

When considering AMPO with f θ ( s, a ) = θ s,a ∈ R , from (17), we obtain that for all s, a ∈ S , A , we have

With the above expression, we have the AMPO's updates for the policies rewritten as

which recovers (34). In particular, the summation in the second line is element-wise and the third line is obtained because of π t s ∝ exp( η t -1 f t s ) , as shown in Appendix C.1.

In Q-NPG with log-linear policies [1, 98, 2], at time t , the updates for the policies have

where ϕ : S × A → R d is a feature map, and

Like in the tabular case, when considering AMPO with θ ∈ R d , f θ ( s, a ) = ϕ ( s, a ) ⊤ θ and v t = d t µ , from (17), we obtain that for all s, a ∈ S , A , we have

Compared the above form with (36), we obtain that

Input: vector to project x ∈ R |A| , parameter ϵ .

1 Initialize y = exp( x ) element-wisely.

2 Let y ( i ) be the i -th smallest element of y .

3 Let i ⋆ be the smallest index for which

Set

Return: the projected vector ( σ ( - ϵ + y a /λ )) a ∈A .

So, the AMPO's updates for the policies can be rewritten as

where the last line is obtained because of π t s ∝ exp( η t -1 f t s ) , as shown in Appendix C.1, and we recover (35).

We next present the squared ℓ 2 -norm and the ϵ -negative entropy. For these two mirror maps, the Bregman projection can be computed exactly but has no closed form.

Squared

ℓ

2

-norm.

Let up to a constant term, as

The associated Bregman divergence becomes the squared Euclidean distance, i.e., D h ϕ ( π s , ¯ π s ) = 1 2 ∥ π s -¯ π s ∥ 2 2 , and ∇ h ∗ ( · ) is the identity function. The update in (15) follows immediately and the Bregman projection step with the Euclidean distance becomes, for all s ∈ S ,

In the projected-Q descent for tabular policies developed by Xiao [93], at time t , the updates for the policies are

When considering AMPO with f θ ( s, a ) = θ s,a and Θ = R |S|×|A| , (15) is solved with

Plugging the above expression into (38), we recover (39).

Notice that the Euclidean projection onto the probability simplex can be obtained exactly, as shown by Wang and Carreira-Perpinán [91].

ϵ -negative entropy [50]. Let ϵ ≥ 0 and define the ϵ -exponential potential as ϕ ( x ) = exp( x -1) -ϵ . The mirror map h ϕ becomes

ϕ

be the identity function. The mirror map

h

ϕ

becomes the squared

ℓ

2

-norm,

Input: vector to project x ∈ R |A| , ω -potential ϕ , precision ε .

1 Initialize

where, for y > 0 ,

̸

2 Define ˜ x ( ν ) = ( σ ( ϕ ( x a + ν ))) a ∈A while ∥ ˜ x (¯ ν ) -˜ x ( ν ) ∥ > ε do

8 Return ˜ x (¯ ν

An exact solution to the associated projection can then be found in ˜ O ( |A| ) computations using Algorithm 2, which has been proposed by Krichene et al. [50, Algorithm 4]. Additionally, following (14), the regression problem in Line 1 of Algorithm 1 becomes

where λ t s can be obtained through Algorithm 2.

The Bregman projection for generic mirror maps can be computed approximately in ˜ O ( |A| ) computations through a bisection algorithm. Krichene et al. [50] propose one such algorithm, which we report in Algorithm 3 for completeness. We next provide two mirror maps that have appeared before in the optimization literature, but do not lead to an exact solution to the Bregman projection. We leave them as object for future work.

Negative Tsallis entropy [74, 57]. Let q > 0 and define ϕ as

The mirror map h ϕ q becomes the negative Tsallis entropy, that is

If q = 1 and following (14), the regression problem in Line 1 of Algorithm 1 becomes

Hyperbolic entropy [34]. Let b > 0 and define ϕ as

The mirror map h ϕ b becomes the hyperbolic entropy, that is

and, following (14), the regression problem in Line 1 of Algorithm 1 becomes

Hyperbolic-tangent entropy. Inspired by the hyperbolic entropy, we consider ϕ as

The mirror map h ϕ becomes

which, to the best of our knowledge, is not equivalent to any mirror map studied in the literature. Following (14), the regression problem in Line 1 of Algorithm 1 becomes

Regarding the limitations of the ω -potential mirror map class, we are aware of two previously used mirror maps that cannot be recovered using ω -potentials: h ( x ) = 1 2 x ⊤ A x [32], for some positivedefinite matrix A , which generates the Mahalanobis distance, and p -norms, i.e. h ( x ) = ∥ x ∥ 2 p [74]. Note that the case where h ( x ) = ∥ x ∥ p p can be recovered by setting ϕ ( x ) = ( px ) p/ (1 -p ) .

We note that tuning the mirror map and the step-size can lead AMPO to encompass the case of deterministic policies, which can be obtained when using softmax policies by sending the step-size to infinity, effectively turning the softmax operator into a max operator. Another simple way of introducing deterministic policies in our framework is to choose the mirror map to be the Euclidean norm and to choose the step-size large enough. Doing so will cause the Bregman projection to put all the probability on the action that corresponds to the maximum value of f θ s . Our results hold in this setting because our analysis does not use the policy gradient theorem (3), which has a different expression for deterministic policies [83].

D Deferred proofs from Section 4.1

D.1 Proof of Lemma 4.1

Here we provide the proof of Lemma 4.1, an application of the three-point descent lemma that accommodates arbitrary parameterized functions. Lemma 4.1 is the key tool for our analysis of AMPO. It is a generalization of both Xiao [93, Equation (44)] and Yuan et al. [98, Equation (50)] thanks to our two-step PMD framework. First, we recall some technical conditions of the mirror map [14, Chapter 4].

Suppose that Y ⊂ R |A| is a closed convex set, we say a function h : Y → R is a mirror map if it satisfies the following properties:

(i) h is strictly convex and differentiable;

(ii) h is essentially smooth, i.e., the graident of h diverges on the boundary of Y , that is lim x → ∂ Y ∥∇ h ( x ) ∥ → ∞ ;

(iii) the gradient of h takes all possible values, that is ∇ h ( Y ) = R |A| .

To prove Lemma 4.1, we also need the following rather simple properties, i.e., the three-point identity and the generalized Pythagorean theorem, satisfied by the Bregman divergence. We provide their proofs for self-containment.

Lemma D.1 (Three-point identity, Lemma 3.1 in Chen and Teboulle [19]) . Let h be a mirror map. For any a, b in the relative interior of Y and c ∈ Y , we have that:

Proof. Using the definition of the Bregman divergence D h , we have

Subtracting (44) and (45) from (46) yields (43).

Lemma D.2 (Generalized Pythagorean Theorem of Bregman divergence, Lemma 4.1 in Bubeck [14]) . Let X ⊆ Y be a closed convex set. Let h be a mirror map defined on Y . Let x ∈ X , y ∈ Y and y ⋆ = Proj h X ( y ) , then which also implies

we have

which implies (47) by applying the definition of Bregman divergence and rearranging terms.

Now we are ready to prove Lemma 4.1.

Lemma D.3 (Lemma 4.1) . Let Y ⊂ R |A| be a closed convex set with ∆( A ) ⊆ Y . For any policies π ∈ ∆( A ) S and ¯ π in the relative interior of ∆( A ) S , any function f θ with θ ∈ Θ , any s ∈ S and for η > 0 , we have that,

where ˜ π is induced by f θ and η according to Definition 3.1, that is, for all s ∈ S ,

Proof. For clarity of exposition, let p s = ∇ h ∗ ( ηf θ s ) . Plugging a = ¯ π s , b = p s and c = π s in the three-point identity lemma D.1, we obtain

Similarly, plugging a = ¯ π s , b = p s and c = ˜ π s in the three-point identity lemma D.1, we obtain

From (49), we have

By rearranging terms, we have

From the Generalized Pythagorean Theorem of the Bregman divergence in Lemma D.2, also known as non-expansivity property, and from the fact that ˜ π s = Proj h ∆( A ) ( p s ) , we have that

Plugging the above inequality into the left hand side of (51) yields which concludes the proof with ∇ h ( p ) = ηf

We also provide an alternative proof of Lemma 4.1 later in Appendix F.

Proof. From the definition of y ⋆ , which is

and from the first-order optimality condition [14, Proposition 1.3], with

D.2 Bounding errors

In this section, we will bound error terms of the type

where ( d π µ , π ) ∈ { ( d ⋆ µ , π ⋆ ) , ( d t +1 µ , π t +1 ) , ( d ⋆ µ , π t ) , ( d t +1 µ , π t ) } . These error terms appear in the forthcoming proofs of our results and directly induce the error floors in the convergence rates.

In the rest of Appendix D, let q t : S × A → R such that, for every s ∈ S ,

So (52) can be rewritten as

To bound it, let ( v t ) t ≥ 0 be a sequence of distributions over states and actions. By using CauchySchwartz's inequality, we have

where the last line is obtained by Assumption (A2) . Using the concavity of the square root and Assumption (A1) , we have that

D.3 Quasi-monotonic updates - Proof of Proposition 4.2

In this section, we show Proposition 4.2 with its proof that the AMPO updates guarantee a quasimonotonic property, i.e., a non-decreasing property up to a certain error floor due to the approximation error, which allows us to establish an important recursion about the AMPO iterates next. First, we recall the performance difference lemma [41] which is the second key tool for our analysis and a well known result in the RL literature. Here we use a particular form of the lemma presented by Xiao [93, Lemma 1].

Lemma D.4 (Performance difference lemma, Lemma 1 in [93]) . For any policy π, π ′ ∈ ∆( A ) S and µ ∈ ∆( S ) ,

For clarity of exposition, we introduce the notation

Proposition 4.2 characterizes the non-decreasing property of AMPO. The error bound (54) in Appendix D.2 will be used to prove the the result.

Proposition D.5 (Proposition 4.2) . For the iterates of Algorithm 1, at each time t ≥ 0 , we have

Proof. Using Lemma 4.1 with ¯ π = π t , f θ = f t +1 , η = η t , thus ˜ π = π t +1 by Definition 3.1 and Algorithm 1, and π s = π t s , we have

By rearranging terms and noticing D h ( π t s , π t s ) = 0 , we have

Then, by the performance difference lemma D.4, we have

which concludes the proof after dividing both sides by (1 -γ ) . The last line follows from

where both terms are upper bounded by √ C v ε approx through (54) with ( d π µ , π ) = ( d t +1 µ , π t +1 ) and ( d π µ , π ) = ( d t +1 µ , π t ) , respectively.

D.4 Main passage - An important recursion about the AMPO method

In this section, we show an important recursion result for the AMPO updates, which will be used for both the sublinear and the linear convergence analysis of AMPO.

For clarity of exposition in the rest of Appendix D, let

For two different time t, t ′ ≥ 0 , let D t t ′ denote the expected Bregman divergence between the policy π t and policy π t ′ , where the expectation is taken over the discounted state visitation distribution of the optimal policy d ⋆ µ , that is,

Similarly, let D ⋆ t denote the expected Bregman divergence between the optimal policy π ⋆ and π t , that is,

Let ∆ t := V ⋆ ( µ ) -V t ( µ ) be the optimality gap.

We can now state the following important recursion result for the AMPO method.

Proposition D.6 (Proposition 4.4) . Consider the iterates of Algorithm 1, at each time t ≥ 0 , we have

Proof. Using Lemma 4.1 with ¯ π = π t , f θ = f t +1 , η = η t , and thus ˜ π = π t +1 by Definition 3.1 and Algorithm 1, and π s = π ⋆ s , we have that

which can be decomposed as

Taking expectation with respect to the distribution d ⋆ µ over states and with respect to the randomness of AMPO and dividing both sides by η t , we have

We lower bound the two terms on the left hand side of (59) separately. For the first term, we have that

where ( a ) follows from Lemma D.4. For the second term, we have that

where ( b ) follows from Lemma D.4 and ( c ) follows similarly to (57), i.e., by applying (54) twice with ( d π µ , π ) = ( d ⋆ µ , π ⋆ ) and ( d π µ , π ) = ( d ⋆ µ , π t ) .

Plugging the two bounds in (59), dividing both sides by (1 -γ ) and rearranging, we obtain

From Proposition 4.2, we have that ∆ t +1 -∆ t -τ ≤ 0 . Consequently, since ν t +1 ≤ ν µ by the definition of ν µ in Assumption (A3) , one can lower bound the left hand side of the above inequality by replacing ν t +1 by ν µ , that is,

which concludes the proof.

D.5 Proof of the sublinear convergence analysis

In this section, we derive the sublinear convergence result of Theorem 4.3 with non-decreasing step-size.

Proof. Starting from Proposition D.6

Summing up from 0 to T -1 and dropping some positive terms on the left hand side and some negative terms on the right hand side, we have

Notice that ∆ 0 ≤ 1 1 -γ as r ( s, a ) ∈ [0 , 1] . By dividing T on both side, we yield the proof of the sublinear convergence

D.6 Proof of the linear convergence analysis

In this section, we derive the linear convergence result of Theorem 4.3 with exponentially increasing step-size.

Proof. Starting from Proposition D.6 by dropping D t +1 t (1 -γ ) η t on the left hand side, we have

Dividing ν µ on both side and rearranging, we obtain

If the step-sizes satisfy η t +1 ( ν µ -1) ≥ η t ν µ with ν µ ≥ 1 , then

Now we need the following simple fact, whose proof is straightforward and thus omitted.

Suppose 0 < α < 1 , b > 0 and a nonnegative sequence { a t } t ≥ 0 satisfies

The proof of the linear convergence analysis follows by applying this fact with a t = E [ ∆ t + D ⋆ t (1 -γ ) η t ( ν µ -1) ] , α = 1 -1 ν µ and b = ( 1 + 1 ν µ ) τ .

Then for all t ≥ 0 ,

E Discussion of the first step (Line 1) of AMPO - the compatible function approximation framework

Starting from this section, some additional remarks about AMPO are in order. In particular, we discuss in detail the novelty of the first step (Line 1) and the second step (Line 2) of AMPO in this and the next section, respectively. Afterwards, we provide an extensive justification of the assumptions used in Theorem 4.3 in Appendices G to I.

As mentioned in Remark 3.3, Agarwal et al. [1] study NPG with smooth policies through compatible function approximation and propose the following algorithm. Let { π θ : θ ∈ Θ } be a policy class such that log π θ ( a | s ) is a β -smooth function of θ for all s ∈ S , a ∈ A . At each iteration t , update

with where W > 0 and A t ( s, a ) = Q t ( s, a ) -V t ( s ) represents the advantage function. While both the algorithm proposed by Agarwal et al. [1] and AMPO involve regression problems, the one in (61) is restricted to linearly approximate A t with ∇ θ log π t , whereas the one in Line 1 of Algorithm 1 is relaxed to approximate A t with an arbitrary class of functions F Θ . Additionally, (61) depends on the distribution d t µ , while Line 1 of Algorithm 1 does not and allows off-policy updates involving an arbitrary distribution v t , as v t is independent of the current policy π t .

F Discussion of the second step (Line 2) of AMPO - the Bregman projection

As mentioned in Remark 3.4, we can rewrite the second step (Line 2) of AMPO through the following lemma.

Lemma F.1. For any policy ¯ π , for any function f θ ∈ F Θ and for η > 0 , we have, for all s ∈ S ,

Equations (12) and (20) are obtained by choosing ˜ π s = π t +1 s and ¯ π s = π t s for all s ∈ S , η = η t , θ = θ t +1 , and by changing the sign of the expression on the right in order to obtain an argmax .

Proof. Starting from the definition of ˜ π , we have

where the second and the last lines are obtained using the definition of the Bregman divergence, and the third line is obtained using (4) ( ∇ h ( ∇ h ∗ ( x ∗ )) = x ∗ for all x ∗ ∈ R |A| ).

Lemmas F.1 and B.1 share a similar result, as they both rewrite the Bregman projection into the MD updates. However, the MD updates in Lemma B.1 are exact, while the MD updates in AMPO involve approximation (Line 1).

Next, we provide an alternative proof for Lemma 4.1 to show that it is the direct consequence of Lemma F.1. The proof will involve the application of the three-point descent lemma [19, Lemma 3.2]. Here we adopt its slight variation by following Lemma 6 in Xiao [93].

Lemma F.2 (Three-point decent lemma, Lemma 6 in Xiao [93]) . Suppose that C ⊂ R m is a closed convex set, f : C → R is a proper, closed 6 convex function, D h ( · , · ) is the Bregman divergence generated by a mirror map h . Denote rint dom h as the relative interior of dom h . For any x ∈ rint dom h , let

Then x + ∈ rint dom h ∩ C and for any u ∈ dom h ∩ C ,

We refer to Yuan et al. [98, Lemma 11] for a proof of Lemma F.2.

Lemma 4.1 is obtained by simply applying the three-point descent lemma, Lemma F.2, to (62) with x + = ˜ π s , f ( u ) = ⟨-ηf θ s + ∇ h (¯ π s ) , u ⟩ , u = π and x = ¯ π s and rearranging terms.

In contrast, it may not be possible to apply Lemma F.2 to (11), as Π(Θ) is often non-convex.

G Discussion on Assumption (A1) - the approximation error

The compatible function approximation approach [1, 61, 15, 21, 2, 98] has been introduced to deal with large state and action spaces, in order to reduce the dimension of the problem and make the computation feasible. As mentioned in the proof idea in Page 8, this framework consists in upper-bounding the sub-optimality gap with an optimization error plus an approximation error. Consequently, it is important that both error terms converge to 0 in order to achieve convergence to a global optimum.

Assumptions similar to Assumption (A1) are common in the compatible function approximation literature. Assumption (A1) encodes a form of realizability assumption for the parameterization class F Θ , that is, we assume that for all t ≤ T there exists a function f θ ∈ F Θ such that

When F Θ is a class of sufficiently large shallow neural networks, this realizability assumption holds as it has been shown that shallow neural networks are universal approximators [39]. It is, however, possible to relax Assumption (A1) . In particular, the condition

can replace Assumption (A1) and is sufficient for the sublinear convergence rate in Theorem 4.3 to hold. Equation (63) shows that the realizability assumption does not need to hold for all t < T , but only needs to hold on average over T iterations. Similarly, the condition

can replace Assumption (A1) and is sufficient for the linear convergence rate in Theorem 4.3 to hold. Additionally, requiring, for all t < T ,

is sufficient for Equation (64) to hold. Equation (65) shows that the error floor in the linear convergence rate is less influenced by approximation errors made in early iterations, which are discounted by the term ( 1 -1 ν µ ) . On the other hand, the realizability assumption becomes relevant once the algorithm approaches convergence, i.e., when t ≃ T and Q t ≃ Q ⋆ , as the discount term ( 1 -1 ν µ ) is applied fewer times.

Finally, although Assumption (A1) holds for the softmax tabular policies and for the neural network parameterization, it remains an open question whether Assumption (A1) is necessary to achieve the global optimum convergence, especially when the representation power of F Θ cannot guarantee a small approximation error.

H Discussion on Assumption (A2) - the concentrability coefficients

In our convergence analysis, Assumptions (A2) and (A3) involve the concentrability coefficient C v and the distribution mismatch coefficient ν µ , which are potentially large. We give extensive discussions on them in this and the next section, respectively.

As discussed in Yuan et al. [98, Appendix H], the issue of having (potentially large) concentrability coefficient (Assumptions (A2) ) is unavoidable in all the fast linear convergence analysis of approximate PMD due to the approximation error ε approx of the Q -function [17, 100, 54, 16, 93, 20, 2, 98]. Indeed, in the fast linear convergence analysis of PMD, the concentrability coefficient is always along with the approximation error ε approx under the form of C v ε approx , which is the case in Theorem 4.3. To not get the concentrability coefficient involved yet maintain the linear convergence of PMD, one needs to consider the exact PMD in the tabular setting [see 93, Theorem 10]. Consequently, the PMD update is deterministic and the full policy space ∆( A ) S is considered. In this setting, at each time t , it exists θ t +1 such that, for any state-action distribution v t ,

and C v is ignored in the convergence analysis thanks to the vanishing of ε approx . We note that the PMD analysis in the seminal paper by Agarwal et al. [1] does not use such a coefficient, but a condition number instead. The condition number is controllable to be relatively small, so that the error term in their PMD analysis is smaller than ours. However, their PMD analysis has only a sublinear convergence rate, while ours enjoys a fast linear convergence rate. It remains an open question whether one can both avoid using the concentrability coefficient and maintain the linear convergence of PMD.

Now we compare our concentrability coefficient C v with others used in the fast linear convergence analysis of approximate PMD [17, 100, 54, 15, 93, 20, 2, 98]. To the best of our knowledge, the previously best-known concentrability coefficient C v was the one used by Yuan et al. [98, Appendix H]. As they discuss, their concentrability coefficient involved the weakest assumptions on errors among Lan [54], Xiao [93] and Chen and Theja Maguluri [20] by using the L 2 -norm instead of the ℓ ∞ -norm over the approximation error ε approx . Additionally, it did not impose any restrictions on the MDP dynamics compared to Cayci et al. [15], as the concentrability coefficient of Yuan et al. [98] was independent from the iterates.

Indeed, Yuan et al. [98] choose v t such that, for all ( s, a ) ∈ S × A ,

where ν is an initial state-action distribution chosen by the user. In this setting, we have

From the above lower bound of v t , we obtain that

where the finite upper bound is independent to t .

As mentioned right after Assumption (A2) , the assumption on our concentrability coefficient C v is weaker than the one in Yuan et al. [98, Assumption 9], as we have the full control over v t while Yuan et al. [98] only has the full control over the initial state-action distribution ν . In particular, our concentrability coefficient C v recovers the previous best-known one in Yuan et al. [98] as a special case. Consequently, our concentrability coefficient C v becomes the 'best' with the full control over v t when other concentrability coefficients are infinite or require strong assumptions [77].

In general, for the ratio E ( s,a ) ∼ v t [ ( d π µ ( s ) π ( a | s ) v t ( s,a ) ) 2 ] to have a finite upper bound C v , it is important that v t covers well the state and action spaces so that the upper bound is independent to t . However, the

upper bound 1 (1 -γ ) min ( s,a ) ∈S×A ν ( s,a ) in Yuan et al. [98] is very pessimistic. Indeed, when π t and π t +1 converge to π ⋆ , one reasonable choice of v t is to choose v t ∈ { d ⋆ µ · π ⋆ , d t +1 µ · π t +1 , d ⋆ µ · π t , d t +1 µ · π t } such that C v is close to 1 .

We also refer to Yuan et al. [98, Appendix H] for more discussions on the concentrability coefficient.

I Discussion on Assumption (A3) - the distribution mismatch coefficients

In this section, we give further insights on the distribution mismatch coefficient ν µ in Assumption (A3) . As mentioned right after (A3) , we have that

which is a sufficient upper bound for ν µ . As discussed in Yuan et al. [98, Appendix H],

The upper bound 1 / ((1 -γ ) min s µ ( s )) of ν ′ µ is very pessimistic and the lower bound ν ′ µ = 1 / (1 -γ ) is often achieved by choosing µ = d ⋆ µ .

Furthermore, if µ does not have full support on the state space, i.e., the upper bound 1 / ((1 -γ ) min s µ ( s )) might be infinite, one can always convert the convergence guarantees for some state distribution µ ′ ∈ ∆( S ) with full support such that

Then by the linear convergence result of Theorem 4.3, we only transfer the original convergence guarantee to V ⋆ ( µ ′ ) -E [ V T ( µ ′ )] up to a scaling factor sup s ∈S µ ( s ) µ ′ ( s ) with an arbitrary distribution µ ′ such that ν ′ µ is finite.

Finally, if d t µ converges to d ⋆ µ which is the case of AMPO through the proof of our Theorem 4.3, then sup s ∈S d ⋆ µ ( s ) d t µ ( s ) converges to 1. This might imply superlinear convergence results as discussed in Xiao [93, Section 4.3]. In this case, the notion of the distribution mismatch coefficients ν µ no longer exists for the superlinear convergence analysis.

We also refer to Yuan et al. [98, Appendix H] for more discussions on the distribution mismatch coefficient.

J Sample complexity for neural network parameterization

We prove here Corollary 4.5 through a result by Allen-Zhu et al. [3, Theorem 1 and Example 3.1]. We first give a simplified version of this result and then we show how to use it to prove Corollary 4.5.

Consider learning some unknown distribution D of data points z = ( x, y ) ∈ R d ×Y , where x is the input point and y is the label. Without loss of generality, assume ∥ x ∥ 2 = 1 and x d = 1 / 2 . Consider a loss function L : R k ×Y → R such that for every y ∈ Y , the function L ( · , y ) is non-negative, convex, 1 -Lipschitz continuous and L (0 , y ) ∈ [0 , 1] . This includes both the cross-entropy loss and the L 2 -regression loss (for bounded Y ).

Let g : R → R be a smooth activation function such that g ( z ) = e z , sin( z ) , sigmoid ( z ) , tanh( z ) or is a low degree polynomial.

Define F ⋆ : R d → R k such that OPT = E D [ L ( F ⋆ ( x ) , y )] is the smallest population error made by a neural network of the form F ⋆ = A ⋆ g ( W ⋆ x ) , where A ⋆ ∈ R k × p and W ⋆ ∈ R p × d . Assume for simplicity that the rows of W ∗ have ℓ 2 -norm 1 and each element of A ∗ is less or equal than 1 .

Define a ReLU neural network F ( x, W 0 ) = A 0 σ ( W 0 x + b 0 ) , where A 0 ∈ R k × m , W 0 ∈ R m × d , the entries of W 0 and b 0 are i.i.d. random Gaussians from N (0 , 1 /m ) and the entries of A are i.i.d. random Gaussians from N (0 , ε A ) , for ε A ∈ (0 , 1] . We train the weights W of this neural network through stochastic gradient descent over a dataset with N i.i.d. samples from D , i.e., we update W t +1 = W t -ηg t , where E [ g t ] = ∇ E D [ L ( F ( x, W 0 + W t ) , y )] .

Theorem J.1 (Theorem 1 of Allen-Zhu et al. [3]) . Let ε ∈ (0 , O (1 /pk )) , choose ε A = ε/ ˜ Θ(1) for the initialization and learning rate η = ˜ Θ ( 1 εkm ) . SGD finds a set of parameters such that

with probability 1 -e -c log 2 m over the random initialization, for a sufficiently large constant c , with

Theorem J.1 shows that it is possible to achieve the population error OPT by training a two-layer ReLU network with SGD, and quantifies the number of samples needed to do so.

We make the following assumption to address the population error in our setting.

Assumption J.2. Let g : R → R be a smooth activation function such that g ( z ) = e z , sin( z ) , sigmoid ( z ) , tanh( z ) or is a low degree polynomial. For all time-steps t , we assume that there exists a target network F ⋆,t : R d → R k , with

where w ⋆,t 1 ,i ∈ R d , w ⋆,t 2 ,i ∈ R d , and a ⋆,t r,i ∈ R , such that

We assume for simplicity ∥ w ⋆,t 1 ,i ∥ 2 = ∥ w ⋆,t 2 ,i ∥ 2 = 1 and | a ⋆,t r,i | ≤ 1 .

Assumptions similar to Assumption J.2 have already been made in the literature, such as the bias assumption in the compatible function approximation framework studied by [1]. The term OPT represents the minimum error incurred by a target network parameterized as F ⋆,t when solving the regression problem in Line 1 of Algorithm 1.

We are now ready to prove Corollary 4.5, which uses Algorithm 4 to obtain an unbiased estimate of the current Q-function. We assume to be in the same setting as Theorem J.1

Proof of Corollary 4.5. We aim to find a policy π T such that

Suppose the total number of iterations, that is policy optimization steps, in AMPO is T . We need the bound in Assumption (A1) to hold for all T with probability 1 -e -c log 2 m , which means that at each iteration the bound should hold with probability 1 -T -1 e -c log 2 m . Through Algorithm 4, the expected number of samples needed to obtain an unbiased estimate of the current Q-function is (1 -γ ) -1 . Therefore, using Theorem J.1, at each iteration of AMPO we need at most

samples for SGD to find parameters that satisfy Assumption (A1) with probability 1 -T -1 e -c log 2 m . To obtain (66), we need

Solving for T and ε approx and multiplying them together, we obtain the sample complexity of AMPO, that is

Due to the statement of Theorem J.1, we cannot guarantee the approximation error incurred by the learner network to be smaller than OPT . Consequently, we have that

A similar bound can be applied to any proof that contains the bias assumption introduced by [1].

We can obtain an improvement over Corollary 4.5 by using the relaxed assumptions in Appendix G, in particular using the condition in (65).

Corollary J.3. In the setting of Theorem 4.3, replace Assumption (A1) with the condition

for all t < T . Let the parameterization class F Θ consist of sufficiently wide shallow ReLU neural networks. Using an exponentially increasing step-size and solving the minimization problem in Line 1 with SGD as in (19) , the number of samples required by AMPO to find an ε -optimal policy with high probability is ˜ Θ( C 2 v ν 4 µ /ε 4 (1 -γ ) 6 ) .

Proof. The proof follow that of Corollary 4.5. Using Theorem J.1, at each iteration t of AMPO, we need at most

samples for SGD to find parameters that satisfy condition (68) with probability 1 -T -1 e -c log 2 m . Summing over T total iterations of AMPO we obtain that the total number of samples needed is

Replacing T and ε approx with the solutions of (67) gives the result.

At this stage, it is important to note that choosing a method different from the one proposed by Allen-Zhu et al. [4] to solve Line 1 in Algorithm 1 of our paper with neural networks can lead to alternative, and possibly better, sample complexity results for AMPO. For example, we can obtain a sample complexity result for AMPO that does not involve a target network using results from [39] and [16], although this requires introducing more notation and background results compared to Corollary 4.5 (since in [16] they employ a temporal-difference-based algorithm, that is Algorithm 3 in their work, to obtain a neural network estimate ̂ Q t of Q t , while in [39] they provide a method based on Fourier transforms to approximate a target function through shallow ReLU networks). We outline below the steps in order to do so (and additional details including the precise statements of the results we use and how we use them are provided thereafter for the sake of completeness).

Step 1) We first split the approximation error in Assumption (A1) into a critic error E [ √ ∥ ̂ Q t -Q t ∥ 2 L 2 ( v t ) ] ≤ ε critic and an actor error E [ √ ∥ f t +1 -̂ Q t -η -1 t ∇ h ( π t ) ∥ 2 L 2 ( v t ) ] ≤ ε actor . In this case, the linear convergence rate in our Theorem 4.3 becomes

[We can obtain this alternative statement by modifying the passages in Appendix D.2. In particular, writing f t +1 -Q t -η -1 t ∇ h ( π t ) = ( f t +1 -̂ Q t -η -1 t ∇ h ( π t )) + ( Q t -̂ Q t ) and bounding the two terms with the same procedure in Appendix D.2 leads to this alternative expression for the error.] We will next deal with the critic error and actor error separately.

Step 2) Critic error. Under a realizability assumption that we provide below along with the statement of the theorem (Assumption 2 in [16]), Theorem 1 from [16] gives that the sample complexity required to obtain E [ √ ∥ ̂ Q t -Q t ∥ 2 L 2 ( d t µ · π t ) ] ≤ ε is ˜ O ( ε -4 (1 -γ ) -2 ) , while the required network width is ˜ O ( ε -2 ) .

Step 3) Actor error. Using Theorem E.1 from [39], we obtain that E [ √ ∥ f t +1 -̂ Q t -η -1 t ∇ h ( π t ) ∥ 2 L 2 ( v t ) ] can be made arbitrarily small by tuning the width of f t +1 , without using further samples.

Step 4) Replacing Equation (67) with the sample complexity of the critic, we obtain the following corollary on the sample complexity of AMPO, which does not depend on the error made by a target network.

Corollary J.4. In the setting of Theorem 4.3, let the parameterization class F Θ consist of sufficiently wide shallow ReLU neural networks. Using an exponentially increasing step-size and using the techniques above to update f θ , the number of samples required by AMPO to find an ε -optimal policy with high probability is ˜ O ( C 2 v ν 5 µ /ε 4 (1 -γ ) 7 ) .

To the best of our knowledge, this result improves upon the previous best result on the sample complexity of a PG method with neural network parameterization [16], i.e., ˜ O ( C 2 v /ε 6 (1 -γ ) 9 ) . We now provide the statements of the aforementioned results we used.

Recalling Theorem 1 in [16] and its assumptions. Consider the following space of mappings:

and the function class:

Consider the following realizability assumption for the Q-function.

t

Assumption J.5 (Assumption 2 in [16]) . For any t ≥ 0 , we assume that Q ∈ F ¯ ν for some ¯ ν > 0 . Theorem J.6 (Theorem 1 in [16]) . Under Assumption 2 in [16], for any error probability δ ∈ (0 , 1) , let

and R > ¯ ν . Then, for any target error ε > 0 , number of iterations T ′ ∈ N , network width

Algorithm 3 in [16] yields the following bound:

where A 2 holds with probability at least 1 -δ over the random initializations of the critic network ̂ Q t .

As indicated in [16], a consequence of this result is that in order to achieve a target error less than ε > 0 , a network width of m ′ = ˜ O ( ¯ ν 4 ε 2 ) and iteration complexity O ( (1+2¯ ν ) 2 ¯ ν 2 (1 -γ ) 2 ε 4 ) suffice.

The statement of Theorem 1 in [16] can be readily applied to obtain the sample complexity of the critic.

Recalling Theorem E.1 in [39] and its assumptions Let g : R n → R be given and define the modulus of continuity ω g as

If g is continuous, then ω g is not only finite for all inputs, but moreover lim δ → 0 ω g ( δ ) → 0 .

Denote ∥ p ∥ L 1 = ∫ | p ( w ) | dw . Define a sample from a signed density p : R n +1 → R with ∥ p ∥ L 1 < ∞ as ( w,b, s ) , where ( w,b ) ∈ R is sampled from the probability density | p | / ∥ p ∥ L 1 and s = sign ( p ( w,b ))

Theorem J.7 (Theorem E.1 in [39]) . Let g : R n → R , δ > 0 and ω g ( δ ) be as above and define for x ∈ R n

Let G α be a gaussian distribution on R n with mean 0 and variance α 2 I . Define the Gaussian convolution l = g | δ ∗ G α with Fourier transform ̂ l satisfying radial decomposition ̂ l ( w ) = | ̂ l ( w ) | exp(2 πiθ h ( w )) . Let P be a probability distribution supported on ∥ x ∥ ≤ 1 . Additionally define and step-size

and for convenience create fake (weight, bias, sign) triples

( w,b, s ) m +1 := (0 , c, m sign ( c )) , ( w,b, s ) m +2 := ( a, 0 , m ) , ( w,b, s ) m +3 := ( -a, 0 , -m ) . Then

and with probability at least 1 -3 λ over a draw of (( s j , w j , b j )) j m =1 from p

We can then characterize the error of the actor by choosing x = ( s, a ) , g = ̂ Q t + η -1 t ∇ h ( π t ) , and f t +1 = 1 m ∑ m +3 j =1 s j σ ( ⟨ w j , x ⟩ + b j ) . We can then make the actor error arbitrarily small by tuning the network width m and δ (note that, since both ̂ Q t and f t are continuous neural networks, g is a continuous function).

Peer reviews:
comments: SUMMARY:
This paper proposes a policy update scheme based on mirror descent for general function approximation. The paper also provides corresponding theoretical analysis on the sub-optimality gap and computation cost. Overall, the reviewer find the theory part not strong enough to prove the advantage of the proposed new algorithm or contribute any new insight to the community.

STRENGTHS:
-

WEAKNESSES:
1. The main results in Theorem 4.3 is not really a meaningful upper bound of the sub-optimality gap of the learned policy since, for a general function approximation class, we do not know how large the term involving $\epsilon_{approx}$, $v_\mu$, and $C_v$ is. More detailed characterization of their scales (how fast do they converge to zero) is necessary to make the current analysis meaningful.
2. It is weird to make assumptions on $\pi^t$ in A1, A2, A3 when you are exactly analyzing the convergence of $\{\pi^t\}$, since we do not know where the policy iterate $\pi^t$ will go to. A more natural way to do this is to assume the conditions hold for all policies in the policy class. (But then the authors need to be more cautious about the assumptions.)
3. No new insights found in the current framework. The main techniques (mirror decent for policy optimization) are already widely used in previous reinforcement learning theory studies. The proposed algorithm seems another form or a very close variant of proximal policy gradient (PPO).

QUESTIONS:
-

LIMITATIONS:
-
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-07-03

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nWhile policy mirror descent methods have been theoretically analyzed in many past works to establish strong convergence guarantees, questions remain regarding the implementation of function approximation in this class of algorithms. The work analyzes general approximators with PMD, obtaining linear convergence guarantees up to standard approximation error. The theoretical approach also yields past algorithms as special cases, and the authors analyze the neural network approximation case as a corollary of their result. The theory also incorporates slightly weaker assumptions/constants in the bounds compared to past work.\n\nSTRENGTHS:\nI think the theoretical contributions of the paper are interesting.\n\n- The paper contributes to understanding how theoretically sound function approximation could be used in the case of policy mirror descent based methods, which already admit strong linear convergence guarantees.\n- The paper is well-written, and the main statements are rigorous and well-defined.\n- The theory unifies many past policy-based approaches under a framework, which is a meaningful contribution to RL theory.\n- AMPO, while having theoretical guarantees, seems to be practically relevant and straightforward to implement.\n- The authors provide a comprehensive comparison with existing results, and their algorithm is well-motivated with examples. The assumptions they introduce seem to be weaker than certain previous results,\n\nWEAKNESSES:\n- Certain parts of the theoretical results (Sections 3.1, 4.2) are adaptations of existing results, and in the case of Section 4.2, it is not clear how the paper uses certain assumptions/quantities in the referred work in their context.\n- The particular assumption (A1) and the final result incorporating a $\\mathcal{O}(\\sqrt{\\varepsilon_{approx}})$ bias could be compared to existing results on approximate PMD, see Q4 below.\n- No experimental results: I think this is a minor weakness. As the main contribution of the paper is theoretical, benchmarking scores with well-tuned existing algorithms could be in my opinion irrelevant. However, since the starting point of the paper was the practical success of PO with function approximation, at least the practical implementation of Algorithm 1 could be demonstrated in some cases.\n\nQUESTIONS:\nQ1: The $\\omega$-potential mirror maps seem to induce a wide range of mirror maps that have a sum of functions of components structure. Are there relevant cases where the $\\omega$-potential mirror map condition is not satisfied?\n\nQ2: In the case of two-layer neural network parameterization (Section 4.2), I found it difficult to compare the assumptions in the paper and the reference work by Allen-Zhu et al, 2019a. The referred work includes a term $OPT$ that depends on the population error (i.e., best achievable error with a finite width neural network). Why can this term be ignored in the Section G, when quantifying $\\varepsilon_{approx}$? While neural networks are universal approximators, I couldn’t follow how large the width $m$ would need to be to control $OPT$, which in turn affects the sample complexity. For example, could it be the case that achieving $OPT < \\delta$ requires $m \\sim \\exp{\\delta^{-1}}$, introducing an additional polynomial term in Corollary 4.4? Or is Corollary 4.4 stated up to an error of $OPT$? A related assumption that was not clear to me was the realizability of the target Q-values (potentially in the infinite width limit), which might need to be introduced to be able to control $OPT$ to be arbitrarily small. If this assumption is indeed necessary, it should be stated clearly.\n\nQ3: As a related question, does the assumption that $F^*$ admitted $A^*, W^*$ with $W^*$ having row norms less than 1 introduce a smoothness assumption on the target Q function?\n\nQ4: As a general question, PMD has been analyzed in the case where policy evaluation or the optimization subproblem can not be computed exactly but approximately, for instance in [Wenhao Zhan, 2021], Section 3.2, Theorem 2, and potentially others. These results also potentially incorporate approximation error, similar to (A1), is there a fundamental difference beyond the incorporation of the particular mirror map? As far as I understand, Theorem 2 of [Wenhao Zhan, 2021] also incorporates linear convergence up to a non-vanishing approximation error.\n\nI am happy to change my score in case the authors clarify the above points.\n\nMinor comments:\nProposition 4.2 could include the (stronger) positive term stated in the appendix as well to ease reading, since (as stated) it seems improvement is not necessarily guaranteed.\n\nLine 985-986, missing $h$ in the equation.\n\nLIMITATIONS:\nThe paper makes the assumptions of the theory clear. Some comments regarding assumptions that were not clear are mentioned in Q1-3 above.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper extends the recently introduced policy mirror descent method of tabular setting to function approximation setting. The essential development seems to reside in exploiting the mirror descent update in the dual space and projecting the dual representation of the policy into a \"realizable\" representation that is consistent with the function approximation used. The author recovers both sublinear and linear convergence when such projection error can be controlled. The resulting method is thus implementable with general parameterization of the policy.\n\nSTRENGTHS:\n1. I enjoy the simplicity of this idea. It is well known that when using general function approximation of the policy it becomes unclear how basic three-point lemma holds as we are performing descent in the parameter space, not the policy space. This paper uses the well-known interpretation of the mirror descent method and suggests to first project the dual representation of the updated policy into the one that is representable by the function approximator. Consequently, one is directly doing approximate descent directly in the policy space and the existing framework of policy mirror descent can be adapted. \n\n2. I also appreciate discussions on the potential choices of distance-generating functions. Example 3.5 and 3.6 appears to be new to me. Theorem 4.5 on how to control the policy projection error with neural network parameterization seems to be useful. \n\n3. Minor comment: using Tsallis divergence as the distance-generating function has been discussed explicitly in arXiv preprint arXiv:2303.04386, 2023.\n\nWEAKNESSES:\n1. It appears to me that Lemma 4.1 can be understood as using the projected dual variables $f^{\\theta}_s - \\nabla h(\\overline{\\pi}_s)/\\eta$ instead of $Q_s$ when doing the policy mirror descent update, especially in the analysis $\\overline{\\pi}$ is set to be the last policy.  In this sense, Lemma 4.1 seems to be a fairly standard result. This can be seen from \n\\begin{align}\n \\tilde{\\pi}_s =  \\mathrm{argmin}  <- \\eta f^{\\theta}_s + \\nabla h(\\overline{\\pi}_s), \\pi> + h(\\pi) - <\\nabla h(\\overline{\\pi}_s),\\pi>\n\\end{align}\nand consequently \n\\begin{align}\n \\tilde{\\pi}_s =  \\mathrm{argmin}  <- \\eta f^{\\theta}_s + \\nabla h(\\overline{\\pi}_s), \\pi> + D(\\pi, \\overline{\\pi}_s).\n\\end{align}\nGiven the above observation standard descent lemma can be applied and yields Lemma 4.1.\n\n2. I personally would much appreciate the perspective of this paper -- the projection step of the policy -- instead of the technicality/importance of Lemma 4.1 that the paper is currently promoting. Kindly correct me if I have missed anything.\n\nQUESTIONS:\nI enjoy reading this paper with great clarity in presenting its essential idea. I have worked on this topic for a while and therefore I do not have additional technical questions.\n\nLIMITATIONS:\nYes, the authors have adequately addressed the limitations and potential future improvement.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-21"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis is a theoretical work and the main claim is the proof of linear convergence rate for policy mirror descent (PMD) algorithms with general policy parametrization. This contrasts with previous results (Lan 2022, Xiao 2022, Cen et al 2021, Zhan et al. 2021,  Cayci et al. 2021, etc.) that proved convergence rates for either tabular or other specific policy parametrizations. The work develops the proof by creating a framework that incorporates general policy parameterization into PMD: Approximate Mirror Policy Optimization (AMPO). The analysis performs a non-trivial extension of the three-point descent lemma (Chen and Teboulle, 1993, Lemma 3.2) and uses the proof techniques of Xiao et al. (2022) for PMD.\n\n**Post rebuttal update**\n\nThe authors did several changes to answer my review, e.g., they implemented their algorithm and ran some experiments.\nI raised my score by 1 point due to including an implementation of the algorithm. The experiments did not fully convince me, so I did not raise my score further (see my discussion with the reviewers for more details). My impression is that the contribution is mainly theoretical, and it fits into the literature on theoretical works in this topic. Looking at the constants, and discussing with the authors, I also had the impression that there is a gap in the theory and practice in this domain. It seems that none of the previous works performed experiments to validate the theory either, and replacing the constants in the bounds with reasonable values makes the bounds vacuous. This is not really a problem of the current work, as previous theoretical works in this domain seemed similar. I would encourage the authors to critically consider these points, and discuss limitations of the theory, unless the experiments strongly support the theory.\n\nSTRENGTHS:\n- The paper is well-written and explains how the contributions relate to previous work.\n- It appears rigorous.\n- Previous proofs had quite restrictive settings, whereas the submitted paper's theorems allow for neural network parameterizations of the policies, so this is quite important for more realistic theory, and seems like a substantial contribution to this line of work.\n\nWEAKNESSES:\n- There are no experiments. It may be good to have to some illustrative experiments that allow understanding the contribution at a glance.\n- The paper is dense and all of the proofs are in the appendix. Given that reviewers are not obligated to read the appendix, and that no evidence or intuitively clear reasons for why the theory is true is provided in the main paper, it is difficult to judge the correctness of the work without reviewers going beyond what is expected of them (although the work does recover the best known rates in the tabular setting as well as several other previous restricted settings, which is promising). The main related past works were all published at journals, so I wonder whether submitting to a journal may be a better option. I think that conference theory papers benefit from intuitive examples and experiments in the main part of the paper.\n\nQUESTIONS:\nI would suggest adding experiments and figures that demonstrate how the work improves over previous bounds.\n\nLIMITATIONS:\nYes, adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-27"}

comments: SUMMARY:
This paper presents a novel framework for analyzing and understanding policy optimization based on mirror descent that can leverage any general parameterizations of the policy class. The new framework induces a new update rule based on mirror descent, and the authors proved the new framework enjoys sublinear and linear convergence for general policy parameterization.

STRENGTHS:
- The authors did a great job presenting the problem's background and motivation, and the differences between the new proposed novel framework and previous related literature, which is quite intuitive and easy to understand. 
- The authors provided solutions (Sec. 3.1) to solve the new proposed update rule (Line 1 in Algorithm 1) by introducing a new mirror map, which is novel and easy to follow. 
- The theoretical analysis for convergence and sample complexity are provided, which makes the whole framework novel and complete.

WEAKNESSES:
- Not necessarily a weak point, but it would be good to provide some empirical examples to explain the motivation or the algorithm.

- Not sure if the current framework can be extended to other settings such as online or offline RL.

QUESTIONS:
- Can authors illustrate if it is possible to extend the results to the offline RL or other general RL setting? 
- I wonder if the policy is deterministic, will it be necessary to introduce the Line 1 update to analyze the general function approximation results?

LIMITATIONS:
NA
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-08-01