Learning mirror maps in policy mirror descent
Authors: Carlo Alfano
Sebastian Rene Towers
Silvia Sapora
Chris Lu
Patrick Rebeschini
Venue: ICLR 2025 Poster (2025)
Keywords: Policy optimization
mirror descent
Primary area: reinforcement learning

Abstract: Policy Mirror Descent (PMD) is a popular framework in reinforcement learning, serving as a unifying perspective that encompasses numerous algorithms. These algorithms are derived through the selection of a mirror map and enjoy finite-time convergence guarantees. Despite its popularity, the exploration of PMD's full potential is limited, with the majority of research focusing on a particular mirror map---namely, the negative entropy---which gives rise to the renowned Natural Policy Gradient (NPG) method. It remains uncertain from existing theoretical studies whether the choice of mirror map significantly influences PMD's efficacy. In our work, we conduct empirical investigations to show that the conventional mirror map choice (NPG) often yields less-than-optimal outcomes across several standard benchmark environments. Using evolutionary strategies, we identify more efficient mirror maps that enhance the performance of PMD. We first focus on a tabular environment, i.e.\ Grid-World, where we relate existing theoretical bounds with the performance of PMD for a few standard mirror maps and the learned one. We then show that it is possible to learn a mirror map that outperforms the negative entropy in more complex environments, such as the MinAtar suite. Additionally, we demonstrate that the learned mirror maps generalize effectively to different tasks by testing each map across various other environments.

Published as a conference paper at ICLR 2025
LEARNING MIRROR MAPS IN POLICY MIRROR DESCENT
Carlo Alfano" R
Sebastian Towers
Silvia Saporal
Department of Statistics
FLAIR University of Oxford
FLAIR, Department of Statistics
University of Oxford
University of Oxford
Chris Lu FLAIR
Patrick Rebeschini
Department of Statistics University of Oxford
University of Oxford
ABSTRACT
Policy Mirror Descent (PMD) is a popular framework in reinforcement learning, serving as a unifying perspective that encompasses numerous algorithms. These algorithms are derived through the selection of mirror map and enjoy finite-time convergence guarantees. Despite its popularity, the exploration of PMD's full potential is limited, with the majority of research focusing map-namely, the negative entropy- -whch gives rise to the renowned Natural Policy Gradient (NPG) method. It remains uncertain from existing theoretical studies whether the choice of mirror map significantly influences PMD's s efficacy. In our work, conduct empirical investigations to show that the conventional mirror map choice (NPG) often yields less-than-optimee outcomes across several standard benchmark environments. Using evolutionary strategies, we identify more efficient mirror maps that enhance the performance of PMD. We first focus on tabular environment, i.e. Grid-World, where relate existing theoretical bounds with the performance of PMD for few standard mirror maps and the Iearned one. We then show that possible learn mirror map that outperforms the negative entropy in more complex environments, such as the MinAtar suite. Additionally, we demonstrate that the learned mirror maps generalize effectively to different tasks by testing each map across various other environments.
particular mirror
INTRODUCTION
Policy gradient (PG) methods (Williams & Peng, 1991; Sutton et al., 1999; Konda & Tsitsiklis, 2000; Baxter & Bartlett, 2001) are some of the most widely-used mechanisms for policy optimization in reinforcement learning (RL). These algorithms are maddeent-based methods that optimize over a class of parameterized policies and have become a popular choice for RL problems, both in theory (Kakade, 2002; Peters& & Schaal, 2008; Bhatnagar al., 2009; Schulman al., 2015; Mnih et al., 2016; Schulman et al., 2017; Lan, 2022a) and in practice (Shalev-Shwartz et al., 2016; Berner al., 2019; Ouyang al., 2022).
Among PG methods, some of the most successful algorithms are those that employ some form of reg- ularization in their updates, ensuring that the newly updated policy retains some degree of similarity to its predecessor. This principle has been implemented in different ways. For instance, trust region policy optimization (TRPO) Schulman al. (2015) imposes Kullback-Leibler divergence Kullback & Leibler (1951) hard constraint for its updates, while proximal policy optimization (PPO) (Schulman et al., 2017) uses clipped objective to penalize large updates. framework that has recently attracted attention and belongs to this heuristic is that of policy mirror descent (PMD) (Tomar et al., 2022; Lan, 2022a; Xiao, 2022; Kuba et al., 2022; Vaswani et al., 2022; Alfano et al., 2023), which applies mirror descent Nemirovski & Yudin (1983) to RL to regularize the policy updates.
PMD consists of a wide class of algorithms, each derived by selecting a mirror map that introduces distinct regularization characteristics. In recent years, PMD has been investigated through numerical
Coorespondence to carlo.alfano ax..uk........kk
I Equal contribution.

---

