Learning on One Mode: Addressing Multi-modality in Offline Reinforcement Learning
Authors: Mianchu Wang
Yue Jin
Giovanni Montana
Venue: ICLR 2025 Poster (2025)
Keywords: Offline reinforcement learning
weighted imitation learning
multi-modality.
Primary area: reinforcement learning
TL;DR: We propose weighted imitation Learning on One Mode (LOM) to address the multi-modality problem in offline reinforcement learning.

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.

Published as a conference paper at ICLR 2025
LEARNING ON ONE MODE: ADDRESSING MULTI- MODALITY IN OFFLINE REINFORCEMENT LEARNING
Mianchu Wang"
Yue Jin
Giovanni Montana* + 'hee Alan Turing Institute
University of Warwick (mianchu. wang, yue. jin. 3, g g.montana @warwick. ac. uk
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 the data. Existing methods often assume unimodal behaviour policies, leading suboptimal performance when this assumption is violated. We propose weighted imitation Learning on One Mode (LOM), 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 complex, multi-modal scenarios.
INTRODUCTION
Offline reinforcement learning (RL) enables policy learning from static datasets, without active environment interaction, making it ideal for high-stakes applications like autonomous driving and robot manipulation (Levine al.) 2020 Ma Ma et al.] 2022] Wang et al. 2024a). key challenge in offline RL is managing the discrepancy between the learned policy and the behaviour policy that generated the dataset. Small discrepancies can hinder policy improvement, while large discrepancies push the learned policy into uncharted areas, causing significant extrapolation errors and poor generalisation (Fujimoto et al. 2019 Yang et al. 2023). Addressing these challenges, existing research has proposed various solutions. Conservative approaches penalise actions that stray into out-of-distribution (OOD) regions Yuu et al. 2020 Kumar et al. 2020), while others regularise the policy by minimising its divergence from the behaviour policy, ensuring better fidelity to the dataset (Fujimoto & Gu 2
Wu et al. 2019). Another solution is weighted imitation learning, which aims to replicate actions from the dataset through supervised learning techniques (Mao al. 2023 Peng al. 2019).
Many real-world datasets introduce an additional challenge: multi-modal action distributions (Wang et al. 2024b) Chen et al. [022)] Zhou et al., 2020). These datasets are common in practice because they often integrate data from diverse sources, such as multiple policies, human demonstrations, or distinct exploration strategies. This diversity arises naturally in domains where the same task can be approached in various ways, leading to states with multiple valid but potentially conflicting actions. For instance, in autonomous driving, different driving styles conservative versus aggressive may lead to different but equally valid ways of navigating a road. Similarly, in robotic manipulation, an object can be grasped in various ways depending the robot's approach, the object's position, and environmental constraints. This multi-modality is not an exception but rather a frequent occurrence in complex, real-world decision-making tasks, as systems often integrate experience from various sources to handle different scenarios. Thus, it becomes crucial to model and manage these multi- modal action distributions effectively.
Most offline RL approaches implicitly assume a unimodal action distribution, which can force policies to converge toward an average action that may not exist in the dataset, leading degraded performance. This limitation is particularly evident scenarios where policies fail to capture complex multi-modai distributions, instead collapsing into suboptimal or invalid averages, studied by Cail jet al. (2023); Wang et al. (2024b): Yang et al. (20226) and illustrated in Appendix A Recent work has attempted to address this by modelling the full multi-mode modal action distribution using expressive

---

Published as a conference paper at ICLR 2025
generative models such GANs, VAEs, and diffusion models (Zhou et al. 2020 Chen et al. 2022 Wang et al. 2023). However, these models often overcomplicate the learning process by capturing the entire action distribution, which is unnecessary when only subset of the modes is relevant for optimal decision-making.
We propose simpler yet effective approach: Weighted Imitation Learning on One Mode (LOM). Our insight is that learning from a single mode the one with the highest expected return is sufficient to generate optimal actions. Instead of modelling the full multi modlll distribution, LOM identifies and focuses on the most promising mode for each state.
LOM operates through three key steps (illustrated in Figure []. First, it models the behaviour policy as a Gaussian mixture model (GMM), capturing the inherent multi-modality in the action space. Each mode in the GMM represents a distinct cluster of actions associated with a state. Second, novel hyper Q-function is introduced to evaluate the expected re- turn of each mode, enabling the dynamic selection of
Behaviour Policy Modelling
the most advantageous one. Finally, LOM performs weighted imitation learning on the actions from the selected mode, ensuring that the learned policy fo- cuses on the most beneficial actions while retaining the simplicity of unimodal policies. This targeted, mode-specific learning strategy simplifies the policy learning process while maintaining even enhancing performance in multi-modal scenarios. By dynami- cally selecting the optimal mode for each state, LOM achieves robust results with reduced complexity.
Mode Selection
Lots, a)l E (O(s,a))
3. Weighted Imitation Learning
This paper has four key contributions: (1) we propose LOM, a novel weighted imitation learning method designed to address the multi-modality problem in of- fline RL; (2) we introduce hyper Q-functions and hyper-policies for evaluating and selecting action modes; (3) we provide theoretical guarantees of con- sistent performance improvements over both the be- haviour policy and the optimal action mode; and (4) empirically demonstrate that LOM outper- forms state-of-the-art (SOTA) offline RL methods across various benchmarks, particularly in multi- modal datasets.
Figure I: The three steps of LOM. (1) Learn network producing the parameters of a GMM to model the behaviour policy. (2) Evaluate each mode via the expected return of its ac- tions and then select the optimal mode Φ² (3) Sample actions from ø² for weighted imita- tion learning.
RELATED WORK
Offline RL The primary challenge in offline RL lies in managing the distribution shift between the behaviour policy, which generated the offline dataset, and the learned policy. This shift can cause the learned policy to produce actions that are not well-represented in the offline dataset, leading to inaccurate value function estimations and degraded performance (Levine et al. 2020 Fujimoto et al.) 2019). To mitigate the risks associated with OOD actions, one approach involves using conservative value functions, which penalise actions that deviate from the behaviour policy (Fujimoto et al. 2019
Kumar al., 20200 Bai et al.) 2022)) Sun et al.) 2022] An et al.) [202]). Another strategy focuses on regularising the learned policy to ensure its proximity to the behaviour policy. This regularisation can be measured using metrics such as mean squared error (MSE) (Beeson Montanal 2024] FFjjimoto
& Gu (221)) more sophisticated metrics like the Wasserstein distance, Lensen-Shamme divergence
(Yang et al. 2022b Wu al.) 2019), and weighted MSE (Ma et al.) 2024). These methods aim to keep the learned policy within safer, well-understood operational space, reducing the likelihood of selecting OOD actions. For a comprehensive review of value penalties and policy regularisation techniques, we refer readers to (Wu et al., 2019). In contrast to these approaches, LOM tackles the distribution shift problem by focusing on a single, optimal mode of the behaviour policy, thereby reducing the risk of OOD actions while still allowing for significant policy improvement.
2

