A Novel Framework for Policy Mirror Descent with General Parameterization and Linear Convergence
Authors: Carlo Alfano
Rui Yuan
Patrick Rebeschini
Venue: NeurIPS 2023 poster (2023)
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
Carlo Alfano
Rui Yuan Stellantis, France"
Patrick Rebeschini
Department of Statistics
Department of Statistics
University of Oxford
rui yuan Ontellantia.com
University of Oxford
car aal.aafssttat
patrick wwweechhnniiie chinesstattttt
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, introduce framework for policy optimization based on mirror descent that naturally accommodates general parameterizations. The policy class induced by scheme recovers known classes, e.g., softmax, and generates new ones depending on the choice of mirror map. Using our framework, obtain the first result that guarantees linear convergence for a policy-gradient-be 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 represents an improvement upon the previous best results, and empirically validate the effectiveness of our theoretical claims on classic control tasks.
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 parameterized policy class and have emerged as popular class of algorithms for [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 games [9] and autonomous driving [81]. In particular, class of PG algorithms that has proven successful in practice consists of building updates that include a hard constraint (e.g., 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 Kulbbeck-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 variant of TRPO that actually a special case of PMD, thus linking TRPO and PMD.
From theoretical perspective, motivated by the empirical success of PMD, there is 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
concerted
"The work was done when the author was affiliated with Télécom Paris.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

entropy mirror map in the linear function approximation regime, i.e., for log-linear policies, either by adding entropy regularization [15], or by using geometrically increasing step-size [20, 2, 98]. The proofs of these results are based on specific policy parameterizaatins,, 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 Optmmization (AMPO), a novel framework designed incorporate general parameterization into PMD in theoretically sound manner. In summary, AMPO is an 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 tabular, linear, or neural network parameterizations. We then develop an update procedure for this policy class based on MD and PG.
method that recovers PMD in different settings, such as tabular
We provide an analysis of AMPO and establish theoretical guarantees that hold for any parameteriza- tion class and any mirror map. More specifically, we show that our algorithm enjoys quasi- monoooni 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 -----rrdienttt-eee method with linear convergence that can
our knowledge, AMPO is the first 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 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).
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. 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 large class of mirror maps, i.e., the w-potential mirror maps in Definition 3.5, AMPO can be implemented in O(IAl) 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., Ö(E-t) (Corollary 4.5), and show how it improves upon previous results.
2 Preliminaries
Let M - (S,A, 14) be a discounted Markov Decision Process (MDP), where S is a possibly infinite state space, A is finite action space, P(s' 8, a) is the transition probability from state s to under action a, r(s,a) E [0, 1] is a reward function, 7 is discount factor, and μ is a target state distribution. The behavior of an agent on an MDP is then modeled by policy π E (A(A))S, where a ~ 7(- s) is the density of the distribution over actions at state s E S, and A A(A) is the probability simplex over A.
Given a policy n, let V* : S + R denote the associated value function. Letting s, and a be the current state and action at time t, the value function /a is defined as the expected discounted cumulative reward with the initial state Su s, namely.
En I'r(s,, at) n, so = s].
V (s) == Ea, ao (-1s) se.ar)
Now letting (H) ;= Ex [V= (s)], our objective is for the agent to find an optimal policy
x E argmax V=(H).
(I)
2

---

As with the value function, for each pair (s,a) E S x A, the state-action value function, or O-function, associated with a policy is defined as
Q (s. a) = Ba (:) ((-ss,a)))
Σ 1=
r ar) 5, = 3,00 = a
We also define the discounted state visitation distribution by
E = - (1- E.
P(s = 60)
where P(s, following policy π and starting from so- The probability (5((s) represents the time spent on state when following policy π.
80) represents the probability of the agent being in state at time when
The gradient of the value function V (n)) with respect to the policy is given by the policy gradient theorem (PGT) [87]:
OV" (p)
1
V,V*(H) Oml (1) = 1-7
())) (8,
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 RRA be convex, continuously differentiable and essentially smooth function' such that Vh())) R|AI, The convex conjugate of h, denoted by h" is given by
set. A mirror map h y R is strictly
h" (x sup suppeyy (x" x) h(x), x¹ E'ERIA
The gradient of the mirror map V/ y - R allows to map objects from the primal space y to its dual space RIAI, x + Vh(z), and viceversa for Vh i... x² + Vh (x*). In particular, from Vh()) = glA|, we have: for all (z, (x,x*)= E ARA
x Vh (Vh(z)) and
Vh(Vh"(=")
Furthermore, the mirror map h induces a Bregman divergence [13, 18] defined as
Dr (x.3) 3) = l1(x) h(w)) (Vh(y), y),
where Da(x, y) 2 0 for all x, y V. We can now present the standard MD algorithm [70, 14]. Let X C y be convex set and V X' IR be differentiable function. The MD algorithm can be formalized as the following iterative procedure in order to solve the minimization problem mingex Vir): for all / 2 0,
y +1 Vh(x*) 176 VV))))
(5)
x ++ Projk((V1) (z/""( )),
(6)
where 7/ set according to step-size schedule (Mx)22 and Proj'x( (.) is the Bregman projection
Proj (v) argmin Dz(2, y).
Precisely, at time t, xf X is mapped to the dual space through Vh(-), where gradient step is performed in (5) to obtain y The next step is to map ff back in the primal space using V.)((). In case Vh (y++)) does not belong to X, projected as in (6).
3 Approximate Mirror Policy Optimization
The starting point of our novel framework is the introduction of novel parameterized policy class based on the Bregman projection expression recalled in (7).
'h is essentially smooth limg-ay [[Vh(r)ll +oo, where 0y denotes the boundary y.
Seee different formulation of MD in (11) and in Appendix (Lemma B.

---

Algorithm 1: Approximate Mirror Policy Optimization
Input:
Initial policy x mirror map h, parameterization class Fe, iteration number T, step-size schedule (Nt 12000 state-action distribution sequence (v*) (2))
for 1 = 0 to T do
Obtain 0++1 11 E θ such that 0*+1 argmingge || Q' Q²-ni -
Update ; E argmin (p, Vh (m.f3 () = (A) (n.ftt1)), (Vh Vs E.
2
PEA(A)
Definition 3.1. Given a parameterized function class _ {f A θ}, mirror map h : y + IR, where V S RRA is convex set with ((A) - V. and 1j > 0, the Bregman projected policy class associated with F and h consists of all the policies of the form:
1T²² n where for all S. T2 E RIAA denote vectors (a s)) and [ a))GEA, respectively.
(nf2)) s S; A E O)
In this definition, the policy is induced by a mirror map h and a parameterized function f", and is ob- tained by mapping f to Y with the operator (-). which may not result in well-defined probabil- ity distribution, and is thus projected on the probability simplex ((A). Note that the choice of is key in deriving convenient expressions for x² 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].
IR! and h is the negative entropy mirror map, i.e.,
Example 3.2 (Negative entropy). If y M(7(11))==> -> (a]s)) (o((=(((s)), then the associated Bregman projected policy class becomes
{
exp(nf?) 8 S S; BE6).
n π [exp(nf%)ll
(8)
where the exponential and the fraction are element-wise and ||-1| the l norm. particular, when f (s, a) = a..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.I for its derivation.
We now construct 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. to our setting. We define the following shorthand: at each time t, let n. == x" f8r V = V*' Q' = Q"', and d', := di'' Further, for any function y S x A + R and distribution v over S x A. let y/ := -(s,-) E RI-I1 and Ilvlli = (o a)) Ideally, we would like to execute the exact MD algorithm: for all 2 0 and for all E S.
f Vh(a,) + 7c(1 V) (n)/d' @ vh(=))+ + n.
(9)
_(+1 Proj" Vh (n((((
(10)
Here, (10) reflects our Bregman projected policy class 3.1. However, we usually cannot perform the update (9) exactly. In general, if fo belongs to a parameterized class F there may not be any gi-l θ such that (9) satisfied for all E 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 I is a standard regression problem where we try to approximate Q' + 7% Vh(a'') with ft+1, and has been studied extensively when f is a neural network [3, 35]. We can then readily use (10) to update "++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 how the two steps of AMPO relate and improve upon previous works.
Thee update is (5) up to a scaling (1 - -d) of ne- With slight abuse of notation, denote Vh( a) [Vh(=)))..

---

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 @ the solution to regression problem aimed at approximating the current Q Q-function Q'. A crucial difference is that, Agarwal et al. [1] approximate Q' linearly with respect to Vo log n' (see (61)), while in Line we approximate Q' and the gradient of the mirror map of the previous policy with any function fo. This generality allows our algorithm to achieve approximation guarantees for a wider range of assumptions on the structure of Q²', Furthermore, the regression problem proposed by Agarwal et al. [I] depends on the distribution dz while ours has no such constraint and allows off-policy updates involving an arbitrary distribution v². 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
πt E af Re [Tx (Q., m (A x!)]
(11)
where II(@) is given parameterized policy class, while the Bregman projection step of AMPO can be rewritten as
++ argmax (n- Vh(T.).p) - (p. T.), E
(12)
We formulate this result as Lemma F.l in Appendix F with a proof. When the policy class II(O) is the entire policy space (AA)S, (1I) equivalent to the two-step procedure (9)-(1) 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 (I) in (11) often a non-convex set, thus the three-poine 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 a convex problem and thus allows the application of the strre-point-desceee lemma [19]. We refer to Appendix F for the conditions of the three-poing punee-point-dsceen lemma in details.
3.1 w-potential mirror maps
In this section, we provide a class of mirror maps that allows to compute the Bregman projection in Line 2 with ([-A) operations and simplifies the minimization problem in Line 1. Definition 3.5 (w-potential mirror map [50]). For u E (-o0, +oo], w S 0, let an w-potential be an increasing C -ddiffcommrphismm o (-00, u) (wo, +-0) such that
L. (=)' S
lim ((x)= w, lim d(x) = +00,
For any w-potential o, the associated mirror map has called w-potential mirror map, defined as r(als) Φ (
(T(TT))= = = E
aEAA
Thanks Krichene et al. [50, Proposition 2] (see Proposition C.1 as well), the policy nh+1 I in Line 2 in- duced by the w-potential mirror map can be obtained with ((AA) computations and can be written as
+++ (a s) ( (s. + X**1),, Vs E A,
(13)
where X+1 E R is a normalization factor to ensure EeEA ( | s) 1 for all s E S, and a(2) max 0) for R. We call this policy class the w-potential policy class. By using (13) and the definition of the w-potential mirror map hg, the minimization problem in Line is simplified to be
gi-l E argmingee ||50 Q' mi max(n--f' -00 -xxx (
(14)
where max(.,)) :) is applied clement-wisely. The w-potential policy class allows AMPO to generate wide range of applications by simply choosing an w-potential 6. fact, it recovers existing approaches to policy optimization, as we show in the next two examples.
Example 3.6 (Squared (2-norm). If = IglAl and ø is the identity function, then hy is the squared l norm, that is (T,))) ||= /2, and Vh is the identity function. So, Line I in Algorithm becomes
g++1 E argmingee || 0 Q' - -/ 711 La(ee))
(15)
5

---

The Vhg also becomes the identity function, and the policy update is given for all E S by
the Proj (n1ss f
argmin
|| - -t 2
(16)
T.EA(A)
which is the Euclidean projection on the probability simplex. In the tabular setting, where S and A are finite and (s. = Os.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, generalize the projected Q-descent algorithm from the tabular setting to general parameterization class F which novel algorithm in the RL literature.
Example 3.7 (Negative entropy). I/ R and d(x) exp(x 1), then he is the negative entropy mirror map from Example 3.2 and Line in Algorithm becomes
g++l E argmingce | - Q'
1t-1 fet T
2
(17)
L2(u')
Consequently, based on Example 3.2, we have rt+1 x exp(n. +(++)) for all 8 E S. In this example, AMPO recovers tabular NPG [82] when f" a) = Ss.a. and Q-NPG with log-linear policies [98] when f and Q' are linear functions for all t 2 0.
We refer to Appendix C.2 for detailed derivations of the examples in this section and efficient implementation of the Bregman projection step. In addition to the ég-norm and the negative entropy, several other mirror maps that have been studied in the optimization literature fall into the class of w-potential mirror maps, such as the Tsallis entropy [74, 57] and the hyperbolic entropy [34], as well as generalization of the negative entropy [50]. These examples illustrate how the class of w-potential mirror maps recovers known methods and can be used to explore new algorithms in policy optimiza- tion. 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 S F and for v >> O, we have
(n12 -l((a)) F. -.)) S Dr. (T Rs) (n( (TT Rs), Vs E
where t is the Bregman projected policy induced by f and h according to Definition 3.1, that is n. argMmn [p, Vh (nf")) for all E S.
The proof of Lemma 4.1 is given in Appendix D.I. Lemma 4.1 describes 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 application of the three-point descent lemma [19], while it cannot be applied to algorithms based on the update in (1I) [88, 89] due the non-convexity of the optimization problem (see also Appendix F). Notice that Lemma 4. 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. similar result for general parameter- ization has been obtained by Lan [55, Proposition 3.5] in the setting of strongly convex mirror maps.
Lemma 4.1 becomes useful when we के n², f f*+1, 1) 1/t and π n² or x². In particular, when / - Vh(m!) v n Lemma 4. allows us to obtain telescopic sums and recursive relations, and handle error terms efficiently, show in Appendix D.
4.1 Convergence for general policy parameterization
In this section, we consider the parameterization class 7 and the fixed but arbitrary mirror map h. We show that AMPO enjoys quasi-monotonic improvement and sublinear or linear convergence, depending the step-size schedule. The first step is to control the approximation error of AMPO.

---

(A1) (Approximation error). There exists Eapprox 2 such that, for all times t 2 0,
[ Q' 1 (V(('') (a(1t)] ( S apprrxx
where (u") is sequence of distributions over states and actions and the expectation is taken over the randomness the algorithm that obtains ft+1.
Assumption (A1) is common in the conventional compatible function approximation approach" [1]. It characterizes the loss incurred by Algorithm in solving the regression problem in Line 1. When the step-size 1)t is sufficiently large, Assumption (A1) measures how well fi+1 approximates the current Q-function Q'. Hence, Eapprox depends both the accuracy of the policy evaluation method used to obtain an estimate of Q' [86, 79, 28] and the error incurred by the function I E F that best approximates Q', that is the representation power of _ Later in Section 4.2, we show how to solve the minimization problem in Line 1 when _ is 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 discussion on technical novelty and Appendix G for relaxed version of the assumption.
As mentioned in Remark 3.3, the distribution v' does not depend the current policy x' for all times & 2 O. 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) (Concentrabillity coefficient). There exists C. 2 such that, for all times t,
Blo.a)mos (( ( 2
CCu.
v((8,a) a)
whenever (dj, T) is either (d) T²), (d",t T'), or (d) (d)
The concentrability coefficient C, quantifies how much the distribution vf overlaps with the distributions (d x*), (d++1, x(+1), (d),=') and (d++1,m) ,7)). It highlights that the distribution of order to avoid large values of C.,, Assumption
should have full support over the environment, (A2) is weaker than the previous best-known concentrability coefficient [98, Assumption 9]. in the sense that we have the full control over o'. We refer to Appendix H for a more detailed discussion. We present our first result the performance of Algorithm 1.
Proposition 4.2 (Quasi-monotues monotonic updates). Let (A1). (A2) be true. We have, for all 2 O,
( m²) + DN 2 2cCooappro
E [vV (m) - v'(w]] 2 E
E. B
The - 7)
1-7
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 cannot lead to performance degradation, up to an error term. The next assumption concerns the coverage of the state space for the agent each time t.
(A3) (Distribution mismatch coefficient). Let d; == d=² There exists Vy 2 such that
sup di(() .ES (s)
S Vus for all times t 2 0.
Since (s) 2 ->)s(s) for all s E S, obtained from the definition of dy in (2), we have that
sup (3)
sup (s)
SES d' (s) 1-7
sES m(s)
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.. 15, 93]. We refer to Appendix I for details, It is worth mentioning that the quasi- monotoio improvement in Proposition 4.2 holds without (A3).
We define the weighted Bregman divergence between the optimal policy " and the initial policy T as D = Ex-dy (Da(T; =²)]- We then have our main results below.
An extended discussion of this approach is provided in Appendix G.
7

---

Theorem 4.3. Let (A1), (A2) and (A3) be true. If the step-size schedule is non-decreasing, i.e., Tt S Mh+1 for all 2 the iterates of Algorithm I satisfy: for every ' O.
()) T
1 (
Di (1-7)m
U/ 11-7)
2(1 + MN) Vp)yCCeEpppr Cappproxx
1-9
teT
Furthermore, if the step-size schedule is geometrically increasing, i.e, satisfies
1/1
Mt++1
2
V2 0,
(18)
Up
-1
we have: for every? 1220 O,
1
Dj V/u To (2 - 1)
T
2(1 + BNV Eappprox approx
V (*()) EE E2 [V (n)]
1-7
1 1-7
Theorem 4.3 is, to the best of 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 20 normm, 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 consid- ering a more general setting, Theorem 4.3 matches or improves upon the convergence rate of previous work 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], impose bounded updates or smoothness of the policy [1, 61], restrict the analysis to the case where the mirror map h is the negative entropy [60, 36]. or to make for noor assumptions
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 - 7)) (27) Al), and can slow down the learning process.
When is a finite state space, a sufficient condition for v4, in (A3) to be bounded is requiring yu to have full support on S. If does not have full support, one still obtain linear convergence for V*(p') (I()",, for an arbitrary state distribution jı' with full support, and relate this quantity to V*(H) (T()). We refer Appendix for 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)
argmin (-ft+1 + m
T Vh(7()) p) m (D m(), Vs S.
Secondly, solving Line 1 of Algorithm leads to ft+1 ", Vh(TT) Q. When the step-size N co, that is n, O, the above viewpoint of the AMPO policy update becomes
E argmin (-Q'.p) =
n E argmax (Q.,p), Es
which is the Pl algorithm. Here ignore the Bregman divergence term Dr (T, (.). as it is multiplied by 1/m,- which goes to O. So AMPO behaves more and more like PI with a large enough step-size and thus is able converge linearly like PI.
Proof idea. We provide a sketch the proof here; the full proof is given Appendix D. In nutshell, the convergence rates of AMPO are obtained by building 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 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.
Step I. Using Lemma 4.1 with = f f 1f nr मे x and T, n² we obtain
(nefg-11 - Vh (T T T 2 0,
which characterizes the improvement of the updated policy.
Step 2. Assumption (A1), Step 1, the performance difference lemma (Lemma D.4), and Lemma frtt 17 = Tt. मे = xf+1, and π. T permit us to obtain the following.
4.1 with न = T f =
Proposition 4.4. Let A, = (*) (+)).. For all t Z 0, we have
B [
E, Ea-d= (T).))))))
2 CoEppeoxx
E [M/ (A.+1 A + Ad <
+ (1+ Uy))
(1-7)m
(1-7)m
1-7
Step 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,
A ++1
Eendy E. (T],
S
(1- My B|A
Ba-d; [Pa())+))] (1-))m(M,-1) + (1+ upp)
2.)
2 /C Cuppprox
(1-9)=+1(4661)
1-y
4.2 Sample complexity for neural network parameterization
Neural networks are widely used in RI. due to their empirical success in applications [67, 68, 84). However, few theoretical guarantees exist for using this parameterization class in policy optimization 160, 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 j from Definition 3. 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) E (S x A) C Rd define (s, a) c (W(ss.) a) + b) with 0 = (c. W, b). where (1) max y. 0) for all y IR is the ReLU activation function and is applied element-wisely. Th W TR. mxd and b m
At each iteration of AMPO, we set vf d' and solve the regression problem in Line of Algorithm I through stochastic gradient descent (SGD). In particular, we initialize entry-wise Wo and b as ii.d. random Gaussian variables from (O0,, 1/m). and c as i.i.d random Gaussian variables from A/0, <a) with сд 1). Assuming access simulator for the distribution v`, we run SGD for K steps on the matrix W, that is, for = 0, O,...,K ----, -1,
- W a (f(() (s, a) Q' a) n. V (s.
(19)
where f(* cT a ((W + (s, (s + b). (s, a) rf and (s( a) is an unbiased estimate of Q'(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 l].
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 I with SGD as in (19), the number of samples required by AMPO to find e-optimal policy with high probability is OCC3v//e -7)"), where has be larger than a fixed and non non-vanishing error floor.
We provide proof of Corollary 4.5 and an explicit expression for the ervor floor in Appendix J. Note that the sample complexity Corollary 4.5 might be impacted by an additional poly(e term. We refer to Appendix for more details and an alternative result (Corollary J.4) which does not include an additional poly(E-) term, enabling comparison with prior works.
Numerical experiments
We provide 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 offte-of-the- art 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.
AMPO performance on CartPole-vl
AMPO performance on Acrobot-v1
500
100
400
300
5.
-00
200
AMPO 66et
with 60g
AMPO with dest sichir)
AMPO with eizi anake))
AMPO with d6e) mi
entropy (e- 6.5)
AMPO with Tsalls entropy (c-03)
AMPO with Thailss entmayy (e- 1.5)
AMPO with T2alls entropy (s-La)
-500
500
300
200
500
Iteration
Iteration:
Figure I: Averaged performance over 50 runs of AMPO in CartPole and Acrobot environments. Note that the maximum values for CartPole and Acrobot are 500 and -80, respectively.
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 w-potential mirror maps. In particular, we choose: (xx) e", which corresponds to the negative entropy (Example 3.7); <(x) = sinh(2)) which corresponds to the hyperbolic entropy [34, see also (41)): ((x) = 2, which corresponds the Euclidean norm (Example 3.6); and the Tsallis entropy for two values of the entropic index [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 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 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 Dij 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 the best choice for AMPO, and that exploring different mirror maps is 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 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 thorough discussion of future work. We believe that the main contribution of AMPO is to provide general framework with theoretical guarantees that can help the analysis of specific algorithms and MDP structures. AMPO recovers and improves several convergence rate guarantees 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.
10

---

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 NeurIPS 2023.
References
[1] Alckh 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.)
[2] Carlo Alfano and Patrick Rebeschini. Linear convergence for natural policy gradient with log-linear policy parametrization. arXiv preprint arXiv:2209. 2022. (Cited on pages 2, 8, 20, 21, 22, 25, 28, 38, and 39.)
[3] Zeyuan Allen-Zhu, Yuanzhi Li, and Yingyu Liang. Learning and generalization in overparame- terized neural networks, going beyond two layers. Advances in Neural Information Processing Systems, 2019. (Cited pages 4, 9, 40, and 41.)
[4] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-paramettrrizaton.. In International Conference on Machine Learning, 2019. (Cited on pages 9 and 43.)
[5] Shun-ichi Amari. Natural gradient works efficiently in learning. Neural Computation, 1998. (Cited on page 21.)
[6] Jonathan Baxter and Peter L. Bartlett. Infinite-horizon policy-gradient estimation. Journal of Artificial Intelligence Research, 2001. (Cited on page 1.)
[7] Amir Beck. First-Order Methods in Optimization. SIAM-Society for Industrial and Applied Mathematics, 2017. (Cited on pages and 23.)
[8] Amir Beck and Mare Teboulle. Mirror descent and nonlinear projected subgradient methods for convex optimization. Operations Research Letters, 2003. (Cited on pages 21 and 26.)
[9] 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
2019, (Cited on page I.)
[10] Jalaj Bhandari and Daniel Russo. Global optimality guarantees for policy gradient methods. arXiv preprint arXiv:/906.00 09766,,786, 2019. (Cited on page 23.)
[11] 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 pages 22, 23, and 25.)
[12] Shalabh Bhatnagar, Richard S. Sutton, Mohammad Ghavamzadeh, and Mark Lee. Natural actor-critic algorithms. Automatica, 2009. (Cited on pages and 22.)
[13] 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.)
[14] Sébastien Bubeck. Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning, 2015. (Cited on pages 3, 25, 31, and 32.)
[15] Semih Cayci, Niao He, and R Srikant. Linear convergence of entropy-regularizzed natural policy gradient with linear function approximation. arXiv preprint urXiv:2106.04090 2021. (Cited pages 2, 7, 8, 22, 25, 38, and 39.)
[16] Semih Cayci, Niao He, and R Srikant. Finite-time analysis of engropy-regularized neural natural actor-critic algorithm. arXiv preprint XXX.:222000008333, 2022. (Cited on pages 8, 9, 20, 22, 24, 39, 43, and 44.)

---

[17] 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.)
[18] Yair Censor and Stavros Zenios. Parallel Optimization: Theory, Algorithms, and Applica- tions. Oxford University Press, USA, 1997. (Cited on page 3.)
[19] Gong Chen and Mare Teboulle. Convergence analysis of a proximal Likke minimization algorithm using bregman functions. SIAM Journal on Optimization, 1993. (Cited on pages 2,
5, 6, 21, 31, and 37.)
[20] Zaiwei Chen and Siva Theja Maguluri. Sample complexity of policy-based methods under off- policy sampling and linear function approximation. International Conference on Artificial Intelligence and Statistics, 2022. (Cited on pages 2, 22, 25, and 39.)
[21] Zaiwei Chen, Sajad Khodadadian, and Siva Theja Maguluri. samie- analysis of off- policy natural actor-critic with linear function approximation. IEFE Control Systems Letters, 2022. (Cited on page 38.)
[22] Ashok Cutkosky and Francesco Orabona. Momentum-based variance reduction in non-convex sgd. Advances Neural Information Processing Systems, 2019. (Cited on pages 22 and 23.)
[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.)
[24] 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.)
[25] 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.)
[26] 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.)
[27] Simon S. Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh. Gradient descent provably optimizes over peer-paaametrized neural networks. In International Conference on Learning Representations, 2019. (Cited on page 9.)
[28] Lasse Espeholt, Hubert Soyer, Remi Munos, Karen Simonyan, Vlad Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. International conference on machine learning, 2018. (Cited on page 7.)
[29] Ilyas Fatkhullin, Jalal Etesami, Niao He, and Negar Kiyavash. Sharp analysis of stochastic optimization under global Kurdyka-lojasiewicz inequality. In Advances in Neural Information Processing Systems, 2022. (Cited on page 23.)
[30] Ilyas Fatkhullin, Anas Barakat, Anastasia Kireeva, and Niao He. Stochastic policy gradient methods: Improved sample complexity for fishee-non-deggeneaate policies. arXiv preprint arXiv:2302.01734 2023. (Cited on page 23.)
[31] 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.)
[32] 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.)
12

