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

Thank you sincerely for the effort and time you spent reviewing our manuscript. We have taken care to address each point you raised, making revisions highlighted in blue within the revised paper. We believe these adjustments will effectively alleviate your concerns.

W1. We agree that several methods enhance policy expressiveness to handle multi-modal action distributions; however, our approach is fundamentally different. Instead of aiming to model or express the entire multi-modal distribution, we explicitly decompose the behaviour policy into its constituent modes and actively select the best mode to generate actions at each state.

The methods mentioned by the reviewer focus on capturing multi-modality for expressive policies, while our goal is to simplify the learning process by isolating and leveraging the most relevant mode for decision-making. As such, these methods are not in direct competition with our approach but could provide additional baselines. Some of the strong baselines had already been included in the original submission; specifically, our original Table 1 included LAPO, WCGAN, and WCVAE. Now, in Table 1, we have also included DMPO as you suggested.

W2. We have added the hyperparameter details to the Appendix for clarity and completeness.

Q1. In the original Section 6.2, we focused on demonstrating that learning on one mode improves upon behaviour cloning (BC and MDN) and weighted imitation learning (AWAC and LOM). To address the reviewer's request, in the revised version, we have added WCGAN and WCVAE as baselines.

The results show that both methods handle multi-modality better than a single Gaussian learned by AWAC but still underperform compared to LOM. This is because WCGAN and WCVAE cannot entirely eliminate actions from sub-optimal modes, which LOM addresses explicitly.

Q2. The benchmark involves tasks with highly multi-modal action distributions, where the policy must select from diverse behaviours to achieve specific predefined goals. While we acknowledge the relevance of LAPO, Diffusion QL, and DMPO, and plan to include these comparisons in future revisions, we believe that the current strong baselines already cover a broad range of different approaches. Our empirical evidence provides strong support for the effectiveness of the idea we propose. The comparisons, though interesting, are unlikely to change the fundamental claim of our paper—that focusing on mode selection has the potential to improve policy performance in multi-modal scenarios.

Q3. Thanks for spotting this! The log-likelihood in a GMM is $\log \sum_i w_i p_i$, where $w_i$ is the mixing coefficient and $p_i$ is the Gaussian component. In the code, we maximised $\sum_i (\log w_i + \log p_i)$ but this was a mistake. Fortunately, we found that the two loss functions produce comparable results and do not affect our overall conclusions. We have updated all the experimental results accordingly in the revised manuscript.

Q4. In the original paper, we have provided an ablation study in Section 6.3, where we evaluated LOM with different numbers of Gaussian components $M = \\{1, 5, 10, 15, 20\\}$ and analysed the effects of these choices. The study demonstrated that $M$ influences the domain of the imitated actions, with improvements achieved at moderate values of $M$. This highlights the importance of balancing expressiveness and generalisation when selecting $M$. In the reply to reviewer N1vf, we have also added additional experiments showing the performance when $M$ is tuned fully offline using cross-validation to optimise the goodness-of-fit metric for the mixture of Gaussians. This demonstrates a practical approach to parameter selection in a purely offline setting.

Please note that optimal tuning of offline RL parameters in a fully offline manner remains an open problem. LOM has only a main parameter to tune, making it relatively simple compared to other methods that involve a larger number of hyperparameters.

Finally, we thank you again for raising important questions about LOM. Have we adequately addressed the main concerns? Please feel free to let us know if there are additional concerns or questions.

Hi Reviewer XhoW,

Thank you for your valuable suggestions on our manuscript. We have incorporated your feedback and made improvements to the paper. We kindly ask if you would consider revising the review score. Your further input would also be greatly appreciated.

Thank you sincerely for the effort and time you have spent reviewing our manuscript. We have carefully analysed and discussed your valuable question.

We believe that LOM can yield better results theoretically under certain conditions. Let us denote the policy that performs only offline policy improvement in the second step (without mode selection) as $\pi_2$. Since $\pi_2$ averages over conflicting and suboptimal modes, its expected value $V^{\pi_2}(s)$ may be lower than the Q-values of high-performing actions selected by $\pi_L$. We assume the following condition: for any states $s$ from the initial state distribution, we have:

$

E_{a \sim \pi_L(\cdot \mid s)} [Q^{\pi_2}(s, a)] \geq V^{\pi_2}(s).

$

This implies:

$

E_{a \sim \pi_L(\cdot \mid s)} [A^{\pi_2}(s, a)] \geq 0,

$

where $A^{\pi_2}(s, a) = Q^{\pi_2}(s, a) - V^{\pi_2}(s)$ is the advantage function under policy $\pi_2$. With the expected advantage being non-negative, the performance difference between $\pi_L$ and $\pi_2$ becomes:

$

