MePoly: Max Entropy Polynomial Policy Optimization

TL;DR

MePoly introduces a polynomial energy-based policy model that explicitly represents complex, multi-modal distributions, enabling effective entropy maximization and outperforming existing methods in decision-making tasks.

Contribution

It proposes a novel polynomial energy-based policy parameterization that provides explicit density and leverages the moment problem for universal approximation.

Findings

Effectively captures complex non-convex distributions

Outperforms baseline methods in diverse benchmarks

Enables exact entropy maximization

Abstract

Stochastic Optimal Control provides a unified mathematical framework for solving complex decision-making problems, encompassing paradigms such as maximum entropy reinforcement learning(RL) and imitation learning(IL). However, conventional parametric policies often struggle to represent the multi-modality of the solutions. Though diffusion-based policies are aimed at recovering the multi-modality, they lack an explicit probability density, which complicates policy-gradient optimization. To bridge this gap, we propose MePoly, a novel policy parameterization based on polynomial energy-based models. MePoly provides an explicit, tractable probability density, enabling exact entropy maximization. Theoretically, we ground our method in the classical moment problem, leveraging the universal approximation capabilities for arbitrary distributions. Empirically, we demonstrate that MePoly effectively captures complex non-convex manifolds and outperforms baselines in performance across diverse benchmarks.