Learning Intractable Multimodal Policies with Reparameterization and Diversity Regularization

TL;DR

This work tackles the challenge of learning expressive multimodal policies in continuous control by reframing intractable multimodal actors as stochastic-mapping policies that can be trained via policy-gradient through reparameterization. It introduces a distance-based diversity regularization that does not require explicit decision probabilities, enabling simultaneous optimization of return and diversity. The proposed DrAC framework unifies amortized and diffusion actors under the stochastic-mapping formulation and demonstrates strong multimodal expressivity and few-shot robustness in multi-goal and generative RL tasks, with competitive MuJoCo performance. The amortized actor emerges as a particularly effective and efficient policy class for multimodal RL, while diffusion actors show slower training, suggesting a preference for amortized architectures in practical settings.

Abstract

Traditional continuous deep reinforcement learning (RL) algorithms employ deterministic or unimodal Gaussian actors, which cannot express complex multimodal decision distributions. This limitation can hinder their performance in diversity-critical scenarios. There have been some attempts to design online multimodal RL algorithms based on diffusion or amortized actors. However, these actors are intractable, making existing methods struggle with balancing performance, decision diversity, and efficiency simultaneously. To overcome this challenge, we first reformulate existing intractable multimodal actors within a unified framework, and prove that they can be directly optimized by policy gradient via reparameterization. Then, we propose a distance-based diversity regularization that does not explicitly require decision probabilities. We identify two diversity-critical domains, namely multi-goal achieving and generative RL, to demonstrate the advantages of multimodal policies and our method, particularly in terms of few-shot robustness. In conventional MuJoCo benchmarks, our algorithm also shows competitive performance. Moreover, our experiments highlight that the amortized actor is a promising policy model class with strong multimodal expressivity and high performance. Our code is available at https://github.com/PneuC/DrAC

Figures (23)

• Figure 1: Motivational examples. A multimodal policy with maximized diversity can achieve multiple goals and enable robustness against environmental changes.
• Figure 2: Key ingredients of our proposed method.
• Figure 3: Average (Avg.) and geometric mean (GM) of pairwise distances for two synthesized data distributions. Average pairwise distance overestimates the diversity of distribution $Y$.
• Figure 4: Evaluation trajectories of all tested algorithms in the simple maze.
• Figure 5: Learning curves in multi-goal PointMaze. Solid lines represent mean performance, and shaded regions indicate standard deviation. All curves are smoothed by the exponential moving average with a coefficient of $0.5$.