---

[33] 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.)
[34] 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.)
[35] Ian Goodfellow, Yoshua Bengio, and Aaron Courville. Deep learning. MIT press, 2016. (Cited on page 4.)
[36] 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.)
[37] Feihu Huang, Shangqian Gao, and Heng Huang. Bregman gradient policy optimization. In International Conference on Learning Representations, 2022. (Cited on pages 22 and 24.)
[38] Arthur Jacot, Franck Gabriel, and Clement Hongler. Neural tangent kernl:: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, 2018. (Cited on page 9.)
[39] Ziwei Ji, Matus Telgarsky, and Ruicheng Xian. Neural tangent kernels, transportation map- pings, and universal approximation. In International Conference on Learning Representations. 2019. (Cited on pages 9, 38, 43, and 44.)
[40] 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 pages 22 and 23.)
[41] Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In International Conference on Machine Learning, 2002. (Cited on page 33.)
[42] Sham M. Kakade. A natural policy gradient. Advances in Neural Information Processing Systems, 2002. (Cited on pages 1, and 21.)
[43] William Karush. Minima of functions of several variables with inequalities as side conditions. Master's 's thesis, Department of Mathematics, University of Chicago, Chicago, II., USA, 1939. (Cited on page 26.)
[44] Michael J. Kearns and Daphme Koller. Efficient reinforcement learning in factored mdps. In International Joint Conference on Artificial Intelligence, 1999. (Cited page 23.)
[45] Sajad Khodadadian, Zaiwei Chen, and Siva Theja Maguluri. Famile analysis of off- policy natural actor-critic algorithm. In International Conference on Machine Learning, 2021. (Cited on pages 22 and 24.)
[46] Sajad Khodadadian, Prakirt Raj Jhunjhunwala, Sushil Mahavir Varma, and Siva Theja Magu- luri. On the linear convergence of natural policy gradient algorithm. In IEEE Conference on Decision and Control, 2021. (Cited on pages 22 and 25.)
[47] Sajad Khodadadian, Prakirt Raj Jhunjhunwala, Sushil Mahavir Varma, and Siva Theja Magu- luri. On linear and super-linear convergence of natural policy gradient algorithm. Systems and Control Letters, 2022. (Cited on pages 22 and 25.)
[48] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference Learning Representations, 2015. (Cited page 10.)
[49] Vijay Konda and John Tsitsiklis. Actor-critic algorithms. In Advances in Neural Information Processing Systems, 2000. (Cited on page 1.)
[50] 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.)
13

---

[51] 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.)
[52] Harold W. Kuhn and Albert W. Tucker. Nonlinear programming. In Proceedings of the Second Berkeley Symposium on Mathematical Statistics and Probubility, 1951. (Cited on page 26.)
[53] Solomon Kullback and Richard A, Leibler. On Information and Sufficiency. The Annals of Muthematical Statistics, 1951. (Cited on page 1.)
[54] 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.)
[55] Guanghui Lan. Policy optimization over general state and action spaces. arXiv preprint arXiv:2211.16115, 2022. (Cited on pages 6, 8, 21, and 24.)
[56] 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.)
[57] Yan Li and Guanghui Lan. Policy mirror descent inherently explores action space. arXiv preprint 2rX3v:2303.04386, 2023. (Cited on pages 6, 10, and 30.)
[58] Yan Li, Tuo Zhao, and Guanghui Lan. Homotopic policy mirror descent: Policy convergence, implicit regularization, and improved sample complexity. arXiv preprint 1iXiv:22011.9455 2022. (Cited on pages 22 and 25.)
[59] 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.)
[60] 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.)
[61] Yanli Liu, Kaiqing Zhang, Tamer Basar, and Wotao Yin. An improved analysis of (variance- reduced) policy gradient and natural policy gradient methods. Advances in Neural Information Processing Systems, 2020. (Cited pages 8, 22, 23, 24, and 38.)
[62] Stanislaw Lojasiewicz. Une propriété topologique des sous-ensembles analytiques réels. Les équations aux dérivées partielles, 1963. (Cited on page 22.)
[63] 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.)
[64] Saeed Masiha, Saber Salehkaleybar, Niao He, Negar Kiyavash, and Patrick Thiran. Stochastic second-order methods improve known sample complexity of SGD for gradient-dominated functions. In Advances in Neural Information Processing Systems, 2022. (Cited on page 23.)
[65] Jincheng Mei, Chenjun Xiao, Csaba Szepesvari, and Dale Schuurmans. On the global con- vergence rates of softmax policy gradient methods. International Conference on Machine Learning, 2020. (Cited on page 22.)
[66] Jincheng Mei, Yue Gao, Bo Dai, Csaba Szepesvari, and Dale Schuurmans. Leveraging non- uniformity in first-order non-convex optimization. In International Conference Machine Learning, 2021. (Cited on page 23.)
[67] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, loannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint 1rXiv:1312.5000 2013. (Cited page 9.)
14

