SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Carlo Alfano
carlo.alfano@stats.ox.ac.uk
Rui Yuan
Stellantis, France*
rui.yuan@stellantis.com
Patrick Rebeschini
patrick.rebeschini@stats.ox.ac.uk
Modern policy optimization methods in reinforcement learning, such as TRPO and PPO, owe their success to the use of parameterized policies. However, while theoretical guarantees have been established for this class of algorithms, especially in the tabular setting, the use of general parameterization schemes remains mostly unjustified. In this work, we introduce a framework for policy optimization based on mirror descent that naturally accommodates general parameterizations. The policy class induced by our scheme recovers known classes, e.g., softmax, and generates new ones depending on the choice of mirror map. Using our framework, we obtain the first result that guarantees linear convergence for a policy-gradient-based method involving general parameterization. To demonstrate the ability of our framework to accommodate general parameterization schemes, we provide its sample complexity when using shallow neural networks, show that it represents an improvement upon the previous best results, and empirically validate the effectiveness of our theoretical claims on classic control tasks.
Policy optimization is one of the most widely-used classes of algorithms for reinforcement learning (RL). Among policy optimization techniques, policy gradient (PG) methods [e.g., 92, 87, 49, 6] are gradient-based algorithms that optimize the policy over a parameterized policy class and have emerged as a popular class of algorithms for RL [e.g., 42, 75, 12, 69, 78, 80, 54].
The design of gradient-based policy updates has been key to achieving empirical success in many settings, such as games [9] and autonomous driving [81]. In particular, a class of PG algorithms that has proven successful in practice consists of building updates that include a hard constraint (e.g., a trust region constraint) or a penalty term ensuring that the updated policy does not move too far from the previous one. Two examples of algorithms belonging to this category are trust region policy optimization (TRPO) [78], which imposes a Kullback-Leibler (KL) divergence [53] constraint on its updates, and policy mirror descent (PMD) [e.g. 88, 54, 93, 51, 89], which applies mirror descent (MD) [70] to RL. Shani et al. [82] propose a variant of TRPO that is actually a special case of PMD, thus linking TRPO and PMD.
From a theoretical perspective, motivated by the empirical success of PMD, there is now a concerted effort to develop convergence theories for PMD methods. For instance, it has been established that PMD converges linearly to the global optimum in the tabular setting by using a geometrically increasing step-size [54, 93], by adding entropy regularization [17], and more generally by adding convex regularization [100]. Linear convergence of PMD has also been established for the negative
*The work was done when the author was affiliated with Télecom Paris.
entropy mirror map in the linear function approximation regime, i.e., for log-linear policies, either by adding entropy regularization [15], or by using a geometrically increasing step-size [20, 2, 98]. The proofs of these results are based on specific policy parameterizations, i.e., tabular and log-linear, while PMD remains mostly unjustified for general policy parameterizations and mirror maps, leaving out important practical cases such as neural networks. In particular, it remains to be seen whether the theoretical results obtained for tabular policy classes transfer to this more general setting.
In this work, we introduce Approximate Mirror Policy Optimization (AMPO), a novel framework designed to incorporate general parameterization into PMD in a theoretically sound manner. In summary, AMPO is an MD-based method that recovers PMD in different settings, such as tabular MDPs, is capable of generating new algorithms by varying the mirror map, and is amenable to theoretical analysis for any parameterization class. Since the MD update can be viewed as a two-step procedure, i.e., a gradient update step on the dual space and a mapping step onto the probability simplex, our starting point is to define the policy class based on this second MD step (Definition 3.1). This policy class recovers the softmax policy class as a special case (Example 3.2) and accommodates any parameterization class, such as tabular, linear, or neural network parameterizations. We then develop an update procedure for this policy class based on MD and PG.
We provide an analysis of AMPO and establish theoretical guarantees that hold for any parameterization class and any mirror map. More specifically, we show that our algorithm enjoys quasi-monotonic improvements (Proposition 4.2), sublinear convergence when the step-size is non-decreasing, and linear convergence when the step-size is geometrically increasing (Theorem 4.3). To the best of our knowledge, AMPO is the first policy-gradient-based method with linear convergence that can accommodate any parameterization class. Furthermore, the convergence rates hold for any choice of mirror map. The generality of our convergence results allows us not only to unify several current best-known results with specific policy parameterizations, i.e., tabular and log-linear, but also to achieve new state-of-the-art convergence rates with neural policies. Tables 1 and 2 in Appendix A.2 provide an overview of our results. We also refer to Appendix A.2 for a thorough literature review.
The key point of our analysis is Lemma 4.1, which allows us to keep track of the errors incurred by the algorithm (Proposition 4.2). It is an application of the three-point descent lemma by [19, Lemma 3.2] (see also Lemma F.2), which is possible thanks to our formulations of the policy class (Definition 3.1) and the policy update (Line 2 of Algorithm 1). The convergence rates of AMPO are obtained by building on Lemma 4.1 and leveraging modern PMD proof techniques [93].
In addition, we show that for a large class of mirror maps, i.e., the  $\omega$ -potential mirror maps in Definition 3.5, AMPO can be implemented in  $\widetilde{\mathcal{O}}(|A|)$  computations. We give two examples of mirror maps belonging to this class, Examples 3.6 and 3.7, that illustrate the versatility of our framework. Lastly, we examine the important case of shallow neural network parameterization both theoretically and empirically. In this setting, we provide the sample complexity of AMPO, i.e.,  $\widetilde{\mathcal{O}}(\varepsilon^{-4})$  (Corollary 4.5), and show how it improves upon previous results.
Let  $\mathcal{M} = (\mathcal{S},\mathcal{A},P,r,\gamma ,\mu)$  be a discounted Markov Decision Process (MDP), where  $\mathcal{S}$  is a possibly infinite state space,  $\mathcal{A}$  is a finite action space,  $P(s^{\prime}\mid s,a)$  is the transition probability from state  $s$  to  $s^\prime$  under action  $a$ ,  $r(s,a)\in [0,1]$  is a reward function,  $\gamma$  is a discount factor, and  $\mu$  is a target state distribution. The behavior of an agent on an MDP is then modeled by a policy  $\pi \in (\Delta (\mathcal{A}))^{S}$  where  $a\sim \pi (\cdot \mid s)$  is the density of the distribution over actions at state  $s\in S$ , and  $\Delta (\mathcal{A})$  is the probability simplex over  $\mathcal{A}$ .
Given a policy  $\pi$ , let  $V^{\pi}: S \to \mathbb{R}$  denote the associated value function. Letting  $s_t$  and  $a_t$  be the current state and action at time  $t$ , the value function  $V^{\pi}$  is defined as the expected discounted cumulative reward with the initial state  $s_0 = s$ , namely,
$$
V ^ {\pi} (s) := \mathbb {E} _ {a _ {t} \sim \pi (\cdot | s _ {t}), s _ {t + 1} \sim P (\cdot | s _ {t}, a _ {t})} \Big [ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r (s _ {t}, a _ {t}) \mid \pi , s _ {0} = s \Big ].
$$
Now letting  $V^{\pi}(\mu)\coloneqq \mathbb{E}_{s\sim \mu}[V^{\pi}(s)]$ , our objective is for the agent to find an optimal policy
$$
\pi^ {\star} \in \operatorname {a r g m a x} _ {\pi \in (\Delta (\mathcal {A})) ^ {s}} V ^ {\pi} (\mu). \tag {1}
$$
As with the value function, for each pair  $(s,a)\in S\times \mathcal{A}$ , the state-action value function, or  $Q$ -function, associated with a policy  $\pi$  is defined as
$$
Q ^ {\pi} (s, a) := \mathbb {E} _ {a _ {t} \sim \pi (\cdot | s _ {t}), s _ {t + 1} \sim P (\cdot | s _ {t}, a _ {t})} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} r (s _ {t}, a _ {t}) \mid \pi , s _ {0} = s, a _ {0} = a \right].
$$
We also define the discounted state visitation distribution by
$$
d _ {\mu} ^ {\pi} (s) := (1 - \gamma) \mathbb {E} _ {s _ {0} \sim \mu} \left[ \sum_ {t = 0} ^ {\infty} \gamma^ {t} P \left(s _ {t} = s \mid \pi , s _ {0}\right) \right], \tag {2}
$$
where  $P(s_{t} = s \mid \pi, s_{0})$  represents the probability of the agent being in state  $s$  at time  $t$  when following policy  $\pi$  and starting from  $s_{0}$ . The probability  $d_{\mu}^{\pi}(s)$  represents the time spent on state  $s$  when following policy  $\pi$ .
The gradient of the value function  $V^{\pi}(\mu)$  with respect to the policy is given by the policy gradient theorem (PGT) [87]:
$$
\nabla_ {s} V ^ {\pi} (\mu) := \frac {\partial V ^ {\pi} (\mu)}{\partial \pi (\cdot \mid s)} = \frac {1}{1 - \gamma} d _ {\mu} ^ {\pi} (s) Q ^ {\pi} (s, \cdot). \tag {3}
$$
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  $\mathcal{Y} \subseteq \mathbb{R}^{|\mathcal{A}|}$  be a convex set. A mirror map  $h: \mathcal{Y} \to \mathbb{R}$  is a strictly convex, continuously differentiable and essentially smooth function such that  $\nabla h(\mathcal{Y}) = \mathbb{R}^{|\mathcal{A}|}$ . The convex conjugate of  $h$ , denoted by  $h^*$ , is given by
$$
h ^ {*} (x ^ {*}) := \sup _ {x \in \mathcal {Y}} \langle x ^ {*}, x \rangle - h (x), \quad x ^ {*} \in \mathbb {R} ^ {| \mathcal {A} |}.
$$
The gradient of the mirror map  $\nabla h: \mathcal{V} \to \mathbb{R}^{|A|}$  allows to map objects from the primal space  $\mathcal{V}$  to its dual space  $\mathbb{R}^{|A|}$ ,  $x \mapsto \nabla h(x)$ , and viceversa for  $\nabla h^{*}$ , i.e.,  $x^{*} \mapsto \nabla h^{*}(x^{*})$ . In particular, from  $\nabla h(\mathcal{V}) = \mathbb{R}^{|A|}$ , we have: for all  $(x, x^{*}) \in \mathcal{V} \times \mathbb{R}^{|A|}$ ,
$$
x = \nabla h ^ {*} (\nabla h (x)) \quad \text {a n d} \quad x ^ {*} = \nabla h \left(\nabla h ^ {*} \left(x ^ {*}\right)\right). \tag {4}
$$
Furthermore, the mirror map  $h$  induces a Bregman divergence [13, 18], defined as
$$
\mathcal {D} _ {h} (x, y) := h (x) - h (y) - \langle \nabla h (y), x - y \rangle ,
$$
where  $\mathcal{D}_h(x,y)\geq 0$  for all  $x,y\in \mathcal{V}$ . We can now present the standard MD algorithm [70, 14]. Let  $\mathcal{X}\subseteq \mathcal{V}$  be a convex set and  $V:\mathcal{X}\to \mathbb{R}$  be a differentiable function. The MD algorithm can be formalized as the following iterative procedure in order to solve the minimization problem  $\min_{x\in \mathcal{X}}V(x)$ : for all  $t\geq 0$
$$
y ^ {t + 1} = \nabla h \left(x ^ {t}\right) - \eta_ {t} \nabla V (x) | _ {x = x ^ {t}}, \tag {5}
$$
$$
x ^ {t + 1} = \operatorname {P r o j} _ {\mathcal {X}} ^ {h} \left(\nabla h ^ {*} \left(y ^ {t + 1}\right)\right), \tag {6}
$$
where  $\eta_t$  is set according to a step-size schedule  $(\eta_t)_{t \geq 0}$  and  $\operatorname{Proj}_{\mathcal{X}}^h(\cdot)$  is the Bregman projection
$$
\operatorname {P r o j} _ {\mathcal {X}} ^ {h} (y) := \operatorname {a r g m i n} _ {x \in \mathcal {X}} \mathcal {D} _ {h} (x, y). \tag {7}
$$
Precisely, at time  $t$ ,  $x^t \in \mathcal{X}$  is mapped to the dual space through  $\nabla h(\cdot)$ , where a gradient step is performed as in (5) to obtain  $y^{t + 1}$ . The next step is to map  $y^{t + 1}$  back in the primal space using  $\nabla h^{*}(\cdot)$ . In case  $\nabla h^{*}(y^{t + 1})$  does not belong to  $\mathcal{X}$ , it is projected as in (6).
3 Approximate Mirror Policy Optimization
The starting point of our novel framework is the introduction of a novel parameterized policy class based on the Bregman projection expression recalled in (7).
${}^{1}h$  is essentially smooth if  $\mathop{\lim }\limits_{{x \rightarrow  \partial \mathcal{V}}}\begin{Vmatrix}{\nabla h}\left( x\right) {\rVert }_{2} =  + \infty$  ,where  $\partial \mathcal{V}$  denotes the boundary of  $\mathcal{V}$  .
See a different formulation of MD in (11) and in Appendix B (Lemma B.1).
Algorithm 1: Approximate Mirror Policy Optimization
Input: Initial policy  $\pi^0$ , mirror map  $h$ , parameterization class  $\mathcal{F}^\Theta$ , iteration number  $T$ , step-size schedule  $(\eta_t)_{t \geq 0}$ , state-action distribution sequence  $(v^t)_{t \geq 0}$ .
for  $t = 0$  to  $T - 1$  do
1 Obtain  $\theta^{t + 1}\in \Theta$  such that  $\theta^{t + 1}\in \mathrm{argmin}_{\theta \in \Theta}\left\| f^{\theta} - Q^{t} - \eta_{t}^{-1}\nabla h(\pi^{t})\right\|_{L_{2}(v^{t})}^{2}.$  4   
2 Update  $\pi_s^{t + 1}\in \operatorname *{argmin}_{p\in \Delta (\mathcal{A})}\mathcal{D}_h(p,\nabla h^* (\eta_tf_s^{\theta^{t + 1}})) = \operatorname{Proj}_{\Delta (\mathcal{A})}^h (\nabla h^* (\eta_tf_s^{t + 1}))$  ,  $\forall s\in S$
Definition 3.1. Given a parameterized function class  $\mathcal{F}^{\Theta} = \{f^{\theta}:S\times \mathcal{A}\to \mathbb{R},\theta \in \Theta \}$ , a mirror map  $h:\mathcal{V}\rightarrow \mathbb{R}$ , where  $\mathcal{V}\subseteq \mathbb{R}^{|\mathcal{A}|}$  is a convex set with  $\Delta (\mathcal{A})\subseteq \mathcal{V}$ , and  $\eta >0$ , the Bregman projected policy class associated with  $\mathcal{F}^{\Theta}$  and  $h$  consists of all the policies of the form:
$$
\left\{\pi^ {\theta}: \pi_ {s} ^ {\theta} = \operatorname {P r o j} _ {\Delta (\mathcal {A})} ^ {h} \left(\nabla h ^ {*} \left(\eta f _ {s} ^ {\theta}\right)\right), s \in \mathcal {S}; \theta \in \Theta \right\},
$$
where for all  $s\in \mathcal{S}$ $\pi_s^\theta ,f_s^\theta \in \mathbb{R}^{\lfloor \mathcal{A}\rfloor}$  denote vectors  $[\pi^{\theta}(a\mid s)]_{a\in \mathcal{A}}$  and  $[f^{\theta}(s,a)]_{a\in \mathcal{A}}$  , respectively.
In this definition, the policy is induced by a mirror map  $h$  and a parameterized function  $f^{\theta}$ , and is obtained by mapping  $f^{\theta}$  to  $\mathcal{V}$  with the operator  $\nabla h^{*}(\cdot)$ , which may not result in a well-defined probability distribution, and is thus projected on the probability simplex  $\Delta (\mathcal{A})$ . Note that the choice of  $h$  is key in deriving convenient expressions for  $\pi^{\theta}$ . The Bregman projected policy class contains large families of policy classes. Below is an example of  $h$  that recovers a widely used policy class [7, Example 9.10].
Example 3.2 (Negative entropy). If  $\mathcal{V} = \mathbb{R}_+^{|\mathcal{A}|}$  and  $h$  is the negative entropy mirror map, i.e.,  $h(\pi(\cdot|s)) = \sum_{a\in \mathcal{A}}\pi(a|s)\log(\pi(a|s))$ , then the associated Bregman projected policy class becomes
$$
\left\{\pi^ {\theta}: \pi_ {s} ^ {\theta} = \frac {\exp \left(\eta f _ {s} ^ {\theta}\right)}{\| \exp \left(\eta f _ {s} ^ {\theta}\right) \| _ {1}}, s \in \mathcal {S}; \theta \in \Theta \right\}, \tag {8}
$$
where the exponential and the fraction are element-wise and  $\| \cdot \| _1$  is the  $\ell_1$  norm. In particular, when  $f^{\theta}(s,a) = \theta_{s,a}$ , the policy class (8) becomes the tabular softmax policy class; when  $f^{\theta}$  is a linear function, (8) becomes the log-linear policy class; and when  $f^{\theta}$  is a neural network, (8) becomes the neural policy class defined by Agarwal et al. [1]. We refer to Appendix C.1 for its derivation.
We now construct an MD-based algorithm to optimize  $V^{\pi^{\theta}}$  over the Bregman projected policy class associated with a mirror map  $h$  and a parameterization class  $\mathcal{F}^{\Theta}$  by adapting Section 2.1 to our setting. We define the following shorthand: at each time  $t$ , let  $\pi^t \coloneqq \pi^{\theta^t}$ ,  $f^t \coloneqq f^{\theta^t}$ ,  $V^t \coloneqq V^{\pi^t}$ ,  $Q^t \coloneqq Q^{\pi^t}$ , and  $d_{\mu}^t \coloneqq d_{\mu}^{\pi^t}$ . Further, for any function  $y: S \times A \to \mathbb{R}$  and distribution  $v$  over  $S \times A$ , let  $y_s \coloneqq y(s, \cdot) \in \mathbb{R}^{|\mathcal{A}|}$  and  $\| y \|_{L_2(v)}^2 = \mathbb{E}_{(s,a) \sim v}[(y(s,a))^2]$ . Ideally, we would like to execute the exact MD algorithm: for all  $t \geq 0$  and for all  $s \in S$ ,
$$
f _ {s} ^ {t + 1} = \nabla h \left(\pi_ {s} ^ {t}\right) + \eta_ {t} (1 - \gamma) \nabla_ {s} V ^ {t} (\mu) / d _ {\mu} ^ {t} (s) \stackrel {(3)} {=} \nabla h \left(\pi_ {s} ^ {t}\right) + \eta_ {t} Q _ {s} ^ {t}, \tag {9}
$$
$$
\pi_ {s} ^ {t + 1} = \operatorname {P r o j} _ {\Delta (\mathcal {A})} ^ {h} \left(\nabla h ^ {*} \left(\eta_ {t} f _ {s} ^ {t + 1}\right)\right). \tag {10}
$$
Here, (10) reflects our Bregman projected policy class 3.1. However, we usually cannot perform the update (9) exactly. In general, if  $f^{\theta}$  belongs to a parameterized class  $\mathcal{F}^{\Theta}$ , there may not be any  $\theta^{t + 1} \in \Theta$  such that (9) is satisfied for all  $s \in S$ .
To remedy this issue, we propose Approximate Mirror Policy Optimization (AMPO), described in Algorithm 1. At each iteration, AMPO consists in minimizing a surrogate loss and projecting the result onto the simplex to obtain the updated policy. In particular, the surrogate loss in Line 1 of Algorithm 1 is a standard regression problem where we try to approximate  $Q^{t} + \eta_{t}^{-1}\nabla h(\pi^{t})$  with  $f^{t+1}$ , and has been studied extensively when  $f^{\theta}$  is a neural network [3, 35]. We can then readily use (10) to update  $\pi^{t+1}$  within the Bregman projected policy class defined in 3.1, which gives Line 2 of Algorithm 1.
To better illustrate the novelty of our framework, we provide below two remarks on how the two steps of AMPO relate and improve upon previous works.
3The update is (5) up to a scaling  $(1 - \gamma) / d_{\mu}^{t}(s)$  of  $\eta_t$
4With a slight abuse of notation, denote  $\nabla h(\pi^t)(s,a)$  as  $[\nabla h(\pi_s^t)]_a$
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  $\theta^{t + 1}$  as the solution to a regression problem aimed at approximating the current  $Q$ -function  $Q^t$ . A crucial difference is that, Agarwal et al. [1] approximate  $Q^t$  linearly with respect to  $\nabla_{\theta} \log \pi^t$  (see (61)), while in Line 1 we approximate  $Q^t$  and the gradient of the mirror map of the previous policy with any function  $f^\theta$ . This generality allows our algorithm to achieve approximation guarantees for a wider range of assumptions on the structure of  $Q^t$ . Furthermore, the regression problem proposed by Agarwal et al. [1] depends on the distribution  $d_\mu^t$ , while ours has no such constraint and allows off-policy updates involving an arbitrary distribution  $v^t$ . See Appendix E for more details.
Remark 3.4. Line 2 associates AMPO to previous approximations of PMD [89, 88]. For instance, Vaswani et al. [89] aim to maximize an expression equivalent to
$$
\pi^ {t + 1} \in \operatorname {a r g m a x} _ {\pi^ {\theta} \in \Pi (\Theta)} \mathbb {E} _ {s \sim d _ {\mu} ^ {t}} \left[ \eta_ {t} \left\langle Q _ {s} ^ {t}, \pi_ {s} ^ {\theta} \right\rangle - \mathcal {D} _ {h} \left(\pi_ {s} ^ {\theta}, \pi_ {s} ^ {t}\right) \right], \tag {11}
$$
where  $\Pi (\Theta)$  is a given parameterized policy class, while the Bregman projection step of AMPO can be rewritten as
$$
\pi_ {s} ^ {t + 1} \in \operatorname {a r g m a x} _ {p \in \Delta (\mathcal {A})} \left\langle \eta_ {t} f _ {s} ^ {t + 1} - \nabla h \left(\pi_ {s} ^ {t}\right), p \right\rangle - \mathcal {D} _ {h} (p, \pi_ {s} ^ {t}), \quad \forall s \in \mathcal {S}. \tag {12}
$$
We formulate this result as Lemma F.1 in Appendix F with a proof. When the policy class  $\Pi(\Theta)$  is the entire policy space  $\Delta(\mathcal{A})^S$ , (11) is equivalent to the two-step procedure (9)-(10) thanks to the PGT (3). A derivation of this observation is given in Appendix B for completeness. The issue with the update in (11), which is overcome by (12), is that  $\Pi(\Theta)$  in (11) is often a non-convex set, thus the three-point-descent lemma [19] cannot be applied. The policy update in (12) circumvents this problem by defining the policy implicitly through the Bregman projection, which is a convex problem and thus allows the application of the three-point-descent lemma [19]. We refer to Appendix F for the conditions of the three-point-descent lemma in details.
3.1  $\omega$ -potential mirror maps
In this section, we provide a class of mirror maps that allows to compute the Bregman projection in Line 2 with  $\widetilde{\mathcal{O}}(|\mathcal{A}|)$  operations and simplifies the minimization problem in Line 1.
Definition 3.5 ( $\omega$ -potential mirror map [50]). For  $u \in (-\infty, +\infty]$ ,  $\omega \leq 0$ , let an  $\omega$ -potential be an increasing  $C^1$ -diffeomorphism  $\phi : (-\infty, u) \to (\omega, +\infty)$  such that
$$
\lim  _ {x \to - \infty} \phi (x) = \omega , \quad \lim  _ {x \to u} \phi (x) = + \infty , \quad \int_ {0} ^ {1} \phi^ {- 1} (x) d x \leq \infty .
$$
For any  $\omega$ -potential  $\phi$ , the associated mirror map  $h_{\phi}$ , called  $\omega$ -potential mirror map, is defined as
$$
h _ {\phi} (\pi_ {s}) = \sum_ {a \in \mathcal {A}} \int_ {1} ^ {\pi (a | s)} \phi^ {- 1} (x) d x.
$$
Thanks to Krichene et al. [50, Proposition 2] (see Proposition C.1 as well), the policy  $\pi^{t + 1}$  in Line 2 induced by the  $\omega$ -potential mirror map can be obtained with  $\widetilde{\mathcal{O}}(|\mathcal{A}|)$  computations and can be written as
$$
\pi^ {t + 1} (a \mid s) = \sigma \left(\phi \left(\eta_ {t} f ^ {t + 1} (s, a) + \lambda_ {s} ^ {t + 1}\right)\right), \quad \forall s \in \mathcal {S}, a \in \mathcal {A}, \tag {13}
$$
where  $\lambda_s^{t + 1} \in \mathbb{R}$  is a normalization factor to ensure  $\sum_{a \in \mathcal{A}} \pi^{t + 1}(a \mid s) = 1$  for all  $s \in S$ , and  $\sigma(z) = \max(z, 0)$  for  $z \in \mathbb{R}$ . We call this policy class the  $\omega$ -potential policy class. By using (13) and the definition of the  $\omega$ -potential mirror map  $h_\phi$ , the minimization problem in Line 1 is simplified to be
$$
\theta^ {t + 1} \in \operatorname {a r g m i n} _ {\theta \in \Theta} \left\| f ^ {\theta} - Q ^ {t} - \eta_ {t} ^ {- 1} \max  \left(\eta_ {t - 1} f ^ {t}, \phi^ {- 1} (0) - \lambda_ {s} ^ {t}\right) \right\| _ {L _ {2} \left(v ^ {t}\right)} ^ {2}, \tag {14}
$$
where  $\max (\cdot ,\cdot)$  is applied element-wisely. The  $\omega$  -potential policy class allows AMPO to generate a wide range of applications by simply choosing an  $\omega$  -potential  $\phi$  . In fact, it recovers existing approaches to policy optimization, as we show in the next two examples.
Example 3.6 (Squared  $\ell_2$ -norm). If  $\mathcal{V} = \mathbb{R}^{|A|}$  and  $\phi$  is the identity function, then  $h_\phi$  is the squared  $\ell_2$ -norm, that is  $h_\phi(\pi_s) = \|\pi_s\|_2^2 / 2$ , and  $\nabla h_\phi$  is the identity function. So, Line 1 in Algorithm 1 becomes
$$
\theta^ {t + 1} \in \operatorname {a r g m i n} _ {\theta \in \Theta} \left\| f ^ {\theta} - Q ^ {t} - \eta_ {t} ^ {- 1} \pi^ {t} \right\| _ {L _ {2} \left(v ^ {t}\right)} ^ {2}. \tag {15}
$$
The  $\nabla h_{\phi}^{*}$  also becomes the identity function, and the policy update is given for all  $s \in S$  by
$$
\pi_ {s} ^ {t + 1} = \operatorname {P r o j} _ {\Delta (\mathcal {A})} ^ {h _ {\phi}} \left(\eta_ {t} f _ {s} ^ {t + 1}\right) = \underset {\pi_ {s} \in \Delta (\mathcal {A})} {\operatorname {a r g m i n}} \left\| \pi_ {s} - \eta_ {t} f _ {s} ^ {t + 1} \right\| _ {2} ^ {2}, \tag {16}
$$