Non-convex entropic mean-field optimization via Best Response flow

TL;DR

This paper investigates non-convex entropy-regularized optimization on probability measures using the Best Response flow, establishing conditions for convergence and applying results to reinforcement learning policy optimization.

Contribution

It extends the application of Best Response flow to non-convex functionals by identifying conditions for contraction and global optimality, relaxing convexity assumptions.

Findings

Best Response flow can be made a contraction with appropriate regularizer.

Unique fixed point corresponds to the global minimizer.

Applications demonstrated in reinforcement learning policy optimization.

Abstract

We study the problem of minimizing non-convex functionals on the space of probability measures, regularized by the relative entropy (KL divergence) with respect to a fixed reference measure, as well as the corresponding problem of solving entropy-regularized non-convex-non-concave min-max problems. We utilize the Best Response flow (also known in the literature as the fictitious play flow) and study how its convergence is influenced by the relation between the degree of non-convexity of the functional under consideration, the regularization parameter and the tail behaviour of the reference measure. In particular, we demonstrate how to choose the regularizer, given the non-convex functional, so that the Best Response operator becomes a contraction with respect to the $L^1$-Wasserstein distance, which ensures the existence of its unique fixed point that is then shown to be the unique global minimizer for our optimization problem. This extends recent results where the Best Response flow was applied to solve convex optimization problems regularized by the relative entropy with respect to arbitrary reference measures, and with arbitrary values of the regularization parameter. Our results explain precisely how the assumption of convexity can be relaxed, at the expense of making a specific choice of the regularizer. Additionally, we demonstrate how these results can be applied in reinforcement learning in the context of policy optimization for Markov Decision Processes and Markov games with softmax parametrized policies in the mean-field regime.