Title: Imitation Beyond Expectation Using Pluralistic Stochastic Dominance
Venue: NeurIPS 2025
Accepted: yes
Keywords: imitation learning
Abstract: Imitation learning seeks policies reflecting the values of demonstrated behaviors. Prevalent approaches learn to match or exceed the demonstrator's performance in expectation without knowing the demonstrator’s reward function. Unfortunately, this does not induce pluralistic imitators that learn to support qualitatively distinct demonstrations. We reformulate imitation learning using stochastic dominance over the demonstrations' reward distribution across a range of reward functions as our foundational aim. Our approach matches imitator policy samples (or support) with demonstrations using optimal transport theory to define an imitation learning objective over trajectory pairs. We demonstrate the benefits of pluralistic stochastic dominance (PSD) for imitation in both theory and practice.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper identifies and addresses a key deficiency in imitation learning (IL): the tendency for algorithms to learn policies that match the *average* performance of demonstrators, potentially leading to behaviors that are a compromise between different expert strategies and are not truly preferred by any single demonstrator. The authors propose a more robust and intuitive objective: **Pluralistic Stochastic Dominance (PSD)**. The goal is to learn an imitator policy whose distribution of rewards is provably better than the demonstrators' distribution, not just in expectation, but across all reward quantiles and for an entire class of plausible reward functions (conical sum rewards).\n\nTo make this concept practical, the work forges a strong connection between PSD and optimal transport (OT). It shows that achieving PSD is equivalent to minimizing a worst-case OT cost. Since solving this directly is intractable, the authors derive a computationally efficient, convex upper bound they term \"Matched Subdominance Minimization,\" which can be solved as a linear program. The paper presents algorithms for both a tabular \"fully realizable\" setting and a deep reinforcement learning setting where the objective serves as a loss for policy gradient updates. The experimental results on several environments convincingly show that the PSD approach successfully preserves distinct behavioral modes from the demonstration set while improving upon them, outperforming a suite of strong IL baselines on novel and well-suited evaluation metrics.\n\nSTRENGTHS AND WEAKNESSES:\n#### Strength\n\nThis shift from matching moments to dominating distributions is a significant conceptual contribution to the field.\n\nThe theoretical development is rigorous and elegant. The authors do an excellent job of connecting the high-level goal of PSD to a practical algorithmic solution through the lens of optimal transport. The relaxation from the intractable worst-case OT problem to the tractable \"Matched Subdominance Minimization\" is clever and provides the key that makes this framework usable. The presentation is exceptionally clear, with Figures 1 and 2 being particularly effective at providing intuition for the core concepts. The experimental validation is also a high point; the authors not only compare against relevant and strong baselines but also propose new evaluation metrics (stochastic and Pareto dominance rates) that directly measure what the paper claims to achieve, which is a hallmark of thoughtful research.\n\n#### Weaknesses\n\nDespite these significant strengths, the work has a few limitations that warrant discussion. The most critical is the framework's reliance on pre-defined, hand-engineered reward features. This assumption, while common in the classic IRL literature, is a major departure from many modern IL approaches (like GAIL) that can learn directly from high-dimensional states (e.g., pixels). This limits the immediate applicability of the method to domains where such features are not readily available. The authors rightly acknowledge this as an area for future work. Additionally, the methodology relies on solving an optimal transport problem in the learning loop, which may present scalability challenges for very large numbers of demonstration and imitator trajectories compared to methods that rely on simpler adversarial training.\n\nQUESTIONS:\n1. A point of clarification regarding the implementation details for the fully-realizable setting. Line 268 mentions that the training demonstrations are matched with candidate trajectories by \"solving a quadratic program based on Eq. (10)\". However, there is no Equation (10) in the paper. Is this referring to the optimization problem in Definition 3.7 (line 203), perhaps with the L2 regularization on the transport plan `γ` making it a quadratic program?\n\n2. Second, concerning the policy model optimization (Algorithm 1), the OT problem is solved in the inner loop to generate training signals. Could the authors comment on the practical computational overhead of this step, especially as the number of demonstration $N$ and imitator $M$ samples grows?\n\nLIMITATIONS:\nThe choice of experimental environments is logical and serves to clearly illustrate the core properties of the PSD framework. However, to fully establish the scalability and robustness of the policy optimization algorithm, it would be compelling to see results on more standard, high-dimensional continuous control benchmarks. For instance, tasks from the MuJoCo suite (e.g., Humanoid) would challenge the policy optimization with more complex dynamics and action spaces, providing stronger evidence of its general applicability in deep RL settings.\n\nFurthermore, the true power of pluralistic imitation would be most apparent on large-scale, heterogeneous datasets that are becoming standard in offline RL and IL research. While the paper's experiments use synthetically generated demonstrations with clear modes, real-world datasets are often more \"messy.\" It would be highly impactful to evaluate PSD on such data. For instance, the **D4RL benchmark** contains mixed-quality data from diverse sources (e.g., scripted policies, human players, and partially trained agents), which is a perfect use-case for a method designed to improve upon, but not average over, demonstrations. More recent efforts like **MimicGen** or the **IKEA ASM** datasets provide demonstrations from multiple humans with inherent variability in strategy and skill. Testing PSD in these settings would provide powerful evidence of its ability to handle the complexity of real-world demonstration data, moving beyond the clean modes in the current experiments and truly validating its \"pluralistic\" promise.\n\nFu, J., Kumar, A., Nachum, O., Tucker, G., & Levine, S. (2020). D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219.\n\nMandlekar, A., et al. (2023). Mimicgen: A data generation system for scalable robot learning using human demonstrations. arXiv preprint arXiv:2310.17596.\n\nBen-Shabat, Y., Yu, X., Saleh, F., Campbell, D., Rodriguez-Opazo, C., Li, H., & Gould, S. (2021). The ikea asm dataset: Understanding people assembling furniture through actions, objects and pose. In Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (pp. 847-859).", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes an approach to imitation learning by introducing Pluralistic Stochastic Dominance (PSD) as an alternative to traditional expectation-based or risk-sensitive objectives. Rather than matching average performance, PSD ensures that the learned policy provides a better distribution over reward outcomes than the demonstrator for all conical sum reward functions. This is achieved via an optimal transport-based formulation that matches demonstration and imitation trajectories while enforcing a form of reward dominance. \nThe paper provides a theoretical formulation of PSD and its connection to stochastic dominance, a practical relaxation of the intractable optimization using matched subdominance minimization. The paper also presents learning algorithms for both fully realizable and policy-model-based settings and a suite of empirical evaluations (grid world, Point Bot, Reacher) showing that PSD-based methods outperform or match baseline approaches in both Pareto and stochastic dominance metrics.\n\nSTRENGTHS AND WEAKNESSES:\nQuality: Strong theoretical grounding: PSD is rigorously defined and connected to prior work in optimal transport and stochastic dominance. Derivation of tractable approximations via convex optimization is well-justified and novel. Experiments are well-controlled and convincingly demonstrate the benefits of the approach across multiple settings and metrics.\nWhile the theory is sound, the computational burden of the proposed methods is only partially addressed, particularly the outer maximization over reward functions. The policy optimization algorithm relies on handcrafted reward features, which limits applicability in more general, high-dimensional settings.\n\nClarity: The paper is clearly written and well-organized, with useful diagrams (e.g., Figures 1, 2, 4, and 8) illustrating core concepts. The step-by-step exposition of theoretical constructs, optimization strategy, and experimental setup is effective.\nSome key terms (e.g., “supporting behaviors”, “matched subdominance”) are introduced with heavy notation before intuitive descriptions are provided. The transition from full PSD definitions to the practical relaxation could be clarified with a higher-level overview or example before diving into math-heavy formulations.\n\nSignificance: The concept of pluralistic dominance fills a critical gap in imitation learning by addressing the preservation of diverse high-quality behaviors rather than average performance. Potential for significant impact in applications requiring robustness or multimodal imitation, such as robotics, autonomous driving, and human-agent interaction.\nGeneralization beyond hand-crafted reward features is mentioned as future work but limits current real-world applicability. Experimental validation is compelling but relatively limited in scale; larger-scale benchmarks would strengthen the claim of practical impact.\n\nOriginality: Introduction of pluralistic stochastic dominance as a criterion in imitation learning is novel and conceptually rich. The proposed method differentiates itself clearly from previous approaches like GAIL, MaxEnt IRL, and BROIL in both objective and outcomes.\nSome connections to mixture model-based methods (e.g., InfoGAIL) could be more deeply explored, especially given their relevance to handling multimodal demonstrations.\n\nQUESTIONS:\n1. How does the approach handle or adapt to situations where no explicit reward features are available, as in high-dimensional visual domains? Would it be feasible to integrate learned representations into the PSD framework?\n\n2. While the outer maximization over reward weights is convex, it is noted to be NP-hard. Can the authors provide more insights into how this scales in practice for high-dimensional feature spaces?\n\n3. The method relies on a subdominance margin (e.g., β = 0.5 in experiments). How sensitive are the results to this choice, and what guidance can be offered for tuning this parameter?\n\n4. InfoGAIL* is included as a strong baseline. Could the authors provide further discussion on whether PSD could be augmented with latent variable models for additional gains, or does PSD render such modeling unnecessary?\n\n5. Since policy optimization relies on subdominance-derived signals rather than reward gradients, how stable is learning over time? Any failure cases or instability patterns observed?\n\nEvaluation Criteria Shift:\nA rebuttal addressing questions 1, 2, and 3, especially regarding how PSD could generalize beyond known reward features and scale computationally, would significantly increase confidence in broader applicability and lead to a higher score.\n\nLIMITATIONS:\nThe paper explicitly discusses core limitations:\n- it assumes access to known reward features (limiting generalization).\n- current methods are restricted to deterministic dynamics; extending to stochastic settings is left for future work.\n- the computational hardness of exact PSD verification is acknowledged and addressed via approximations.\nFurther discussion of sample efficiency, especially under high-dimensional state/action spaces, and the feasibility of deployment in complex real-world environments would improve the section.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a method for imitation learning based on stochastic dominance over the distribution of possible rewards linearly induced by the demonstrations. To do so, the authors use an optimal transport relaxation. The authors evaluate their approach on a lava world, point bot and reacher.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n- The proposed approach relies on a unique approach and outperforming the demonstrators across a wide range of reward functions is a powerful result with solid theoretical grounding. \n\nWeaknesses\n- One concern I had was the gap between the proposed motivation of the paper and the experiments. Generally, if we are assuming a mixture of behavior in our demonstrations, then why not simply match the distribution of demonstrations (e.g., with the very popular diffusion policy approach (Chi et al., 2023))? While there are practical cases where we might be interested in multiple reward functions (the authors mention risk-sensitivities), the experiments do not examine these settings. This is of concern, given that the reward function is assumed to be linear with the features. As such, it’s unclear to me whether the paper is actually more effective than the baselines at the types of problems it's attempting to address, rather than just toy constructions with randomly sampled rewards. \n- It’s fairly clear but worth mentioning that the proposed approach introduces a more complex optimization problem than in typical imitation learning. Consequently, it’s unlikely to have immediate practical implications.\n\nQUESTIONS:\nQuestion:\n- Several of the inverse RL baselines rely on interaction with the environment. My understanding from the paper is that PSD does not require further data collection. Would allowing interaction with the environment allow for stronger theoretical guarantees?\n\nScore increase:\n- More extensive experimental analysis, particularly on a less toy example where multiple reward objectives are desirable.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-07-02"}