---

Published as a conference paper at ICLR 2025
Weighted Imitation Learning Our proposed method contributes the growing field of weighted imitation learning. which revolves around two fundamental questions: which actions should be
imitated and how should these actions be weighted (Brandtonbrener et al., [02II) I. Most existing
methods imitate actions directly from the dataset (Peng et al. 2019 Wang et al. [2018] Nair et al.
202TT, Wang et al., 2020), while some approaches (Siegel et al. 2020; Mao al., 2023) suggest imitating actions from the most recent policies. These methods typically weigh actions based on their advantage, conditioned either on the behaviour policy (Peng al., 2019) Wang et al. 2018) or on the most recently updated policy (Nair et al. 2021 Wang et al. 2020] (Siegel et al. 20200 Mao et al. 2023). However, a key limitation of these approaches is their underlying assumption that the imitated actions are drawn from unimodal Gaussian distribution, complex, real-world datasets where multi-modal distributions are common (Lynch et al. 2020 Yang et al. 2022b). LOM addresses this limitation by forgoing the unimodal assumption, instead extracting a highly-rewarded Gaussian mode from the behaviour policy.
oversimplification that often fails in
Multi-modality in Offline RL Offline RL datasets are often collected from multiple, unknown policies, leading to states with several valid action labels, which may conflict with one another. Most existing methods employ unimodal Gaussian policy, which inadequately captures the inherent multi- modality in such datasets. To address this limitation, PLAS (Zhou et al. (020) decodes variables sampled from VAE latent space, while LAPO (Chen al. 2022) leverages advantage-weighted latent spaces for further policy optimisation. Additionally, Yang et al. (2022)) demonstrate that
GANs can model multiple action modes. and GOPlan (Wang al. 2024b) extends this by applying exponentiated advantage weighting to highlight highly-rewarded modes. DAWOG (Wang et al. 2024a) and DMPO (Osa et al. 2023) separate the state-action space and learn a conditioned Gaussian policy or a mixture of deterministic policies to solve the regions individually. Techniques such as normalising flows and diffusion models can also model multi-modal distributions by progressively transforming an initial distribution into the target distribution (Wang et al.) [0233 Akimov et al.) 2022). However, these techniques are computationally intensive, making their inference processes inefficient. Beyond direct policy improvements, approaches like TD3+RKL (Cai et al. (023) exploit the mode-secking property of reverse KL-divergence to narrow action coverage, while BAW (Yang al. 2022a) filters actions with values in the leftmost quantile. Despite these advances, most methods attempt to capture or model the entire multi-modal distribution, leading to increased complexity. In contrast, LOM simplifies the learning process by extracting and focusing on a single, optimal mode from the behaviour policy, effectively addressing the multi-modality challenge without the need to model the entire distribution.
PRELIMINARIES
To lay the foundation for our method, we first introduce the Markov Decision Process (MDP) and the concept of weighted imitation learning. An MDP is defined by the tuple M = (8, A, P, 7),, where S is the state space, A the action space, and P the transition dynamics. The function represents the reward, and y is the discount factor. The objective in RL is to learn a policy π that maximises the expected discounted return:
( Eropa
Σ Entr(sn
where 7 = {so, ag. ..... ST, a is trajectory sampled under policy n. The state-action value function Q?' a.) represents the expected return starting from state st, taking action a,, and following ı thereafter:
T Σ
7'rr(s) an)
Q= 0z) = BEst1. ac
(1)
The value function is defined as (*t(s) E. (ss (st, at)], and the advantage function is A As(s,, a,) = (st, a.) V*()).
In offline RL, the agent learns exclusively from a fixed dataset D collected by one or more behaviour policies, denoted as wh. without further interaction with the environment (Levine et al., 2020). Our
approach is based on weighted imitation learning, where the objective is to optimize
J J(T) = Egndr (-)))))))
exp
A (s, log ((a
3

---

Published as a conference paper at ICLR 2025
where 3 is from D, weighted by the exponentiated advantage of the behaviour policy Ty thereby implicitly constraining the KL-divergence between and ne (Wang et al. 2018). Recent works have extended this approach by imitating actions from the current learned policy (Siegel et al.) 2020; Mao et al.
This formulation encourages learning policy π that imitates actions
2023). weighting actions based the advantage conditioned on the current policy (Nair al.
2021
Wang et al. 20200).
METHOD
In this section, we present the LOM method, specifically designed for offline RL with heterogeneous datasets. We begin by modelling the behaviour policy using a Gaussian Mixture Model (GMM) to capture the inherent multi multi-modality. Following this, formalise the problem using hyper- Markov decision process (H-MDP), where a hyper-policy dynamically selects the most promising mode for each state based expected returns. Next, we propose hyper Q-function the hyper-policy and introduce a greedy hyper-policy for mode selection at each step for weighted imitation learning. Finally, we prove that the resulting policy of our LOM method outperforms both the behaviour policy and a greedy policy.
evaluate
4.1 MOLTI-MODAL BEHAVIOURAL POLICY
To address the challenge of multi-modaa action distributions in offline RL, we model the behaviour policy Gaussian Mixture Model (GMM). GMMs are flexible in representing complex distributions and can approximate any density function with sufficient components (MeLachlan & Basford 1988). making them well-suited for capturing the heterogeneous nature of offline datasets that often arise from multiple policies or exploration strategies. The behaviour policy (a s) is expressed as:
M Σ
(a 8)
(s)
where a' (s) represents the mixing coefficients for each mode i, and Øi (a |) denotes the Gaussian component for mode at state s. Each mode captures distinct cluster of actions associated with the state. By using a GMM, our model captures the diverse modes present in real-world offline RL datasets, where multiple valid strategies may coexist for the same task. This approach preserves the richness of the original data while avoiding the limitations of unimodal policy approximations.
4.2 HYPER-MARKOV DECISION PROCESS
To extend the traditional MDP framework introduced in Section 3 and formulate the offline RL problem with multi-modal behaviour policy, introduce the Hyper- Markov Decision Process (H-MDP). The H-MDP accounts for multiple policy modes, extending the standard MDP to a higher level of abstraction. Formally, the H-MDP is defined as MI (S, A, P,,, r, Y, Ω, An, P rH), where S, A. P, r, and y retain their standard meanings from the MDP. The novel components introduced are:
{Ø ..... &h),, representing the Gaussian modes, and the hyper-action space Aи = {L, M). which indexes these modes.
The transition dynamics Pnr conditioned on the hyper-action E An and the reward function for u are defined as:
Pu((e+1 82, ui) = Ba (P( st, 2.))
Th (St, un) = Eag [r(se, (r(s) at)]
At each time step. the agent selects Gaussian mode based on hyper-pplicy Ç, which maps states probability distribution over modes, Ç(u s) : x Ah + [0, 1]. The hyper-policy determines the likelihood of selecting each mode in given state. The agent then selects action a according to the Gaussian distribution (a 8). Consequently, the action follows a composite policy æç. which corresponds to the hyper- syper-policy& Ç and can be expressed as:
T<(a s)= Σ ((u sa s),
(3)
HEAy
4

---

Published as a conference paper at ICLR 2025
In the H-MDP, the agent's objective is to learn hyper-policy that maximises the expected discounted return:
T
J(C = E Eropp
Σ
Arrr(s ut)
(4)
2==
where P. is the distribution of state-mode trajectories following the hyper-policy Ç.
(st, ues Se41, 1141....) induced by
4.3 HYPER Q-FUNCTION
To evaluate the quality of selecting a mode u in a given state 8, we define the hyper Q-function (x, u). This function quantifies the expected return when choosing mode at state and subsequently following the hyper-policy Ç. Formally, the hyper Q-function defined as:
22 (s, u) Be
Σ 7 (SE, u4) 80 = 8, 14p
4
(5)
This expectation is taken over future states and modes encountered while following the hyper-polige Ç after selecting mode in state The term TH (st, uz) denotes the expected reward for selecting mode u, in state 8t, with y being the discount factor.
Proposition 1. The hyper Q-function can be linked to the standard value function Q" (s, a) via:
23 (s, u) = Ea (+10) 1Q (s,a)]-
(6)
The proof can be found in Appendix B.1
This proposition establishes critical relationship between the hyper Q-function and the standard Q-function. It shows that the expected return for selecting mode u (i.e., the hyper-action) is equivalent to the expected value of actions sampled from the Gaussian component (a s). evaluated under the composite policy тς. This result is particularly important when dealing with multi-modal behaviour policies, allows the agent to effectively compare and evaluate different modes based on their associated action distributions (a ().
4.4 MODE SELECTION
Since the hyper Q-function evaluates the expected return after selecting mode u at state, we now propose greedy hyper-policy that improves upon the behaviour policy by selecting the most advantageous mode at each state based on the hyper Q-function.
In correspondence with the multi-modal behaviour policy m, defined in Eq. 2 we introduce the behavioural hyper-policy, denoted as (.. which reflects the mode-selection strategy implicit in the behaviour policy that generated the offline dataset. Specifically, (o s) = a (s). where a (s)) is the mixing coefficient for the w-th Gaussian component in the GMM state This formulation establishes the connection between the original behaviour policy na and the H-MDP framework. allowing represent the behaviour policy in terms of mode selection.
To improve upon the behavioural hyper-policy So. we define a greedy hyper-policy So- which selects the mode that maximises the hyper Q-function Q (s, u) for each state:
(s) = arg max (s, HEAH
This greedy hyper-policy ensures that the agent selects the mode with the highest expected return at each state, thereby improving the overall policy performance relative to the behaviour policy.
Theorem 1. The composite policy induced by the greedy hyper-policy improves upon the behaviour policy:
vac (s) 2 (s),, Vs E.
(8)
The proof can be found in Appendix B.2
This improvement property guaranteed by Theorem I forms the foundation of our LOM algorithm. ensures that by systematically selecting modes with the highest expected return at each state, the agent's 's composite policy will, over time, perform at least as well as the original behaviour policy.
5

---

Published as a conference paper at ICLR 2025
4.5 WEIGHTED IMITATION LEARNING ON ONE MODE
Based on the greedy hyper-policy Co. we obtain the corresponding greedy policy "c,, which, according to Theorem is not worse than TT. However, rather than simply imitating "s we apply weighted imitation learning to further improve the policy. Specifically, we aim to find policy π that maximises the expected improvement 17( J(TT) The expected improvement ny(x) can be expressed in terms of the advantage A a) (Kakade & Langford, 2002; Schhlman et al. 2015): v(TT) (-)ad=())aaat(--e)))
(s, a)]. However, in offline RI., state samples from the distribution d are unavailable. Instead, we approximate the expected improvement by using estimate (Kakade) & LLagffrdd 2002): f)(x)) = Exwdre, (.).a>=(--l)) [A (s, a)]. This approximation provides a good
estimate of 17((T) if π and ",,, are close in terms of KI--divergence (Schulman et al., 2015). As a result, we can formulate the following constrained policy optimisation problem:
arg max Σ dieep (8)
Σ T(a 8) (s,
(9)
s.t. Edar (3)DKI(T() 1 MT, ( (8) S..
Solving the corresponding Lagrangian leads to the optimal policy n": x- (a 1s) Zea (a
8) exp AT (s, a)) We then learn a policy by minimising the KI.-divergence to this optimal policy, resulting in the weighted imitation learning objective:
arg min Esdd (-) ( (- s))]
s)][[(
(10)
1 A (s) exp
arg max don (.amt, (11s)
log T(a log s)
(11)
In this framework, the actions to be imitated are sampled from "C,' a unimodal Gaussian policy, thereby addressing the multi-modality issue. Furthermore, actions within this mode weighted according to their advantage, encouraging the policy to concentrate on the highly-rewarded actions. As shown in Theorem 2 the policy learned by LOM has theoretical advantage over both TT, and the original behaviour policy Th.
Theorem 2. The LOM algorithm learns a policy nt. that is at least as good both the composite policy induced by the greedy hyper-policy TC, and the behaviour policy no. Specifically, for all states E S:
VAL (s) 2 Va (s) 2 (s).
(12)
The proof is provided in Appendix B.3
This theorem implies that our method introduces a novel two-step improvement process. LOM achieves further policy improvement compared to one-step algorithms while requiring less off-policy evaluation than multi-step algorithms (Brandfonbrener et al., 202T).
Theorem 3.
The LOM algorithm learns policy TL, whose improvement upon the composite policy induced by the greedy hyper-policy TC, is bounded. For all s E S, we have
A Amax
v (s) co (s) 1- ((LL)
1-+V V D (dr,, Idxx,,))
(13)
where Amax- max ir (s, a)}- The proof is provided in Appendix B.4
The bound shows that the learned policy TI is guaranteed to perform at least as well as the baseline policy Tc,, with the performance improvement quantified by the maximal advantage and reduced by a penalty term proportional the divergence in their state visitation distributions emphasizing that maximising advantage while minimising divergence leads to better performance.
5 PRACTICAL ALGORITHM
In this section, we present the practical implementation of the LOM algorithm. First, we model the multi-modal behaviour policy using mixture density network, as described in Section 5.1 In
Section [32 we outline the mode selection process, which involves learning the hyper Q-function for the behavioural hyper-policy. This hyper Q-function is then used to derive the greedy hyper-policy C and the corresponding greedy policy TG,- which selects a single mode. Finally, in Section (.33 we describe how learn policy through weighted imitation learning the selected mode.

---

Published as a conference paper at ICLR 2025
Algorithm 1 Weighted imitation learning one mode (LOM).
Initialise: A MDN behaviour policy Tr with parameter a target policy de with parameter 6, a value function
Qu with parameter 3 and its ulowly-updated copy Qu behaviour hyper Q-function Qo with parameter
(); offline dataset D.
1: # Learn the MDN behaviour policy.
2: for i .,.,, Las do
3: Update P by minimising L(p) = ((ss,ar) (ooge) (aa st]].
Eq. 15
4: end for
5: for i =o, Ja do
6: Sample transitions T {8s, Pt: 8t41, a++1) ~ D.
7: # Learn the value function
8:
Update & by minimising ((*)) = E- ((Q (.) (re 7Q (81+1, (++)))++
Eq. 16
9:
* Learn behavioural hyper Q function..
Update & by minimising ( Unl (( Basma" (s ([))
11:
* Learn the target policy.
12:
Select the optimal mode u arg ma...e (o()).u)
Sample action from the optimal mode âr ~ (st)
14:
Get the mean the mode ā; = (s:)
D The mean of the mode
15:
Estimate the advantage A(st, â.) = Qu ōe) Qulss, kz)
16:
Update 0 by minimising L(0) -Bn [exp(A(st, dz)) log to(ae st)]-
Eq. 18
17:
Update the target network.
if i mode update_dells = 0 then
19:
1 + pd + (1 - p)p
20:
end if
21:
end for
5.1 BEHAVIOUR POLICY MODELLING
At the start of the LOM algorithm, we model the behaviour policy using mixture density network (MDN) (Bishop) T994). The MDN receives a state & and represents the resulting action distribution with GMM, including its mixing coefficients, locations, and scales. We index the Gaussian components as the hyper-actions, ranging from to M.
In the algorithm, we learn neural network Ta(a| 8) parameterised by p to estimate the behaviour ) which are used to estimate the policy. The network produces the parameters {af" , ,,, J probability density function for each policy mode:
2
(2To((s) exp (-Ila- (-1la-m(6)11)) (4()))12)
(a s)
(14)
(s)(s)2
where (4(s) 2 and o (s) exp(2,).
The mixing coefficient as(s) is processed through softmax function: (s = exp(2"") )/ Σ j-1 exp(2*)), ). The parameters of the network are updated
by minimising the negative log-likelihood:
C(p) - [log (a {)]]
(15)
LOM trains the MDN until convergence, after which the parameter p is fixed for the remainder of the algorithm.
5.2 HYPER Q-FUNCTION LEARNING
We learn the hyper Q-function of the behavioural hyper-policy based Proposition | Specifically, we first learn a f-function, QT of the behaviour policy, and then estimate its expectation over actions sampled from mode We instantiate QT with neural network Q₁ parameterised by s and updated to minimising the TD error:
(16) where Qp- is slowly updated copy of Q Following Proposition we learn Qú, a hyper Q-function the behavioural hyper-policy, which is denoted as Qo parameterised by ó. The optimisation objective of learning Qs is:
((b)= E(s. [Q ar) (r. + (st+(ss11.000 (c41)))
((Φ) E, ~ D.lrD,,ww-iiii (n)iteem((A)) [[Qp(st, u) - Eai (11s) ax)))
(17)
With this hyper fuuuction, get greedy hyper- policy (s(s) arg Aa Qo(s, u). Corre- sponding to C (8), we get the greedy policy with one mode, Tc, (a 8) = (( 1s).
7

---

Published as a conference paper at ICLR 2025
Table 1: Averaged normalised scores on D4RL benchmarks over 4 random seeds. 0 represents the performance of a random policy and 100 represents the performance of an expert policy. Standard deviations are provided as subscripts for multi-modal handling algorithms. The highest mean is highlighted in bold.
Dataset
Env
OneStep
AWAC
CQL
TD3BC
LAPO
DMPO
WCGAN
WCVAE
LOM
balfcheetah
50.4
47.9
47.0
48.3
45.9.03
47.3104
48.2g1.
50.5g1.
51.0
medium
hupper
87.5
59.8
53.0
59.3
51.6.9a
71.2
78.6224
89.0225
100.8
walker2d
84.8
83.1
73.3
83.7
80.7g0s
79.414.7
82.4g17
85.011
85.1 85.1
halfcheetah
42.7
2..
45.2
42.3
45.010.
48.8 10.7
medium-replay
hopper
98.5
69.8
88.7
60.9
58.6
89.2
90.3ga1
09.0ga
99.2.1.1
walker2d
21.7
$2.
72.6.4.
74.0aa.
84.8
halfcheetah
75.1
64.9
75.6
90.7
93.0
91.125.
76.52a.1
80.0g1.7
92.7.1.3
medium-expert
hopper
108.6
100.1
105.6
98.0
105. 105.2=4.7
78.42100
110.0.2.4
100.0 100.0+1.5
110.1:1.4
walker2d
111.3
110.0
107.9
110.1
111.1g0a
109.9.04
99.7gt.
111.0g0.
111.3.08
balfcheetah
88.2
81.7
96.3
96.7
95. 95.9
97.0g10
90.7g1.7
80.0g0
95.2.o.
expert
hupper
106.9
109.5
96.5
107.8
10.
93.6215a
107.3g18
108.0g1.3
111.0;0.8
walker2d
110.7
110.1
108.5
1 110.2
112.2.0.1
111.4403
109.3a1.4
I1 111.0
109.3
halfcheetah
64.7
66.6
73.6
71.7
74.2=1.8
71.421a
64.2ga.1
66.023.
78.6
full repl-replay
hupper
69.8
100.1
98.2
76.8
100.0.a.a
101.5259
80.2g5.
86.5444
102.0.27
walker2d
67.1
78.3
02.7
90.2
96.2-2.8
95.4g18
53.0g5.4
72.0ga.
97.9
5.3 POLICY LEARNING
Finally, we learn the policy Ta by maximising the objective:
1 AT (s, log m |s)
J(0) = E Eaw (11) exPetip B
(18)
expetip (.) the exponential function with clipped range (0, C], where C is a positive number for numerical stability. We estimate the advantage using: A (s, a) x Qu a) Baw (-1s) 1Q (s, a)], where we use the Q-function conditioned on the behaviour policy Te rather than Tc, in order reduce extrapolation error and computational complexity. Further, employ the approach used in (Mao et al. 2023) to replace the expectation over a ~ TC, (: 8) with the Q-value of the mean action of the Gaussian policy TT, to reduce computation.
where |8 a
The algorithm starts with training a MDN to model the behaviour policy, and then iteratively update the hyper Q-function Qá. the value function Qu and the target policy πe. The code has been open sourced
Implementation details and the pseudocode can be found in Algorithm I
EXPERIMENTAL RESULTS
We aim to answer the following questions: (1) How does LOM compare SOTA offline RL algorithms, particularly those handling multi-modality? (2) Can performance improvements be attributed learning on one mode? (3) How does the number of Gaussian components affect LOM's performance?
6.1 LOM ACHIEVES SOTA PERFORMANCE IN D4RL BENCHMARK
We evaluate LOM on three MuloCo locomotion tasks from the D4RL benchmark (Fu et al.
2020):
halfcheetah, hopper, and walker2d. Each environment contains five dataset types: (i) medium 1M
samples from a policy trained to approximately one-third of expert performance; (ii) medium-replay the replay buffer of policy trained match the performance of the medium agent (0.2M for halfcheetah, 0.4M for hopper, 0.3M for walker2d); (iii) medium-expert - a 50-50 split of medium and expert data (just under 2M samples); (iv) expert- IM samples from a fully trained SAC policy (Haarnoja et al., 2018); and (v) full-replay - 1M samples from the final replay buffer of an expert policy. Notably, the medium-replay and full-replay datasets are highly multi-modal, states may have multiple action labels from different policies.
Table shows the benchmark results of our method against SOTA offline algorithms: OneStep (Brandfonbrener et al.] 2021), AWAC (Nair et al., 2021), CQL (Kumar et al. 2020), and TD3BC (Fjjimoto & Gu, 2021). b. We also include four algorithms designed for handling
LAPO (Chen et al. 2022). deterministic mixture policy optimisation (DMPO) (Osa et al. 2023).
GitHub repository: https//gghhuukcccm///aaaaaaaaaa
8

---

Published as a conference paper at ICLR 2025
(a) Environment
(b) BC
(c) MDN (M=4)
(d) AWAC
(e) LOM (M=4)
Figure Comparative study the FetchReach task with highly
datasets. (a) The
FetchReach robot is tasked with reaching one of four specified goals using an expert dataset. The robot arm receives a reward of for reaching the goal in the first quadrant and for reaching any of the other three goals. The dataset contains actions directed toward all four goals, with conflicting directions. (b) Action distribution learned by behaviour cloning using a unimodal Gaussian policy model. (c) Action distribution learned by MDN using GMM policy model. (d) Action distribution learned by AWAC, which applies weighted imitation learning over the entire action distribution using unimodal Gaussian policy. (e) Action distribution learned by LOM.
weighted conditioned GAN (WCGAN), and weighted conditioned VAE (WCVAE) (Wang et al. 2024b). The experimental results are shown in Table where the baselines' results are mainly sourced from Mao et al. (2023) or reproduced using their official implementations. The results demonstrate that L.M outperforms the baseline algorithms in 12 out tasks. tasks with multi-modal datasets, LOM surpasses all the baselines, and achieves a performance improvement ranging from 0.2% to 7.9% over the SOTA results.
6.2 LEARNING ON ONE MODE DRIVES IMPROVEMENTS
To explore the source of LOM's improvements, design three tasks with highly multi-modal datasets: FetchReach, FetchPush, and FetchPickAndPlace In FetchReach, as shown in Figure 2a the objective control a robotic arm to reach one of four specified positions, symmetrically distributed on the zyy-llane. The robot receives a reward of 2 for reaching the position in the first quadrant, 1 for reaching any of the other positions, and otherwise. In FetchPush and FetchPickAndPAl robot is tasked with moving cube to one of four positions Each dataset contains 10 transitions collected by an expert policy. Given a state and four goals in different quadrants, the dataset contains four trajectories corresponding to each goal, making the datasets highly multi-modal.
the desk, with the same reward structure.
We compare LOM with behaviour cloning (BC), AWAC, and MDN, introduced in Section 5.1. Specifically, BC and MDN imitate the behaviour policy using a unimodal Gaussian policy and multi-modal Gaussian policy, respectively. AWAC applies weighted imitation learning to imitate the behaviour policy across all modes (Nair et al. (222T). LOM, however, distinguishes multiple behaviour modes using MDN and focuses on learning from the best mode. Table 2 shows that LOM outperforms the baseline algorithms by an average of 53%.
In Figure 2 we further investigate the behaviour of different policies. Given a state, the multi-modal dataset contains four actions, each targeting distinct goals in separate quadrants. Figure 2b shows that BC learns an averaged action distribution. Figure de that MDN effectively reconstructs the action distribution, allowing LOM select the highest-rewarding mode. Figure 2d shows that while AWAC can learn the best mode (in the first quadrant), it tends to generate OOD actions.
In contrast, Figure 2c illustrates that LOM imitates actions from the best single mode, filtering out suboptimal actions and focusing on highly-rewarde that learning on one mode effectively improves policy performance by concentrating on the most rewarding actions.
ones. These experiments demonstrate
Table 2: Average scores on the tasks with extremely multi-modal datasets.
BC
MDN
AWAC
WCGAN
WCVAE
LOM
FetchReach
14.3
17.1
32.4
42.313.
40.8
47.2
FetchPush
11.8
18.944.
26.8+1.4
41.243.7
41.842.
44.3
FetchPick
3.640.8
8.1
22.7+2.3
29.242.8
30.522.1
34.243.6

---

Published as a conference paper at ICLR 2025
(a) M=2
(b) M=4
(c) M=10
(d) M=20
(c) M=50
Figure 4: Action modes learned by MDN with varying numbers of mixtures. Red dots represent samples from the highest-reward mode. The original actions are clustered around (1. 1), (1, -1), (-1,-1), and (-1, 1) but do not extend beyond these points. In (d) and (e), the red dots collapse into a single point in the first quadrant due to the small standard deviation the mode.
6.3 EFFECT OF THE NUMBER OF GAUSSIAN COMPONENTS
We analyse the influence of the number M of Gaussian components in the LOM algorithm. The results, shown in Figure 4 are based on the multi-modal FetchReach environment. The figure illustrates the decomposition of action modes using blue contours, while red dots represent the actions imitated by LOM. When M is too large (Figure 4d 4e), the standard deviation of the imitated Gaussian mode becomes excessively small, leading to overly narrow the domains of action learning. Conversely, smaller values of M results in modes with a large standard deviation (Figure (4). causing LOM to capture actions outside the dataset and imitate OOD actions. This highlights the need to investigate LOM's ' sensitivity to this
Figure SegureBcome compares comparees LOM across M 51,5, 10, 15, 20} in the medium- replay and full-replay datasets from the D4RL benchmark. The per- formance generally improves as the number of mixture components in- creases, particularly in the medium- replay dataset, where larger M val- ues lead to continuous performance gains. However, in the full-replay datasets, performance increases only up to M 10. This indicates that while increasing the number of com- ponents benefits performance, there may be diminishing returns beyond a certain threshold. Our experiments suggest that M primarily influences the domain of the imitated actions, and improvements can be achieved with a moderate, though not highly specific, value of M.
Medium Replay Debasets
Recby Detasees
310
200
194
19-2
H9L5
10h20
7041
Figure 3: Normalised scores for varying numbers of Gaussian components in the medium replay and full replay datasets.
CONCLUSIONS AND DISCUSSIONS
In this paper, we introduced LOM, a novel offline RL method specifically designed to tackle the challenge of multi-modality in offline RI datasets. Unlike existing methods that aim to model the entire action distribution, LOM focuses on imitating the highest-rewarded action mode. Through extensive experiments, we demonstrated that LOM achieves SOTA performance, and attributed these improvements to the central idea of learning on one mode, which simplifies the learning process while maintaining robust performance.
We observed that the hyper Q-function satisfies a Bellman-like equation, opening the possibility for learning through Bellman updates. However, one challenge that emerges is the difficulty in accurately estimating the Bellman target for the hyper Q-function. as it requires computing the expectation over all possible hyper-actions, which may lead to extrapolation errors. Exploring deterministic hyper-policy is promising future direction to mitigate this issue, potentially simplifying the estimation process and further improving performance. Furthermore, the hyperparameter in LOM is intrinsically linked to the multi-modality present in the dataset. This parameter can be estimated using statistical techniques like bump hunting (Friedman Fisher 1999) and peak finding (Du et al.] 2006), eliminating the need for fine-tuning, However, key challenge lies in the fact that each state corresponds to a distinct multi-modal action distribution, making it computationally intensive to identify and count the modes for every individual state.
10

---

Published as a conference paper at ICLR 2025
Acknowledgments We acknowledge support from UKRI AI Turing Acceleration Fellowship (EPSRC EP/VO24868/1).
REFERENCES
Dmitry Akimov, Vladislav Kurenkov, Alexander Nikulin, Denis Tarasov, and Sergey Kolesnikov. Let offline RL flow: Training conservative agents in the latent space of normalizing flows. In 3rd Offline RL Workshop: Offline RL as a Launchpad, 2022.
Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty sncertainty-based based offline reinforcement learning with diversified Q-ensemble. Advances Neural Information Processing Systems, 2021.
Chenjia Bai, Lingxiao Wang, Zhuoran Yang. Zhi-Hong Deng, Animesh Garg, Peng Liu, and Zhao- ran Wang. Pessimistic bootstrapping for uncertainty-driven offline reinforcement learning. In International Conference on Learning Representations, 2022.
Alex Beeson and Giovanni Montana. Balancing policy constraint and ensemble size in uncertainty- based offline reinforcement learning. Machine Learning, 113(1):443 488, Jan 2024. ISSN 1573-0565.
Christopher M. Bishop. Mixture density networks. Technical report, Aston University, Birmingham, 1994.
David Brandfonbrener, William F Whitney, Rajesh Ranganath, and Joan Bruna. Offline RL without off off-policy evaluation. In Advances in Neural Information Processing Systems, 2021.
Yuanying Cai, Chuheng Zhang, Li Zhao, Wei Shen, Xuyun Zhang, Lei Song, Jiang Bian, Tao Qin, and Tie-Yan Liu. TD3 with reverse KL regularizer for offline reinforcement learning from mixed datasets. In IEEE International Conference on Data Mining, 2023.
Xi Chen, Ali Ghadirzadeh, Tianhe Yu, Jianhao Wang, Alex Yuan Gao, Wenzhe Li, Liang Bin, Chelsea Finn, and Chongjie Zhang. LAPO: Latent- advantage-we -eantage-weighted weighted policy optimization for offline reinforcement learning. Advances in Neural Information Prucessing Systems, 2022.
Pan Du, Warren A. Kibbe, and Simon M. Lin. Improved peak detection in mass spectrum by incorporating continuous wavelet transform-based pattern matching. Bioinformatics, 22(17): (959-2065,.0 07 2006. ISSN 1367-4803.
Jerome H. Friedman and Nicholas I. Fisher. Bump hunting in high-dimensional data. Statistics and Computing, 92): 143, Apr 1999. ISSN 1573-1375.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning, 2020.
Scott Fujimoto and Shixiang Gu. A minimalist approach to offline reinforcement learning. In Advances in Neural Information Processing Systems, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning, 2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International Conference on Machine Learning, 2018
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning.. In International Conference on Machine Learning, 2002.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, 2020.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives open problems, 2020.

---

Published as a conference paper at ICLR 2025
Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet. Learning latent plans from play. In Conference on Robot Learning, 2020.
Chengzhong Ma, Deyu Yang, Tianyu Wu, Zeyang Liu, Houxue Yang, Xingyu Chen, Xuguang Lan, and Nanning Zheng. Improving offline reinforcement learning with in-sample advantage regularization for robot manipulation. IEEE Transactions on Neural Networks and Learning Systems, pp. 1 1-13, 2024.
Yecheng Jason Ma, Jason Yan, Dinesh Jayaraman, and Osbert Bastani. Offline goal-conditioned reinforcement learning via f-advantage regression. In Advances in Neural Information Processing Systems, 2022.
Yixiu Mao, Hongchang Zhang, Chen Chen, Yi Xu, and Xiangyang Ji. Supported trust region optimization for offline reinforcement learning. In International Conference on Machine Learning, 2023.
Geoffrey J McLachlan and Kaye E Basford. Mixture models: Inference and applications clustering. volume 38. M. Dekker New York, 1988.
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. AWAC: Accelerating online reinforcement learning with offline datasets, 2021.
Takayuki Osa, Akinobu Hayashi, Pranav Deo, Naoki Morihira, and Takahide Yoshiike. Offline reinforcement learning with mixture of deterministic policies. Transactions on Machine Learning Research, 2023. ISSN 2835-8856.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-wcighted regression: Simple and scalable off-policy reinforcement learning, 2019.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning, 2015.
Noah Siegel, Jost Tobias Springenberg, Felix Berkenkamp. Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner. Nicolas Heess, and Martin Riedmiller. Keep doing what worked: Behavior modelling priors for offline reinforcement learning. In International Conference on Learning Representations, 2020.
Hao Sun, Lei Han, Rui Yang, Xiaoteng Ma, Jian Guo, and Bolei Zhou. Exploit reward shifting in value-based deep-RI.: Optimistic curiosity- exploration and conservative exploitation via linear reward shaping. In Advances in Neural Information Processing Systems, 2022.
Mianchu Wang. Yue Jin, and Giovanni Montana. Goal-conditioned offline reinforcement learning through state space partitioning. Machine Learning, 11365)2435--265, May 2024a. ISSN 1573-0565.
Mianchu Wang, Rui Yang, Xi Chen, Hao Sun, Meng Fang, and Giovanni Montana, GOPlan: Goal- conditioned offline reinforcement learning by planning with learned models. Transactions on Machine Learning Research, 2024b. ISSN 2835-8856.
Qing Wang, Jiechao Xiong, Lei Han, peng sun, Han Liu, and Tong Zhang. Exponentially weighted imitation learning for batched historical data. In Advances in Neural Information Processing Systems, 2018.
Zhendong Wang, Jonathan Hunt, and Mingyuan Zhou. Diffusion policies as expressive policy class for offline reinforcement learning. In International Conference on Learning Representations, 2023.
Ziyu Wang, Alexander Novikov, Konrad Zolna, Josh S Merel, Jost Tobias Springenberg, Scott Reed, Bobak Shahriari, Noah Siegel, Caglar Gulcehre, Nicolas Heess, and Nando de Freitas. Critic regularized regression. In Advances in Neural Information Processing Systems, 2020.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning, 2019.
12

---

Published as a conference paper at ICLR 2025
Rui Yang, Yiming Lu, Wenzhe Li, Hao Sun, Meng Fang, Yali Du, Xiu Li, Lei Han, and Chongjie Zhang. Rethinking goal-conditioned supervised learning and its connection to offline RL. In International Conference Learning Representations, 2022a.
Rui Yang, Lin Yong, Xiaoteng Ma, Hao Hu, Chongjje Zhang, and Tong Zhang. What is essential for unseen goal generalization of offline goal-conditioned RL? In International Conference on Machine Learning, 2023.
Shentao Yang, Zhendong Wang, Huangjie Zheng, Yihao Feng, and Mingyuan Zhou. A behavior regularized implicit policy for offline reinforcement learning, 2022b.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 2020.
Wenxuan Zhou, Sujay Bajracharya, and David Held. PLAS: Latent action space for offline reinforce- ment learning. In Conference on Robot Learning, 2020.
13

---

Published as a conference paper at ICLR 2025
A MOTIVATING EXAMPLE
To illustrate the challenge of capturing complex multi-modal action distributions, we present a toy example in Fig. 5 In this example, we generate a dataset where each state is associated with multiple valid actions, simulating the multi-modal nature of real-world offline RL datasets. We evaluate four models Gaussian, Conditional Variational Auto-Encoder (CVAE) (Zhou al., 2020 Chen et al. 2022), Conditional Generative Adversarial Networks (CGAN) YYngg et al. 2022b Wang ct al. 2024b), b, and Mixture Density Network (MDN) (Bishop 1994) on their ability to model the multi-modal action distribution.
Subfigures (a-c) show the action distributions Iearned by each model. Specifically, we visualise how well each model captures distinct modes in the action space. The Gaussian model fails separate modes, while the CVAE, CGAN, and MDN demonstrate stronger mode separation, with the MDN showing minimal mode overlap and the clearest boundaries between modes.
Additionally, we assess each model's ability to capture high-reward actions. Subfigures 5 (f-j) illustrate the models' performance in learning positively-rewarded multi-modal action distributions. Our results reveal that the MDN with ranked components closely adheres to in in-distribution actions and achieves superior outcomes in terms of reward, compared to the other models.
Mhdt Casts Data Foal Peamrt
BAY LONEOIIIII WEWBROOK
(a)
(b)
(f)
(g)
(h)
Figure 5: An example of modelling the multi-modal behaviour policy in one-step MDP. The z-axis represents the state, and the y-axis represents the corresponding multi-modal actions. (a) shows the action distribution from the offline dataset. (b)-(e) illustrate the action distributions learned by Gaussian model, conditional VAE, conditional GAN, and MDN, respectively. (f) shows the action distribution of offline dataset with rewards, where actions the first and third quadrants receive reward of 1, and others receive 0. (g)-(i) illustrate the action distributions learned by weighted Gaussian, weighted conditional VAE, and weighted conditional GAN models, respectively. (j) shows the action distribution learned by the top 10 of the 20 MDN components, ranked by hyper Q-function.
B THEORETICAL RESULTS
B. PROOF OF PROPOSITION
Proposition The hyper Q-function can be linked to the standard value function (s, a) via:
Qi (8, u) E (a)g=(1)) (Q (s,
(19)
14

---

Published as a conference paper at ICLR 2025
Proof:
24 (s, u) =Es,.us.
TH ($t, (8), 80 8, do
T-1
=T,( Σ
p(5411=8|
+ r=(s
8,
Σ s)r(sa)
Σ p(s6+1 = E ((u 8)7 (( Σ da s)r(s,a)
t=
Σ
Σ p(5441 =s (110)<so ((ul s)o" ( s)g r),
k,a
T-1
= Σ ((als)r(ssa)) +
Σ
Σ p(8+11 = 8 ())(c((a ,))
t=
T-1
Σ (a 8) r(s,a) + Σ
Σ
8) + r(s,
t= 8,0
Σ la r(s,a) + Σ
Σ
8)7 ((8_t)
t=0
(1) (s)
(20)
B.2 PROOF OF THEOREM
II
Theorem The composite policy induced by the greedy hyper-policy improves upon the behaviour policy:
V (s) 2 (s), E.
Proof. Now, we prove that the improved policy *c, is uniformly good better than the behaviour policy mb.
Te (st)
=EoINTE [r(e, ae) + =Beei [Q(*+11, (x41)]]
-Enwne [r(s,, 0z) + 7 miy |r(8++1, +++
=0 21l
= -Burne, (r() 0.) + yE, M [r(s 01+1) +..
-
-1 +
[r(4 )I]]
> Eornns [r(, an) + B 1 (1 ax+1) +...
144 i [r 1 -1 + yB ay ae [r(4 a(+))]l  =Eorne (1(st, 0z) + 1 1(5 +...
ILE ILE H-
7-lBat 1 , [ ZEONNTe (r(8,, 0e) + 1 |r (1 +
(21)
7/Bnn 1891 mQ +++1))]]]
(22)
> Eorrne (r(st, (x) + 1 (TT 1+41)]
> Enrnoo (r(s,, ar) + 1 nQa (s441, (si)
15

---

Published as a conference paper at ICLR 2025
The derivation from Eq. 21 Eq. 22 is based on:
B, Baa 1Q (s,a)]
Eg T (11)))
E, 1Q (8,
(23)
E. Cu Q4 (8, u)]
Esgart (16)
Eggawn 12 (s.a))
B.3 PROOF OF THEOREM 2
Theorem 2
LOM learns policy zL. which is uniformly good better than the improved policy TG and the behaviour policy Ty. That is,
Vs E S, V (s) 2 V (8) 2 (s).
(24)
Proof. The proof structured in two parts. First, we show that the LOM-learned policy Hy is at least good as the improved policy *c, specifically, for all states S, we establish that V (s) 2 y (s). Second, we prove that the improved policy Tc, performs no worse than the behaviour policy Tb. i.e., V/ (s) 2 (T)(s) for all s S.
As established in Wang et al. 20188) the following lemma provides sufficient condition for proving that one policy no worse than another. We restate this result as the following lemma:
Lemma (from Wang et al. 2018)): Suppose two policies and ā satisfy:
g( x)) g(T(a 8)) + h(8, (*(,a)));
(25)
where g(.) is a monotonically increasing function, and h(s,) :) is monotonically increasing for any fixed s. Then we have:
v*() 2 VT()), As
(26)
This means that is uniformly good as better than
Since the LOM-learned policy imitates the optimal policy n², we have:
x" r" 8) = 7(s) Tc, (a sxep L (3)
(8.0))
(27)
which gives:
1
1 A ( log Tc, (a s) +
log TL (a 8) log
(s,
(28)
((s)
where log(-) is a monotonically increasing function and log 2(2) is constant. Therefore, TL is uniformly good better than the improved policy Tea
1
B.4 PROOF OF THEOREM 3
Theorem The LOM algorithm learns a policy TT, whose improvement upon the composite policy induced by the greedy hyper-policy TC, is bounded: for all S:
VoL (s) c0 (s) z 1 1- (mL))
Amax 1
D
(29)
where Amax max. a + (s, a)). The proof is provided in Appendix
B.4
Proof. Following performance difference lemma (akaade & Langford 2002 ). we have
T (s) V (8)
E
Bonda arni (A (s)
(30)
16

---

Published as a conference paper at ICLR 2025
Our goal is to express the right-hand side in terms of (T()) which uses dn, rather than dry. To do this, we define the discrepancy between dx,, and dre, as:
( - dxz (6) - dmco (s).
(31)
We decompose the expectation over dip as:
E B [f())] = E Bonddre [f(b)] + [
(32)
where E [f({]] - -.,
Applying this to the expression:
VTL (5) V Te (s)
(Sondrr aaay [A (s, a)] + Badd [AT
(33)
1 7-7
The first term is exactly ((TL) scaled by T-y" We need to bound the enror term:
1-7 B BawAaam. [A (s, a)l
(34)
Assuming the advantage function Tco (s, a) is bounded:
1A (s, a)) S Amax.
(35)
We can then bound by
Amax
Amax
Amax
|e|
Σ |((1))
1-7 1 dxc, 111
Drv (rv(daaldde))
(36)
1-
By applying Pinsker's inequity, have
1 D
DTV (dreldel
V
(37)
Finally, we have
1 ((TL)
Amax
(* 2 Vkc (5) 1-7
1- V Dru DDe(darlldre
(38)
n
C FURTHER EXPERIMENTAL DETAILS
C.I BASELINES
In the experiments discussed in Section 6 we compare LOM with other multi-modal modelling approaches, such WCGAN and WCVAE, in capturing complex action distributions in offline reinforcement learning. These models are evaluated on their ability to reconstruct and learn from the multi-modal behaviour policy in the FetchReach, FetchPush, and FetchPickAndPlace tasks, and across various D4RL benchmarks.
Weighted Conditional GAN (WCGAN): The WCGAN learns a discriminator D and a generator to optimise the following adversarial objectives:
min E(s,a)~D [w[ss a) log D(s, a)] (x-Da" [log(1 D(ss
(39)
where wis a) exPetsp
(s, Here, Th represents the advantage of the behaviour policy The WCGAN aims to generate action distributions that are weighted by this advantage, prioritising actions that yield higher returns.
Weighted Conditional VAE (WCVAE): The WCVAE learns conditional variational autoencoder with an encoder F. and decoder D. The objective to optimize:
max E.D
[w(s, a) log ((a 2, 8) w(s, a) Dxi, s, || p(2 s))], (40) B(s.a))D,e
where p(2 |s) is the unit Gaussian prior distribution of the latent variable conditioned on The weighting function w(s a) encourages the WCVAE to focus on high-reward actions.
C.2 HYPERPARAMETERS
The hyperparameters used in LOM were largely the same for all of the experiments reported. In Table we give a list and description of them, as well as their default values. The hyperparameters M used in Table are selected from (2, 5, 10, 15, 20).
17

---

Published as a conference paper at ICLR 2025
Symbol paper
Description
Default values
Number of Gaussian components
Please check Figure
Advantage weights
5
C
Weight clips
50
Ployak averaging coefficient
0.995
update_delay _eelly
Target update delay
2
Pp
Gaussian mixture model
sttat__dim, 512, 512, * ac_ sc_dim + M]
ma
Policy learned by LOM Q-network
[state_dime 512, 512, ac ac_dim] Sstat_ddim ac_dim, 512, 512, 1] [state dim M, 512, 512, M]
Q
Q
Behaviour hyper Q-function
Table 3: Hyperparameters used in the experiments.
18


Reviews:
[{"review_id": "GpDVtRbdjU", "reviewer": "Reviewer_N1vf", "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.", "questions": "See weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper is well-organized and easy to follow.\n2. 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.\n3. 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?\n2. 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "4DDo3sJ6oE", "reviewer": "Reviewer_Eg5z", "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.", "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?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is well-written and easy to follow. \n\n2. 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. \n2. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gAROgomv8P", "reviewer": "Reviewer_HzTN", "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.", "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?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "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.\n\nThe 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6KQXoEcVGY", "reviewer": "Reviewer_XhoW", "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.\n\nThe 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.", "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.\n\n- 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.\n\n- 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.\nHowever, 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.\n\n- 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.\n\n=== comments after the rebuttal ===\n\nThe concerns raised in the initial review were appropriately addressed in the rebuttal.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "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.\n\n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)