V^{\pi_L}(s) - V^{\pi_2}(s) = \frac{1}{1 - \gamma} E_{s' \sim d_{\pi_L},\ a \sim \pi_L(\cdot \mid s')} [A^{\pi_2}(s', a)] \geq 0.

$

This result follows from the performance difference lemma (Kakade and Langford., 2002; Schulman et al., 2015).

Therefore, under the condition, LOM can produce a better policy than the one learned from performing only offline policy improvement without mode selection. We recognise that this is an interesting problem, and we will investigate it further—specifically, we will determine the exact conditions under which the policy improvement holds.

We would like to thank you again for raising the important question about LOM. Please feel free to let us know if you have any additional concerns or questions.

References

Kakade, S., & Langford, J. (2002). Approximately Optimal Approximate Reinforcement Learning. ICML.

Schulman, J., Levine, S., Abbeel, P., Jordan, M., & Moritz, P. (2015). Trust Region Policy Optimization. ICML.

We are grateful for your time in evaluating our manuscript. The manuscript has been accordingly updated in blue to reflect these changes. We hope that our responses will effectively resolve your concerns.

W1. We appreciate your observation regarding the use of established techniques. While LOM employs tools like GMMs and weighted imitation learning, its novelty lies in reframing the problem as a mode-selection problem. To the best of our knowledge—and as acknowledged by the reviewers—this approach is original.

Unlike methods that model the entire multi-modal action distribution (e.g., GANs or VAEs) or assume unimodality, LOM focuses exclusively on the highest-reward mode. This reframing enables a simpler, more efficient policy learning process that avoids unnecessary complexity while maintaining robust theoretical guarantees (Theorems 1-3). By selecting and imitating a single, most promising mode, LOM addresses multi-modality in a way that is both practical and novel. More detailed comparisons to existing work are presented in the Multi-modality in offline RL, Section 2.

W2. In Figure 4, when we increase the number of Gaussian components, the standard deviation of the highest-reward mode becomes smaller and smaller (the area of the red dots are gradually collapse into a single point). Since the policy $\pi_L$ is learned by imitating the actions from the highest-rewarded mode (i.e., the red dots), we say increasing the number of Gaussian components narrows domains of action learning. It is as narrow as a dot, reducing the potentially useful actions. You may zoom in the Figure to check the highest-reward mode in Figure 4.d and Figure 4.e.

Q1. The conservative nature of the hyper Q-function and the weighted imitation learning method make us use a greedy policy safely. The hyper Q-function conditioned on the behaviour policy inherently regularizes action evaluation, maintaining a conservative stance (Kumar et al., 2020). The imitated action mode is derived from a fixed dataset, so the learned policy is constrained towards the behaviour policy in terms of KL-divergence (Equation 9).

Following the suggestion, we modify the definition in Equation 7 with $\zeta_g(s) = \arg \max_{u \in \mathcal{A}} \alpha^u(s)$. Here $\zeta_g$ is a conservative hyper-policy which select the mode that the behaviour policy most likely to take. We report the experiment results in the following table. It shows the conservative mode selection method is worse than a greedy selection, indicating that overly constraining to the behaviour policy limits policy improvement.

We would like to thank you again for raising these questions about LOM. Have we sufficiently addressed the main concerns? Please feel free to let us know if there are additional concerns or questions.

Reference

Kumar, A., Zhou, A., Tucker, G., & Levine, S. (2020). Conservative Q-Learning for Offline Reinforcement Learning. NeurIPS.

We sincerely appreciate the time and effort you dedicated to reviewing our paper. We have addressed each of these points individually and revised the paper in blue accordingly. We hope that our responses will effectively resolve your concerns.

W1. Our original Theorem 2 demonstrates that the proposed algorithm learns a policy $\pi_L$ that is at least as good as both the composite policy induced by the greedy hyper-policy $\pi_{\zeta_g}$ and the behaviour policy $\pi_b$. Although we did not provide lower bounds in our original submission, we believe that this result already adds theoretical support to the strong empirical performance.

We have taken your comment into consideration, and have attempted to derive such bounds. This further analysis has been included in the manuscript as Theorem 3, and we summarise these additions in what follows.

There are two inequalities in the theorem: $V^{\pi_L}(s) \geq V^{\pi_{\zeta_g}}(s)$ and $V^{\pi_{\zeta_g}}(s) \geq V^{\pi_b}(s)$. We have been able to derive a lower bound for the first inequality and have added it to the revised manuscript as Theorem 3. This result states that, $\forall s \in \mathcal{S}$,

$

V^{\pi_L}(s) - V^{\pi_{\zeta_g}}(s) \geq \frac{1}{1-\gamma} \hat{\eta}(\pi_L) - \frac{A_{\max}}{1-\gamma} \sqrt{\frac{1}{2} D_{KL}(d_{\pi_L} || d_{\pi_{\zeta_g}})},

$

where $A_{\max} = \max_{s, a} |A^{\pi_{\zeta_g}}(s, a)|$, $\hat{\eta}(\cdot)$ is the approximated expected improvement (defined in Section 4.5), $\pi_L$ is the policy optimising the LOM objective (Equation 11). The equality holds when $\pi_L(a \mid s) = \pi_{\zeta_g}(a \mid s)$ for all $s \in \mathcal{S}$, as the policy improvement $\hat{\eta}(\pi_L)$ equals zero and the KL-divergence $D_{KL}(d_{\pi_L} || d_{\pi_{\zeta_g}})$ is also zero. The bound shows that the performance of the learned policy $\pi_L$ is guaranteed with the expected performance improvement quantified by the advantage and reduced by a penalty term proportional to the divergence in their state visitation distributions—emphasizing that maximizing advantage while minimizing divergence leads to better performance.

To establish a lower bound for the second inequality, we need to relate $V^{\pi_{\zeta_g}}$ and $V^{\pi_b}$. This can be achieved by constructing the following equation:

$

V^{\pi_b}(s_t) = \sum_a \pi_b(a_t \mid s_t) [r(s_t, a_t) + \gamma E_{s_{t+1} \sim \mathcal{P}(\cdot \mid s_t, a_t)} [V^{\pi_b}(s_{t+1})]] = \sum^M_{i=1} \alpha^i(s_t) \sum_a \phi_i(a_t \mid s_t) [r(s_t, a_t) + \gamma E_{s_{t+1} \sim \mathcal{P}(\cdot \mid s_t, a_t)} [V^{\pi_b}(s_{t+1})]]

$

Here we define $\alpha_{\max}=\max_{s} \max_i \alpha^i(s)$, and we have

$

V^{\pi_b}(s_t) \leq \alpha_{\max} \sum^M_{i=1}\sum_a \phi_i(a_t \mid s_t) [r(s_t, a_t) + \gamma E_{s_{t+1} \sim \mathcal{P}(\cdot \mid s_t, a_t)} [V^{\pi_b}(s_{t+1})]] = \alpha_{\max} \text{First Term} +  \alpha_{\max} \text{Second Term}

$

where,

$

\text{First Term} = \alpha_{\max} E_{a_t \sim \pi^{\zeta_g}(s_t)} [r(s_t, a_t) + \gamma E_{s_{t+1} \sim \mathcal{P}(\cdot \mid s_t, a_t)} [V^{\pi_b}(s_{t+1})]]

$

$

\text{Second Term} = \alpha_{\max} \sum_{i \neq \zeta_g(s)} \sum_a \phi_i(a_t \mid s_t) [r(s_t, a_t) + \gamma E_{s_{t+1} \sim \mathcal{P}(\cdot \mid s_t, a_t)} [V^{\pi_b}(s_{t+1})]]

$

Our goal is to extract the term $V^{\pi_{\zeta_g}}(s_t) = E_{{a_{t+i}} \sim \pi^{\zeta_g}(s_{t+i})} [\sum^T_{i=t} \gamma^{i-t} r(s_{t+i}, a_{t+i})]$. It seems to appear in $\text{First Term}$. However, each expansion on $E_{s_{t+1} \sim \mathcal{P}(\cdot \mid s_t, a_t)} [V^{\pi_b}(s_{t+1})]$ leads an additional term, like the second term, inside the expectation of the $\text{First Term}$, making it challenging to induce $V^{\pi_{\zeta_g}}(s_t)$. Consequently, we provide only a non-worse guarantee in Theorem 1.

W2. We appreciate your suggestion to explore automatic model selection techniques to tune LOM's hyper-parameter. We have now implemented cross-validation to select the number of Gaussian components $M$ from $\\{1, 5, 10, 15, 20\\}$. This process optimised negative log-likelihood as a measure of goodness of fit.

In the cross-validation process, we employed 3-fold cross-validation, training the GMM model until convergence for $M=\\{1, 5, 10, 15, 20\\}$, and selected the $M$ that minimised the negative log-likelihood on the validation set. We run $3$ different random seeds and report the mean and standard deviation of the model's performance. As shown in the following table, this approach allows us to identify a near-optimal $M$ that enhances model performance while keeping the selection process efficient.

It is important to notice that optimal tuning of offline RL parameters in a fully offline manner is still an open problem. All the baselines we have used in our comparison feature one or more hyper-paramaters, which are usually optimised in a similar way, i.e. by on-line tuning. For this reason, we believe that the approach we have taken aligns with current practice and the results we have reported provide a fair comparison. As such, we don't believe that this aspect is a weakness specific to our methodology. Your suggestion has been valuable: as seen in the previous table, optimising a goodness of fit metric produces performance results close to optimal.

We believe that the revised manuscript has become much stronger after incorporating your valuable feedback. We would like to thank you again for raising important questions about LOM. Have we sufficiently addressed the main concerns? Please feel free to let us know if there are additional concerns or questions.

Hi Reviewer N1vf.

Thank you for your valuable suggestions on our manuscript. We have incorporated your feedback and made improvements to the paper. We kindly ask if you would consider revising the review score. Your further input would also be greatly appreciated.

Thank you very much for increasing the score! We are glad that we could address your concerns.