---

[68] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep rein- forcement learning. Nature, 2015. (Cited on pages 7 and 9.)
[69] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lilli- crap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International Conference on Machine Learning, 2016. (Cited on page 1.)
[70] Arkadi Nemirovski and David B. Yudin. Problem Complexity and Method Efficiency in Optimization. Wiley Interscience, 1983. (Cited on pages 1, 3, and 21.)
[71] Yurii E. Nesterov and Boris T. Polyak. Cubic regularization of Newton method and its global performance. Mathematical Programming, 2006. (Cited on page 23.)
[72] Gergely Neu, Anders Jonsson, and Vicenç Gómez. A unified view of entropy-regularian markov decision processes. arXiv preprint arXiv: 1705.07798, 2017. (Cited on page 22.)
[73] Lam M. Nguyen, Jie Liu, Katya Scheinberg, and Martin Takáč. 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.)
[74] Francesco Orabona.
modern introduction online learning, 2020. URL httttp/pp////p (Cited on pages 6, 10, 30, and 31.)
bu.
[75] Jan Peters and Stefan Schaal. Natural actor-critic. Neurocomputing, 2008. (Cited on pages and 22.)
[76] Boris T. Polyak. Gradient methods for the minimisation of functionals. USSR Computational Mathematics and Mathematical Physics, 1963. (Cited on page 22.)
[77] Bruno Scherrer. Approximate policy iteration schemes: A comparison. In International Conference on Machine Learning, 2014. (Cited on page 39.)
[78] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference Machine Learning. 2015. (Cited on pages 1, 19, and 22.)
[79] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. In International Conference on Learning Representations, 2016. (Cited on pages and 10.)
[80] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint (707.00777.06444 2017. (Cited on pages 1, 19, and 22.)
[81] Shai Shwarz-Shwartz, Shaked Shammah, and Amnon Shashua. Safe, multi-agent, reinforce- ment learning for autonomous driving. arXiv preprint 1rXiv:16,0.03295, 2016. (Cited on page 1.)
[82] 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 pages 1, 6, 8, 21, 22, and 24.)
[83] 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.)
[84] David Silver, Julian Schritwieser, 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.)
15

---

[85] Wen Sun, Nan Jiang, Akshay Krishnamurthy, Alekh Agarwal, and John Langford. Model- based in contextual decision processes: Pac bounds and exponential improvements over model-free approaches. In Conference on Learning Theory, 2019. (Cited on page 23.)
[86] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press,
2018. (Cited on pages 7 and 18.)
[87] 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 3, and 5.)
[88] 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.)
[89] 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 Artificial Intelligence and Statistics, 2022. (Cited on pages 5, 6, 8, 19, and 23.)
[90] Lingxiao Wang. Qi Cai, Zhuoran Yang, and Zhaoran Wang. Neural policy gradient meth- ods: Global optimality and rates of convergence. In International Conference on Learning Representationss 2020. (Cited on pages 7, 8, 9, 22, and 24.)
[91] 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.)
[92] Ronald J. Williams and Jing Peng. Funetion optimization using connectionist reinforcement learning algorithms. Connection Science, 1991. (Cited on page 1.)
[93] Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 2022. (Cited pages 1, 2, 6, 7, 8, 9, 19, 20, 21, 22, 24, 25, 28, 29, 31, 33, 37, 38, 39, and 40.)
[94] 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.)
[95] 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.)
[96] 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.)
[97] 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.)
[98] 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.)
[99] Andrea Zanette, Ching-An An Cheng, and Alekh Agarwal. Cautiously optimistic policy optimiza- tion and exploration with linear function approximation. In Conference on Learning Theory, 2021. (Cited on pages 20, 22, and 24.)
[100] 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-11) 2023. (Cited on pages 1, 8, 22, 25, and 39.)
16

---

[101] Junyu Zhang, Alec Koppel, Amrit Singh Bedi, Csaba Szepesvari, and Mengdi Wang. Varia- tional policy gradient method for reinforcement learning with general utilities. In Advances in Neural Information Processing Systems, 2020. (Cited on pages 7 and 23.)
[102] Junyu Zhang, Chengzhuo Ni, Zheng Yu, Csaba Szepesvari, and Mengdi Wang. On the convergence and sample efficiency of variance-reduee reducced policy gradient method. In Advances in Neural Information Processing Systems, 2021. (Cited on page 23.)
17

---