Published as a conference paper at ICLR 2025
experiments (Tomar et al., 2022). but it has mainly been analysed from theoretical perspective. To the best of our knowledge, research has been mostly focused either on the particular case of the negative entropy mirror map, which generates the natural policy gradient (NPG) algorithm (Kakade, 2002; Agarwal et al., 2021), or on finding theoretical guarantees for a generic mirror map.
NPG has been proven converge to the optimal policy, up to a difference in expected return (or error in several settings, e.g. using tabular, linear, general parameterization (Agarwal et al., 2021) or regularizing rewards (Cen et al., 2021). It has been shown that NPG benefits from implicit regularization (Hu et al., 2022), that it can exploit optimism (Zanette et al., 2021; Liu et al., 2023), and some of its variants have been evaluated in simulations (Vaswani et al., 2022: 2023). When considering other specific mirror maps investigated in the RL literature, the Tsallis entropy has been noted for enhancing performance in onfline settings (Tomar al., 2022) and for offering improved sample efficiency in online settings (Li & Lan, 2023), when compared to the negative entropy.
There is substantial body of theoretical research focused on the general case of mirror maps. This research demonstrates that PMD achieves convergence to the optimal policy under the same conditions as NPG, as evidenced by various studies (Xiao, 2022; Lan, 2022b; Yuan et al., 2023; Alfano et al., 2023). Except for the setting where we have access to the true value of the policy, convergence guarantees in stochastic settings are subject to an error floor due to the inherent randomness or bias within the algorithm. In the majority of PMD analyses involving generic mirror maps, both the convergence rate and the error floor show mild dependence on the specific choice of mirror map. Typically, the effect of the mirror map appears explicitly multiplicative factor in the convergence rate and it appears implicitly in the error floor. These analyses often rely on upper bounds, meaning they may not accurately reflect the algorithms' actual performance in applications.
In this work, we contribute to the literature with an empirical investigation of PMD, with the objectives of finding a mirror map that consistently outperforms the negative entropy mirror map and of understanding how the theoretical guarantees of PMD relate to simulations. We first consider set of tabular environments, i.c. WriddWWrldd (Oh et al., 2020), where we compare the empirical results to the theoretical guarantees given by Xiao (2022), which we can compute as we have full control over the environment. In this setting, we provide a learned mirror map which outperforms the negative entropy in all tested environments. Our experiments suggest that the error floor appearing in prototypical PMD convergence guarantees not a good performance indicator: our learned mirror map achieves the best value while presenting the worst theoretical error floor, implying that the upper bounds typically considered in the PMD literature are loose with respect to the choice of the mirror map. Additionally, our experiments indicate that having small policy updates leads smoother value improvements over time with less instances of performance degradation, as suggested by the monotonic policy improvement property given by Xiao (2022). We then consider two non-tabular settings, i.e. the Basic Control Suite and the MinAtar Suite, which are more realistic but also more complex, and therefore prevent us from computing the exact theoretical guarantees. We learn mirror map for each of these environments and, also in this case, show that the learned mirror maps lead to a higher performance of PMD than the negative entropy. Moreover, we show that the learned mirror maps generalize well to other tasks, by testing each of the learned mirror maps on all the other environments. Lastly, we tackle continuous control tasks in MuJoCo (Todorov et al., 2012), where we show that the mirror map learned on one environment surpasses the negative entropy across several environments.
To establish our findings, we employ the standard formulation of PMD (Xiao, 2022) for the tabular setting and the continuous control tasks, and we used generalized version, Approximate Mirror Policy Optimization (AMPO) (Alfano et al., 2023), for the non-tabular setting with discrete action spaces. To allow optimization over the space of mirror maps, introduce parameterization schemes for mirror maps, one for PMD and one for AMPO. Specifically, we propose parameterization for w-potentials, which have been shown to induce a wide class of mirror maps (Krichene et al., 2015). We use evolution strategies (ES) to search for the mirror map that maximizes the performance of PMD and AMPO over an environment.
AMPO (Alfano et al., 2023) is a recently-proposed PMD framework designed to integrate general parameterization schemes, in our case neural networks, and arbitrary mirror maps. It benefits from theoretical guarantees, as Alfano et al. (2023) show that AMPO has muasi-monotonic updates as well as sub-linean and linear convergence rates, depending on the step-size schedule. These desirable properties make AMPO particularly suitable for our numerical investigation.
2

---

Published as a conference paper at ICLR 2025
ES are type of population based stochastic optimization algorithm that leverages random noise to generate a diverse pool of candidate solutions, and have been successfully applied to a variety of tasks (Real et al., 2019; Salimans al.. 2017; Such et al., 2018). The main idea consists in iteratively selecting higher-performing individuals, w.r.t. a fitness function, resulting in a gradual convergence towards the optimal solution. ES algorithms are gradient- free and have been shown to be well-suited for optimisation problems where the objective function is noisy or non-differentiabbl and the search space is large complex (Beyer, 2000; Lu al., 2022: 2023). We use ES to search over the parameterized classes of mirror maps we introduce, by defining the fitness of particular mirror map as the value of the last policy outputted by PMD and AMPO, for fixed hyper-parameteen
The rest of the paper is organized as follows. In Section 2, we introduce the setting of RL as well as the PMD and AMPO algorithms. We describe the methodology behind our numerical experiments in Section 3, which are then discussed in Section 4. Finally, we discuss related works from the literature on automatic discovery of machine learning algorithms in Appendix A and give our conclusions in Section 5.
2
PRELIMINARIES
2.1 REINFORCEMENT LEARNING
Define a discounted Markov Decision Process (MDP) as the tuple M (S,A, P.r.,, where S and A are respectively the state and action spaces, P(s' s,a) is the transition probability from state 8 to s' when taking action a, r(s, a) E [o, 1] is the reward function, 7 a discount factor, and is starting state distribution. A policy E (A(A)) (.))), where ((A) is the probability simplex over A, represents the behavior of an agent on an MDP. whereby at state s E S the agents takes actions according to the probability distribution T(. |).
Our objective is for the agent to find a policy that maximizes the expected discounted cumulative reward for the starting state distribution μs. That is, want to find
argmax Es~n(V TE(A(A))S
Here V S -> T denotes the value function associated with policy and is defined
7'r(st, (*((): = E a.)
n. 80 =
where 8t and are the current state and action at time I and the expectation is taken over the trajectories generated by ag ~ क (-11s) and 8141 P(Iss,,;;) a:).
Similarly to the value function, we define the Q-function associated with a policy æ as
(s E
Σ 20
1= 7'r(se: a.) T,, So 8,00 =
where the expectation once again taken over the trajectories generated by the policy When the state and action spaces are finite, the Q-function can be expressed as Q = (I > PP)) rr. where pa a square matrix where the position ((s, a), (s', a')) is occupied by define the discounted state visitation distribution as
a). We also
E (s) = (1 - 7) w Σ P(s = n, s0)
n, so) represents the probability of the agent being in state at time when following policy и and starting from sp- The probability distribution over states d[((s) represents the proportion of time spent state when following policy
where P(s,
2.2 POLICY MIRROR DESCENT
We review the PMD framework, starting from mirror maps (Bubeck, 2015). Let X' C R- be a convex set. A mirror map h X' - IR. is strictly convex, continuously differentiable and essentially smooth function' that satisfies Vh(x') = A In particular, we consider mirror maps belonging to the w-potential mirror map class, which contains most mirror maps used in the literature.
'A function is essentially smooth if lim [Vh(2)] +oo, where 0. denotes the boundary of X.
3

---

Published as a conference paper at ICLR 2025
Definition 2.1 (w-potential mirror map (Krichene et al., 2015)). For E (-o0, tool, w 0, an w- w-potential is defined as an increasing C* -iffeeomorphism 6 (-o0, u) + (w, t00) such that
lim <(x) - w, lim o(x) = +00, L. x (x) S co.
For any w-potential (, we define the associated mirror map ha as
((ass) o (x)dd.
(T(TT)) = E
When Φ(z) et we recover the negative entropy mirror map. which the standard choice of mirror map in the RL literature (Agarwal et al., 2021; Tomar et al., 2022; Hu et al., 2022; Vaswani al., 2022; 2023), while we recover the æ2-norm when ((x) x (see Appendix B.1). Mirror maps belonging to this class are particularly advantageous as they can be defined by a single scalar function 6, without having to account for the dimension of the action space A. The Bregman divergence (Bregman, 1967; Censor & Zenios, 1997) induced by the mirror map J is defined as
(x, v) = h(x) h(y) (Vh(y), y).
where DA y) 2 0 for all z, y X. As further discuss in Appendix C, the Bregman divergence measures how far two points are in geometry induced by the mirror map. Given a starting policy T², learning rate 17 and mirror map h, PMD can be formalized as an iterative algorithm: for all iterations t O,
++| E a Eg (O,, nz) DD(TT; m9)]s
where we used the shorthand: V : V Q' === Q dz == =z and y. == v(s.)) E R-A, for any function y : S x A - TR. PMD benefits from several theoretical guarantees and, in particular, Xiao (2022) shows that PMD enjoys quasi-monotonie updates and convergence the optimal policy. We give here slight modification of the statements of these results. Specifically, we do not upper-bound a term regarding the distance between subsequent policies in the result on quasi monotonic updates, which we use to draw connections between theory and practice. Additionally, to obtain the convergence rate for PMD with constant step-size in the setting where we do not have access to the true Q-function, we combine the analyses on the sublinear convergence of PMD and linear convergence of inexact PMD given by Xiao (2022). We provide a proof in Appendix D. Theorem 2.2 (Xiao (2022)). Following update (2), we have that, for all 2
V (n) - V*()) - 1--
max @1 O'lloo max aES
t+1
(3)
sE8
where 11-110o and ||-||1 represent the loo and the l norms, respectively, and Q is an estimate of the true Quuuccioo, Additionally, at each iteration T > we have
max 11 2% (TT (1 7)
B[V((N)] S
Easdy [Da(TT T0)
(*( Σ KT
Th (1
(1 1 7)3
The statement of Theorem 2.2 similar to many results in the literature on PMD and NPG (Agarwal et al., 2021; Xiao, 2022: Lan, 2022b; Hu et al., 2022). That is, the convergence guarantee involves two terms, i.c. convergence rate, which involves the Bregman divergence between the optimal policy and the starting policy, and an error floor, which involves the estimation error max.es 1 lo Given that by setting To Ex-d; [Dx (T (1 >) we obtain the 2(T(1- convergence rate 2(7(1 - 2) and that the error floor has no explicit dependence on the mirror map, Equation (4) suggests that the mirror map has mild influence on the performance of PMD. The only way the mirror map seems to affect the convergence guarantee is, implicitly, by changing the path of the algorithm and therefore influencing the estimation error. Similar observations can be made for several results in the PMD literature that share the same structure of convergence guarantees. On the other hand, Equation (3) suggests that mirror maps that prevent large updates of the policy cause the PMD algorithm to less prone to performance degradation, as the lower bound is close to O when the policy update distance maxses m!+1 æ ll is small. One of the contributions of
4

---

Published as a conference paper at ICLR 2025
work is to challenge and refute the conclusion that the mirror map has little influence the convergence of PMD, highlighting a gap between theoretical guarantees, which are based on upper bounds, and the actual performance observed in PMD-based methodologies. Our empirical studies reveal that the choice of mirror map significantly influences both the speed of convergence and the minimum achievable error floor in PMD. Additionally, provide evidence that mirror map that prevents large updates throughout training leads to a better performance, as suggested by (3).
When applying PMD to continuous control tasks, we replace the tabular policy in (2) with parametrized one. That is, for all iterations / 2 0, we obtain the updated policy as
1 E argmaxn,:0e d, (- (s" a)) - DA (Te(()s),
(5)
where {TTe E } is a class of parametrized policies.
2.3 APPROXIMATE MIRROR POLICY OPTIMIZATION
AMPO a theoretically sound framework for deep reinforcement learning based on mirror descent, it inherits the quasi-monotonic updates and convergence guarantees from the tabular case (Alfano et al., 2023). Given a parameterized function class F = { S x A + R, E }) an initial scoring function fd, step-size 7, and w-potential mirror map hgs AMPO can be described, for all iterations t, by the two-step update
r'(a 8) = (s. S,EA,
(6)
g E argmin Ex a x 6E0
- n max(nf a) + +2)) (7)
where X3 R normalization factor to ensure o.. ((A) for s E S, ft : fo and (2) = max(2, 0) for E R. Theorem by Krichene et al. (2015) ensures that there always exists a normalization constant X₂ E TR. As shown by Alfano et al. (2023), AMPO recovers the standard formulation of PMD in (2) in the tabular setting. Assuming for simplicity that Φ(x) > 0 for all z E R. the minimization problem in (7) implies that, at each iteration 1, f is an approximation of the sum of the Q-funtioons up to that point, that is ft E.= Q'. Therefore, the scoring function f' serves as an estimator of the value of an action.
METHODOLOGY
Denote by H the class of w-potentials mirror maps. Our objective to search for the mirror map that maximizes the value of the last policy outputted by our mirror descent based algorithms, for fixed time horizon T. That is, we want to find
h" E argmax E [V (14)
(8)
NEW
where the expectation is taken over the randomness of the policy optimization algorithm and the policy updates are based on the mirror map h. Depending on the setting, we parameterize the mirror map by parameterizing either Φ=1 or & monotonically increasing functions. One of the primary objectives of this work it to motivate further research into the choice of mirror maps, as influenced by the choice of the function o, moving beyond the conventional use of Φ(x) = e. This is achieved by examining how different choices of & influence the trajectory training and demonstrating that, in many cases, there is a mirror map that outperforms the negative entropy by large margin.
3.1 POLICY MIRROR DESCENT
We parameterize 6-1 non-negative and the activation functions are equally split among the following convex and concave 1/3 log((x)+ + 10-3) and e", where (z)+ monotonic mon-linear:ties: x², {x) (x) (x) max(z, 0). To ensure that we are able to recover the negative entropy and the & norm,, we add
layer neural network with 126 hiden units, where all kernels
+ b log(2:) the final output, where a,
5

---

Published as a conference paper at ICLR 2025
To search for the best mirror map within this class, we employ a slight variation of the OpenAI-ES strategy (Salimans et al., 2017), adapted to the multi-task setting (Jackson et al., 2024). Denote by e the parameters of the mirror map and by (1b) the objective function in (8). Given distribution of tasks &, we estimate the gradient V FyFF ((b) as
[Eee
[ - - 20
F
where (00, Ia) is the multivariate normal distribution, is the number of parameters, and a is hyperparameter regulating the variance of the perturbations. To account for different reward scales across tasks, we perform rank transformation of the objective functions, whereby, for each sampled task
return I for the higher performing member between & + de and oe, and O for the other. We note that, when the distribution of tasks & covers a single task, we recover the standard OpenAI-ES strategy.
3.2 APPROXIMATE MIRROR POLICY OPTIMIZATION
As non-tabular environments with discrete action space, consider the Basic Control Suite (BCS) and the MinAtar suite. Given the higher computational cost of simulations on these environments w.r.t. the tabular setting, replace the neural network parameterization for the mirror map with one with fewer parameters, in order to reduce the dimension of the mirror map class we search over. We define the parameterized class & (ov + [0. 1], v E F. with
if 0,
if x S V1,
(x(x)=
Ei= W
if Σ vss < I E/ ti.
if > 1,
where < j n 1. 1. In other words, ø is defined as piece-wise linear function with n steps where, for all j S n, (E', v/s) = j/n and subsequent points are interpolated with straight segment. This is illustrated in Figure We note that the w- w-potentials within & violate some of the constraints in Definition 2.1, as they are non-decreasing instead of increasing and lim (xx) = 1 for all é E 4. In Appendix B.2, show that, if E 4, we can construct an w-potential o' that satisfies the con- straints in Definition 2.1 such that é and o' in- duce the same policies along the path of AMPO.
65 -.
0.13
0.26
0/2
0.02
0.0
To effectively learn the hyperparameters, we employ the Separable Covariance Matrix Adaptation Evolution Strategy (sep-CMA) (Ros & Hansen, 2008), variant of the popular algorithm CMA-ES (Hansen & Ostermeier, 2001). CMA-ES is. essentially, second-order method adapted for gradient free optimization. At every generation, samples n new points from a normal distribution, parameterized by mean vector m& and a covariance matrix Ct. That is, the samples for generation k æ4 ..,..., ..... distributed i.i.d. according to x A(m,, Ck). For each generation k denote the m S n best performing samples as = ....,,",. x.k Then update the mean vector as m = Em w ek where Σ W/ = 1, so that the next generation distributed around the weighted mean of the best r* performing samples. Ck+1 also updated to reflect the covariance structure of æ , in complex way beyond the scope of this text. Due to the need to update Cj based on covariance information, CMA-ES exhibits a quadratic scaling behavior with respect to the dimensionality of the search space, potentially hindering its efficiency in high-dimensional settings. To improve computational efficiency, we adopt sep-CMA, which introduces a diagonal constraint on the covariance matrix and reduces the computational complexity of the algorithm.
Figure 1: A plot visually demonstrating the param- eterization for d.

---

Published as a conference paper at ICLR 2025
Dence
Labryrtth
Sodeen rooms
Small contsor
Standard mobe
Ransom tode
201e
161 Brwwenee.coo chipe
En Dape
Estindoo
Revostie entrogy
ITRD
Figure 2:
Comparison between the learned map and the negative entropy and fe-norm mirror maps across a range of held-out configurations of Grid-World. We display the average over 256 runs and report the standard error shaded region. The column "Random tasks", reports the averaged metrics for 256 randomly sampled configurations of Grid-World.
4 EXPERIMENTS
In this section, discuss the results of our numerical experiments. We start by presenting the tabular setting, where we track errors in order to understand what properties are desirable in a mirror map, and proceed by showing our results in the non-tabular setting.
4.1 TABULAR SETTING: GRID-WORLD
As tabular setting we adopt discounted and infinite horizon version of Grid-World (Oh et al., 2020), which is large class of tabular MDPs.
Model architecture and training We define the tabular policy one-layer neural network with a softmax head, which takes as input one-hot encoding of the environment state and outputs a distribution over the action space. We train the policy using the PMD update in (2), where we solve the minimization problem through stochastic gradient descent and estimate the Q-function through generalized advantage estimation (GAE) (Schulman et al., 2016). We perform simple grid-search over the hyperparameters to maximize the performance for the negative entropy and the 2--norm mirror maps, in order to have fair comparison. We report the chosen hyperparameters in Appendix E. The training procedure is implemented in Jax, using evosax (Lange, 2022a) for the evolution. We run on four A40 GPUs, and the optimization process takes roughly 12 hours.
Environment We adopt the version of Grid- World implemented by Jackson et al. (2024), and adapt to the discounted and infinite horizon setting. We learn a single mirror map by training PMD on a continuous distribution of Grid-World environments, and test PMD with the learned mirror map on five held-out configurations from previous publications (Oh et al., 2020; Chevalier-Boisver et al., 2024) and on 256 randomly sampled configurations. For all PMD iterations we track two quantities that appear in Theorem 2.2, that is the estimation error m Q: loo and the distance between policy updates max.es |m/+11 - T |l To obtain the true Q-function, we compute the transition matrix pr and use the formula specified in Section 2, that is Q' (1 -- 'r. PMD is run for 128 iterations and 218 250k total environment steps for all World configurations.
We show the results of our simulations in Figure 2. As shown by the first row, the learned mirror map, the f norm, and the negative entropy consistently rank first, second and third, respectively, in all tested configurations. in terms of final value. These results advocate the effectiveness of our methodology, as we are able to find a mirror map that outperforms the benchmark mirror maps in all tested environments. We highlight that the first five environments in Figure 2 are part of the task
7

---

Published as a conference paper at ICLR 2025
Actobor-e1
Carthlleot
Adteris-kinttrr
Preemhy-dinaatt
Auerage
12 64
Act 6
Score BcaM
Score
Score acale
Nogative entropy
Noro
Lcamed
Figure 3: Comparison between the learned mirror map, the 22-nrrm and the negative entropy across range standard environments. The top plots present the performance of AMPO for all mirror maps, reporting the average over 100 realizations and a shaded region denoting the standard error around the average. The middle plots report the w-potentials that induce the mirror maps. The bottom plots report the policy distribution according to (6), for each mirror map and score scales. The score scales are obtained by multiplying the vector [1,..., [A] by a variable c E [0,
distribution used during the evolution, demonstrating the generalizability of our approach. Moreover, have that the l 22-norm consistently outperforms the negative entropy, further proving that the negative entropy is not always the best choice of mirror map. Another shared property among all training curves, is that the learned mirror map presents slower convergence in the initial iterations w.r.t. the fg-norm and to the negative entropy, as testified by a lower value, but convergence to higher value in the long run. Lastly, we note that all environments most of the value improvement happens in the first 50k environment steps.
To gain better understanding of why the learned mirror maps outperforms the negative entropy and the &g-no fornrrm, and to draw connections with the theoretical results outlined in Section 2.2, we report the estimation error and the distance between policy updates for mirror maps. The first conclusion that we draw is that smaller estimation error does not seem to be related to a higher performance. the contrary, the second row of Figure 2 shows how the three mirror maps consistently have the same ranking in both value and estimation error, which is exactly the opposite of what Theorem 2.2 would suggest. On the other hand, the first and third rows of Figure 2 show how the lower bound on the performance improvement in (3) brings some valid insight the behaviour of PMD during its first iterations, which are the ones that bring the largest improvement. In all configurations, we have that in the initial iterations of PMD the learned mirror map induces the smallest policy update distances as well as the performance curve with fewer dips in value, while both the negative entropy and the fg-norm induce larger policy update distances and performance curves with several dips in value. This observation confirms the behaviour described by (3), whereby a small distance between policy updates prevents large performance degradation in policy updates.
4.2 NON-TABULAR SETTING: BASIC CONTROL AND MINATAR SUITES
Model architecture and training We define the scoring function as deep neural network, which we train using the AMPO update in (7) and (6), where (7) is solved through Adam and the Q-function is estimated through GAE. We optimize the mirror map for each suite, using the hyper- pprrmeea tuning framework Optuna (Akiba et al., 2019). This is done to ensure we are looking at a fair benchmark of performance when using the negative entropy mirror map. We report the chosen dyper-parametert in Appendix E. We then initialize the
of AMPO for the negative entropy
8

---

Published as a conference paper at ICLR 2025
Table 1: The table contains, for each entry, the value of the final policy outputted by AMPO trained on the environment corresponding to the column with the mirror map learned on the environment corresponding the row. The last row represents the performance of AMPO with the negative entropy for the corresponding column environments. The value is averaged over 100 runs. Green cells correspond a value higher than that associated to the negative entropy.
Acrobot
CartPole
Asterix
Freeway
Spacelnvaders
Acrobot
-88.49
476.41
24.03
56,00
144.01
CartPole
-83.76
499.93
27.26
52.26
100.07
Asterix
-103.55
490.86
30.22
58.56
122.00
Freeway
-82.51
457.47
3.20
58.21
143.93
SpaceInvaders
-78.29
489.56
4.36
22.27
170.24
Negative entropy
-105.63
359.14
17.80
53.69
81.77
parameterized mirror map be an approximation of the negative entropy² and Sep-CMA-ES on each environment separately. The whole training procedure is implemented in JAX, using gymnax environments (Lange, 2022b) and evosax (Lange, 2022a) for the evolution. We run on GTX 1080Ti GPUs, and the optimization process takes roughly 48 hours for a single environment.
We test AMPO the Basic Controle Suite (BCS) and the MinAtar Suite. For BCS, we run the evolution for 600 generation, each with 500k timesteps. For MinAtar, we run the evolution for 500 generation with 1M timesteps, then run 100 more generations with 10M timesteps.
Environments
Our empirical results are illustrated Figure 3, where we show the performance of AMPO for the learned mirror map, for the negative entropy and for the fg-norm. For all environments and mirror maps, we use the syper-parametent returned by Optuna for AMPO with the negative entropy. Our learned mirror map leads better overall performance in all environments, apart from Acrobot, where it ties with the fy-norm. We observe the largest improvement in performance on Asterix and SpaceInvaders, where the average return for the learned mirror map more than double the one for the negative entropy. Figure 3 also suggests that different mirror maps may result in different error floors, as shown by the performance curves Acrobot, Asterix and Spacelnvaders, where the negative entropy converges to lower point than the learned mirror map.
The second and third row of Figure 3 illustrate the properties of the learned mirror map, in comparison to the negative entropy and the f2-norm. In particular, the second row shows the corresponding w-potential, while the third row shows the policy distribution induced by the mirror map according to (6), depending on the scores assigned by the scoring function to each action. A shared property among all the learned mirror maps is that they all lead to assigning 0 probability to the worst actions for relatively small score scales, whilst the negative entropy always assigns positive weights all actions. In more complex environments, where the evaluation of a certain action may be strongly affected by noise where the optimal state may be combination locked (Misra et al., 2020), this behaviour may lead to a critical lack of exploration. However, it appears that this is not the case in these environments, and hypothesise that by setting the probability of the worst actions to 0 the learned mirror maps avoid wasting samples and hence can converge to the optimal policy more rapidly.
Our last result consists in testing each learned mirror map across the other environments we consider. In Table 1, we report the value of the final policy outputted by AMPO for all learned maps, plus the negative entropy, and for all environments, averaged over 100 runs. The table shows that the learned mirror maps generalize well to different environments and to different sets of hyper-parameterd are shared within BCS and MinAtar but not between them. In particular, we have that the mirror maps learned in Acrobot and Asterix outperform the negative entropy in all environments, those learned in CartPole and Freeway outperform the negative entropy in 4 out of 5 environments, and that learned in SpaceInvaders outperforms the negative entropy in out of 5 environments. These results show that our methodology can be useful in practice, as it benefits from good generalization across tasks.
which
We achieve this by assigning vi log(i/(i - 1)); for i 2,.... a n. and setting 3 log(10).

---

Published as a conference paper at ICLR 2025
4.3 CONTINUOUS CONTROL: MuJoCo
Model architecture and training We define the policy as deep neural network and optimize it us- ing the PMD update in (5), which is computed using Adam and the estimate of the Q-function obtained through GAE. As previously discussed, we ensure a fair comparison by tuning the hyper-paramers in
to maximize the performance of the negative entropy. We report the chosen Appendix E. The mirror map is initialized to be close to the negative entropy and is trained for 256 generations using OpenAI-ES on a single MuloCo environment. We run on eight A40 GPUs, and the optimization process takes roughly 24 hours.
Environments We use the brax library (Freeman al., 2021) simulate three MuJoCo environ- ments, i.e. Hopper, Halfcheetah, and Ant. The mirror map is learned on Hopper and is then tested on all environments, using 10 environment steps and 488 PMD update steps.
Hopper
Haltchestah
Ant
265
Value
100
PMAcy
Psiccy MEDICC Steps
Nagathen artrapy
Figure 4: Comparison between the mirror map learned on Hopper and the negative entropy three MuJoCo environments. The plots present the performance of PMD for both mirror maps, reporting the average over 8 realizations and shaded region denoting the standard error around the average.
Figure 4 shows that the mirror map learned on Hopper significantly improves in terms of final perfor- mance upon the negative entropy in both train and test environments. This result confirms the ability of our methodology to learn generalizable mirror maps, even in complex continuous control tasks.
5 CONCLUSION
Our study presents an empirical examination of PMD, where we successfully test the possibility of learning mirror map that outperforms the negative entropy in both the tabular and non-tabular settings. In particular, we have shown that the learned mirror maps perform well on a set of configurations in Grid-World and that they can generalize to different tasks in BCS, in MinAtar, and in MuJoCo. Additionally, we have compared the theoretical findings established in the literature and the actual performance of PMD methods in the tabular setting. highlighting how the estimation error is not a good indicator of performance and validating the intuition that small policy updates lead to less instances of performance degradation. Our findings indicate that the choice of mirror map significantly impacts PMD's effectiveness, an aspect not adequately reflected by existing convergence guarantees.
Our research introduces several new directions for inquiry. From theoretical perspective, obtaining convergence guarantees that reflect the impact of the mirror map is an area for future exploration. On the practical side, investigating how temporal awareness (Jackson al., 2024) or specific environmen- tal challenges, such as exploration, robustness to noise, or credit assignment (Osband et al., 2019), can inform the choice of mirror map to improve performance represents another research focus.
ACKNOWLEDGMENTS
Carlo Alfano is funded by the Engineering and Physical Sciences Research Council. Patrick Rebes- chini was funded by UK Research and Innovation (UKRI) under the UK government's 's Horizon Europe funding guarantee [grant number EP/Y028333/1J. For the purpose of Open Access. the authors have applicd CC BY public copyright licence to any Author Accepted Manuscript (AAM) version arising from this submission.
10

---

Published as a conference paper at ICLR 2025
REPRODUCIBILITY STATEMENT
The implementation of our experiments can be found httptp///ggtttbbbccooooo///////a ceazninggmirroo-mae
REFERENCES
Agarwal, A., Kakade, M., Lee, D., and Mahajan, G. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. Journal of Machine Learning Research, 2021.
Akiba, T., Sano, S., Yanase, T., Ohta, and Koyama, M. Optuna: A next-generation hyperparameter optimization framework, 2019.
Alfano, C., Yuan, R., and Rebeschini, P. novel framework for policy mirror descent with general parametrization and linear convergence. Advances in Neural Information Processing Systems, 2023.
Baxter, J. and Bartlett, P. L. Infinitehhorizon volicy-gradien estimation. Journal of Artificial Intelligence Research, 2001.
Belousov, B. and Peters, f-divergence constrained policy improvement. arXiv preprint 1rXhv:1801000056. 2017.
Berner, C.. Brockman, G., Chan, B., Cheung, V., Debiak, Dennison, C.,, Farhi, D., Fischer, Q.. Hashme, S., Hesse, C., et al. Dota 2 with large scale deep reinforcement learning. arXiv preprint 191.v 2019.
Beyer, H.-G. Evolutionary algorithms in noisy environments: theoretical issues and guidelines for practice. Computer Methods in Applied Mechanics and Engineering, 2000.
Bhatnagar, S., Sutton, R. S., Ghavamzadch, M., and Lee, M. Natural actor-critic algorithms. Automatica, 2009.
Bregman, I.. M. The relaxation method of finding the common point of convex and its application to the solution problems convex programming. USSR Computational Mathematics and Mathematical Physics, 1967.
Bubeck, S. Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning, 2015.
Cen, S., Cheng, C., Chen, Y., Wei, Y., and Chi, Y. Fast global convergence of natural policy gradient methods with entropy regularization. Operations Research, 2021.
Censor, Y. and Zenios, S. A. Parallel Optimization: Theory, Algorithms, and Applications. Oxford University Press, USA, 1997.
Chevalier-Boisvertt, M., Dai, B., Towers, M., Perez-Vicente, R., Willems, L., Lahlou, S., Pal, S., Castro, S., and Terry, J. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. Advances in Neural Information Processing Systems, 2024.
Freeman, C. D., Frey, E., Raichuk, A., Girgin, S., Mordatch, I., and Bachem, O. Brax - a differentiable physics engine for large scale rigid body simulation, 2021. URL http: /tt::///iiti /gtthbbttthuhh/g com/ google/brax.
Hansen, and Ostermeier, A. Completely derandomized self-adaptation evolution strategies. Evolutionary Computation, 2001.
Houthooft, R., Chen, Y., Isola, P., Stadie, B., Wolski, F., Jonathan Ho, O., and Abbeel, P. Evolved policy gradients. Advances in Neural Information Processing Systems, 2018.
Hu, Y., Ji, Z., and Telgarsky, M. cctor-critic implicitly biased towards high entropy optimal policies. In International Conference Learning Representations, 2022.

---

Published as a conference paper at ICLR 2025
Jackson, M. T., Jiang, M., Marker-Holde,, J., Vuorio, R., Lu, C., Farquhar, G., Whiteson, S., and Foerster, Discovering general reinforcement learning algorithms with adversarial environment design. Advances in Neural Information Processing Systems, 2023.
Jackson, M. T., Lu, C., Kirsch, L., Lange, R. T., Whiteson, S., and Foerster, J. N. Discovering temporally-aware reinforcement learning algorithms. In International Conference on Learning Representations, 2024.
Kakade, S. M. A natural policy gradient. Advances in Neural Information Processing Systems, 2002.
Kirsch, I.,, van Steenkiste, S., and Schmidhuber, Improving generalization in meta reinforcement learning using learned objectives. In International Conference on Learning Representations, 2020.
Konda, V. and Tsitsiklis, J. Actor-critic algorithms. In Advances in Neural Information Processing Systems, 2000.
Krichene, W., Krichene, S., and Bayen, A. Efficient bregman projections onto the simplex. In IEEE Conference Decision and Control, 2015.
Kuba, J. G., De Witt, C. A. S., and Foerster, J. Mirror learning: A unifying framework of policy optimisation. In International Conference on Machine Learning, 2022.
Kullback, S. and Leibler, R. A. On Information and Sufficiency. The Annals of Mathematical Statistics, 1951.
Lan, G. Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and generalized problem classes. Mathemutical programming, 2022a.
Lan, G. Policy optimization over general state and action spaces. arXiv preprint arXiv:2211.16715. 2022b.
Lange, R. T. evosax: Jax-based cvolution strategies. arXiv preprint arXiv:2212.04188, 2022a.
Lange, R. T. gymnax: A JAX-based reinforcement learning environment library, 2022b. URL http://s httptttttbbttttt//////////ttt com Robeet /gmmaxx.
Li, Y. and Lan, G. Policy mirror descent inherently explores action space. arXiv:2303.04386, 2023.
arXiv preprint
Liu, Q. Weisz, G., György, A., Jin, C., and Szepesvari, C. Optimistic natural policy gradient: a simple efficient policy optimization framework for online rl. Advances Neural Information Processing Systems, 2023.
Lu, C., Kuba, J., Letcher, A., Metz, L., Schroeder de Witt, C., and Foerster, J. Discovered policy optimisation. Advances in Neural Information Processing Systems, 2022.
Lu, C., Willi, T., Letcher, A., and Foerster, J. N. Adversarial cheap talk. In International Conference Machine Learning. 2023.
Misra, D., Henaff, M., Krishnamurthy. A., and Langford, J. Kinematic state abstraction and provably efficient rich-observation reinforcement learning. In International conference on machine learning, 2020.
Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T., Harley, T., Silver, D.,, and Kavukcuoglu, K. Asynchronous methods for deep reinforcement learning. In International Conference on Machine Learning, 2016.
Nemirovski, and Yudin, D. B. Problem Complexity and Method Efficiency Optimization. Wiley Interscience, 1983.
Oh, J., Hessel, M.. Czameeki, W. M.. Xu, Z., van Hasselt, H. P., Singh, S., and Silver, D. Discovering reinforcement learning algorithms. Advances in Neural Information Processing Systems, 2020.
Orabona, F. A modern introduction to online learning, 2020. URL https: / /open. bu. edu/ handle/ 221 1444/// 2/////144/400 4/40900.
12

---

Published as a conference paper at ICLR 2025
Osband, I., Doron, Y., Hessel, M., Aslanides, J., Sezener, E., Saraiva, A., McKinney, K., Lattimore. T., Szepesvari, C., Singh, S., et al. Behaviour suite for reinforcement learning. In International Conference Learning Representations, 2019.
Ouyang, L.,, Wu, J., Jiang, X., Almeida, D., Wainwright, C., Mishkin, P., Zhang, C., Agarwal, S., Slama, K., Ray, A., et al. Training language models to follow instructions with human feedback. Advances in Neural Information Processing Systems, 2022.
Peters, J. and Schaal, S. Natural actor-critic. Neurocompuring, 2008.
Real, E., Aggarwal, A., Huang, Y., and Le, Q. V. Regularized evolution for image classifier architecture search. Proceedings of the AAAI Conference on Artificial Intelligence, 2019.
Ros, R. and Hansen, N. A simple modification in ema-es achieving linear time and space complexity.
In International conference on parallel problem solving from nature. Springer, 2008.
Salimans. T., Ho, J., Chen, X., Sidor, S., and Sutskever, Evolution strategies as a scalable alternative to reinforcement learning. arXiv preprint arXiv: 7703.33844 2017.
Schulman, J., Levine, S., Abbeel, P., Jordan, M., and Moritz, P. Trust region policy optimization. In International Conference on Machine Learning, 2015.
Schulman, J., Moritz, P., Levine, S., Jordan, M., and Abbeel, P. High-dimensional continuous control
using generalized advantage estimation. In International Conference on Learning Representations, 2016.
Schulman, J.. Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. arXiv preprint 1eX.v:1777.06:6 2017.
Shalev-Shwartz, S., Shammah, S., and Shashua, A. Safe, multi-agent, reinforcement learning for autonomous driving. arXiv preprint arXiv: 1610.0399, 2016.
Such, F. P., Madhavan, V., Conti, E., Lehman, J., Stanley, K. O., and Clune, J. Deep neuroevolution:
Genetic algorithms are competitive alternative for training deep neural networks for reinforcement learning, 2018.
Sutton, R. S., McAllester, D. A., Singh, S. P., and Mansour, Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems, 1999.
Todorov, E., Erez, T., and Tassa, Mujoco: A physics engine for model-based control. 2012 IEEE/RSJ International Conference Intelligent Robots and Systems, 2012.
Tomar, M., Shani. L., Efroni, Y., and Ghavamzadeh, M. Mirror descent policy optimization. In International Conference on Learning Representations, 2022.
Vaswani, S., Bachem, O., Totaro, S., Müller, R., Garg, S., Geist, M., Machado, M. C., Castro, P. S., and Le Roux, N. A general class of surrogate functions for stable and efficient reinforcement learning. International Conference on Artificial Intelligence and Statistics. 2022.
Vaswani, S., Kazemi, A., Babanezhad, R., and Roux, N. L. Decision-aware actor-critic with function approximation and theoretical guarantees. Advances in Neural Information Processing Systems, 2023.
Williams, R. J. and Peng. J. Function optimization using connectionist reinforeement learning algorithms. Connection Science, 1991.
Xiao, L. On the convergence rates of policy gradient methods. Journal of Machine Learning Research, 2022.
Yuan, R., Du, S. S., Gower, R. M., Lazaric, A., and Xiao, L. Linear convergence of natural policy gra- dient methods with lon-linear policies. In International Conference on Learning Representations, 2023.
Zanette, A., Cheng, C.-A., and Agarwal, A. Cautiously optimistic policy optimization and exploration with linear function approximation. In Conference on Learning Theory, 2021.
13

---

Published as a conference paper at ICLR 2025
A RELATED WORKS
Discovering reinforcement learning algorithm has been an active area of research in the last years, where researchers have shown that handcrafted algorithms are not always optimal and can be outperformed by automatically discovered algorithms. Oh et al. (2020) obtained an algorithm capable of solving GridWorld and Atari tasks without relying common notions in the field, such as Q-value functions. Houthooft et al. (2018) used ES to meta-train a policy loss network that outperforms PPO and Kirsch et al. (2020) discovered related to our work are Lu et al. (2022) and Jackson et al. (2024), who employ ES to discover RL algorithms within the Mirror Learning class (Kuba et al., 2022). Given function f, the class of algorithms they consider follows the policy update step
loss function for deterministic policies. Most closely
[[ (() []
Tgge+1 argmax Bswdy 1/1 (el (Q' (8,a)) (-lor
This update is similar the the one (5) but replaces the Bregman divergence with a different penalty term called drift, which recovers the f-divergences when f is convex function. To the best of our knowledge, this class of algorithms has fewer theoretical guarantees than PMD, in particular we are not aware of finite-time convergence guarantees convergence guarantees involving approximation or estimation errors (Belousov & Peters, 2017; Kuba et al., 2022).
B FURTHER DISCUSSION ON 4-POTENTIALS
B. NEGATIVE ENTROPY AND & No-N
If (xx) = ="-1, then the associated mirror map ha is the negative entropy. We have that Lo x (x)de = Lo log(r)de = [x: log(2) - =]] = -11 S +00
The mirror map hs becomes the negative entropy. up to constant, (als)
ho(m))
Σ
log(x)dr
|A| - 1 Σ
x(a ( log( s))
REA"
REA
If ((x)= = x, then the associated mirror map ho is the fz-norm. We have that
f. (x)(x)
f xdz
x
<
400
2 0
The mirror map hs becomes the E--norm, up to a constant, as
hT( Σ
((a)))
rdr
1
Σ
x(a
s)2 1.
2
aEA
B.2 PARAMETRIC &
We show here that the parametric class of w-potentials we introduce in Section 3 results in a well defined algorithm when used for AMPO, even if it breaks some of the constraints in Definition 2.1. In particular, w-potentials within & are not C* dddifffeomorphssss and are only non-decreasing. We can afford not meeting the first constraint as the proof for Theorem by Krichene et (2015), which establishes the existence of the normalization constant in Equation (6), only requires Lipschitz- continuity. As to the second constraint, we build an augmented class &' that contains increasing o--oeentials and show that it results in the same updates for AMPO 4.
Let A, be the n--dimensional probability simplex. क' (d's :R [0, 1), 1h },). with
We define the parameterized class
i < 0, if < x S Vi,
(x)
z-
Eimr Vi n
if Σソ
v
if
14

---

Published as a conference paper at ICLR 2025
where 1 1 = The functions within Φ' are increasing and continuous. We start by showing that 6. E &' and Φv E d are equivalent for Equation (6), for the same parameters v. At each iteration t, have that
π + (a s) of (s, (1))+)) (1))+1))
(N)+t))
= (s, (5+1))
This due to the following two facts. For x S 0, (o'(z) and (o(x)) are the same function. Also, ( (nf ft (s,a) a) (X)++1)) has to be less or equal to
result of the projection, meaning that
nf (s, a) + (X); 1, where (+)((-) and ( are equivalent.
Since and d' induce the same policy, share that same normalization constant, and have the same value at 0, they also induce the same expression for Equation (7).
C FURTHER DISCUSSION ON BREGMAN DIVERGENCE
To provide an intuition on the Bregman divergence, we report here some discussion from Orabona (2020). Given mirror map h, the associated Bregman divergence between two points x, y E X is defined as
Da (x. == h(x)) h(v) (Vh(y).x x
Since h is convex, the Bregman divergence always non-negative for x, X. We build more intuition for the Bregman divergence using Taylor's theorem. Assume that h is twice differentiable in an open ball B around y and æ E Then there exists 0 S a s such that
Bx y) = h(x) h(y) Vh(u) (x v) = (a - u) Vh(2)(x y),
where - orT + a ayy. Therefore, the Bregman divergence can be approximated by squared local norms that depend on the Hessian the mirror map h. This means that the Bregman divergence will behave differently on different areas of X, depending on the value of Hessian of the mirror map.
D PROOF OF THEOREM 2.2
In this section we outline how we use the proof' given by Xiao (2022) in order to obtain Theorem 2.2.
D.1 PRELIMINARY LEMMAS
We start by presenting two preliminary lemmas, which are ubiquitous the literature on PMD. The first characterizes the difference in value between two policies, while the second characterizes the PMD update.
Lemma D.1 (Performance difference lemma, Lemma Xiao (2022)). For any policy n, n² ((A)$ and E A(S),
(T(P) v (1) ---I---
I-n -- Banay [00 n xT3)
Lemma D.2 (Three-point decent lemma, Lemma 6 in Xiao (2022)). Suppose that C R is a closed convex set, f : - is a proper, closed convex function, Dn(-.-) is the Bregman divergence generated by a mirror map h. Denote rint dom the relative interior of dom h. For any æ E rint dom h, let
a arg
min u dom Anc
[f(u) + (D(,,)))
Then t rint dom h CC and for any E domhnn
f(x*) + (x x) f(u) Dn(u, Dn(uxx*).
A convex function is proper if dom is nonempty and for all dom f. f(z) > -00. A convex function is closed, if it is lower semi-continuous.
15

---

Published as a conference paper at ICLR 2025
D.2 QUASI-MONOTONE UPDATES
We first prove that PMD enjoys quasi-monotonic updates (Equation (3)), that is PMD updates have an upper bound on how much they can deteriorate performance.
Proposition D.3 (Lemma 11 in Xiao (2022)). At each time L 2 0, we have
(Th V.. n
1+1
T 2 0.
Additionally, we have that
-v)))((()) 2 11--111
max dES
10 -Qell max |T m
(10)
-
Proof. Using Lemma D.2 with x+ = obtain
n,t1 C = A(A), f(u) (Q: u), x= n² and m²+1,
(n. - $ (T m²) - Di (T , DT +.t1)
(11)
By rearranging terms and noticing D, (T m.) = 0, we have
(ne 7 r+1 - x 2 ( ++1 x + (T - 2 0.
(12)
Equation (10) can be obtained using the performance difference lemma and (12):
V++ (n) V'())=
E B
[1Qt, m -3']
E
snddt
+1
(Q': TT -+41 - -r3))
1
[(0) - 2. .+1 - m33)
E Bandtti
1
[| Q¹ loll | milh1]
7 B +1
D.3 CONVERGENCE GUARANTEE
We can now prove the convergence guarantee reported in Equation (4). The proof of our result is obtained combining the analysis from the sublinear convergence guarantee with the analysis for the inexact linear convergence guarantee given by (Xiao, 2022, Theorems 8 and 13). For two different time t, t' 2 0, denote the expected Bregman divergence between the policy x' and policy x² where the expectation is taken over the discounted state visitation distribution of the optimal policy di by
D Bandy (Di(T m)))
Similarly, denote the expected Bregman divergence between the optimal policy n² and of by D = Esndy
Theorem D.4 (Theorems 8 and 13 in Xiao (2022)). Consider the PMD update in (2), at each iteration 2 O, we have
1 E[Db]]
1
4
(*) [M)]
S T
mh 7) + (1 (1- (1
max
I Qi
7)2
$ES
KTT
Proof. Using Lemma D.2 with x = - C = A(A), f(u) = (Q', u), x n" and = 11 we
have that
(m :3 - < Du(a", m²) - D m(11)- - -
16

---

Published as a conference paper at ICLR 2025
which can be decomposed
(ne - - (n ",- =5) S DA(TT ,,,",,,, x²) ) Da(== ( (T ,-),
Taking expectation with respect to the distribution d' over states and with respect to the randomness of PMD and dividing both sides by Te, we have
Randy [ - =]
+ E [ssd; [ T3 - -
E[; - Di+1 - D.+ (13)
We lower bound the two terms on the left hand side of (13) separately. For the first term, we have that
E
Bandy
[
-
- m
1 E
E
Bawd!!
(2h1 - x
7 E E amdd 110 =5 - m {
1
E
E Band =
[0 - 0z: n:t11
1- 7
E [d' v++ (dt)) 1 7
E
E
snad
[0: - Q2 n n
2 [d( vd-)
1
E
B
Ban
[10: - 1|
1
?
2
(ch))
E
max
11 - 10%-0% 120
EES
where (a) follows from Lemmas D.3 and the fact that 1 (s) 2 (1 (s) Vs S, and (6) follows from D.I. For the second term, we have that
E Bs (0 - +
E [and: [1Q' .. - + Eawdy
10: D2 oे m
(b) ([V*(u) - V'(n)((1 - -7) + Bandy
(Q2 Q' ~ - m5-
> B[V*(u) - v'(w](1 -7) E
max
1@: Q'
sES
where (b) follows from Lemma D.I.
Plugging the two bounds in (13), dividing both sides by (1 7) and rearranging, we obtain
D[D++1] M 7) + E[V v'())] S E[D; Di+1l E (dd - V [dū)]
N 7)
1-7
Q
Qi
max #ES
(1 - 7)2
Summing up from 0 to T - 1 and dropping some positive terms on the left hand side and some negative terms on the right hand side, we have
(V) - Σ E[V) < D
E [V - v2()))
1/c
1 1-7
KT
4
E max 10:: Q'll
+ (1 7)2
Notice that E (T - rs,, a) [0, 1). By dividing T on both side, we yield
the statement.
E TRAINING DETAILS
We give the hyper-parameter parameters we for training in Tables 2, 3 and 4. Hyper- tuning was performed differently for each method. For Gridworld, we performed grid search over the hyper-parameters and selected those that maximized the averaged performance of the negative entropy
17

---

Published as a conference paper at ICLR 2025
and fz-norm over 256 randomly sampled Gridworld environments. For Basic Control Suite and MinAtar, we used the optuna library to search over the hyper-parameters to maximize the average performance of the negative entropy all environments. We used "CmaasSaapler" the sampler. Lastly, we used Weights and Biases to optimize the hyperparameters for MuloCo in order to maximimize the average performance of the negative entropy over all environments.
option for
Table 2: Hyper-parameter settings of PMD for different sets of environments
Parameter
Grid Mordd
MuJoCo
Number of environment steps
2 ~ 2.5e5
lc7
Number of environments
64
2048
Unroll length
32
10
Number of minibatches
1
128
Number of update epochs
32
8
Adam learning rate
le-4
Sgd learning rate
40
Gamma
0.99
0.99
Max grad norm "
1.0
0.1
0.5
Table 3: hyper-parameter settings of AMPO for different sets of environments
Parameter
BCS
MinAtar
Number of environment steps
5e5
le7
Number environments
4
256
Unroll length
128
128
Number of minibatches
4
Number of update epochs
16
Adam learning rate
4e-3
7e-4
Gamma
0.99
0.99
Max grad norm
1.4
AMPÓ learning rate
0.9
0.9
Table 4: Hyper-parameter settings of OpenAI-ES, which we used in the tabular setting, and of Sep-CMA-ES, which used in the non-tabular case.
OpenAI-ES
Sep-CMA-ES
Population Size
512
128
Number of generations
512 (256 for MuJoCo)
600
Sigma init
0.5
2
Sigma Decay
0.995
Learning rate
0.01
18

---

Published as a conference paper at ICLR 2025
F COMPUTATIONAL COSTS OF ES
The computational costs associated with ES are inevitable in this field of research that focuses on automatically discovering optimizers. In the literature, the most popular methods to discover optimizers are meta-gradients (Oh et al., 2020; Jackson et al., 2023) and ES (Lu et al., 2022). Discovering algorithms using meta-gradients consists in introducing some meta- meta--aaameters that influence the agent training procedure, training batch of agents, and differentiating through the training procedure w.r.t. the meta-parameters to maximize the final performance of the agents. ES consists in running several training procedures in parallel and estimating the meta-gradient as in (9), therefore avoiding the need for differentiation. This property particularly desirable in settings, like ours, where the agent is updated many times and differentiating through the whole training procedure is unfeasible due to memory constraints, meaning that it is necessary to limit the differentiation to the final updates. We decided employ ES due to this property and because has been found to be more successful in discovering optimizers (Jackson et al., 2024). Additionally, our experiments implemented in JAX, which through parallelism and just in time compilation renders the computational costs of ES feasible.
19


Reviews:
[{"review_id": "LsUFQi9Sxw", "reviewer": "Reviewer_WJaW", "summary": "This paper studies the impact of different mirror graphs within the Policy Mirror Descent (PMD) framework in reinforcement learning, traditionally dominated by negentropy graphs. Through empirical testing, the authors identify alternative mirror mappings that significantly improve performance and generalization in different environments, challenging the default choice of mirror mappings in PMD.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The author presents a novel approach by replacing the conventional mirror map with a learned mirror map, offering valuable insights for reinforcement learning research.\n2. The experiments cover both a tabular setting and a real-world application, enhancing the potential for this work to be applied to real-world problems.", "weaknesses": "1. There are some typos in this paper, for example, in equation (5), it should be $\\pi_{\\theta_t}$ instead of $\\pi_{t}heta$?\n2. The author didn't provide related works of this paper, making it hard to understand and follow. And it is better to provide more training details in Appendix C.\n3. It seems that the author only provides comparisons of their method with the negative entropy and d $\\ell_2$-norm mirror map. Can the author provide comparisons with other methods?\n4. The computational cost of the proposed method appears to be relatively high. It would be beneficial for the author to include a discussion comparing the computational cost of this method to that of existing approaches. This comparison would provide a valuable disscussion for understanding the trade-offs in efficiency and practicality.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bnUNxL2a6Z", "reviewer": "Reviewer_4HAM", "summary": "This paper explores the role of mirror maps in the Policy Mirror Descent (PMD) framework in reinforcement learning (RL), which includes various algorithms for policy optimization. Traditionally, the negative entropy mirror map, giving rise to the Natural Policy Gradient (NPG), is the most studied, but the authors argue that this choice is often sub-optimal. Through empirical experiments, they show that alternative, learned mirror maps can enhance PMD's performance across different RL environments, starting with tabular settings like Grid-World and scaling to more complex scenarios like MinAtar and MuJoCo. Their findings suggest that mirror map choice has a more substantial impact on PMD performance than theoretical convergence guarantees indicate. The manuscript highlights the potential for learned mirror maps to generalize across tasks and calls for further theoretical research to understand the influence of mirror maps on PMD's effectiveness.", "questions": "1. How does the mirror map learned on e.g. Hopper perform on very different environments (other than Mujoco's)? Is performance/training dynamics significantly different from using the negative entropy mirror map?\n2. I find it surprising that using negative entropy or l2 norm does not solve Asterix-MinAtar under 10M steps. I personally recall experimenting with Asterix and rather easily solving it. Could you provide a more detailed description of how hyperparameters were selected for each method, or include ablation studies showing the impact of key hyperparameters?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 3, "strengths": "- Empirical validation: the authors provide empirical evidence that challenges the reliance on the negative entropy mirror map in PMD.\n- Comparison with theory: the paper highlights a gap between theoretical bounds and actual empirical performance.\n- Algorithm: the use of ES makes sense here and is well justified.\n- Clarity: the background and methodology sections are well written.", "weaknesses": "- Reliance on ES which may not be very scalable (to higher dimensional mirror map parameterization).\n- Experiments only include small environments where finding better mirror maps may not lead to significantly superior learning curves. Testing on a broader array of environments, especially those with different reward structures or high-dimensional state spaces would enhance the robustness and applicability of the findings.\n- Experiments do not seem statistically significant and may hinder hyper-parameter optimization for some methods and not e.g. for the negative entropy mirror map baseline.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jHqlpOMGAW", "reviewer": "Reviewer_CN9i", "summary": "This paper investigates the Policy Mirror Descent in RL, which is known for its finite-time convergence guarantees and encompasses various algorithms by selecting a mirror map. However, the study explores the potential of alternative mirror maps (instead of negative entropy mirror map, which leads to the Natural Policy Gradient method), demonstrating that the conventional choice often results in suboptimal performance. In empirical studies, the authors use evolutionary strategies to identify more effective mirror maps, showing improved performance in tabular and non-tabular environments. The research highlights that the choice of mirror map significantly affects PMD's effectiveness, challenging existing theoretical assumptions and suggesting new directions for future research.", "questions": "Q1: Although the experiments demonstrate that the learned map is more effective, I believe this does not necessarily conflict with the theoretical assertion that these methods share the same bound. Could the authors provide further explanation on this point?\n\nQ2: In lines 165-166, what is the relationship between Bregman divergence and the mirror map $h$? Could this be further clarified?\n\nQ3(minor): There is a typo in Equation 5 (\\theta).", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The research effectively highlights the limitations of existing convergence guarantees in capturing the impact of mirror maps. It provides robust empirical evidence that challenges the conventional choice of mirror maps, offering a fresh perspective on the potential of PMD. \n\n2. The experimental design is well-conceived and rigorous, aligning closely with the preceding analysis.\n\n3. The paper is well-written, with clear definitions and precise descriptions. Despite covering a substantial amount of theory, it remains accessible and easy to read.", "weaknesses": "From my perspective, this paper makes a satisfactory and engaging contribution, though there is always room for improvement.\n\n1. The paper offers compelling empirical evidence but does not provide a comprehensive theoretical framework to elucidate why certain mirror maps outperform others. Although developing such a framework is challenging, it represents a promising avenue for future research. \n\n2. the study's focus on specific benchmark environments may not fully encompass the diversity and complexity of real-world scenarios in reinforcement learning, potentially limiting the generalizability of the findings.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)