Learning on One Mode: Addressing Multi-modality in Offline Reinforcement Learning

Authors: Mianchu Wang, Yue Jin, Giovanni Montana
Primary area: reinforcement learning
Keywords: Offline reinforcement learning, weighted imitation learning, multi-modality.

Abstract:
Offline reinforcement learning (RL) seeks to learn optimal policies from static datasets without interacting with the environment. A common challenge is handling multi-modal action distributions, where multiple behaviours are represented in the data. Existing methods often assume unimodal behaviour policies, leading to suboptimal performance when this assumption is violated. We propose weighted imitation Learning on One Mode (LOM), a novel approach that focuses on learning from a single, promising mode of the behaviour policy. By using a Gaussian mixture model to identify modes and selecting the best mode based on expected returns, LOM avoids the pitfalls of averaging over conflicting actions. Theoretically, we show that LOM improves performance while maintaining simplicity in policy learning. Empirically, LOM outperforms existing methods on standard D4RL benchmarks and demonstrates its effectiveness in complex, multi-modal scenarios.

Decision: accept (poster)
Reviewer scores: 6, 6, 8, 6

Reviews and discussion:
<thread>
creation date:
1730176857335

thread_type:
official_review

<thread_object>
creator: reviewer_xhow

<content>
summary:
This paper presents an offline reinforcement learning algorithm that accounts for the multimodal distribution of actions within a given dataset. In the proposed algorithm, the behavior policy is modeled as a Gaussian mixture policy, and the mode yielding the highest expected return is identified based on an approximated Q-function associated with the behavior policy. The policy is then learned using an advantage-weighted actor-critic approach, guided by actions generated by the mode of the behavior policy that maximizes the expected return.

The proposed method was evaluated on standard D4RL tasks and custom tasks based on Fetch tasks with multimodal datasets. Experimental results show that the proposed method outperformed baseline methods.

strengths:
- The idea is interesting and reasonable. While existing methods often focus on enhancing the flexibility of the policy model, the proposed approach aims to extract a unimodal distribution from a multimodal dataset.

- In the experiments, the proposed method outperformed the baseline methods.