Appendix
Table of Contents
A Related work
18
B Equivalence of (9)-(10) and (11) in the tabular case
25
AMPO for specific mirror maps
26
Deferred proofs from Section 4.1
31
E Discussion of the first step (Line 1) of AMPO the compatible function approximation framework
37
F Discussion the second step (Line 2) AMPO the Bregman projection
37
G Discussion Assumption (A1) the approximation error
38
H Discussion Assumption (A2) the concentrability coefficients
39
Discussion on Assumption (A3) the distribution mismatch coefficients
40
Sample complexity for neural network parameterization
40
Here we provide the related work discussion, the deferred proofs from the main paper and some additional noteworthy observations.
A Related work
We provide extended discussion for the context of work, including comparison of different PMD frameworks and comparison the convergence theories of PMD in the literature. Furthermore, discuss future work, such extending 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, give 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 8 E S,
π
argmax(m. p++1 m,EA(A)
Vh(=T), n.) (T m.),
(20)
where f/
( x Q2 following Line of Algorithm 1. The proof of (20) can found in Lemma F. in Appendix F.
f++1
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 the estimated Q-functio.. That is, for all E S,
π
argmax (Q', Te). R,EA(A)
(21)
AMPO behaves like GPI when we perfectly approximate f!+1 to the value Q¹ (e.g. when we consider the tabular case) and f)t - +oo (or ne 0) which the case with the use of geometrically increasing step-size schedule.
18

---

Mirror Descent Modified Policy Iteration (MD-MPI) [33]. Consider the full policy space II (AA).. The MD-MPI's update as follows:
T ++1 E argmax (0' T.) - (Ts, T.). Vs S.
(22)
x,EA(A)
In this case, the PMD framework of Xiao [93], which special case of AMPO, recovers MD-MPI with the fixed siep-size 1)t = 1. Consequently, Assumption (A1) holds with Eapprox = 0, and we obtain the sublinear convergence of MD-MPI through Theorem 4.3, which is
(n) T
Σ [V() KT
S 1 Dj (1 (1-1))
V/ )
As explained later in Appendix H that the distribution mismatch coefficient 2/ is upper bounded by ( - ). we obtain an average regret of MD-MPI as O (1- (1-))))) 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:
π ++
argmnax Es [(A5, n.>]
(23)
TEI
such that Eg [Dx (T π.)] S 6,
where A = Q - Vt represents the advantage function, the negative entropy and 5 0. Like GPI, TRPO is equivalent to AMPO when at each time I, the admissible policy class is II {Tr ((4)S; Eswdy Da TT) s 5}, and we perfectly approximate Q' with т)c - +00.
Proximal Policy Optimization (PPO) [80]. The PPO's update consists of maximizing a surrogate function depending the policy gradient with respect to the new policy. Namely,
++1 argmax Eswdy [L(T,; m()]
(24)
#ell
with
L(T,,, m²) = Eowx [min (r" ()), (1..
where r (s. a) [ss, a) / (s,, a) is the probability ratio between the current policy x' and the new one, and the function clip(r" (s, a), 1= clips the probability ratio r" a) to be more than 1 + and no less than - e. PPO has also a KL variation [80, Section 4], where the objective function is defined as
L(T,, m = (A - (TT nT),
where h is the negative entropy. In an exact setting and when II = ((A)S, the KL variation of PPO still differs from AMPO because inverts the terms in the Bregman divergence penalty.
Mirror Descent Policy Optimization (MDPO.) [88]. π A-1
The MDPO's update is follows:
argmax Basa (n,44 n.) (T m.)]
(25)
sEll
where II is parameterized policy class. While it is equivalent to AMPO in exact setting and when II = A (AA),, 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:
++1 argmax Wegwed [V* (n) (ne V V*( ( - (m m))l
(26)
n²: BEE
E argmax Exna; Q' T (TT T)),, t:: AE@
The second line is obtained by the definition of V and the policy gradient theorem (3). The discussion the same as the previous algorithm.
19

---

Mirror Learning [51]. The on on-policy version of the algorithm consists of the following update:
π [+1 argmax B,nd;, (00s , Ta) D(T +[))
(27)
where (T(=") is policy class that depends on the current policy x¹ and the drift functional D is defined as a map D: : (AA) x ((A) + R such that D(T,, *a) 2 and ( # = 0. The drift functional D recovers the Bregman divergence as a particular case, in which case Mirror
Learning equivalent to AMPO in an exact setting and when IT A(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 framework for policy optimization AMPO. For AMPO, we establish in Theorem 4.3 both (1/TT) convergence guarantee by using a non-decreasing step-size and linear convergence guarantee by using 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. 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 e-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 general parameterization class J
As discussed in Section 4.1, the results of Theorem 4.3 hold for 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 sample complexity analysis of AMPO which could improve that of Zanette et al. [99] and Hu al. [36]. As we discuss in Appendix A.3, this 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 I simplifies the policy optimization step into a regression problem; and Lemma 4.1 establishes a key result for policies belonging 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 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 con- sidering the log-linear policy parameterization and the entropy mirror map. On the contrary, AMPO solves 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 I of Algorithm 1. This innovation is key to the generality of the algorithm, as allows AMPO to employ any mirror map and any parameterization class. Moreover, AMPO is computationally efficient for 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 neural network whose
20

---

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 as- sumptions. Xiao [93] requires an -nom on the approximation error of Q', i... 1''' - Q' I εapprox. for all t I T. Alfano and Rebeschini [2] and Yuan et al. [98] require an Ly-norm bound on the error of the linear approximation of Q', i.e. | Φ ('(liz("') I Eupprox for some feature mapping o S x A - I and vector w E R", for all = T. Our approximation error Sapprox 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 repre- sentation 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 iter- ations. In addition, we improve the concentrability coefficients of Yuan et al. [98] by defining the expectation under an arbitrary state-action distribution of 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 parameteriization 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-powe 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 sample complexity analysis of AMPO when employing shallow neural networks that improves upon previous state-of-the thet-arrt 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 comparison wih Lan [55]. Our diffences can be outlined in two points.
Lan [55] propose 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, 'V((T') and the second step is to find the policy through 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 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 -nom over the action space. Let
Edct Egwe'
[|[ - Q¹ - 7/ VV((m)))
Esto Esnet
|| - -f-f!!++1! 1111113]
where the expectation E[ is taken w.r.t. the stochasticity of the algorithm employed obtain ft+1, Lan [55] assume that both Edet and Esto are bounded for all iterations In contrast, our assumptions are weaker as they are required to hold for the define in Section 3. Additionally, Lan [55] establishes O(1/VT) convergence rate for their algorithm without regularization and a O(ITT) convergence rate the regularized case, in both cases using bounded step-sizes. We improve upon these results by obtaining a (1/T)) convergence rate without regularization and a linear convergence rate.
we
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 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
21

---

Neu et al. [72], is a special case of PMD. For instance, PMD and NPG are the building blocks of the -tate-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 (1² EE [V" (4)] S c 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 al. [82] establish a ((IVVT) convergence rate for unregularized NPG and (L/TT) for regularized NPG. Agarwal et al. [1]. Khodadadian et al. [45] and Xiao [93] improve the convergence rate for unregularized NPG and NAC to (1)TT) 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, 7anette et al. [99] and Hu et al. [36] achieve ((IvVT) convergence rate by developing variants of PMD methods for the linear MDP [40] setting. The same O(1/VT) 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 (1/TT) for log-linear policies. For smooth policies, the convergence rate is later improved (1/TT) either by adding an extra Fisher-non canher-non-deee non-degeneracy condition on the policies [61] or by analyzing NAC under Markovian sampling [94]. Yang al. [95] and Huang et al. [37] consider Lipschitz and smooth policies [97], obtain ((/VTT) convergence rates for PMD-type methods and faster (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/ VT) 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(L/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 decaying step-size a constant step-size. Under these step-size schemes, our AMPO's ' O(1TT) sublinear convergence rate the state of the art: recovers the best-known convergence rates in the tabular setting [33, 93] without regularization; improves the ((L/VT) 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-linean policies [98]; it matches the OL/TT) sublinear convergence rate for smooth and 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; 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 for an overview of recent sublinear convergence analyses of NPG/PMD.
cahhernoondegencee
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, AMPO's s linear convergence rate is the state of the art: not only it recovers the best- known 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 thethe- art linear convergence rate for PG-based methods with general parameterization, including the neural network parameterizations. We refer to Table for an overview of recent linear convergence analyses of NPG/PMD.
Alternatively, by exploiting a Polyak-LLojsiewvic (PL. condition [76, 62], fast linear convergence results can be achieved for PG methods under different settings, such as linear quadratic control problems [3] I] and softmax tabular policies with entropy regularization [65, 97], The PL condition
22

---

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 PI. 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, 30J when the PL condition is replaced by the weak PL condition [97], which is satisfied by Fisher-non-degenerale 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-degeet -engeneeeee policies combined with variance reduction [64].
Outside of the literature focusing on finite time convergence guarantees, Vaswani al. [89] and Kuha
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 and 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 I of Algorithm 1. For instance, by using the dual averaging algorithm [7, Chapter 4J, it could be possible to replace the term Vh(mT) with f for all s E S, to make the computation of the algorithm more efficient. That is, it could be interesting consider the following variation of Line in Algortihm 1:
p - Q'
76-1 T)
S Capprox.
(28)
Lg(u*)
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 I 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 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- -egeeeerre 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 study the application of AMPO to the offline setting. In the main text, we have discussed how to extend Algorithm and Theorem 4.3 to the offline setting, where 2 can be set as the atatenaction distribution induced by 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 x' and the behaviour policy. Several methods have been introduced to deal with this issue, such as constraining the current policy 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-fiel 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 -funccion
23

---

Table 1:
Overview of sublinear convergence results for NPG and PMD methods with constant step-size in different settings. The dark blue cells contain our new results. The light blue cells contain previously known results that we recover as special cases of our analysis. The pink cells contain previously known results that we improve upon by providing a faster convergence rate. The white cells contain existing results that have already been improved by other literature that could not recover under our general analysis.
Algorithm
Rate
Comparisons to our works
Setting: Softmax tabular policies
Adaptive TRPO [82]
0(1/V)) They employ regularization
O(/TT) We have weaker approximation emror Assumption (A1) with Ly instead &oo norm
Tabular off-policy NAC [45]
Tabular NPG [1]
O(l/T)
MD-MPI [33]
O(1/T) We match their results when (s, a) O.,a-
We recover their results when f (s. a) = 0z
Tabular NPG/
O(1/T) we have a weaker approximation error Assumption (A1) with Ly instead of loo norm.
projected Q-descent [93]
Setting: Log-linear policies
Q-NPG [1]
O(1/VT)
We recover their results when f (s, a) is linear to 6
Q-NPG/NPG [98]
O(1/T)
Setting: Softmax two layer neural policies
Neural PPO [60]
O(I/TT)
Neural NAC [90]
O(I/V))
Regularized neural NAC [16]
O(1/T)
We match their results without regularization.
Setting: Linear MDP
NPG [99, 36]
0(1/VT)
Setting: Smooth policies
NPG [1]
0(1/V))
NAC under Markovian sampling [94]
O(L/T)
NPG with
O(L/T)
Fisher-non-degenerate policies [61]
Setting: Lipschitz and Smooth policies
Variance reduced PMD [95, 37]
O(1/T)
We match their results without variance reduction.
Setting: Bregman projected policies with general parameterization and mirror map
We match their results without regularization;
Regularized PMD [55]
O(I/T)
we have a weaker approximation error Assumption (A1) with La instead of & norm.
AMPO (Theorem 4.3, this work)
O(/T))
24

---

Table 2: Overview of linear convergence results for NPG and PMD methods in different settings. The darker cells contain our new results. The light cells contain previously known results that we recover as special cases of our analysis, and extend the permitted concentrability coefficients settings. The white cells contain existing results that we could not recover under our general analysis.
Algorithm
Reg.
C.S.
A.I.S.
N.I.S.
Error assumption
Sctting: Soltmax tabular policies
NPG [17]
PMI [TO0]
NPG [54]
Eno
NPG [58]
800
NPG [TT]
NPG [46, 47]
t,o
NPG Projeroeected Q-descent [93]
8.00
Setting: Log-linear policies
NPG [15]
L22
Off-policy NAC [20]
foo
Q-NPG
L2
Q-NPG/NPG [98]
L
Setting: Bregman projected policies with general parameterization and mirror map
AMPO (Theorem 4.3, this work)
La
Reg.: regularization; C.S.: constant step-size; A.I.S.: Adaptive increasing step-size; N.L.S: Non-adaptive increasing step-size.
** Error assumption. to means that the approximation error assumption uses the -nomm;; and L2 means that the approximation error assumption the weaker Lę norm.
as accurate as possible, such as using offline data for training, and to construct adaptive representation learning for F to closely approximate Q-function, that Capprox 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 found in Bubeck [14, Chapter 4.2]. However, for the sake of completeness, 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 function V(),, thut is,
y 11 Vh(x') - VV
(29)
x 1+ Proj
(30)
Then the mirror descent update can rewritten as
1+1 E argmin (a, (V(x) Dn(a,
(31)
=EX
25

---

Proof. From definition of the Bregman projection step, starting from (29) we have
x²+1
(Vh') (Vh"(y+t1) argmin D (z, Vh (/(())
argmin Vh(x) - Vh(Vh' (
-V---hh-0 -
SEX
argmin Vh(x) - - - x
-- -
E
SEX
argmin Vh(x) - (1,
aEX
(29)
argmin (h - (a, VA(T*) - T3 V=))
nEX
argmin (z, VV (V(x))==e) + Vh(x) - Vh(x') - (Vh(x"),
=EX
argmin
n/ (x, VV (2))))==z) (xr,
xEX
where the second and the last lines are both obtained by the definition the Bregman divergence.
The one-step update in (31) is often taken as the definition of mirror descent [8], which provides proximal view point of mirror descent, i.e., gradient step in the primal space with a regularization of Bregman divergence.
C AMPO for specific mirror maps
In this section, we give the derivations for Example 3.2, which is based on the Krash-KKunnTuuun (KKT) conditions [43, 52], and then provide details about the w-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
h(TT)
Σ Taa log(z(a s)),
VT, A(A) and Vs S.
aeA
For every state a E S, we solve the minimization problem
n E argmin (n(Ms, (n(()
TEE((A)
through the KKT conditions. We formalize it as
argmin
DA(TT, Vh Vh" (nf"))
F.ERIAl
subject to (T 1) = 1,
(a 8) 2 Va
where 1 denotes vector in R with coordinates equal to 1 clement-wisely. The conditions then become
(stationarity)
log(=?) n5 (A, + 1)1 -- <a s) = Va A,
(complementary slackness)
(primal feasibility)
(Tf,1)=11 1, 7 (s) 20
AaEA.
(dual fcasibility)
e2 VaE
where log(m,) is applied element- X, R and cg R are the dual variables, and c, devotes the vector [ch] aEAA- It is easy to verify that the solution
exp(ns?)
T
[cxp(1f2)
26

---

with X, = log ( exp(nf((s.a)) - -1 and c = 0 for all a E A. satisfies all the conditions.
When (s, a) = 0s.a we obtain the tabular softmax policy aa s) & exp(nos.a) When (s, a) = H Ø(s.a) a) is a linear function, for f E R² and for a feature function o ::S. S x A + R. we obtain the log-linear policy (a 8) & exp(10" d(),a)). When f :S x A + IR is neural exp(nf"'(s,a)))
network, we obtain the softmax neural policy (a s) x exp(nf((s,a)).
C.2 More on w-potential mirror maps
In this section, we provide details about the w-potential mirror map class from Section 3.1, including the derivation of (14), several instantiations of w-potential mirror map mentioned Section 3.1 with their derivations, and an iterative algorithm to find approximately the Bregman projection induced by w-potential mirror map when an exact solution
available.
We give a different but equivalent formulation of Proposition 2 of Krichene et al. [50]. Proposition C.1. For (-o0, +oo and S O, an increasing C c--diffeomorphism @ (-oc, u) + (w, +o0) is called an w-potential if
lim o(x) too, [ (x)(x) 00.
lim
((x) w,
Let the mirror map h be defined as
(((as)
he (T(TT))
Σ
o 5- (x)dx.
aEAA"
We have that n². is solution to the Bregman projection
min
Proj ( (T(())112))
xEd,
and only if there exist normalization constant X R such that
T' = (f a) + x)),, Va E A,
(32)
and EnE a'(a ()) ===,, where for all E S and (=) = max(z, for E ER.
We can now use Proposition C. to derive (14).
Consider an w-potential mirror map hs associated with an w-potential ø. By definition, we have
Vho(T() 16 (a' (||aa))sall s))
(33)
Plugging (32) into (33), we have
11 (TT)
o ( s))LaEA
+ X.)))
mma (0- (61--1 +;) x (0)( aed
mmax (M, + N4
[max (n--1
(0) 15)] + X
where the third line is obtained by using the increasing property of Φ1 as Φ is increasing. Finally, plugging the above expression of Vho(T()) into Line 1, we obtain (14), which is
gi+1 argmingee B(s.))~or [(f (s, a) (', a) .. if 0)) [
where the term X; 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 constant term are equivalent and generate the same algorithm. We start with the negative entropy, which leads to closed solution for X and therefore for the Bregman projection.
27

---

Negative entropy.
Let x(x) = exp(x 1), which is an w-potential with w = 0, +oo, and
[ f (
S log(x) Idr [xr [ogg(())) = S +00.
The mirror map he becomes the negative entropy, as
((las) (log(2) 1)dr Σ T(a 8) log(TT(a) s)),
(T Σ
and the associated Bregman divergence becomes the KL divergence, i.e, Dhg (T)) = (L( Equation (17) follows from Equation (14) by the fact that
(0) log(0) + -1-
which means that max(M:-1f Φ (0) X1) 7/ element-wisely,
As we showed in Appendix C.I. the Bregman projection (Line of Algorithm 1) for the negative entropy has closed form which is x exp(n.fa*1).
In NPG with softmax tabular policies [1, 93], at time t, the updates for the policies have
=!+1 nh @ exp(n) E
(34)
When considering AMPO with (s, a) = Os.a E R, from (17), we obtain that for all 8, E S,A. have
g++1 argmin
8- Q' 1/-1 grl
= f = (s. 7h- fs
DERÍSIXIA
n
L2(0*)
mf = a) n m-
Q'
With the above expression, we have the AMPO's updates for the policies rewritten
m exp(n 1
exp(n Q' n-
n. o (x)(m Es
which recovers (34). In particular, the summation in the second line element-wise and the third line is obtained because of n³ a exp(n=-1f5), as shown in Appendix C.I.
In Q-NPG with log-linear policies [1, 98, 2], at time f, the updates for the policies have
1+ (a s) x x (a s) exp(n)(ss a) w²).
(35)
where :: x A - R is feature map, and
w E argmin (s,a)))dd [() - (ss,u) Tw)]
(36)
wERd
Like in the tabular case, when considering AMPO with A E Rdd. 1 a) = a(s,a) TO and o² = dr
from (17), we obtain that for all s, a E S, A, we have
27
a++1 argmii B Bee,aymag DER
Mh-1
(1 a) a)To² TOt)
((8,
1/.
2
(
Th-1 0e
argmin E (s,a)ndd
BeRa
Compared the above form with (36), obtain that
- = gt+1
7h-1 gt
n1 mww + 1
(37)
17.
28

---

Algorithm 2: Bregman projection for e-negative entropy
Input: vector to project z E IR -AA,, parameter e.
Initialize y cxp(z) element-wisely.
2 Let y( be the i-th smallest element of y.
Let 2 be the smallest index for which
(1 c(A) - 1+ 1))
Ev ~00> >0.
521
Set
Σ 2>1 (()
A = 1 c(1 A| - 1)
Return: the projected vector (a (-c - ya/ A))
So, the AMPO's updates for the policies can be rewritten as
π +++ (a s)
x axp(no(ss.a) T
(37) explo(s, a) (new + +-100)) x z'(a 5) exp(n)o(s, a) w²),
where the last line obtained because of n exp(m)-1f)) as shown in Appendix C.I, and recover (35).
We next present the squared fo-norm and the e-negative entropy. For these two mirror maps, the Bregman projection can be computed exactly but has no closed form.
Squared norm. Let be the identity function. The mirror map hi becomes the squared n--nrmm up to constant term, as
h (T(
-E acA"
Σ
(als)
x dx
1 (7)
aEA
The associated Bregman divergence becomes the squared Euclidean distance, i.e., Dr. (T,, R,)
|| * and Vh (-) the identity function. The update in (15) follows immediately and the Bregman projection step with the Euclidean distance becomes, for all 8 E S,
n
Proj hy ((A)
(Vh' (nif++)) Proj (neft f argmin ||p nt.ft+111
(38)
PEΔ(A)
In the projected-Q descent for tabular policies developed by Xiao [93], time t, the updates for the policies are
A ++1 argmin PEA(A)
| + +n -pll E.
(39)
When considering AMPO with f 0) = 0s.a and = RISIi|A (15) is solved with
ss a) = 1 = Q', a) + ''(a s).
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].
c-negative entropy [50]. Let 2 0 and define the c-exponential potential as (bx) exp(x - 1) The mirror map his becomes
(a|s)
hT(T)))= Σ
(log(x+e++I)dz Σ [TT(a 8) In( 8) + -)) +e)]
29

---

Algorithm 3: Bregman projection for w-potential mirror maps
Input: vector to project a E IR AA,, w-potential o, precision e. Initialize
== (1) max Ta aEA 2 == -
max Fa atA
Define k(v) = (a(Ø (I, + +))))
while |(((D) - -(())ll > do
Let ot t (1 + 22/2
EoEA Ja (ut) > then
else
D+
V
8 Return ((D)
An exact solution the associated projection then be found in Ō(.Al) computations using Algorithm 2, which has been proposed by Krichene et al. [50, Algorithm 4]. Additionally, following (14), the regression problem in Line of Algorithm becomes
g++1 E argmin || - Q' - m max(m=11f 1 + log(e) >X)1
Age
where X can be obtained through Algorithm 2.
The Bregman projection for generic mirror maps can be computed approximately in (IIA) compu- tations 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 the optimization literature, but do not lead to an exact solution to the Bregman projection. We leave them object for future work.
Negative Tsallis entropy [74, 57J.
Let q > and define e as
exp(x 1)
1f
(x)x)
(0 (0=1)=)
else.
The mirror map has, becomes the negative Tsallis entropy, that is
hag (Tx) Σ (a s) log, (I(a s));
(40)
where, for y > 0.
log(y) if = 1,
log, (v)
else.
q # 1 and following (14), the regression problem in Line of Algorithm becomes
g++l argmin eEe
Q¹ 7 fr
17
L2(er)
Hyperbolic entropy [34]. Let b > 0 and define o as
(v(x) = biih((r)
The mirror map hag becomes the hyperbolic entropy, that
has (Tx) = Σ T[a s) arcsinh(T(a s)/b) V m[aa 8)2 +
(41)
acA
and, following (14), the regression problem in Line 1 of Algorithm I becomes
g++1 E argmin
f Q' me max(m--1 1"" -> (x(()))
0e0
30

---

Hyperbolic-tangent entropy.
Inspired by the hyperbolic entropy, we consider @ as
(xx) (anh(x)/22+000
The mirror map hi becomes
1
hT(T,) E( s) 1) arctanh(2T(a
-) 2
+ log <a s) T(a s)), (42)
aEA
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 of Algorithm becomes
0²+1
aargmin 0E
5 Q' 1
T)t
L2(e²)
Regarding the limitations of the w-potential mirror map class, mirror maps that cannot be recovered using w- o--poenntils: h(x) = 2 Ar [32], for some positive- definite matrix A, which generates the Mahalanobis distance, and p-norms, i.e. h(x) = || [74].
aware of two previously used
Note that the case where h(x) |I>|| can be recovered by setting ((x) (pz))/((-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 this setting because our analysis does not use the policy gradient theorem (3), which has a different expression for deterministic policies [83].
the action that corresponds to the maximum value of fg Our results hold in
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 c A is a closed convex set, we say a function y + IR is a mirror map if satisfies the following properties:
(i) à is strictly convex and differentiable:
(ii) f is essentially smooth, i.e., the graident of h diverges the boundary of V, that is lim [[h(x)] + oo;
(iii) the gradient takes all possible values, that Vh())) = R
To prove Lemma 4.1, we also need the following rather simple properties, i.c., the three- three-point identity and the generalized Pythagorean theorem, satisfied by the Bregman divergence. We provide their proof' for celf-containment.
Lemma D.1 (Three-poie identity, Lemma 3.1 in Chen and Teboulle [19]). Let h be mirror map. For any a, b the relative interior of and E V, we have that:
Da(c, a) + D)(a,b) Dí, (c, (Vh(b) Vh(a),
(43)
Proof. Using the definition of the Bregman divergence Dr., we have Vh(aa,, e a) h(()) h(a) (c,c,a), (Vh(b), b) -h(a) h(b) (a,a,
(44)
(45)
(Vh(b), c b) -((c) h(b) Da(c, Subtracting (44) and (45) from (46) yields (43).
(46)
31

---

Lemma D.2 (Generalized Pythagorean Theorem of Bregman divergence, Lemma 4.1 in Bubeck [14]). Let X cy be closed convex set. Let h be mirror map defined on Y. Let æ X, y E y and 1¹" Projx()) then
(Vh(v*) vh(y).y" x) s O,
which also implies
Dh (2,y") + D (2' .1) - (x.
(47)
Proof.
From the definition of y", which is
y argmin Da (y' v): y'EX
and from the first-order optimality condition [14, Proposition 1.3), with D'''''' Vh(y') Vh(y), for all y/ e e
we have
(Vh(y") (h(y))-(w)) which implies (47) by applying the definition of Bregman divergence and rearranging terms.
V Day'' .o''''')vv=> "y - <> 0
Now we are ready to prove Lemma 4.1.
Lemma D.3 (Lemma 4.1). Let y C. RlA be a closed convex set with ((A) V. For any policies
π E ((A)S and π in the relative interior of ((A)s, any function fr with E O, any E S and for 7 we have that,
(ns Vh((a), n, - n.) S DA(Ts; Ta) DDi (T,(T,
where t is induced by 1 and according to Definition 3.1, that for all E S.
n. Proji (Vh (nf2)) = argmin Da(T', (nff))
(48)
x(EM(A)
Proof. For clarity of exposition, let p. = Vh^(n Plugging a = R, b = and c == m, in the threc-point identity lemma D.I, we obtain
Du(T,, F.) (T. P&) + Dr ( Ps) = (V)(T.) Vh(p,), - nz) Similarly, plugging 7., b Ps and
(49)
म. in the three-point identity lemma D.I, we obtain
Da R8) (=t Po) (T.. (Vh(H,) - Vh(ps), R. #.)
(50)
From (49), we have
DA(Te, TT)) Dr( Po) + Dí (TT.
(Vh(TT)) Vh(pp), F, - TT))
(Vh(TT)) Vh(pp), To - (Vh(T..) - Vh(p.), T. -Ta)
DA F.) DA Pa) D + (Vh(F,) Vh(p,), न. - TT))
By rearranging terms, have
(T(TT: R.) (,, R.) - (T., P&) ) + ( P.) (Vh(#.) Vh(p.), म. - -g)- (51) (T.,
From the Generalized Pythagorean Theorem of the Bregman divergence in Lemma D.2, also known as non-expansivity property, and from the fact that म. Proj (p,), we have that
DA(TT. #.) + ( P.) < DA (TesPs)
-D, p.) +
S - DD 7.).
Plugging the above inequality into the left hand side of (51) yields
Dn(T., R.) - (m( F.) (ms 2 (Vh(a.) - Vh(ps), - a.
-m)),
which concludes the proof with Vh(ps) = ns..
We also provide alternative proof of Lemma 4. later in Appendix F.
32

---

D.2 Bounding errors
In this section, we will bound error terms of the type
E Eg [Q" a) + Ti [V
(52)
where (d TT) {(dy, T" (dft) )) (d
, (d+ x }')} These error terms appear in the forthcoming proof' of our results and directly induce the error floors the convergence rates.
In the rest Appendix D, let g² :S A R such that, for every & e S,
d = I 111 V((TT) E ERA!
So (52) can be rewritten
E Berdy away, (s, + VVh(=)) (
Bowdy awa, [Q'(s, dsss (53)
To bound it, let (u)) be sequence of distributions states and actions. By using Cauchy- Schwartz's 's inequality, we have
E Egndy [Q* as) (())((aa))
/ SES,REA V)((s,a)
(s)
/ SESS
(d)(s){(a|<)))
f JSES,OEA
('ss,))(('(ssa)) -
a((sa)
(1
Blandmor
B(ara)err (() as)
('(80)
Coll ((Q"(s. (1,
where the last line is obtained by Assumption (A2). Using the concavity of the square root and Assumption (A1), we have that
E
Bandyyaan (s) g'(s S Cyeaaprro
(54)
D.3 Quasi-monotonie updates - Proof of Proposition 4.2
In this section, we show Proposition 4.2 with proof that the AMPO updates guarantee a quasi- monotonic property, i.e., non-decreasing property up certain error floor due 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 analysis and well known result in the RL literature. Here we use a particular form the lemma presented by Xiao [93, Lemma 1].
Lemma D.4 (Performance difference lemma, Lemma in [93]). For any policy n, T' E A(A) and 1 ASS).
(p) V (1) 1 E Bondy [ .T - -
For clarity of exposition, we introduce the notation
T =:= 2 CoEaprox
Proposition 4.2 characterizes the non-decreasing property of AMPO. The error bound (54) Ap- pendix D.2 will be used to prove the the result.
33

---

Proposition D.5 (Proposition 4.2). For the iterates of Algorithm 1, at each time & 2 0. we have
( m + +Di ( x(+11 Te 77
([ v'(n]] 2 E Et
Proof. Using Lemma 4.1 with n², f pett - r/., thus कं n + by Definition 3.1 and Algorithm 1, and n. = T we have
(ne _ - D/((TT (T =3) Dn (= m²) (ITTT +3
(55)
By rearranging terms and noticing D) (T m²)) = 0, we have
(n R / 2 Da (T ,T;) Dr (T TT + )
2 O.
(56)
Then, by the performance difference lemma D.4, we have
(1 (1-7)8(((+(()) V'O]] =
E
snut
[10% n, 1+1 x
E
[10' T3 T
snuth
+E
E
110's - 43 1. TT>?
(M [ + DDi (IT 12
E
E
snattt1
E E
smatth
[1Q2 =
1 t-1
T
(T+++ , + Din (T mlt1)
E
E
7( 7))
176
which concludes the proof after dividing both sides by (1 7).
The last line follows from
E
E smatt
[1Q'
d. m
r
E R 1
[Q((s, a) (s.
(57)
E E Boad
+1
aaw; (Q'(s, d'(s,a))
(44)
2 CiEenor 7(1 7),
(58)
where both terms are upper bounded by / CaEaprox through (54) with (d T) = (d/t x(1) and
(dī TT = (d++1 T²), respectively.
D.4 Main passage An important recursion about the AMPO method
In this section, we show 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
d sup
Ft
Los
EES
(s)
For two different time t, 2 let Dj, denote the expected Bregman divergence between the policy n' and policy n²', where the expectation taken over the discounted state visitation distribution of the optimal policy di that is,
D Eand; [Di(T's Tt))
Similarly, let Dj denote the expected Bregman divergence between the optimal policy nt and n² that is,
== = Bawd; m()]
Let A, = V (p) - V (14) be the optimality gap.
We now state the following important recursion result for the AMPO method.
34

---

Proposition D.6 (Proposition 4.4). Consider the iterates of Algoritim 1, at each time t 2 0, we have
D.++1 (1 - 7)me
Dj
Di+1
E
+ v/ (A++1 - - + 4
S (1 (1->)m (1 7)-
+ (1 + 1/)TT
Proof. Using Lemma 4.1 with π = n². 1 n n/, and thus t = nlt1 by Definition 3.1 and
Algorithm 1, and n, = n. we have that
(ng -o=- و D x²) D x xxx1 1.--py -Du(='t),at
which can be decomposed
(neg': n. x = + (ne x; - n S DA(x", (n , ,
Dís (1T n²)
Taking expectation with respect to the distribution di over states and with respect to the randomness of AMPO and dividing both sides by Th. we have
E Eandit [(u's TT' x
[1]
+ E Ea-ap [V mT --)]]]]]]
LED:
Di+1 DD (59)
We lower bound the two terms on the left hand side of (59) separately. For the first term, have that
[ssdd [(gh nT. x [
(56)
d'
[1gh nt. - x 25tt>]]
dy
++1
smatt
Loo
Mx+
E
E
Fandet [1Qts T 7T. nイtt
+1/+1
E
1- [igl Qu m3 - n [1+1
Ve+1 (1 - 7) [V'(M) (4) - v [w)]
[E [Band [1g Q2 o. - m []
+V++1
Ve+1 - 7 [V' (1) v(+ [n)] -1417(11 7)
M - 7E [A++1 - A.J (1 - 7);
where (a) follows from Lemma D.4. For the second term, we have that
E E Esndy [(gh nT [
E Eandy [1Q', -. [ + E Exnd; [9g Q2 ग =
B - 7) + E Bowdy 14 - Q2 x5= - -
B] - 7) -(1 7),
where (b) follows from Lemma D.4 and (c) follows similarly (57), i.e., by applying (54) twice with (d% T) (d x²) and (d .) = (d) T²).
Plugging the two bounds in (59), dividing both sides by (1 (1)7) and rearranging, we obtain
D
D
Di+1
+ 22+1 (A++1 - - - A S (1-7)m
+ T.
7)me
(1-7)m
From Proposition 4.2, we have that A4t1 A. < O. Consequently, since 1/+1 S 1/2 by the definition of y in Assumption (A3), one can lower bound the left hand side of the above inequality by replacing V++1 by uy, that is,
D
D
Di+1
1
E
+ Mu (De+1 - A. +
SE
+
7)
(1
7)th
(1
7)
which concludes the proof.
35

---

D.5 Proof of the sublinear convergence analysis
In this section, we derive the sublinear convergence result of Theorem 4.3 with non-decraaiing step-size.
Proof. Starting from Proposition D.6
D.+1
DJ
Di+1
E
+ U/ (Ae+1 - A.) + A
S
+ (1+ v/)T
(1 7)me
(1 7)
(1 7)
7/ S 11t11
DJ+1 [1 (1-
Dj
Di+1
E
Up (A141 A.) A
S
E
+ (1 21)
(60)
(1
7)m
(1
m)++1
Summing up from 0 to T and dropping some positive terms on the left hand side and some negative terms on the right hand side, we have
Σ E KT
< (1 - 7)70
Dj
+ Dn + T(1 (11++u)T 3
Dj
V/ 1-
+ +T(I + +1)).
(1 777o
Notice that Do - as a(s,a)= [0,1]. By dividing T on both side, we yield the proof of the sublinear convergence
(M) 1 ΣE S 1 7 ((1- - 7010
Dj
1 1-7/ --
) + +1++))) V/)T.
KT
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.
D (1-7)% on the left hand side, we have
Proof. Starting from Proposition D.6 by dropping
Di
D
E [u (A,+1 - A.) + A.] S E
+ (1 Un)T.
(1
- 7)m
(1-7)m1
(1-
Dividing Up on both side and rearranging, we obtain
Di+1
1
+(+ 11
1 1/
E Aut + (1 (unmee
<
E Aet
+
V/
(1 - 7) (vp- 1)
If the step-sizes satisfy 71+1 (P/
1) 2 n with 1/2 2 then
A141 + (1 - 7) 1 1) (1 (1- 1 - A. - ( 1) + (1+ ) T.
Diai
Dj
Now we need the following simple fact, whose proof straightforward and thus omitted. Suppose 0 < a < 1,b >>b>00 0 and nonnegative sequence {arhe2o satisfies
ac+1 S aa, + b
Vt 2 0.
Then for all / 2 0,
0, S a'an 0g
1
The proof of the linear convergence analysis follows by applying this fact with
Dr A + (1- (1- a== and b 1 + + T.
36

---

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 detail the novelty of the first step (Line 1) and the second step (Line 2) of AMPO this and the next section, respectively. Afterwards, we provide an extensive justification of the assumptions used Theorem 4.3 in Appendices G to
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 policy class such that log n² (a s) is 3-smooth function of for all E S, A. At each iteration t, update
g++1 = 0t + nw',
with
w' argmin ||A' - -w- Vo Volog T' (x(d) 1
(61)
V*(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' with V, log whereas the one in Line 1 of Algorithm is relaxed to approximate A with an arbitrary class of functions F Additionally, (61) depends on the distribution d' while Line I of Algorithm does not and allows off- off-policy updates involving an arbitrary distribution v as of independent of the current policy x².
where W and
F
Discussion of the second step (Line 2) of AMPO - -he Bregman projection
As mentioned in Remark 3.4, rewrite the second step (Line 2) of AMPO through the following lemma.
Lemma F.1. For any policy n. for any function f E F for 0, we have, for all S,
F argmin (p, (nf()) = #s E arggminnnff" + Vh(*s),p) + (D R.). PEA(A)
PEΔ(A)
Equations (12) and (20) are obtained by choosing म, m and n. मं for all 8, nl 01+1, and by changing the sign of the expression on the right order to obtain an argmax.
Proof. Starting from the definition of म. we have
7. E argmin (p. Vh (nf())
PEA(A)
E argmin h(p) ((Vh" (nf.)) (VI(V) vn"()!)) rEA(A)
E argmin h(()) (nf%.11 PEA(A)
E argmin (-n + Vh( p) + h(p) - h(T.)) (Vh(F.), -x) PEΔ(A)
E argmin(-nff + vh( p) + (D, pEA(A)
(62)
where the second and the last lines obtained using the definition of the Bregman divergence, and the third line is obtained using (4) (Vh(Th* (z"))" = x for all x e R(A)).
Lemmas F.1 and B. share a similar result, as they both rewrite the Bregman projection into the MD updates. However, the MD updates in Lemma B.1 exact, while the MD updates in AMPO involve approximation (Line 1).
Next, we provide alternative proof for Lemma 4.1 to show that it is the direct consequence of Lemma F. 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].
37

---

Lemma F.2 (Three-point decent lemma. Lemma 6 in Xiao [93]). Suppose that C C. R is closed convex set. f C + E is a proper, closed convex function, Du(:,)) is the Bregman divergence generated by a mirror map h. Denote rint dom h the relative interior of dom h. For any dom h, let
xt Earg
min u uEdombnne
[f(u)+ Dx(,x)))
Then x* rint dom NC and for any E dom h nc,
f(x*) + (x x)) S f(u) + Pa(u,x) - Dí (u,,
We refer to Yuan et al. [98, Lemma 11] for proof of Lemma F.2
Lemma 4.1 is obtained by simply applying the three-point descent lemma, Lemma F.2, (62) with xt *.. f(u) = (-nfa + Vh(T,), u), u and n. and rearranging terms.
In contrast, it may not be possible to apply Lemma F.2 to (11), as (I(@) 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 com- putation feasible. As mentioned in the proef idea in Page 8, this framework consists in upper-boundee the sub-optimality gap with an optimization eπror plus approximation error. Consequently. it is important that both error terms converge
in order to achieve convergence to global optimum.
Assumptions similar to Assumption (A1) are common in the compatible function approximation literature. Assumption (A1) encodes form of realizability assumption for the parameterization class T that is, we assume that for all 1 s T there exists a function f E T such that
||
n. - ( )F S Capprox-
When e 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]. is, however, possible to relax Assumption (A1). In particular, the condition
Σ E
E
g++ Q' ni- Vh=>"(("")) (l
Eaapprx
(63)
T
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. but only needs hold average T iterations. Similarly, the condition \
Σ (1- IST
1 U/
Q' - - V ([..))) S
Eapprox
(64)
can replace Assumption (A1) and is sufficient for the linear convergence rate in Theorem 4.3 to hold. Additionally, requiring, for all T,
-(T-()
E
- Vh(T -'
E
-
" ( [ S
1
Eapprox
(65)
T
1/4
sufficient for Equation (64) to hold. Equation (65) shows that the error floor in the linear conver- gence rate is less influenced by approximation eπrors made in early iterations, which are discounted
by the term (1-11 -.,)
On the other hand, the realizability assumption becomes relevant once the
algorithm approaches convergence, i.e., when ~ T' and Q* ~ ==,, as the discount term (1 applied fewer times.
is
Finally, although Assumption (A1) holds for the softmax tabular policies and for the neural network parameterization, it remains 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.
"A convex function is proper if dom f is nonempty and for all E dom f. f(x) > -o0. A convex function closed, if is lower semi-continuous.
38

---

H Discussion on Assumption (A2) - -the concentrability coefficients
In our convergence analysis, Assumptions (A2) and (A3) involve the concentrability coefficient C. and the distribution mismatch coefficient Vμ, which are potentially large. We give extensive discussions on them in this and the next section, respectively.
As discussed in Yuan et al. 198, Appendix H]. the issue of having (potentially large) concentrability coefficient (Assumptions (A2)) unavoidable in all the fast linear convergence analysis of approxi- mate 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 eπrоr Eppprxx under the form of Eapppoox 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 ((A)s is considered. In this setting, at each time t, it exists 6++1 such that, for any state-action distribution v²,
g++ Q' - Th V=(T'))))
MI La(tt) = 0 = Eapprox:
and C, is ignored in the convergence analysis thanks to the vanishing of Sapprox. 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 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 compare our concentrability coefficient C₁ 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, 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 Lg-norm instead of the fo norm over the approximation error Epproxx Additionally, it did not impose any restrictions on the MDP dynamics compared to Cayci al. [15], as the concentrability coefficient of Yuan et al. [98] was independent from the iterates.
Indeed, Yuan al. [98] choose of such that, for all (s,a) ESx (
DQ
(o0 (1-)
as a) (1
7" P(sv s, ae a n²,,, so: an)
where v is an initial state-action distribution chosen by the user. In this setting, have 0'(s,a) 2 (1 - ())
From the above lower bound of v², we obtain that
E(s,a)oor (
[
2 9
(T())² 8)²
(3(3,a)
(In))SSA (sssaa)
(sa))SSXA as(s S (1-7) min v(s,a))
where the finite upper bound is independent to t.
As mentioned right after Assumption (A2), the assumption our concentrability coefficient C, weaker than the one in Yuan et al. [98, Assumption 9], as we have the full control over v' while Yuan et al. [98] only has the full control over the initial state-action distribution concentrability coefficient C. recovers the previous best-known one in Yuan et al. [98] as a special case. Consequently, our concentrability coefficient C. becomes the "best" with the full control over y² when other concentrability coefficients are infinite or require strong assumptions [77].
In particular, our
(((())(())) 2 ('(aa))
In general, for the ratio ((
(
to have finite upper bound C., it is important
that vf covers well the state and action spaces so that the upper bound is independent to However, the
39

---

upper bound (1-7) min M(s,)) in Yuan et al. [98] is very pessimistic. Indeed, when T' and x converge to ", one reasonable choice of of is to choose v² E {d) x²" d b++ d , r²} such that C. 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 uy in Assumption (A3). As mentioned right after (A3), we have that
(x) SE (8)
(s) 1/ 10 SE (sss
sup
sup
1-y
which is sufficient upper bound for V/2 As discussed in Yuan et al. [98, Appendix H],
1/(1 7) - 2/2 s 1/((1 - min u(s)).
The upper bound 1/((1-)) min, ju(s) of / is very pessimistic and the lower bound v/, 1/(1-1) is often achieved by choosing di
Furthermore, if j does not have full support on state space, i.e., the upper bound 1/((1 7) min, (u())) might be infinite, one can always convert the convergence guarantees for some state distribution yl E A(S) with full support such that
V*(N) - EVV = E I uls)) ES 1e(s) ul(s (V*
pls)
< sup #ES n'(s)
-V((
Then by the linear convergence result of Theorem 4.3, we only transfer the original convergence guarantee to V*()') 2(VT('²') up to scaling factor sup a(e) with an arbitrary distribution μ² such that v/ is finite.
sES [()
Finally, if d; converges to di which is the case of AMPO through the proof of our Theorem 4.3, then d=)(a) suPses (x) converges to This might imply superlinear convergence results discussed in Xiao [93, Section 4.3]. In this case, the notion of the distribution mismatch coefficients V/ 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.
Sample complexity for neural network parameterization
We prove here Corollary 4.5 through a result by Allen-Zhu al. [3, Theorem I and Example 3.1]. We first give a simplified version of this result and then we show how to use to prove Corollary 4.5.
Consider learning some unknown distribution D of data points z (x, y) Rd y, where x is the input point and y is the label. Without loss of generality, assume ||^||| a loss function T, R x y - IR such that for every y E Y, the function L(:.y) is non-negative, convex, 1-Lipschitz continuous and (00 y) E [0, 1]. This includes both the cross-entropy loss and the Lgrregression loss (for bounded )?).
and œa = 1/2. Consider
Let RR + IR be a smooth activation function such that g(z) = = sin(z), sigmoid(z), tanh(=) or is low degree polynomial.
Define F* Rd - R such that OPT Ep(L(F"(2), v)] is the smallest population error made by a neural network of the form F A"")W" where A E Rkxp and W e IRxx Assume for simplicity that the rows of W* have 4--norm 1 and each element of A" is less equal than 1.
40

---

Define a ReLU neural network F(z, Wo) ( + bo), where Au E R Wo E Rexd the entries of Wo and by are i..d. random Gaussians from (0,)/m) and the entries of A are i..d. random Gaussians from (O,, EA), for εл E (0, 1.. We train the weights W of this neural network through stochastic gradient descent over a dataset with N i..dd samples from D, i.e., we update W++1 = W Mit, where E g V (F(( Wo + W,),)),
Theorem J.1 (Theorem I of Allen-Zhu et al. [3]. Let E (0, O(L/pk)), choose E = = the initialization and learning rate
θ (e SGD finds a set of parameters such that
1 J-1
Σ
E Ba.0)> [L (F(e; wi).
OPT +E
J
with probability 1 e eoloo
over the random initialization, for sufficiently large constant c, with
poly(k,p) and sample complexity min min(N, J} = poly(k, log m)
size m poly(E)
E²
Theorem J.1 shows that 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 address the population error in our setting.
Assumption J.2. Let g R R be smooth activation function such that g(z) et, sin(2). sig gg)) tanh(z) or is a low degree polynomial. For all time-steps f, we assume that there exists a target network F= IRd + R&,, with
1........ Sbt) and (x) F (11 1....
where wiy E RRd 22 E Rdd and o E R, such that
[[| Q' " "''''(') La(w') OPT.
We assume for simplicity || lla = | 112 = and a 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 F when solving the regression problem in Line of Algorithm
We are now ready prove Corollary 4.5, which uses Algorithm 4 to obtain unbiased estimate of current Q-function. We assume be in the same setting as Theorem J.1
Proof of Corollary 4.5. We aim to find policy x² such that
V*(4) E (VT(())) Se.
(66)
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 clog" which means that each iteration the bound should hold with probability 1 1 --e--lc m. Through Algorithm the expected number of samples needed to obtain an unbiased estimate of the current Q-function is - -7)- Therefore, using Theorem J.1, at each iteration of AMPO we need most
poly(k, p. log I")
2 2pproxx (1 - 7)
samples for SGD to find parameters that satisfy Assumption (A1) with probability -1 -cloo m. To obtain (66), we need
1 1111 I-
1))²
1 vyy
''(1+
Dj
)
s E
and 2(1 + vu) CoEappron
S 5
(67)
To - 1)
1-7
41

---

Algorithm 4: Sampler for unbiased estimate (' a) of (s,,a
Input: Initial state-action couple (5o, 0o), policy x², discount factor E [0,
Initialize Q'(so, ao) r(50; an), the time step n.
2 while True do
With probability Y+ Sample Sn+11 P(. 8n, an)
Sample an+1 ~ (- sn+1)
Q'( (o) 5- (* 0n) + 1(8n+1, An+1)
4
Otherwise with probability (1 (1-19 -y)) break
Accept Q (0)
Output: (1ss, aa)
Solving for T and Eapprox and multiplying them together, we obtain the sample complexity of AMPO, that is
poly(ki, log C? e²(1 - 76
Ö
Due to the statement Theorem J.1, we cannot guarantee the approximation error incurred by the learner network to be smaller than OPT. Consequently, have that
4(1 + (1++))VC.OPT
&
1 1-y
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, particular using the condition (65).
Corollary J.3. the setting of Theorem 4.3, replace Assumption (A1) with the condition
-2
E E p Q²= (on-m-Vhh VV (()
(2( T
1
Eapprox:
(68)
V/
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 Line with SGD as in (19), the number of samples required by AMPO to find an e-optimal policy with high probability is O(C3y
Proof. The proof follow that of Corollary 4.5. Using Theorem J.1, at each iteration i of AMPO, we need most
T
1
-)
poly(k, (k. log log T)
1
v2
U/
22p (1 7)
42

---

samples for SGD to find parameters that satisfy condition (68) with probability 1 - 1-T-le e- e- Summing over T total iterations of AMPO we obtain that the total number of samples needed is
1
2(
poly(k, p. log m, log T)
1
V/
2 eapprox 7)
1
T
1
27
poly (k: p. log m, log T)
1 -t
1
1/
Up
aapprox (1 - 7)
V/u
1ST
(1-
1
-2(T+1)
T
1
2T
poly (k, p. log m, log T)
U/
1
1/4
Vpp
Spprox (1 -?)
-2
-1)
T-
2 poly(k; (k:, p, log m, log T)
0
22 2upprxx (1 - 7)
Replacing T and Eupprox with the solutions of (67) gives the result.
At this stage, it is important note that choosing a method different from the one proposed by Allen-Zhu et al. [4] to solve Line 1 in Algorithm 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
algorithm, that is Algorithm 3 in their work, to obtain neural network estimate Q' of Q', while in [39] they provide a method based on Fourier transforms to approximate a target function through shallow ReL.U 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 Ely y II - Q'I S Écritic and an actor error E /IIet1 - Q mh VV(T7)12 S Eactor. In this case, the linear convergence rate in our Theorem 4.3 becomes
(*()) - E [vT(W]] x T (1 11-1)11)1 (1
1/ ))
Dī nol%, + 2(1 + v/ Vn)VCollennd + Eactor)
1-7
[We can obtain this alternative statement by modifying the passages in Appendix D.2. particular, writing ft+1 - Q' - T/ 'Vh(=') = (ft+1 @' nh VV(('') + (Q' - Q²) 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 from [16] gives that the sample complexity required to obtain El BI/IiQ= Q2Il2210 < E is Ö(E *(1 -)), while the required network width is Ō(==
Step 3) Actor error.
Using Theorem E.1 from [39]. we obtain that Ely [|ftt1 - n. ( (19)) can be made arbitrarily small by tuning the width of
f+ 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.A. 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 E-optimal policy with high probability is OCC3, ///s/ (1)))
43

---

To the best of our knowledge, this result improves upon the previous best result on the sample complexity of PG method with neural network parameterization [16], i... (1-)))
We now provide the statements of the aforementioned results used.
Recalling Theorem 1 in [16] and its assumptions. Consider the following space of mappings:
Ho = {o IRd + B. sup S
wER
and the function class:
Fo = {9( fo=(
Bun--rio.aalluuuua (II ( 0ff v E Ho!
Consider the following realizability assumption for the Q-function.
Assumption J.5 (Assumption [16]). For any 2 we assume that Q' Fg for some 0. Theorem J.6 (Theorem 1 in [16]) Under Assumption 2 in [16], for any enror probability 5 E (0, 1). let
t(m', ,)) 6) and R > D. Then, for any larget error > 0, number of iterations T E N, network width
4 4 + + (T/o((//),
16 + (R+ + E(ml, .8) (1 R) )
m
(1- 7222
and step-size
<2 7)
ac (1 + 21122
Algorithm 3 in [16] yields the following bound:
[
Q'11i1ttar - xt) As < (1+ 2R) (1+2R)²
E
E(1 ((1~))VT 3e,
where Ag holds with probability least 1 8 over the random initializations the critic network Q¹.
As indicated in [16], a consequence of this result is that in order to achieve a target error less than
E > 0, a network width m Ō g+ and iteration complexity O (1+20) (E 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 R + R be given and define the modulus of continuity Wy as
(g( = sup ERt'
(g(x)) g(x') max((]= | < 1 + ö, || - x r'112 S 6}.
g is continuous, then wy is not only finite for all inputs, but moreover lims-o (g( - 0.
Denote ||p]l,, = (p(w) ldw. Define sample from signed density RB++1 - IR with ||p|I < oo as (w, 6, s), where (w, b) is sampled from the probability density |pl/ and sign(p(w, b))
Theorem J.7 (Theorem E.I in [39]). Let Rd + RR 8 > 0 and (8) be as above and define for 1
5
M sup 19(x)), 110
()(x)== = f(x) s 87.
V8+ 22 W (8))
Let G. be a gaussian distribution Rn with mean O and variance a?I. Define the Gaus- sian convolution I = g)5 * Ga with Fourier transform i satisfying radial decomposition I(w) [(w) exp(2mi0n(w))) Let P be probability distribution supported |||||| S 1. Additionally define
44

---

c ((0)a(0) / [ - [|w]l2)) 2 2/||w||2 s (0:(() - ||w]]2))] dw
a / w wWW))ldw
r vn+ 2/ log 24T 7)
(6) p p: 4n [w)] cos( (2s b))[[}6| & ||||| S sr),
and for convenience create fake (weight, bias, sign) triples
(w, b, f) = (0, c, m sign(c)), (w, b, m) = (a, 0, m), (w,b,))m+33 = (-a, 0, -m).
Then
le| s M 1 (2T0²)) -/22
(2
Iplr, S 211y,olea V (2ma2)nt1²
and with probability at least
3λ over a draw of ((8j, Wj, by) b from
1
m+3 Σ
s)o((W),)) 2) + by) L(P) S Swug(d) + r|pp|e
1+ /2log(1/A)
m j=1
mm
We can then characterize the error of the actor by choosing x = a). Q + 1/2 VV((''), and f++ 1 Eye1 ()0((w,,2)) 2) + bj). We can then make the actor error arbitrarily small by tuning the network width m and ô (note that, since both Q' and ft are continuous neural networks, is continuous function).
45


Reviews:
[{"review_id": "hTJbYQFjum", "reviewer": "Reviewer_qVps", "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. \n\n", "questions": "- Can authors illustrate if it is possible to extend the results to the offline RL or other general RL setting? \n- 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", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 4, "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. \n- 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. \n- 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.\n\n- Not sure if the current framework can be extended to other settings such as online or offline RL. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive and encouraging remarks. Please find the answers to your questions in Experimental results for AMPO and below.\n\n**Application of AMPO to online/offline RL and other settings.**\nAMPO can be applied to online and offline RL settings, and has a great potential to be applied to other RL settings. As of now, and provided that we can evaluate the Q-function, we have obtained that AMPO can be applied to a general discounted MDP using the Bregman projected policy class induced by any parametrized functions, including tabular softmax, log-linear, and neural policy classes (Line 136-138 and Appendix C.1).\nFor the online setting, the policy can be evaluated through roll-out (Algorithm 1 in [1]) or temporal difference learning [2], while the distribution $v^t$ in Algorithm 1 and in Assumptions (A1) and (A2) can be chosen as the discounted state visitation distribution of the current policy $d^t_\\mu$. The experiments that we provide in the general comment correspond to this setting.\n\nRegarding the offline setting, the policy evaluation step can be done offline, e.g.\\ using importance sampling (Section 3.1 in [3]). Moreover, the update in Line 1 of Algorithm 1 and the results in Theorem 4.3 can be applied to the offline setting, which is mentioned in Lines 165 and 264 of our paper, as the distribution $v^t$ is independent of the current policy $\\pi^t$. In practice, one can choose $v^t$ to be the state-action distribution induced by an arbitrary behavior policy that generates the data. \nHowever, adapting AMPO to the offline setting 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 $\\pi^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 [2]. We leave introducing offline RL techniques in AMPO as future work and will include the above discussion in Appendix A.5.\n\nAdditional settings that have been addressed using the policy mirror descent (PMD) framework are mean-field games [4] and constrained MDP [5]. We hope to build on the existing literature for these settings and see whether our results can bring any improvements. We thank the reviewer for pointing out this direction of research, which we will include in the future work section in Appendix A.5.\n\n**Deterministic policies.** The PMD framework that encompasses AMPO is based on the policy gradient theorem that we report in Equation (3). This expression for the gradient only holds when the policy is stochastic, while in the case of deterministic policies the gradient has a different expression, as studied by [6]. For this reason, AMPO in its current formulation is not suited for the deterministic policy setting.\n\nHowever, one may ensure that the policies outputted by AMPO are deterministic by tuning the mirror map and the step-size. This is possible with softmax policies, induced by the negative entropy mirror map, by sending the step-size to infinity, effectively turning the softmax operator into a max operator.\nAnother 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^\\theta_s$. The policy can be deduced by the expression of the closed form update between Lines 193 and 194 with $\\phi (x) = x$. In both cases, it is still necessary to perform Line 1 of Algorithm 1 in order to obtain the updated parametrization function.\n\n1. 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.\n2. Gerald Tesauro. Temporal difference learning and TD-Gammon. Communications of the ACM, 1995.\n3. Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems, 2020.\n4. 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.\n5. Dongsheng Ding, Kaiqing Zhang, Jiali Duan, Tamer Başar, and Mihailo R. Jovanović, Natural Policy Gradient Primal-Dual Method for Constrained Markov Decision Processes. Advances in Neural Information Processing Systems, 2020.\n6. David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic Policy Gradient Algorithms. In International Conference on Machine Learning, 2014.\n7. Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 2022.\n8. Rui Yuan, Simon Shaolei Du, Robert M. Gower, Alessandro Lazaric, and Lin Xiao. Linear convergence of natural\npolicy gradient methods with log-linear policies. In International Conference on Learning Representations, 2023.\n9. Yanli Liu, Kaiqing Zhang, Tamer Basar, and Wotao Yin. An improved analysis of (variance- reduced) policy gradient and natural policy gradient methods. Advances in Neural Information Processing Systems, 2020. \n10. Semih Cayci, Niao He, and Rayadurgam Srikant. Linear convergence of entropy-regularized natural policy gradient with linear function approximation, 2021.\n11. Zaiwei Chen, Sajad Khodadadian, and Siva Theja Maguluri. Finite-sample analysis of off-policy natural actor–critic with linear function approximation. IEEE Control Systems Letters, 2022."}, {"review_id": "RwRBVfLS56", "reviewer": "Reviewer_rns2", "summary": "This 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.", "questions": "I would suggest adding experiments and figures that demonstrate how the work improves over previous bounds.", "limitations": "Yes, adequately addressed.", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for your review and remarks. Please find the answers to your questions in Experimental results for AMPO and below.\n\n**Interpretation of the theory.**\nAn interpretation of our theory can be provided by connecting AMPO to the Policy Iteration algorithm, which shares the linear convergence of AMPO. To see this, we first rewrite the Bregman projection step of AMPO (Line 2 of Algorithm 1) as\n$$\n\\pi\\^{t+1}\\_s= \\arg\\min_{\\pi \\in \\Delta(\\mathcal{A})} \\langle -\\eta_t f_s^{\\theta^{t+1}} + \\nabla h(\\pi^t_s), \\pi \\rangle + \\mathcal{D}_h(\\pi, \\pi^t_s), \\quad \\forall s \\in \\mathcal{S},\n$$\n\n$$\n\\qquad = \\arg\\min_{\\pi \\in \\Delta(\\mathcal{A})} \\langle - f_s^{\\theta^{t+1}} + \\frac{1}{\\eta_t}\\nabla h(\\pi^t_s), \\pi \\rangle + \\frac{1}{\\eta_t}\\mathcal{D}_h(\\pi, \\pi^t_s), \\quad \\forall s \\in \\mathcal{S},\n$$\n\nas suggested by Reviewer **oNAx** in Equation ($*$) for the **alternative proof of Lemma 4.1**. Secondly, we have$\nf^\\{t+1}\\_s - \\frac{1}{\\eta_t}\\nabla h(\\pi\\^t\\_s) \\approx Q\\^{\\pi\\^t}\\_s\n$\n by solving Line 1 of Algorithm 1 (see Line 239 as well).\nWhen the step-size $\\eta_t \\rightarrow \\infty$, and  $1/\\eta_t \\rightarrow 0$, the above AMPO policy update viewpoint thus becomes\n\n$$\n\\pi\\^{t+1}\\_s = \\arg\\min_{\\pi \\in \\Delta(\\mathcal{A})} <-Q^{\\pi^t}_s, \\pi>, \\quad \\forall s \\in \\mathcal{S},\n$$\n\n$$\n\\qquad= \\arg\\max_{\\pi \\in \\Delta(\\mathcal{A})} < Q^{\\pi^t}_s, \\pi>, \\quad \\forall s \\in \\mathcal{S},\n$$\n\nwhich is the Policy Iteration algorithm.\nHere we ignore the Bregman divergence term $\\mathcal{D}_h(\\pi, \\pi^t_s)$ as it is multiplied by a term that goes to 0. So AMPO behaves more and more like Policy Iteration with large enough step size and is able to converge linearly as Policy Iteration does. We will add this interpretation in the revised version of the paper.\n\n**Proof sketch.** Thank you for your suggestion, we will provide a proof sketch in the revised version of the main paper. In a nutshell, the convergence rates of AMPO are obtained by building on Lemma 4.1 (please refer to Appendix D.1 for its proof) and leveraging the policy mirror descent (PMD) proof techniques of [1]. More specifically, following the conventional compatible function approximation approach in [2], the idea is to write the global optimum convergence results in an additive form, that is\n$$\\mbox{sub-optimality gap} \\leq \\mbox{optimization error} + \\mbox{approximation error}.$$ \nThe separation between the two errors is allowed by Lemma 4.1, while the optimization error is bounded through the proof technique of [1] and the approximation error is characterized by Assumption (A1). Once Lemma 4.1 is established, the proof can be decomposed into three main steps.\n\n*Step 1.* Using Lemma 4.1 with $\\bar{\\pi} = \\pi^t$, $f^\\theta=f^{t+1}$, $\\eta =\\eta_t$, $\\tilde{\\pi} = \\pi^{t+1}$ , and $\\pi_s = \\pi^t_s$, we obtain that\n\n$$\n\\langle \\eta_t f\\^{t+1}_s - \\nabla h(\\pi\\^t\\_s),\\pi\\^{t+1}_s - \\pi\\^t\\_s\\rangle \\geq 0,\n$$\n\nwhich characterizes the improvement of the updated policy.\n\n*Step 2.* Assumption (A1), Step 1, the performance difference lemma (Lemma 1 in [1]), and Lemma 4.1 with $\\bar{\\pi} = \\pi^t$, $f^\\theta=f^{t+1}$, $\\eta =\\eta_t$, $\\tilde{\\pi} = \\pi^{t+1}$, and $\\pi_s = \\pi^\\star_s$ permit us to obtain the following proposition.\n\n*Proposition D.6.*\nLet $\\Delta_t := V^\\star(\\mu) - V^t(\\mu)$. For the iterates of Algorithm 1, at each time $t\\geq0$, we have\n$$\n\\mathbb{E}\\left[\\nu\\_\\mu\\left(\\Delta\\_{t+1}-\\Delta\\_t\\right)+\\Delta\\_t\\right]\\leq\\mathbb{E}\\left[\\frac{\\mathbb{E}\\_{s \\sim d\\_\\mu\\^\\star}\\left[\\mathcal{D}\\_h(\\pi\\^\\star\\_s,\\pi\\^t\\_s)\\right]}{(1-\\gamma)\\eta\\_t} -\\frac{\\mathbb{E}\\_{s \\sim d\\_\\mu\\^\\star}\\left[\\mathcal{D}\\_h(\\pi\\^\\star\\_s,\\pi\\^{t+1}\\_s)\\right]}{(1-\\gamma)\\eta_t}\\right]+(1+\\nu\\_\\mu)\\frac{2\\sqrt{C\\_v \\varepsilon\\_\\mathrm{approx}}}{1-\\gamma}.\n$$\n\n*Step 3.* Proposition D.6 leads to the sublinear convergence rate using a telescoping sum, and to the linear convergence rate by  rearranging into\n\n$$\n\\mathbb{E}\\left[\\Delta\\_{t+1} + \\frac{\\mathbb{E}\\_{s \\sim d\\_\\mu\\^\\star}\\left[\\mathcal{D}\\_h(\\pi\\^\\star\\_s,\\pi\\^{t+1}\\_s)\\right]}{(1-\\gamma) \\nu\\_\\mu\\eta\\_t}\\right] \\leq \\left(1-\\frac{1}{ \\nu\\_\\mu}\\right)\\mathbb{E}\\left[\\Delta\\_t+\\frac{\\mathbb{E}\\_{s \\sim d\\_\\mu\\^\\star}\\left[\\mathcal{D}\\_h(\\pi\\^\\star\\_s,\\pi\\^t\\_s)\\right]}{(1-\\gamma)\\eta\\_t( \\nu\\_\\mu-1)}\\right] + \\left(1+\\frac{1}{\\nu\\_\\mu}\\right)\\frac{2\\sqrt{C_v \\varepsilon\\_\\mathrm{approx}}}{1-\\gamma},\n$$\nwhich then induces a recursion.\n\n1. Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 2022.\n2. 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."}, {"review_id": "MWk0BbaPZw", "reviewer": "Reviewer_oNAx", "summary": "This 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.", "questions": "I 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. ", "limitations": "Yes, the authors have adequately addressed the limitations and potential future improvement.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. 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. ", "weaknesses": "1. 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\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for your positive and insightful remarks. Please find the answers to your questions in Experimental results for AMPO and\nbelow.\n\n**Tsallis divergence.** We thank the reviewer for pointing out the missing prior work [1] regarding the use of Tsallis entropy. We include it in Line 218 in the revised version of the paper.\n\n**Alternative proof of Lemma 4.1.** We are indebted to the reviewer for providing the alternative insight on $\\tilde{\\pi}_s$ in Lemma 4.1. We agree with the reviewer and provide the complete derivation here.\n\n$$\n\\begin{aligned}\n\\tilde{\\pi}\\_s &= \\arg\\min\\_{\\pi\\' \\in \\Delta(\\mathcal{A})} \\mathcal{D}\\_h(\\pi\\', \\nabla h^*(\\eta f\\_s\\^\\theta)) \\\\\\\\\n&= \\arg\\min\\_{\\pi' \\in \\Delta(\\mathcal{A})} h(\\pi') - h(\\nabla h^*(\\eta f\\_s\\^\\theta)) - <\\nabla h(\\nabla h^*(\\eta f\\_s\\^\\theta)), \\pi' - \\nabla h^*(\\eta f\\_s\\^\\theta)> \\\\\\\\ \n&= \\arg\\min\\_{\\pi' \\in \\Delta(\\mathcal{A})} h(\\pi') - <\\eta f\\_s\\^\\theta, \\pi'> \\\\\\\\\n&= \\arg\\min\\_{\\pi' \\in \\Delta(\\mathcal{A})} <-\\eta f\\_s\\^\\theta + \\nabla h(\\bar{\\pi}\\_s), \\pi'> + h(\\pi') - h(\\bar{\\pi}\\_s) - <\\nabla h(\\bar{\\pi}\\_s), \\pi'> \\\\\\\\ \n&= \\arg\\min\\_{\\pi' \\in \\Delta(\\mathcal{A})} <-\\eta f\\_s\\^\\theta + \\nabla h(\\bar{\\pi}\\_s), \\pi'> + \\mathcal{D}\\_h(\\bar{\\pi}\\_s, \\pi'), \\quad \\quad \\quad (*)\n\\end{aligned}\n$$\n\n\n\n\nwhere the second and the last lines are obtained by the definition of the Bregman divergence, and the third line is obtained by Equation (4): $\\nabla h(\\nabla h^*(x^*)) = x^*$ for all $x^* \\in \\mathbb{R}^{|\\mathcal{A}|}$. Consequently, we can apply the three-point descent lemma in [2] to prove Lemma 4.1.\nWe will include this proof in the revised version of the paper in Appendix D.1.\n\n**Emphasis of the projection step.** Based on the previous point, we agree with the reviewer that the key contribution of the paper is the idea of performing the parametrization function update in the dual space rather than in the policy space and obtaining the updated policy through the projection step.\nWe will replace the comment about extending the three-point descent lemma from convex to non-convex functions, which corresponds to Lines 233-237, to better emphasize the contribution in the revised version of the paper. \n\nPlease note that the comments made in Lines 174-184 still hold. In particular, our formulations of the update and the projection step are necessary to apply Lemma 4.1, which on the contrary cannot be applied to previous updates such as Equation (11). This improvement represents an important technical contribution that allows to extend the analysis from the tabular ([2]) and linear parametrization ([3]) settings to the general policy parametrization setting and we will make sure to highlight it in our revision.\n\n1. Yan Li and Guanghui Lan. Policy Mirror Descent Inherently Explores Action Space, 2023.\n2. Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 2022.\n3. 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."}, {"review_id": "F3tp1QVN1H", "reviewer": "Reviewer_zS6s", "summary": "While 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.", "questions": "Q1: 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.", "limitations": "The paper makes the assumptions of the theory clear. Some comments regarding assumptions that were not clear are mentioned in Q1-3 above.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "I 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,", "weaknesses": "- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive and insightful remarks. Please find the answers to your questions in Experimental results for AMPO and\nbelow.\n\n**Q1: $\\omega$-potential class.** We are aware of two previously used mirror maps that cannot be recovered using $\\omega$-potentials: $h(x) = \\frac{1}{2}x^\\top A x$, for some matrix $A$, which generates the Mahalanobis distance, and $p$-norms, i.e. $h(x) = \\|x\\|_p^2$. Note that the case where $h(x) = \\|x\\|_p^p$ can be recovered. We will add these observations in our revision.\n\n**Q2: Appendix G.** We agree with the reviewer that the exposition in Appendix G, especially regarding how to use the results from [2], could have been clearer. We will add the following clarifications in the proof of Corollary 4.4 in the revised version of the paper.\n\n*Assumption* Let $g:\\mathcal{R}\\rightarrow\\mathcal{R}$ be a smooth activation function such that $g(z) = e^z,~ \\sin(z)$, $\\mathtt{sigmoid}(z)$, $\\tanh(z)$ or is a low degree polynomial. For all time-steps $t$, we assume that there exists a target network $F^{\\star,t}:\\mathcal{R}^d\\rightarrow\\mathcal{R}^k$, with\n\n$$\nF\\^{\\star,t} = (f\\^{\\star,t}\\_1, \\dots, f\\^{\\star,t}\\_k)\\quad \\text{and} \\quad f\\^{\\star,t}\\_r(x)= \\sum\\_{i=1}\\^p a\\^{\\star,t}\\_{r,i}g(\\langle w\\^{\\star,t}\\_{1,i},x\\rangle)\\langle w\\^{\\star,t}\\_{2,i},x\\rangle\n$$\n\nwhere $w\\^{\\star,t}\\_{1,i} \\in \\mathbb{R}\\^d$, $w\\^{\\star,t}\\_{2,i} \\in \\mathbb{R}\\^d$, and $a\\^{\\star,t}\\_{r,i} \\in \\mathbb{R}$, such that \n$$\\mathbb{E}\\big[\\left\\lVert F\\^{\\star,t}- Q\\^t-\\eta\\_t\\^{-1}\\nabla h(\\pi\\^t) \\right\\rVert\\_{L\\_2(v\\_t)}\\^2\\big] \\leq OPT.$$\nWe assume for simplicity $\\|w\\^{\\star,t}\\_{1,i}\\|\\_2 = \\|w\\^{\\star,t}\\_{2,i}\\|\\_2 = 1$ and $|a\\^{\\star,t}\\_{r,i}| \\leq 1$.\n\nCompared to [2], this assumption represents a particular case of the class of loss functions for which their results hold. In the RL setting, assumptions of this type 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 parametrized as $F^{\\star,t}$ when solving the regression problem in Line 1 of Algorithm 1.\n\nIn the proof of Corollary 4.4, we will indicate that, due to the statement of Theorem 1 by [2], which we provide in Appendix G, we cannot guarantee the approximation error incurred by the learner network to be smaller than $OPT$. Consequently, since in the proof of Corollary 4.4 we require\n$$\n\\frac{2(1+\\nu_\\mu)\\sqrt{C_v \\varepsilon_\\mathrm{approx}}}{1-\\gamma}\\leq\\frac{\\varepsilon}{2},\n$$\nwhere $\\varepsilon$ upper-bounds the difference between the values of the optimal policy and the last policy, we also have that\n$$\n\\varepsilon\\geq\\frac{4(1+\\nu_\\mu)\\sqrt{C_v OPT}}{1-\\gamma}.\n$$\nA similar bound can be applied to any proof that contains the bias assumption introduced by [1].\n\nRegarding how to tune the width of the target network $F^{\\star, t}$ in order to control $OPT$, [3] have shown that polynomial width suffices to approximate target functions arbitrarily well, given some regularity conditions on the target function. However, the result form [3] holds for ReLU networks, while we are not aware of results regarding the approximation properties of neural networks with smooth activation functions, that is the class to which $F^{\\star, t}$ belongs. We leave addressing this gap in the literature as future work.\n\n**Q3: Assumption on bounded norms of the parameters.** The condition $\\|w\\^{\\star,t}\\_{1,i}\\|\\_2 = \\|w\\^{\\star,t}\\_{2,i}\\|\\_2 = 1$ and $|a\\^{\\star,t}\\_{r,i}| \\leq 1$ in the previous comment is only stated for simplicity and can be relaxed. We summarize here the discussion from footnote 6 in page 5 of [2] on this matter. For general $\\|w^*\\_{1,i}\\|\\_2 \\leq B$, $\\|w^*\\_{2,i}\\|\\_2 \\leq B$, $|a^*\\_{r,i}| \\leq B$, the scaling factor $B$ can be absorbed into the activation\nfunction $g'(x) = g(Bx)$. For the activation functions that we consider in Appendix G, that is the ones considered in Example 2.1 of [2], the sample complexity result described in Theorem G.1 of our paper is not affected by changes on the scaling factor.\n\n**Incorporating Assumption A1 in previous results.** Using the AMPO formulation of the parametrization function update in the dual space and the policy projection step, we can take advantage of Lemma 4.1 and the proof technique from [4] to make Assumption (A1) a mean square error assumption. If we were to use the proof technique from [Wenhao Zhan, 2021], we would need Assumption (A1) to hold in $\\ell_\\infty$-norm, due to the passages in their Appendix B.2. That is, we would have to assume that there exists $\\varepsilon_\\mathrm{approx}\\geq 0$ such that, for all times $t\\geq 0$,\n\n$$\n\\mathbb{E}\\big[\\left\\lVert f\\^{t+1}- Q\\^t-\\eta\\_t\\^{-1}\\nabla h(\\pi\\^t)\\right\\rVert\\_\\infty\\big] \\leq \\varepsilon\\_\\mathrm{approx}\n$$\n\nwhich is a stronger assumption than our Assumption (A1), especially in contexts where the state and action spaces are very large.\n\n**Proposition 4.2.** We will write the tighter bound for Proposition 4.2 in the main body of revised version of the paper. We thank the reviewer for this suggestion.\n\n\n\n1. 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.\n2. 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.\n3. Ziwei Ji, Matus Telgarsky, and Ruicheng Xian. Neural tangent kernels, transportation mappings, and universal approximation. In International Conference on Learning Representations, 2019.\n4. Lin Xiao. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 2022."}, {"review_id": "21R6fhC7IL", "reviewer": "Reviewer_WBF5", "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.", "questions": "-", "limitations": "-", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "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.\n2. 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.)\n3. 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Please find below an answer to the issues you have raised regarding the significance and soundness of our work. Due to space constraints, please find the bibliography for this reply in the reply to Reviewer qVps.\n\n**Characterization of $\\epsilon_\\mathrm{approx}$, $\\nu_\\mu$ and $C_v$ in Theorem 4.3.**\nUpper bounds involving terms similar to $\\epsilon_\\mathrm{approx}$, $\\nu_\\mu$ and $C_v$ have been established for the linear convergence of natural policy gradient (NPG) with log-linear policies in Theorem 1, 3 and 4 in [8], and for the linear convergence of NPG with tabular softmax policies in Theorem 10 and 14 in [7]. The particularity of Theorem 4.3 is that it unifies and extends (in terms of applications to general parametrization and arbitrary mirror maps) the current theory of policy mirror descent (PMD) analysis. Notably, Theorem 4.3 recovers  the best known linear convergence results of [7] and [8], as indicated in Lines 298-300, in Lines 304-306 and in Table 2.\n\nIn order to shed a clearer light on how to interpret Theorem 4.3, we provide below the characterizations that we have obtained for the approximation error $\\epsilon_\\mathrm{approx}$, the distribution mismatch coefficient $\\nu_\\mu$, and the concentrability coefficient $C_v$ and put them into perspective with the existing literature. For additional discussion, please refer to Lines 252-261, Lines 270-273, Lines 280-284, Remark 3.3, and Appendices A.3, E, and F of our paper.\n\n**1.** The approximation error $\\epsilon_\\mathrm{approx}$, which is defined in Assumption (A1) as the error incurred when solving the regression problem in Line 1 of Algorithm 1, characterizes how well the function $f^\\theta$ approximates the $Q$-function. This approximation error follows from the conventional compatible function approximation analysis by [1], which is, to the best of our knowledge, one of the current best known proof techniques for policy gradient analysis. Together with $\\epsilon_\\mathrm{approx}$, this proof technique has been extensively analysed in the literature [8,9,10,11].\n\nAssumption (A1) for $\\epsilon_\\mathrm{approx}$ is in fact weaker than Assumption 6.1 made in the seminal paper [1], which is a very common assumption in the literature and can be expressed as\n$$\n\\mathbb{E}[\\|Q\\^t-w\\^\\top\\nabla\\_\\theta\\log\\pi\\^t\\|\\_{L_2(d\\^t\\_\\mu\\cdot\\pi\\^t)}] \\leq \\epsilon\\_\\mathrm{approx},\n$$\nbecause we do not contrain the approximation of the $Q$ function to be linear in $\\nabla_\\theta\\log\\pi^t$.\n\nFurthermore, we show that Assumption (A1) holds when the parametrization function $f^\\theta$ is parametrized as a shallow ReLU network in Appendix G. This result allows us to obtain the sample complexity of AMPO for the important case of neural network parametrization.\n\n**2.** As to the concentrability coefficient $C_v$, terms of this nature are unavoidable in all the fast linear convergence analysis of approximate PMD (see Appendix F for citations). Moreover, our concentrability coefficient $C_v$ is weaker than the one in Assumption 9 in [8] which was previously the ''best'' known concentrability coefficient among the works cited in Lines 1128-1130.\nThis is because we have full control over $v^t$ in Assumption (A2), while [8] only has full control over the initial state-action distribution.\nLastly, we provide an upper bound for $C_v$, that is\n\n$$\\mathbb{E}\\_{(s,a)\\sim v\\^t}\\bigg[\\left(\\frac{d\\_\\mu\\^\\pi(s)\\pi(a|s)}{v\\^t(s,a)}\\right)\\^2\\bigg] \\leq \\frac{1}{(1-\\gamma)\\min\\_{(s,a) \\in \\mathcal{S} \\times \\mathcal{A}}\\nu(s,a)},$$\n\nwhere $\\nu$ is the initial state-action distribution. However, this upper bound is very pessimistic. When $\\pi^t$ and $\\pi^{t+1}$ converge to $\\pi^\\star$, we can choose $v^t \\in \\{d_\\mu^\\star\\cdot\\pi^\\star, d_\\mu^{t+1}\\cdot\\pi^{t+1}, d_\\mu^\\star\\cdot\\pi^t, d_\\mu^{t+1}\\cdot\\pi^t\\}$ such that $C_v$ is close to $1$.\n\n**3.** As to the distribution mismatch coefficient $\\nu_\\mu$, it is a standard term among the works cited in Lines 281-283. We provide an upper bound for $\\nu_\\mu$, that is\n$$\\max_{s\\in\\mathcal{S}}\\frac{d_\\mu^\\star(s)}{d_\\mu^t(s)}\\leq\\frac{1}{1-\\gamma}\\max_{s\\in\\mathcal{S}}\\frac{d^\\star_\\mu(s)}{\\mu(s)}:= \\nu_\\mu'.$$\nWe also provide lower and upper bounds for $\\nu_\\mu'$, i.e.\n$$1/(1-\\gamma) \\leq \\nu_\\mu' \\leq 1/((1-\\gamma)\\min_s\\mu(s)).$$\nThe upper bound $1/((1-\\gamma)\\min_s\\mu(s))$ of $\\nu_\\mu'$ is very pessimistic, while the lower bound $\\nu_\\mu' = 1/(1-\\gamma)$ can be achieved by choosing $\\mu = d_\\mu^\\star$. In Appendix F, we also cover the case where $\\mu$ does not have full support on the state space. Overall, $\\nu_\\mu$ is finite and can be reasonably small.\n\n**Time-step dependence of the assumptions.** As we argued in our above reply **Characterization of $\\epsilon_\\mathrm{approx}$, $\\nu_\\mu$ and $C_v$ in Theorem 4.3**, assumptions on $\\pi^t$ regarding the approximation error, the concentrability coefficient, and the distribution mismatch coefficient are widely adopted in previous literature [1,7,8,9,10,11]. Our Assumptions (A2) and (A3) are already the weakest among them and we provide upper bounds that do not depend on the time-step $t$.\n\n**Literature.** The application of mirror descent to policy optimization is a vast and growing body literature regarding the design of new algorithms and the study of their theoretical properties. Within this literature, our work distinguishes itself by being the first to proposes performing the update in the dual space rather than in the policy space and obtaining the policy through a Bregman projection. Thanks to this idea, we are able to provide, to the best of our knowledge, the first result that establishes linear convergence for a policy gradient-based method involving general policy parametrization. Our works paves the way for a new direction of research that investigates the properties of specific mirror maps in the general parametrization setting."}]

Author rebuttal:

**Thank you**  We would like to thank the reviewers for taking the time to review our paper and for their careful remarks.
In particular, we thank Reviewer **qVps**, **rns2**, **oNAx** and **zS6s** for recognizing the importance of our contributions, considering that our work is substantial and promising. We also thank Reviewer **rns2** and **zS6s** for acknowledging the technical soundness of our analysis. In this rebuttal, we address a common point raised by the reviewers regarding the empirical validation of our work first and reviewer-specific remarks below. We will integrate all reviewers' suggestions to improve the clarity of the paper. Please note that, following the reviewers' comments, we have implemented AMPO and added a PDF file that contains new empirical results for AMPO. We will refer to these empirical results in our response (Figures 1 and 2).

### Answer for Reviewers **qVps**, **rns2** and **zS6s**: Experimental results for AMPO

As Reviewer **zS6s**  pointed out, the main contribution of our paper is theoretical and comparing the performance of AMPO with fine-tuned existing algorithms is not the focus of our work. However, we also agree with the suggestions of Reviewers **qVps**, **rns2** and **zS6s** to provide a practical implementation of AMPO and to validate our theoretical findings with experiments. As a result, we have implemented AMPO with neural network parametrization and evaluated its performances in the well-known CartPole environment for several hyperparameter choices. We compare these performances with the convergence rates obtained in the AMPO paper and, where possible, with existing convergence rates from the literature such as  [1].



**Implementation details** As Reviewer **oNAx** and **zS6s** mentioned, AMPO is straightforward to implement. Our implementation is based upon the PPO code from PureJaxRL [2], which obtains the estimates of the $Q$-function through generalized advantage estimation (GAE) and performs the policy update using ADAM optimization and mini-batches. To implement AMPO, we (i) replaced the PPO loss with the expression to minimize in Equation (12) from our paper, (ii) replaced the softmax projection with the Bregman projection, (iii) saved the constants $\lambda$ along the sampled trajectories in order to compute Equation (12), and (iv) implemented the case of increasing step-sizes. The implementation details will be included in an additional section dedicated to the experimental results for AMPO in the appendix.

**Empirical evaluation** In the attached pdf, we provide empirical evaluation for 3 different settings of AMPO:
1. constant step-size, 2 hidden layers NN as parametrization function, negative entropy mirror map;
2. increasing step-size, 1 hidden layer NN as parametrization function, negative entropy mirror map;
3. increasing step-size, 2 hidden layers NN as parametrization function, negative entropy and hyperbolic entropy mirror map.

Here, the negative entropy mirror map (Example 3.6) corresponds to choosing $\phi(x) = e^x$ in Equation (12), and the hyperbolic entropy mirror map (Appendix C.2) corresponds to choosing $\phi(x) = \sinh(x)$ in Equation (12). Their derivations are both provided in Appendix C.2. Note that the case of 1 hidden layer NN as parametrization function in Setting (2) corresponds exactly to the theoretical case of shallow NN studied in our work in Section 4.2, and the case of 2 hidden layers NN as parametrization function in Settings (1) and (3) corresponds to the typical case used in practice for the task considered. For each setting, we provide both the performance of a sampled run of AMPO (Figure 1) and the average performance over 20 runs of AMPO (Figure 2). Additionally, we provide the performance of PPO compared against Setting (3), which was found to be the most successful. We note that the value of the optimal policy for the CartPole environment is 500.

From Setting (1), we observe that the error of AMPO with respect to the optimal policy can be bounded by our predicted convergence rate $O(1/T)$ (which is plotted in log-scale as $- \log(t) + c$ for a hand-picked constant $c$ in Figures 1(a) and 2(a)). We also observe that our $O(1/T)$ convergence rate is tighter than the $O(1/\sqrt{T})$ convergence rate previously obtained in the literature [1] (and plotted as $-0.5 \log(t) + c'$ for a hand-picked constant $c'$).

From Settings (2) and (3), we observe that the error of AMPO with respect to the optimal policy can be bounded by the linear convergence rate we predicted, for varying types of parametrization and mirror maps, i.e., expressions of $\phi$. To the best of our knowledge, this is the first time the hyperbolic entropy mirror map has been applied in RL. Lastly, we observe that the performance of AMPO, although not fine-tuned, is competitive when compared to PPO for the Cartpole environment.

We thank the reviewers for suggesting that we look into the empirical aspects of AMPO, as the empirical evaluation we have provided corroborates our theoretical findings on the performance of AMPO, and in our view, justifies further investigations with regards to its fine-tuning. We will include this empirical evaluation in the revised version of the paper. In particular, due to the limited space, we will include the first row of plots (a) and (c) in Figure 2 in the main body of the revised version of our paper, to show the efficiency of AMPO both with constant and increasing step-sizes.

1. 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.
2. Chris Lu, Jakub Grudzien Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered Policy Optimisation. Advances in Neural Information Processing Systems, 2022.

Meta-review: The paper provides the first linear convergence result for a policy gradient-based method under general policy parametrization going beyond the standard tabular or log-linear setting. The contributions are mainly on the theoretical side.  The paper was thoroughly considered by five reviewers, who are knowledgeable and familiar with the topic. Most reviewers agree with the soundness and clarity of the results after the extensive discussion, albeit some doubts on the assumptions made.  There were also some concerns with the lack of empirical validation, which has been adequately addressed by the authors with preliminary experiments provided during the rebuttal. 

Therefore, I am recommending for acceptance and ask authors to incorporate reviewers' suggestions and enrich the numerical results in the final version.
Decision: Accept (poster)