Learning mirror maps in policy mirror descent

Authors: Carlo Alfano, Sebastian Rene Towers, Silvia Sapora, Chris Lu, Patrick Rebeschini
Primary area: reinforcement learning
Keywords: Policy optimization, mirror descent

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.

Decision: accept (poster)
Reviewer scores: 6, 6, 8

Reviews and discussion:
<thread>
creation date:
1730641239513

thread_type:
official_review

<thread_object>
creator: reviewer_cn9i

<content>
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.

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. 

2. The experimental design is well-conceived and rigorous, aligning closely with the preceding analysis.

3. 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.

1. 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. 

2. 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.

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?

Q2: In lines 165-166, what is the relationship between Bregman divergence and the mirror map $h$? Could this be further clarified?

Q3(minor): There is a typo in Equation 5 (\theta).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730658287630

thread_type:
official_review

<thread_object>
creator: reviewer_4ham

<content>
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.

strengths:
- Empirical validation: the authors provide empirical evidence that challenges the reliance on the negative entropy mirror map in PMD.
- Comparison with theory: the paper highlights a gap between theoretical bounds and actual empirical performance.
- Algorithm: the use of ES makes sense here and is well justified.
- 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).
- 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.
- 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.

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?
2. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730686689666

thread_type:
official_review

<thread_object>
creator: reviewer_wjaw

<content>
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.

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.
2. 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$?
2. 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.
3. 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?
4. 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.

questions:
See weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732387485125

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General comment

comment:
We would like to thank the reviewers for their careful remarks and for taking the time to review our paper. In particular, we thank Reviewers *CN9i* and *4HAM* for recognizing that our work challenges conventional choices in PMD and that our paper is well written. We also thank Reviewer *WJaW* for acknowledging that our work offers valuable insights. We address here a common concern on the computational costs of ES and reply to reviewer specific questions below. 


**[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 [1,2] and ES [3]. Discovering algorithms using meta-gradients consists in introducing some meta-parameters that influence the agent training procedure, training a 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 estimating this gradient by running several training procedures in parallel and estimating the gradient as in Line 263 of our paper, therefore avoiding the need for differentiation. The two methods have a similar computational cost, as both require training multiple agents in parallel to update the meta-parameters. We decided to employ ES as it has been found to be more successful in discovering optimizers [4]. Additionally, our experiments are implemented in JAX, which through parallelism and just in time compilation renders the computational costs of ES feasible. As to the scalability of ES to higher dimensional mirror maps, we highlight that all maps belonging to the $\omega$-potential mirror map class can be efficiently parametrized by a single scalar function, i.e. the $\omega$-potential, regardless of the dimension of the action space.

[1] Oh, Junhyuk, et al. "Discovering reinforcement learning algorithms." Advances in Neural Information Processing Systems 33 (2020).

[2] Jackson, Matthew T., et al. "Discovering general reinforcement learning algorithms with adversarial environment design." Advances in Neural Information Processing Systems 36 (2023).

[3] Lu, Chris, et al. "Discovered policy optimisation." Advances in Neural Information Processing Systems 35 (2022).

[4] Jackson, Matthew Thomas, et al. "Discovering Temporally-Aware Reinforcement Learning Algorithms." The Twelfth International Conference on Learning Representations (2024).
</content>
</thread_object>

</thread>


Metareview:
metareview:
In this paper, the authors argue that the well-studied negative entropy mirror map, which results in natural policy gradient, is often sub-optimal. They empirically show that learning a proper mirror map using extensive search methods, such as evolutionary strategies,   
can result in improved performance in tabular and non-tabular environments. This work highlights the importance of the choice of mirror map in the performance of policy mirror descent algorithms and calls for further research to understand the influence of mirror maps on the effectiveness of these algorithms.

(+) The paper is well-written and easy to follow. The authors used clear definitions and precise descriptions, which helps their work to be accessible to a large audience.
(+) The experimental design is well-conceived and rigorous, and it is closely aligned closely with the analysis.
(+) It properly challenges the conventional choice of mirror maps and offers a fresh perspective on the potential of policy mirror descent algorithms.

Despite some issues (mainly minor) raised by the reviewers, which I strongly recommend that the authors address them, the paper could be useful for the community.