Title: Learning Intractable Multimodal Policies with Reparameterization and Diversity Regularization
Venue: NeurIPS 2025
Accepted: yes
Keywords: reinforcement learning

multimodal policies

diversity
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


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a policy gradient method for multimodal actors, and distance-based diversity regularization to boost exploration for diverse policies. Experiments are mainly conducted on two diversity-critical benchmarks, to show the improved diversity of the policies and learning performances. Experiments on conventional benchmark MuJoCo show the worse performances of the proposed methods against unimodal baseline SAC.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nThe paper proposes a new sampling-based diversity measure for regularization purpose.\n\nThe two diversity-critical benchmarks are more proper to evaluate the multimodal policy representation.\n\nThe experiment analysis is quite comprehensive.\n\nWeakness:\n\nFig. 2 is hard to interpret, I cannot understand how Fig. 2(a) unifies the two types of actors.\n\nThe unification of amortized and diffusion actors with reparameterized policy gradient seems trivial.\n\nThe paper mentions that the sampling-based diversity measure can be parallelized over GPUs for generating multiple samples, but this inevitably introduces additional computational costs during training.\n\nOn traditional MoJoCo benchmarks, the proposed methods perform worse than unimodal policy SAC.\n\nMissing references and baselines. The multimodal policy representation in RL has been studied for a long time, and I think there are quite a lot of papers missing as references and baselines from this paper. Discussion and comparisons against these baselines would be necessary to prove the advantages of the proposed methods.\n\nFor mixture-of-Gaussian:\n\n[1] Peng, X. B., Chang, M., Zhang, G., Abbeel, P., and Levine, S. MCP: learning composable hierarchical control with multiplicative compositional policies.\n\n[2] Ren, Jie, et al. \"Probabilistic mixture-of-experts for efficient deep reinforcement learning.\n\nFor gradient optimization on diffusion policy:\n\n[3] Ren, Allen Z., et al. \"Diffusion policy policy optimization.\n\n[4] Chen, Yuhui, Haoran Li, and Dongbin Zhao. \"Boosting continuous control with consistency policy.\"\n\n[5] Ding, Zihan, and Chi Jin. \"Consistency models as a rich and efficient policy class for reinforcement learning.\"\n\nQUESTIONS:\nIn Fig. 5 (a), why are there a drop in performances in other baselines for simple and medium environments?\n\nHow much additional computation is introduced by distance-based diversity measure during training?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 1, "ORIGINALITY": 2, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes an actor-critic RL algorithm, Diversity-regularized actor-critic (DrAC), to train multimodal policies. The authors motivate their approach by noting that while previous works have employed generative models such as diffusion models as policies, they are limited by the fact that computing policy likelihoods for such models is intractable. To solve this issue, DrAC considers a parameterization that maps current states and a latent noise vector to the action, which also includes existing works on diffusion models. The authors propose training these actors through PGRep, formalized in Proposition 1. Experiments are performed in a modified PointMaze environment to study multimodal behaviours, game content generation to study the effectiveness of DrAC for generative tasks, and MuJoCo benchmark for continuous control tasks.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n- This work unifies amortized actors and diffusion actors under a single framework, which could be useful for studying such parameterizations.\n- The paper contains diverse experiments, demonstrating that DrAC elicits multimodal behaviours, generates diverse latent vectors for a generative model, and performs competitively with other methods on continuous control tasks.\n- The authors have taken care to use the same network architecture across all methods for comparison. They also tune SAC properly on the MuJoCo benchmark, matching my personal observation that it is indeed a very strong baseline and most methods do not actually outperform SAC on these tasks.\n- The authors have provided a detailed hyperparameter table for DrAC and baselines in the appendix, as well as provided the code in the supplementary material. While I did not study the code, the details in the appendix are quite thorough, and the authors’ efforts in promoting reproducibility are appreciated.\n\n**Weaknesses**\n\n- Proposition 1 asserts that any actor expressible as a deterministic mapping $a = f_\\theta(s,z)$, where $z$ is sampled from some noise distribution, can be trained with the gradient $\\nabla_a Q \\nabla_\\theta f_\\theta$. This seems to be a restatement of the deterministic policy gradient theorem [1], except that the policy takes an additional input. As noted by the authors, existing methods like DACER (and also Diffusion QL [2] in the offline setting) use DPG to train actors. Because this practice of taking gradients through the Q-function is well understood and widely used, could the authors clarify what conceptual or algorithmic advance is provided by PGRep?\n- The assumption in Proposition 1 that $\\nabla_z f_\\theta(s, z)$ is full rank for all $s \\in \\mathcal{S}$ is a very strong assumption. It is well known that the gradients of a neural network develop a low-rank structure [3], which means the conditions for Proposition 1 will almost never be satisfied in practice.\n- The paper adopts a very specific operational definition of diversity and makes two claims that merit reconsideration.\n    - Lines 46–50 and 106–108 state that DACER is *unique* in accounting for decision diversity because it injects Gaussian noise into the sampled action. This overlooks the fact that any actor that samples from a multimodal policy (e.g., diffusion actors trained to sample approximately from the Boltzmann distribution of $Q$ [4,5,6,7]) inherently produces action diversity without additional noise, and DACER’s unimodal Gaussian perturbation does not necessarily translate into richer behavioural diversity. The current wording, therefore, understates prior work that achieves diversity implicitly.\n    - Diversity is measured via the geometric mean of pairwise action distances, which favours a single, broadly dispersed cloud over well-separated clusters. In many RL settings, separated clusters better capture multimodal behaviour, with each cluster corresponding to a distinct strategy, whereas a diffuse unimodal cloud may just reflect stochastic deviation around one mode. The paper would benefit from explaining why a “balanced spread” is preferable to explicit mode coverage. Since this regularizer is intended as a proxy for entropy, an empirical check would help: e.g., compute the actual entropy of the two distributions in Figure 3 - does distribution X indeed have higher entropy than distribution Y? Demonstrating that the geometric mean aligns more closely with entropy than alternative measures would strengthen the motivation for this choice.\n- Some closely related works [6,7] also learn a policy that samples from the Boltzmann distribution of the Q-function, and are not discussed. In particular, DQS [6] also exhibits multimodal behavior in a 2D navigation setting with sparse rewards, very similar to the experiment in Section 5.1\n\n*[1] Silver, David, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. \"Deterministic policy gradient algorithms.\" In International conference on machine learning, pp. 387-395. Pmlr, 2014.*\n\n*[2] Wang, Zhendong, Jonathan J. Hunt, and Mingyuan Zhou. \"Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning.\" In The Eleventh International Conference on Learning Representations, 2023.*\n\n*[3] Baker, Bradley T., Barak A. Pearlmutter, Robyn Miller, Vince D. Calhoun, and Sergey M. Plis. \"Low-rank learning by design: the role of network architecture and activation linearity in gradient rank collapse.\" arXiv preprint arXiv:2402.06751 (2024).*\n\n*[4] Psenka, Michael, Alejandro Escontrela, Pieter Abbeel, and Yi Ma. \"Learning a diffusion model policy from rewards via q-score matching.\" arXiv preprint arXiv:2312.11752 (2023).*\n\n*[5] Yang, Long, Zhixiong Huang, Fenghao Lei, Yucun Zhong, Yiming Yang, Cong Fang, Shiting Wen, Binbin Zhou, and Zhouchen Lin. \"Policy representation via diffusion probability model for reinforcement learning.\" arXiv preprint arXiv:2305.13122 (2023).*\n\n*[6] Jain, Vineet, Tara Akhound-Sadegh, and Siamak Ravanbakhsh. \"Sampling from Energy-based Policies using Diffusion.\" arXiv preprint arXiv:2410.01312 (2024).*\n\n*[7] Ishfaq, Haque, Guangyuan Wang, Sami Nur Islam, and Doina Precup. \"Langevin Soft Actor Critic: Efficient Exploration through Uncertainty-Driven Critic Learning.\" arXiv preprint arXiv:2501.17827 (2025).*\n\nQUESTIONS:\n- It is somewhat surprising that DrAmort works better than DrDiffus, since one would expect diffusion models to be more powerful samplers compared to a simple neural network conditioned on a noisy latent. The authors offer one possible explanation, that the number of denoising steps is too low. Could the authors try with more diffusion steps to get a sense of how well DrDiffus (and the other diffusion methods) can perform if we disregard the extra compute?\n- The parameter $\\beta$ seems like a crucial component of DrAC since it controls the trade-off between diversity and optimizing the Q-function. Could the authors discuss how sensitive DrAC is to $\\beta$ and how to tune this parameter in practice? Including an ablation study would also strengthen the empirical contribution.\n- Another parameter that warrants a discussion is the number of diversity estimation pairs. How does increasing the number of pairs affect the performance of DrAC?\n- From Table 2, the authors use the target entropy for SAC and DACER to be $0.5|\\mathcal{A}|$, whereas the standard practice is to use $|\\mathcal{A}|$. Could the authors explain this particular choice?\n- Also from Table 2, the authors use a fixed entropy coefficient for DACER in the multi-goal and game generation experiments, whereas the paper uses automatic temperature tuning like SAC. Could the authors explain this particular choice?\n- The authors mention (line 318) that DACER uses a more advanced critic learning technique. Could the authors elaborate on this? From their paper, it looks like they use standard double Q-learning.\n- Line 220: the authors state that determining the range of $D_\\theta$ is trivial, and they introduce $\\beta$ to “ease the triviality”. Could the authors explain what they mean here?\n\n**Minor errors and suggestions:**\n\n- Line 132 should also include the citation for Ziebart et al. [X] when discussing maximum entropy RL.\n- Line 134: the definition of $Q_\\text{soft}$ should include entropy for all future states $\\mathcal{H}(\\cdot | s_{t+k})$\n- Line 160: in Equation 6, the model $\\epsilon_\\theta$ does not predict the mean, it predicts the noise at each step.\n- Line 189: “The geometric mean is sensitive to small values as production is taken instead of summation” is grammatically incorrect - “production” is not a valid word to refer to the product of several terms.\n\nLIMITATIONS:\nThe authors discuss avenues for future work, including improved architectures, different diversity metrics, and temperature schedules. However, there is no discussion on the limitations of their approach. Some points to include here could be the strong full-rank assumption in Proposition 1, the fact that the diversity surrogate may not align with the true entropy of the policy, and the dependence of DrAC on the parameter $\\beta$.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces PGRep, a reparameterized policy gradient method that works with a wide range of policies, including amortized and diffusion actors. It avoids using log-probabilities by adding a simple distance-based diversity regularizer, making it easier to train complex policies. The ideas build on existing methods like SAC and SQL, but the paper puts them together into a clear and useful framework. The experiments cover navigation, level generation, and continuous control tasks, showing strong results across settings.\n\nSTRENGTHS AND WEAKNESSES:\n**Strength**\n- The paper provides a clear, general gradient formula that works for any stochastic-mapping actor, such as amortized or diffusion policies. This makes the method broadly applicable.\n- The method works well on a range of tasks, including navigation, level generation, and continuous control. It often outperforms strong baselines, showing that the approach is robust.\n\n**Weaknesses**\n- The reparameterization trick is standard in SAC and SQL-style methods. While some diffusion-based approaches use similar ideas [1] for trajectory optimization, typical diffusion policies for imitation learning do not involve policy gradients or reparameterization.\n- The benchmark for diversity measure is narrow. It only uses the reachable goal count as the diversity metric in evaluation. However, I think some diversity measures in conventional benchmarks are also important and insightful (see questions below).\n\n[1] Huang, Zhiao, et al. \"Reparameterized policy learning for multimodal trajectory optimization.\" International Conference on Machine Learning. PMLR, 2023.\n\n**Typos**\n- **Line 77** There is an extra period “.” at the end of the sentence—please remove it.\n- **Line 142 (Eq. 3)** The first term inside the argmax should read\n  $\n  \\alpha\\,\\varrho\\bigl(\\pi(\\cdot|s)\\bigr)\n  $\n  instead of\n  $\n  \\alpha\\,\\varrho\\bigl(\\pi(\\cdot|S_{t+k})\\bigr).\n  $\n- **Line 197** Replace “i.d.d.” with the standard abbreviation “i.i.d.”\n\nQUESTIONS:\n- For Figure 7(b), would you be able to provide **quality-diversity curves** for each algorithm by sweeping the diversity-regularisation coefficient (e.g., $\\alpha$ or $\\beta$) and plotting the resulting final return versus diversity?  This would reveal (i) how each method moves along the return–diversity frontier, and (ii) whether DrAC’s Pareto front consistently dominates those of the baselines.\n- Related to the previous point, how sensitive is DrAC’s performance to the choice of $\\alpha$/$\\beta$? \n- Figure 8 reports only episodic return on MuJoCo tasks.  To support the claim that DrAC delivers *both* quality and diversity, could you include an explicit diversity measure or visualization in this setting as well (e.g., action entropy, pairwise action distance, trajectory-level variation, or heat map)?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a class of diversity-regularized actor-critic (DrAC) algorithms for learning multimodal policies in RL with continuous actions. DrAC assumes a reparameterizable policy class (including amortized and diffusion actors/policies) and uses the reparameterization trick for calculating the policy gradient. The diversity of actions given a state is promoted by a geometric-based distance loss for a pair of generated actions. The policy optimization loss and this distance loss are balanced by a coefficient that is tuned in a similar way to SAC’s temperature. Two instances of DrAC algorithms based on amortized and diffusion actors, respectively, are empirically evaluated against some common baselines. In multi-goal environments and generative tasks, the amortized variant is shown to be generally better, achieving both high return and more diverse behaviors, while the diffusion variant is shown to be comparable to a diffusion baseline. In a conventional benchmark (MuJoCo), both variants and other amortized/diffusion-based baselines do not outperform SAC and are sometimes worse.\n\nSTRENGTHS AND WEAKNESSES:\nThe strengths of the paper include the following:\n1. The paper provides some valuable insights about the utilities of more expressive policies in RL with continuous actions. On the one hand, unlike recent works that mostly focus on a simple 2D multi-goal toy problem [11, 28, 41], the paper demonstrates the utility of multimodal policies in two settings that are understudied in relevant work. On the other hand, it is surprising that, despite consistent claims of improvement for more expressive policies, they do not outperform the standard Gaussian policy in MuJoCo under the same neural network (NN) structure and hyperparameters. Both sets of results are interesting to many RL researchers who are curious about more expressive policies and call for further investigation on when they are useful.\n2. The paper proposes a novel loss to encourage diversity for multimodal policies and demonstrates its effectiveness. The automatic coefficient adjustment also reduces the need to tune the relevant hyperparameters.\n3. The paper is well-written, clear, and very easy to follow. \n\nThe paper also has some weaknesses:\n1. The current results are suggestive and not thorough.\n    - In the MuJoCo experiments, multimodal policies, being evaluated under the same neural network structure and hyperparameters as standard Gaussian policies, didn't outperform them. It's unclear whether this is due to insufficient hyperparameter tuning or if multimodal policies simply aren't necessary in these environments.\n    - The process of tuning the temperature hyperparameters in the multi-goal experiment is not elaborated, potentially hindering the validity and rigorousness of the results.\n2. There is no analysis on the sensitivity of the temperature hyperparameter $\\beta$. There appear to be two values used in different settings (0.8 and 0.15), and little is known about their sensitivity. Further analysis on this could improve the understanding and applicability of the work.\n3. One of the claimed contributions, specifically the derivation of policy gradient for reparameterizable policies, appears to be highly incremental to prior work (for example, Theorem 3 in Lan et al., 2022) and does not introduce substantial novelty beyond that. A more thorough review and discussion of existing work is necessary to accurately establish its place within the field. Considering other contributions, this is minor but still a weak point of the paper. \n\nLan, Q., Tosatto, S., Farrahi, H., & Mahmood, R. (2022). Model-free Policy Learning with Reward Gradients. In International Conference on Artificial Intelligence and Statistics (pp. 4217-4234). PMLR.\n\nQUESTIONS:\nQuestions that may impact the evaluation:\n1. For the multi-goal achieving experiment, could you elaborate on the process of tuning and selecting the temperature hyperparameter? What hyperparameters are tuned for each method?\n2. How is $\\beta$ tuned in the MuJoCo experiment?\n3. Are DrAmort and DrDiffus sensitive to $\\beta$ in different settings, respectively?\n\nQuestions or suggestions that do not impact the rating:\n1. What is used as the ground truth density for Figure 3?\n2. If I understand correctly, the actor in S$^2$AC should also be considered an amortized actor. However, I’m not sure if the current definition is broad enough to include this instance. Could you shed some light on this? If possible, extending the current definition to be more general would be helpful.\n3. In Line 132, the introduction of the entropy-regularized RL framework is attributed to Haarnoja et al. (2018). However, at least Toussaint (2009) introduced such a framework under the name of probabilistic inference. The discussion around Line 132 should be updated to assign credit properly.\n\nToussaint, M. (2009, June). Robot trajectory optimization using approximate inference. In Proceedings of the 26th annual international conference on machine learning (pp. 1049-1056).\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}