MePoly: Max Entropy Polynomial Policy Optimization

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.

💡 Research Summary

MePoly (Maximum‑Entropy Polynomial Policy Optimization) introduces a novel policy parameterization that bridges the gap between expressive multimodal policies and the need for explicit, tractable probability densities in maximum‑entropy reinforcement learning (MaxEnt RL) and imitation learning (IL). Traditional Gaussian or Gaussian‑mixture policies are limited in expressivity and often suffer from mode collapse, while recent diffusion‑based or flow‑matching policies can represent complex supports but lack a closed‑form likelihood, making entropy and KL terms difficult to compute accurately.

The core idea of MePoly is to model the policy as an energy‑based model (EBM) whose energy function is a polynomial in the action variables. For a given state (s), a neural network predicts natural parameters (\lambda(s)). These parameters are combined with a feature vector (T(a)) that contains all monomials (or, more stably, orthogonal Legendre polynomials) up to a chosen degree (K). The policy density is then defined as

\