Learning Intractable Multimodal Policies with Reparameterization and Diversity Regularization

AI summary

Overview

• Research area: Deep reinforcement learning (RL) for continuous control, specifically multimodal policy learning, decision diversity, and intractable actor models such as diffusion and amortized actors.
• Technical level: Advanced. The paper assumes familiarity with policy gradient methods, actor-critic architectures, maximum-entropy RL, the reparameterization trick, diffusion models, and Stein variational gradient descent (SVGD).
• Scope: The paper proposes a unified formulation of intractable multimodal actors and a diversity-regularized actor-critic algorithm (DrAC), evaluated on multi-goal PointMaze, a game level generation benchmark, and MuJoCo-v4.

What This Paper Is About

Most continuous deep RL algorithms use deterministic or unimodal Gaussian actors, which cannot represent complex multimodal decision distributions that matter when a task has several valid solutions. Multimodal actors based on diffusion or amortized models can represent such distributions, but their decision probabilities have no closed-form expression, so standard entropy regularization and direct policy gradient do not apply, and existing methods struggle to balance performance, decision diversity, and efficiency at the same time. The paper's goal is to show that these intractable actors can be optimized directly by policy gradient through reparameterization, and to add a diversity objective that never needs explicit decision probabilities.

Key Contributions

1. A unified "stochastic-mapping actor" formulation. The authors define actors as π_θ = {f_θ, p_z}, a parameterized mapping f_θ(s, z) combined with a fixed latent distribution p_z, and show that both amortized actors, f^Amort_θ(s,z) ≡ g_θ(s ⊕ z), and diffusion actors, f^Diffus_θ(s,z) ≡ x_0, fit this definition. They prove that the policy gradient of any such actor can be estimated by the reparameterization trick (PGRT), without accessing gradients of decision probabilities.
2. A distance-based diversity regularization. Instead of entropy, the paper regularizes on pairwise distances between sampled actions, using the geometric mean of pairwise distances on a log scale, which requires no explicit access to decision probabilities.
3. The DrAC algorithm. Diversity-regularized Actor Critic combines regularized policy gradient via PGRT with SAC-style double critics, double target critics, and an automatic coefficient adjustment that drives the measured diversity toward a target diversity value D̂ = log(β√|A|).
4. Two diversity-critical domains plus standard benchmarks. The paper identifies multi-goal achieving and generative RL as domains where multimodality is essential, tests DrAmort and DrDiffus against SQL, DACER, S²AC, and SAC, and concludes that the amortized actor is a promising policy model class.

Main Findings

• Amortized and diffusion actors share one structure: Both can be written as a mapping from state and latent noise to actions, which makes the reparameterization trick applicable to both, and lets Q-function gradients be backpropagated directly into the actor's weights.
• Mean pairwise distance overestimates diversity: When a distribution forms small clusters far apart, the mean pairwise distance metric reports high diversity artificially; the paper uses the geometric mean of pairwise distances instead, on a log scale, and reports that the mean-based version led to underwhelming performance.
• Multi-goal PointMaze trajectories: In the simple maze, DrAmort learns the most diverse and uniform trajectories. Despite using the same diffusion actor, DrDiffus learns to reach different goals while DACER does not. DrAmort and SQL consistently reach the most goals in all mazes, with DrAmort learning faster than SQL.
• Success rates in PointMaze: DrAmort, DrDiffus, and DACER learn fast and stably, while SQL, SAC, and S²AC show a decrease in success rate in the early stage and reach optimal success rate more slowly.
• Few-shot robustness: DrAmort shows the best robustness in all mazes under both perturbations (Removal of half the goals, and Obstacle placements blocking the shortest paths), with SQL outperforming the remaining algorithms. DrDiffus outperforms DACER on robustness against removal. SAC struggles to learn multiple goals and to generalize despite a high temperature.
• Diffusion actors underperform in multimodality: The underwhelming multimodality of diffusion actors may be due to the small number of diffusion steps of 20 used as suggested by DACER; even with that small number, training and inference speed is much slower than amortized actors.
• Generative RL results: DrAmort shows the best episodic returns for both the MarioPuzzle and MultiFacet reward styles. S²AC performs worst on all tasks. DrDiffus and DACER do not outperform SAC. In the return-diversity space, DrAmort dominates other algorithms under MarioPuzzle, and dominates under MultiFacet except for SQL, which shows the best diversity but poor episodic return. DrDiffus, DACER, and SAC are non-dominated with each other.
• SVGD as a gradient estimator: SQL uses SVGD to approximate soft policy iteration and performs badly; the authors suggest that reparameterization trick may be a better gradient estimator than SVGD for RL.
• MuJoCo performance: DrAmort is competitive, achieving the best performance in three out of six locomotion tasks and surpassing SQL in a total of five. DrDiffus is comparable to DACER, but DACER and DrDiffus do not outperform SAC, even though the DACER paper uses deeper networks, more advanced activations, and smaller learning rates than the default SAC settings.
• Not reported: The paper does not report numerical wall-clock or memory measurements; speed comparisons are stated qualitatively. Hyperparameters are said to be detailed in Appendix A.2, which is not included in the available content.