weaknesses:
- The comparison with related methods is insufficient. While the proposed approach is interesting for its aim to extract a unimodal action distribution from a given dataset, several methods enhance policy expressiveness to handle multimodal action distributions. Using a Gaussian mixture model is a straightforward and intuitive approach. A recent study explores Gaussian mixture policies and mixtures of deterministic policies based on AWAC (https://openreview.net/forum?id=zkRCp4RmAF). Additionally, diffusion models offer a promising approach (https://openreview.net/forum?id=AHvFDPi-FA). Comparisons with these methods should be included in the experiments.

- Hyperparameters used in the experiments are not provided. Please provide a table or appendix section detailing the hyperparameters used for each experiment, including any tuning process.

questions:
- Why are LAPO, WCGAN, and WCVAE absent from the experiments with extremely multimodal datasets? These methods should be compared to evaluate their capability in handling multimodal action distributions.

- Is it possible to include Diffusion QL and AWAC with a Gaussian policy for comparison? These methods should also be included in the experiments with extremely multimodal datasets.

- In the attached code, I noticed that the loss function for the GMM might be incorrect. The provided code shows the loss as $\mathcal{L} = \sum_i ( \log w_i + \log p_i$ ) where $p_i$ is a Gaussian component.
However, the likelihood in a GMM is given by $\sum w_i p_i $ and its log likelihood is $\log ( \sum w_i p_i )$, which differs from $ \sum_i ( \log w_i + \log p_i$ )$.  Was this modification made intentionally or by mistake? If it was intentional, please provide the rationale.

- How sensitive is the proposed algorithm to the number of Gaussian components? Please provide an ablation study or sensitivity analysis showing how performance changes with different numbers of Gaussian components across various tasks or datasets.

=== comments after the rebuttal ===

The concerns raised in the initial review were appropriately addressed in the rebuttal.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730358601590

thread_type:
official_review

<thread_object>
creator: reviewer_hztn

<content>
summary:
This paper addresses the multi-modality problem in offline reinforcement learning (RL). It proposes to first perform model selection and then further improve the greedy hyper policy that captures the best mode.

strengths:
I enjoy reading this paper. I think its key contribution is learning on one mode. By doing this, the KL-divergence term in offline policy improvement would not burden the policy to cover the entire support distribution generated by the behavior policy.

The key problem lies in the fact that, in the behavior data, the action distribution has some clear modes that should be represented by discrete values and some continuously distributed probability mass within each mode. The paper proposes dealing with these two aspects using different approaches. For the first aspect, it involves selecting between different modes to ensure the learned policy remains within the offline data distribution. For the second aspect, it involves using offline policy improvement with KL-divergence to constrain the continuous action distribution.

weaknesses:
No clear weaknesses.

questions:
I am curious about the theoretical aspect. Clearly, this paper employs a two-step improvement process. Empirically, it seems reasonable that this approach would perform better than doing just the second step without the mode selection. However, I am wondering if there is any theoretical analysis comparing the two-step improvement with only performing offline policy improvement in the second step. Which approach do you think would yield better results theoretically?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730526088889

thread_type:
official_review

<thread_object>
creator: reviewer_eg5z

<content>
summary:
The paper presents LOM, an offline imitation learning approach designed to address the multi-modality of batch datasets. The method employs a Gaussian mixture model to capture the behavior policy, selecting the optimal mode based on expected returns. Subsequently, it estimates the advantage by leveraging information from the selected mode and incorporates this advantage into a weighted imitation learning objective using an exponential weighting scheme. The approach demonstrates competitive performance against baselines on the D4RL benchmark.

strengths:
1. The paper is well-written and easy to follow. 

2. It effectively leverages the multimodal nature of the data to design a practical method.

weaknesses:
1. While the method appears broadly applicable and versatile, it closely resembles existing techniques, such as weighted advantage imitation learning and multi-modal policy networks. The approach mainly combines established methods, making the contribution appear incremental. 
2. Furthermore, although the author analyzes how the number of mixture components impacts the results, the experimental findings require additional clarification. For instance, in Section 6.3, the statement “When M is too large, it results in overly narrow domains of action learning, which may exclude potentially useful actions from the dataset” would benefit from a more thorough explanation, including potential reasons behind this effect. Without such elaboration, the analysis may feel incomplete.

questions:
1. In mode selection, the greedy policy is applied. Is this the optimal strategy for choosing the mode? I wonder if a more conservative approach might yield better results. Could the author elaborate more about this choice?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730692473752

thread_type:
official_review

<thread_object>
creator: reviewer_n1vf

<content>
summary:
This paper introduces Weighted Imitation Learning on One Mode (LOM), a novel approach for offline reinforcement learning that addresses the challenges of multi-modal action distributions by focusing on the single mode with the highest expected return. Using a Gaussian mixture model to identify modes and a hyper Q-function for mode selection, LOM significantly improves performance in multi-modal environments, outperforming other methods on standard benchmarks.

strengths:
1. This paper is well-organized and easy to follow.
2. The paper proposes an innovative approach to handling multi-modal action distributions in offline reinforcement learning, which provides valuable insights for research of muti-modality in RL, by focusing on the most promising action mode and using a Gaussian mixture model to identify modes.
3. This paper presents extensive experiments validating the effectiveness of LOM across diverse benchmarks. The results demonstrate that LOM outperforms state-of-the-art offline RL algorithms in multi-modal settings.

weaknesses:
1. For the theoretical perspective, the result that $V^{\pi^L}(s) \geq V^{\pi^{\zeta_g}}(s) \geq V^{\pi^b}(s)$ is relatively weak, Is it possible to provide a lower bound for the gap between the value functions, or could the author explain the difficulties in doing so?
2. The effectiveness of LOM depends on the choice of the number of Gaussian components (M) in the Gaussian Mixture Model. Have the author tried any automatic model selection techniques for GMMs, such as the Bayesian Information Criterion or cross-validation, and how these might be applied to LOM?

questions:
See weakness.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733156522899

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We would like to sincerely thank all the reviewers and the area chair for their valuable time, insightful comments, and constructive feedback during the review process. Your thoughtful suggestions have greatly contributed to improving the quality of our work, and we deeply appreciate your effort and dedication.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This is a strong paper that tackles a relevant problem in offline reinforcement learning. The reviewers found the approach of selectively applying constraints based on mode structure in the behavior data to be novel and effective. The key strength lies in the paper's insightful observation that learning on a single mode can alleviate issues caused by full distribution matching in offline policy improvement. The method of addressing discrete modes and continuous distributions within modes separately is technically sound. While further empirical validation could strengthen the work, the core idea is compelling.  Based on the innovative approach and insightful analysis, I recommend accepting this paper.