Methodology in Plain English

The authors start by noting that a wide range of expressive policies share a simple structure: draw random noise z from a fixed distribution, then pass the state s and z through a neural network to produce an action. Because the action is a deterministic function of the noise, the chain rule lets you push the critic's Q-value gradient straight back into the network weights, which gives an estimate of the policy gradient without ever computing the probability of an action. This is the reparameterization trick applied to RL.

For diversity, the authors avoid entropy, since entropy needs a probability density that these actors do not have. Instead they measure how far apart actions sampled from the same state tend to be. They find that simply averaging pairwise distances can be fooled by data that forms a few tight clusters far apart, so they use the geometric mean of pairwise distances and take a log, which is more sensitive to small distances and prevents that overestimation. The L2 distance is the metric used.

Putting these together, the critic is trained with a bootstrapped target that includes the actor's current diversity at the next state, and the actor is trained to maximize Q plus a scaled diversity term. The diversity weight is tuned automatically: a learnable scalar is optimized so that measured diversity matches a target value derived from a temperature hyperparameter and the action space dimension, an approach borrowed from SAC's automatic temperature adjustment. Two variants are produced by swapping in different actor models: DrAmort uses an amortized actor that concatenates state and noise, and DrDiffus uses a diffusion actor that iteratively denoises noise over 20 steps.

Experiments include a multi-goal PointMaze from D4RL with three maze maps of increasing difficulty, where reward is sparse (+100 per goal reached, ending the episode), plus robustness tests that remove half the goals or insert obstacles, measured by a five-episode success rate. A game level generation benchmark with MarioPuzzle and MultiFacet reward styles measures episodic return and diversity via average pairwise Hamming distance. Finally, six MuJoCo-v4 locomotion tasks test general performance. All training uses five seeds, the same network architectures and common hyperparameters as SAC's default settings, and curves are smoothed with an exponential moving average with coefficient 0.5.

Why This Matters

For research, the paper provides a general lens for understanding and designing multimodal RL: any actor that maps state and fixed latent noise to actions can be trained by backpropagating Q gradients, which separates the question of actor expressivity from the question of whether the actor is tractable. It also argues, against the current interest in diffusion policies, that the amortized actor is both more diverse in practice and much faster, and it replaces entropy with a probability-free diversity objective that can be attached to any such actor.

Real-world applications the paper motivates or cites:

• Navigation with fallback routes: In a maze with multiple valid paths, a unimodal policy converges to a single shortest route and fails catastrophically when that route is obstructed during deployment; a diverse multimodal policy keeps alternative routes available.
• Zero-sum games: Strategic diversity that requires multimodal policies is described as critical for maintaining robustness against adaptive opponents.
• Game content generation: The game level generation benchmark uses RL to propose new level segments with a pretrained decoder, where diversity of generated levels matters alongside quality.
• Generative science and media tasks: The paper cites RL applications in code generation, music generation, artificial intelligence for science, and drug design, where generating diverse new chemical structures is treated as an essential need besides quality.

Industry relevance: the work targets domains where deployed policies must keep several acceptable behaviors available, such as robotics navigation and procedural content generation, and it releases code at https://github.com/PneuC/DrAC. Because DrAmort's training and inference speed is comparable to traditional actors and significantly faster than diffusion actors, the amortized variant is the more practical choice for latency-sensitive deployment, while diffusion actors are described as needing more powerful networks and more elaborate hyperparameter tuning to realize their potential.

Future Directions

• Scheduling the temperature rather than using the current automatic target-matching scheme.
• Designing better stochastic-mapping actors that fall within the proposed unified formulation.
• Exploring other distance metrics and other diversity metrics beyond log pairwise L2 distance.
• Integrating more learning tricks into DrAC, since the paper states it only investigates a basic implementation.
• Determining whether diffusion actors can match amortized actors with better neural network architectures and more careful hyperparameter tuning.

Target Audience

Researchers and graduate students working on reinforcement learning, continuous control, and generative policy models, especially those interested in multimodal policies, exploration, and quality-diversity trade-offs. It is also relevant to practitioners who need robust policies with multiple fallback behaviors, such as robotics engineers and procedural content generation developers, though the mathematical content on policy gradients, reparameterization, diffusion processes, and SVGD makes it most accessible to readers with a background in deep RL.

Authors’ 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