Diverse Policies Recovering via Pointwise Mutual Information Weighted Imitation Learning
Authors: Hanlin Yang
Jian Yao
Weiming Liu
Qing Wang
Hanmin Qin
Kong hansheng
Kirk Tang
Jiechao Xiong
Chao Yu
Kai Li
Junliang Xing
Hongwu Chen
Juchao Zhuo
QIANG FU
Yang Wei
Haobo Fu
Venue: ICLR 2025 Poster (2025)
Keywords: Imitation Learning
Policy Diversity
Offline Learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: This paper presents a novel method for recovering diverse policies in imitation learning by weighting state-action pairs based on their significance, with theoretical justifications and empirical evaluations demonstrating the effectiveness.

Abstract: Recovering a spectrum of diverse policies from a set of expert trajectories is an important research topic in imitation learning. After determining a latent style for a trajectory, previous diverse polices recovering methods usually employ a vanilla behavioral cloning learning objective conditioned on the latent style, treating each state-action pair in the trajectory with equal importance. Based on an observation that in many scenarios, behavioral styles are often highly relevant with only a subset of state-action pairs, this paper presents a new principled method in diverse polices recovering. In particular, after inferring or assigning a latent style for a trajectory, we enhance the vanilla behavioral cloning by incorporating a weighting mechanism based on pointwise mutual information.
This additional weighting reflects the significance of each state-action pair's contribution to learning the style, thus allowing our method to focus on state-action pairs most representative of that style.
We provide theoretical justifications for our new objective, and extensive empirical evaluations confirm the effectiveness of our method in recovering diverse polices from expert data.

Published as a conference paper at ICLR 2025
DIVERSE POLICIES RECOVERING VIA POINTWISE MUTUAL INFORMATION WEIGHTED IMITATION LEARNING
Hanlin Yang! Jian Yao",, Weiming Liu², Qing Wang", Hanmin Qin', Hansheng Kong", Kirk Tang', Jiechao Xiong?, Chao Yu Kai Li Junliang Xing', Hongwu Chen', Haobo Haobo Fu
Juchao Zhuo?, Qiang Fu?, Wei Yang",
Sun Yat-sen University. Guangzhou, China
Teneent Al Lab, Shenzhen, China
Institute of Automation, Chinese Academy of Sciences, Beijing, China
Tsinghua University, Beijing, China
ABSTRACT
Recovering a spectrum of diverse policies from a set of expert trajectories is an important research topic in imitation learning. After determining a latent style for a trajectory, previous diverse policies recovering methods usually employ vanilla behavioral cloning learning objective conditioned on the latent style, treating each state-action pair in the trajectory with equal importance. Based on observation that in many scenarios, behavioral styles are often highly relevant with only a sub- set of state-action pairs, this paper presents a principled method in diverse polices recovery. In particular, after inferring or assigning a latent style for a tra- jectory, enhance the vanilla behavioral cloning by incorporating weighting mechanism based on pointwise mutual information. This additional weighting re- flects the significance of each state-action pair's contribution to learning the style, thus allowing our method to focus on state-action pairs most representative of that style. We provide theoretical justifications for our new objective, and exten- sive empirical evaluations confirm the effectiveness of our method in recovering diverse policies from expert data.
INTRODUCTION
Imitation Learning (IL. is about observing expert demonstrations in performing a task and learn- ing to mimic those actions (Hussein et al., 2017: Osa et al.) 2018). Vanilla behavioral cloning (BC) (Pomerleau) 1991) learns mapping from state to actions using expert state-action pairs via supervised learning, which is simple to implement but may have the issue of compounding errors. Generative Adversarial Imitation Learning (GAIL) (Ho & Ermon 2016) mitigates the issue via learning both a discriminator and a policy. Despite their wide applications, these methods lack mechanisms generate diverse policies, which may be essential in certain tasks.
There has been an increase in recent research addressing policy diversity in imitation learning (Li et al. 2017 Wang al., 2017| /Zhan al. 2020 Shafuulaah al. 2022 Mao et al. (023). which can generally divided into two categories. In one category, the latent style 2 of a trajectory is inferred in unsupervised manner, for instance, by an expectation maximization (EM) procedure. In the other category, the style of a trajectory is determined by user-specified function, for instance, programmatic labeling function. No matter in which category, samples that are used to train a style- conditioned policy π [a]ss, =) are treated with equal importance in those methods. However, in many cases, we observe that the relevance of different state-action pairs to the trajectory style can vary significantly. For example, in autonomous driving tasks, the diversity of overtaking policies (from the left or right) is primarily relevant to the overtaking period in the trajectory. In other words, the preceding normal driving period is less relevant to the overtaking diversity. Again, in basketball game, the diversity of shooting position behaviors primarily related to the part leading up to the
"Equal contribution
Work done during an internship at Tencent Al Lab "Equal advising. Correspondence Haobo Fu thaohofu@tteneet.com))

---

Published as a conference paper at ICLR 2025
shot, while it is less relevant other parts in the trajectory. In nutshell, policies with different styles could have significantly different degrees of overlap in different areas of the state-action space. Therefore, when learning style-conditioned policy, the relevance between the stateacttion pairs in the trajectory and the behavioral style should be taken into consideration.
In this paper, we propose new diverse policies recovering method by leveraging the relevance of the state-action pairs with the trajectory styles. We focus on the situation where the style label of a trajectory has been provided, as unsupervised learning (either by EM or mutual information maximization) of the style 2 leads to uncontrolled styles and often only qualitative evaluation. In particular, we introduce an additional importance weight based on Pointwise Mutual Information (PMI) (Church & Hanks 1990; Manning & Schutze 1999] Bouma (209)), in traditional conditional BC, to quantify the relevance of state-action pairs with the conditioned style. Intuitively, state- action pairs with larger posterior of the corresponding style p(=|s, a) are given larger weight. In practical implementation, we utilize Mutual Information Neural Estimation (MINE) (Belghazi et al.) 2018) to estimate the PMI between state-action pairs and style variables. We term the proposed method Behavioral Cloning with Pointwise Mutual Information Weighting (BC-PMI).
Our theoretical analysis indicates that our new weighted learning objective unifies two extreme cases in recovering diverse policies. When the mutual information between the style and the state-action pair is zero, which means there is style diversity in expert data, objective degenerates to vanilla BC, which views the data as generated from one policy. By contrast, when policies with dif- ferent styles have no overlap in the state-action space, our objective degenerates to learning different style polices separately. Empirical results in Circle 2D, Atari games and professional basketball player dataset demonstrate that BC-PMI achieves better performance in recovering diverse policies than the baseline methods.
2 RELATED WORK
Imitation Learning Imitation learning (II.) methods designed to mimic the behaviors of ex-
perts. Behavioral Cloning (BC) (Pomerleau 1991), a well-known II. algorithm, learns a policy by directly minimizing the discrepancy between the agent and the expert in the demonstration data. However, offline learning methods like BC suffer from compounding errors and the inability to han- dle distributional shifts during evaluation (Ross et al. 2011: Fujimoto et al., 2019 Wu et al. 2019
Peng et al. 2019 Kostrikov et al., (021). Inverse Reinforcement Learning (IRL) Ng et al. 2000,
Arora & Doshi (202). another type of II. learns a reward function that explains the expert behavior and then uses this reward function to guide the agent's 's learning process. Popular IRL. approaches
like Generative Adversarial Imitation Learning (GAIL)) (Ho & Ermon (016) and Adversarial In- verse Reinforcement Learning (AIRL) (Fu et al. 2017) use adversarial training learn policy that is similar to the expert policy while being robust to distributional shift. However, these methods are limited to imitating single policy and do not address the issue of promoting diverse policies. When imitating diverse policies, BC approaches using supervised learning tend to learn an average policy, which does not fully capture the range of diverse behaviors (Codevilla al. (018)). GAIL tends to learn a policy that captures only a subset of the expert's control behaviors, which can be viewed as modes of distribution. Consequently, the learned policy fails cover all styles of the expert's diverse behaviors Wang et al. (2017).
Policy Diversity Diversity is crucial in imitation learning algorithms, especially in practical con- trol tasks and multi-player games (Zhu et al., 2018), as diverse policies in control tasks can enhance the robustness of adapting to various environments. In contrast, Al with diverse policies can max- imize the player's experience and the ornamental value games and competitions (Yannakakis & Togelius 2018). Some approaches utilize coammation-thhorete methods to address this issue and learn the behavioral styles of policies. InfoGAlL (Li et al. 2017) and Intention-GAN (Hausman et al. 2017) augment the objective of GAIL with the mutual information between generated trajec- tories and the corresponding latent codes. Wang et al. (2017) use a variational autoencoder (VAE) module to encode expert trajectories into a continuous latent variable. Eysenbach et al. (2088) pro- pose a general method called Diversity Is All You Need (DIAYN) for learning diverse skills without explicit reward functions. DIAYN focuses
discovering skills in online reinforcement learning tasks (Campos et al.) 2020 Sharma et al.) 2019 AAchiam et al. 2018), whereas our method specifi- cally considers pure offline imitation learning scenarios. More recently, Mao et al. (2023) introduced
2

---

Published as a conference paper at ICLR 2025
Stylized Offline RL (SORL), which utilizes unsupervised learning methods cluster styles and fits policies to each cluster separately, optimizing an equal number of policies while using the KL diver- gence constraint. However, these methods focus inferring the latent style clustering the trajectories with different styles without considering the relevance between the state-action pairs in the trajectory and the behavioral style. In contrast, focus the situation where the style label of a trajectory has been provided. By introducing the PMI weights, we quantify the relevance of state-action pairs with the conditioned style, allowing the policy focus on the samples that highly relevant to the style.
3 PRELIMINARY
3.1 PROBLEM SETTING AND VANILLA BELIAIIRR CLONING
We consider the standard Markov Decision Process (MDP) Sutton & Barto 2018) as the mathe- matical framework for modeling sequential decision-making problems, which is defined by a tuple (S, A, P. vr,do, do, T*, where S is finite of states, A is finite of actions, P S A x S IR is the transition probability function, r S + R is the reward function, dg is the initial distribution, [0. 1] maps from state to distribution over ac- tions. Let di and d' = T Σi= d; denote the distribution over states at time step f and the average distribution over T time steps induced by π,, respectively. The vanilla BC loss function is:
and T is an episode horizon. A policy π S x
LBC
E
[ [ogg((ass]]
(s.0)~De
which aims to maximize the probability of selecting action a for a given policy under the state s.
Building upon the basic setup, we focus on an assumption where the expert demonstrations D.. which consist of many diverse trajectories т, are collected by stylized expert policies denoted as (Te
(1) Te. ................................ ................................ ne {. Let Z denotes the variable indicating which stylized policy belongs to, and p(T = i) denotes the probability of 7 sampled under policy me Our objective to learn conditioned policy [[als, 2) such that trajectories generated by x[als, 2) closely match the demonstrations in D. that exhibit the corresponding style 2.
3.2 MUTUAL INFORMATION NEURAL ESTIMATION
Mutual Information Neural Estimation (MINE) is powerful technique for estimating the mutual information between two random variables using neural networks (Belghazi et al. 2018). It has been widely used in various domains, including representation learning (Hjelm al. 2018), gener- ative modeling (Chen et al.) 2016), and imitation learning (Eysenbach et al. 2018). The key idea behind MINE is formulate the estimation of mutual information Given two random variables X and Y, the mutual information I(X; Y) can be expressed as the Bellback--eibler (KL) divergence between the joint distribution PxY and the product of marginal distributions Px & Py:
optimization problem.
I(X;Y) = Drr (Pxy Px @ Pr).
MINE approximates this KL divergence using a lower bound based on the Donsker-Varand rep- resentation (Donsker Varadhan 1983):
I(X; Z sup Exxy [To] TEE
(e=o)),
where F is a class of functions T: neural network Ta. which takes as input samples from the joint distribution Pxy and the product of marginal distributions Px Φ Py. The network is trained to maximize the lower bound, equivalent to estimating the mutual information.
R. In MINE, this function class is parameterized by
BEHAVIORAL CLONING WITH POINTWISE MUTUAL INFORMATION WEIGHTING
As discussed in the introduction, in many real-world applications, the impact of different state-action pairs the style vary greatly, and often only part of the trajectory is highly relevant to the
3

---

Published as a conference paper at ICLR 2025
style. Hence, when training style-conditione policy, would like to assign different weights to different state-action pairs based on their relevance to that style. We aim to develop such a method that can capture the specific influence of each (s, a) pair the style, thereby achieving more generalized imitation objective:
min
E (s.DD
[- (a]s, x) a(s, a, =)],
(4)
where (s, a, 2) a weighting function. Intuitively, state-action pairs that are more exclusive to style should have larger weights in learning that style-conditioned policy, and vice-versa. Drawing inspiration from the information theory, we introduce the PMI (Church Hanks, [990] Manning & Schutze 1999 Bouma 2009) to quantify the contribution of (s, a) when learning style-condttioned policy:
P(2; 8, a) = log p(s, a, 2) = log p(z]s.a) p(s, p(s,a)p(z)
(5)
p(2)
We utilize the exponential of ((2; 8, as the weight, which is the ratio of the posterior given (s, a) and the prior of the style z. This weight has the following properties. When the posterior of style given (s, a) is larger than the prior, i.e., p(z|s, > p(z), this means there is high relevance between (s, a) and style 2, and we should give larger weight. If (s, a) and style = are nearly independent, then p(=ls,a) a) p(z). When we have p(=|s, a) < p(=), this means (s,a) is more likely generated by other styles, and we should give a lower weight here. Accordingly, Eq turns into the following form:
[- log To(als, 2) P(:ss.a)
log To(a (c)s, p(2 a)
CBC-PMI(0) (s,0,e)<D.
E (e,0,))~De
p(z)
4.1 THEORETICAL ANALYSIS
We give some theoretical justifications for our BC-MI objective in G. Considering an extreme case where the mutual information between (s, a) and style approaches zero, which means policies with different styles have nearly the same state-action distribution. In this case, the BC-PMI objec- tive degenerates into vanilla BC, which means there is no difference in training style conditioned policies and an average unconditioned policy. In the opposite extreme case, where all (s, a) pairs exhibit significant style differences, which means trajectories of different styles have minimal over- lap, we find that the BC-PMI objective degenerates into behavior cloning on each style. Formally, have the following proposition:
Proposition 1. (a). When the mutual information 1(2; S, A) equals to O, it indicates that there is no distinction in the trajectory style corresponding to all the state-action pairs. In this case, the BC-PMI objective degenerates to the vanilla behavior cloning objective:
arg min (Mr(() arg min LBc
(b). When the conditional entropy HZ A) equals 0, indicates that there a significant dis- tinction in the trajectory style corresponding all the state-action pairs. In this case, the BC-PMI objective degenerates to the behavior cloning on each style:
K
LBC PMI (0) Σ (O).
(8)
where ClS (0) is the behavior cloning loss the subset of data with style label i.
Proof. Refer Appendix A.I
Equation means when H(Z)S, equals to 0, the BC-PMI objection function can be viewed as separating the (s, a) by the style and optimize it respectively. This is associated with the clustering- based behavior cloning objective, where each trajectory first assigned to a specific style cluster based on its style label, and then behavior cloning is performed within each cluster.
4

---

Published as a conference paper at ICLR 2025
The above analysis provides insights into the mechanisms of BC-PMI. In practice, the mutual in- formation I I( A) often lies between the two extremes and Z By adjusting the weight of a state-action sample in learning style-conditioned policy, BC-PMI can smoothly interpolate between vanilla behavior cloning and clustering behavior cloning, allowing it to use the expert data more effectively than previous methods that treat each state-action sample with equal importance.
4.2 PRACTICAL IMPLEMENTATION
To practically estimate the PMI values in 6Q. we employ the Mutual Information Neural Esti- mation (MINE) (Belghazi et 2018) method. MINE is neural network-based approach that can effectively estimate mutual information between high-dimensional random variables. By leveraging the Donsker- Varadhan representation (Donsker Varadhan, (983) of the Kullback-Leible (KL) divergence, MINE allows for the estimation of mutual information using neural networks.
Let Ta S x Ax + I be neural network parameterized by Φb The mutual information between (s, a) and be estimated as:
[(s,0;2) 2
E [To(s, a, 2)] log (5,0,2)~De
E (e.m)--..R-p()))
The neural network To is trained to maximize the lower bound in Eq. which is equivalent to minimizing the KL divergence between the joint distribution p(s, a, 2) and the product marginals p(s, a=)(=). According to the proof of Theorem 1 in MINE (Belghazi et al. 2018), the optimal solution for To is:
Th (s, 0, 2) log a(s,a,z) log p(zls,a)
(10)
p(s, a(s.a)p(z)
p(z)
which is exactly the PMI aim to estimate.
In practice we can train the MINE network Ta using samples from the expert demonstrations De and the style distribution p(z). The training objective for T is:
[To(s, log
E
e [T(
(11)
(s,0,2)~DD
[a.a)-De..~p(a))
By optimizing Eq. obtain approximation of the PMI values, which be used as the weights the behavioral cloning objective in (6. The weights can be denoted as:
( a, 2) exp[TT as,
(12)
Furthermore, in order to reduce the variance of the gradient in optimizing Eq. we can subtract an optimal baseline from the weight, like what has been done in A3C (Mnih al. 2016). This turns (q.( to the following objective:
min (a
- log (als,,)) a))
(13)
where 5
(s, exp((T(s, as),, In practice, estimate 6 using the moving average.
E
The pseudo-code for the BC-PMI algorithm is shown in Algorithm |
5 EXPERIMENTS
In this experimental section, we aim to address the following questions:
QI. Can our method recover diverse and controllable policies from diverse style trajectorics?
Q2. Does the PMI weight have interpretability and improve policy diversity and style calibration to the algorithm?
Q3. Can our method cope with complex, real-world tasks, particularly those that involve extensive datasets derived from human participants?
5

---

Published as a conference paper at ICLR 2025
Algorithm 1 Behavioral Cloning with Pointwise Mutual Information Weighting (BC-PMI)
Initial the parameters & of the neural network T, the parameters 0 of the policy Given the expert demonstrations Dzi
for i 0, 1, do
Draw b minibatch samples from the joint distribution: (fs, ((s,a)01,211)) (1s,
Do;
Draw b samples from the Z marginal distribution: ,(),...,E())) ~ P(Z);
Update the neural network T using:
max à
b Σ To (00, 6
1
-log
Poll(o,a)) (0.))
(14)
end for
for = 0, 1, do
Sample random minibatch of N state-action pairs from De; Update the policy и with the PMI weighting using:
min E (8,0,2)DD
1 [exp(TT(s.a,z)) - -
(15)
end for
5.1 STYLES AND BASELINES
The dimension of the style needs to be specified if the style needs to be controllable. Our method can handle both style-labeled data and data without style labels. For the latter, getting the style labels can be expensive when relying manual annotations and uncontrollable when using unsupervised approaches. Instead, the programmable labeling functions (Ratner al.) 2016: Zhan et al., 2020 can be used to automatically generate style labels.
We compare PMI-BC the following baselines: (1) BC, which directly imitates expert actions across all styles; (2) CBC, which separates trajectories of different styles and uses BC to imitate each style separately; (3) CTVAE, which is the conditional version of TVAEs (Wang et al. 2017);
(4) IIIoGAIL. (Li et al. 2017), which infers the latent style of trajectories by maximizing the mu- tual information between the latent codes and trajectories; (5) SORL (Mao et al., 2023 + which use Expectation- Expectation-MMxxiiizttti Maximization (EM) algorithm to classify the trajectories from the heterogeneous dataset into clusters where each represents a distinct and dominant motion style. Since CBC, CT- VAE, and our method use the style labels, concatenate the state and style input to the network of InfoGAIL and SORL to ensure a fair comparison. Specifically, we first introduce a toy example called Circle 2D provide a simple visualization and analysis. Then Atari games, detailed anal- ysis and validation were provided regarding the PMI weights. Lastly, we evaluate all these baseline methods in Section 5,4 to illustrate the effectiveness of BC-PMI.
5.2 CIRCLE 2D
The Circle 2D environment is 2D plane where an agent can freely move at a constant velocity by selecting its direction, denoted as Pt, at time step t. For the agent, the observation at time step t includes the state information from time step / - 4 to The offline expert trajectories consist of four different styles, each generated by random expert policy. The expert policy generates trajectories that resemble circular patterns after period of translation (75 time steps). This design aims to introduce partial diversity in the trajectories. In this environment, each episode consists of 300 time steps. If the first loop around the circle is completed before reaching 300 steps, the agent continues circling until the end of the episode. Hence, in this scenario. there is minimal difference in the offline trajectories during the first 75 steps, and the trajectory differences vary as the agent's position progresses after 75 steps. During the imitation learning training process, the expert trajectories used are noisy, meaning there is randomness introduced in both the sampled actions and the environment.
"For the SORL algorithm, we only used the EM clustering part.

---

Published as a conference paper at ICLR 2025
(a) Exphrr( noise)
(b) Vanilla BC
(c) Conditional BC
(d) BC-PMI
Figure I: Visualized comparison of trajectories generated by different policies.
This task toy example that can help the readers intuitively understand our motivation. In this task, the agent's behavioral diversity is highly related to the curvature of the circle moving after 75 steps and has nothing to do with the linear movement part. Therefore, (8, a) that moves according to the curvature of specific style after 75 steps will be assigned higher PMI, and the linear motion part or the movement curvature which are not related to the style will be assigned lower PMI. The visualized results shown in Figure I By introducing PMI, the calibration of the BC policies has been further improved. TableTiis comparison of numerical results. The metrics used are as follows: (1)DTW, which is Dynamic Time Warping, dynamic programming algorithm for calculating the similarity of sequences of different lengths; (2)ED, which is Euclidean Distance; (3)KL, which is used calculate the state-action distribution difference of the trajectory.
Table Comparison style calibration across different metrics in Circle 2D.
Style
Metrics
BC
CBC
BC-PMI
DTW
88.590 + 9.329
8.130 + 0.732
7,511 + 0.545
Class (Red)
ED
133.100 20.210
22.039 + 1.383
21.591
1.091
KL
8,037 0.870
0.044 + 0.006
0.037 + 0.004
DTW
45.729 5.293
7.631 + 0.821
Class 2 (Green)
ED
88.180 + 7.541
35.355 2.691
35.777 + 2.912
KI.
1.111 +0.072
0.043 0.003
0.037 + 0.003
DTW
77.619 7.812
18.972 + 2.323
11.576 + 1.788
Class 3 (Orange)
ED
130.839 + 10.322
107.833 + 8.213
74.267 + 8.002
KI.
16.839 1.022
0.271 + 0.031
0.135 + 0.014
DTW
69.527 + 10.924
7.755 + 1.832
7.577 + 1.051
Class 4 (Blue)
ED
97.066 14.239
27.230 + 2.110
26.749 2.347
36.238 + 2.981 KL
0.401 + 0.061
0.219 + 0.019
5.3
ATARI GAMES
In this experiment, we concentrate on three widely recognized Atari games: Alien, MsPacman and SpaceInvaders. The datasets utilized this study are sourced from Atari-Head (Zhang al. 2018 2020), an extensive collection of human game-play data. Atari-Head are meticulously recorded in cemi-frame-by-fraat manner, ensuring high data quality and granularity, which facilitates in-depth analysis and robust evaluation of our proposed method.
This experiment consists three parts. Firstly, we demonstrate the convergence of the lower bound of mutual information in Eq. which indicates the relevance between state-action pairs and styles. Secondly, we provide interpretability of PMI weights by assessing the extent to which they appropri- ately reflect the influence of the current (s, a) pair on the style. Lastly, we evaluate the controllability of the BC-PMI policy, which refers to the ability of the policy to act according to a given style once it is specified. In Atari and Baskeball benchmark (in next subsection), since the expert policies are not available, is challenging to compute DTW, ED, or KL as we did in Table [] We evaluate the style calibration of the policy by comparing the accuracy of the style of the agent's s actual trajectory with the given style. A higher value indicates better diversity of the policy. There are also other candidate metrics evaluate state diversity action diversity (Belkhale al. 2024

---

Published as a conference paper at ICLR 2025
In the Alien and MsPacman game, employ two styles: the area style and the range style. The former divides the map into four distinct areas, distinguishing the agent's preferences for moving toward each area. The latter models the agent's displacement trajectory distribution, differentiating the variance of the distribution. A higher variance indicates a tendency for the agent to move across areas, while lower variance indicates preference for movement within a single area. In the SpaceInvader game, we also utilize two styles: the firing rate style and the area style. For further details, please refer to Appendix B.4
the map Gaussian
The results in Figure 2 indicate a high rele- vance between the agent's state-action distri- bution and the style. Moreover, as the rele- vance increases, the converged MI values also increase. The MsPPacmnnAR style divides the map into four areas: top-left, top-right, bottom- left, and bottom-right. The agent's trajec- tory exhibits clear distinctions under this style, which is reflected in the larger converged MI values, while the MsPaemanRG style repre- sents the style based the range of movement, which is less distinct compared to the move- ment area style.
1.4
1.2
1.0
0.8
0.6
Alenva
0.4
Spatelmaddesatt
SpacelnadderiR
0.2
MsPacmankG
0.0
200
400
600
800
1000
1200
1400
1600
1800
Step
Figure MI between state-action pairs and styles. FR: Fire Rate style; AR: Movement Area style; RG: Movement Range style.
Game Frame
1.0
0.0
-0.5
100
150
200
250
300
350
400
450
Step
Figure 3: The PMI weight values related to style for each frame along a trajectory, with the corre- sponding style being the bottom-left area style (green color) in the game frame. The agent in the game frame is indicated by white circle, and the corresponding actions are indicated by white ar- rows. No arrow indicates that the action is NOOP
As shown in Figure 3 in the first frame, the agent is located at the edge of the area and does not take any action, resulting in near-zero relevance with the style. In the second and fourth frames, the agent moves to the corresponding area and exhibits a tendency to continue moving towards that area, leading fifth frames, the agent is positioned at the edge or outside of the style area and shows a tendency to move away from the style area, resulting in lower relevance with the style.
higher relevance with the style and higher PMI values. Conversely, in the third and
To verify whether different styles of policies affect the cumulative reward, we conducted experi- ments in three different Atari environments, as shown in Table The results indicate that most styles, such as those related to position and range of movement, do not impact the cumulative re- ward of the agent. However, some styles that are highly correlated with scoring can significantly influence the cumulative reward. For instance, in Space Invaders, the Fire style, characterized by higher firing frequency, significantly increases the agent's 's cumulative reward, while lower firing frequeney reduces the cumulative reward.
Finally, we compared the calibration of several different style policies across these Atari environ- ments, shown in Table 3 In styles related to position and firing, BC-PMI outperformed all baselines. In styles related to range, BC-PMI achieved performance comparable to CTVAE. The reason is that the mutual information distinguishing capability of styles related to position and firing higher (refer to Figure?? 2). Consequently, the trained MINE network can assign more accurate PMI values to state-action pairs, resulting in higher accuracy for the corresponding styles.
8

---

Published as a conference paper at ICLR 2025
Table Comparison of the cumulative reward across different styles in Atari. BC-PMI-x means the reward of the x-th style.
Env
Style
BC
BC-PMI-1
BC-PMI-2
BC-PMI-3
BC-PMI-4
Pos
460.5:1133.4
416.0 160.8
432.0:117..7
500.3+:11.7
496.1=191.5
Alien
Range
460.5:1173.4
440.0::217.0
501.24424.9
423.2+1240.9
Ms
701.5:+:82.4
665.44+216.3
803.5+-211.2
674.3
146.5
666.2++172.2
Pacman
Range
701.5+112.4
636.1 +66.11
759.5++280.6
260.8
Space Invaders
Pos
208.2++73.0
178.9+554.5
187.0+:54.9 207.9+70.3
200.2+ 260.6+86.7
Fire
208.2++73.0
94.8+47.2
Table 3: Comparison style calibration across different styles in Atari.
Env
Style
BC
CBC
CTVAE
InfoGAIL
SORL
BC-PMI
Pos
0.23
0.43:10.04
0.52+10.4
0.3410.11
0.31
0.61:10.08
Alien
Range
0.31
0.51+0.13
0.57-10.09
0.31
0.27
0.55:10.12
MS
Pos
0.17=0.05
0.48=00.07
0.55+1000
0.33
0.27
0.57
Pacman
Range
0.33:10.09
0.53
0.61
0.29
0.30
0.56
Space Invaders
Pos
0.36+0.06 0.71
0.66110.10
0.41
0.37
0.8311006
Fire
0.27+0000
0.41+-0.88
0.47++0.06
0.36+
0.34
0.51
5.4 PROFESSIONAL BASKETBALL PLAYER DATASET
In this experiment, validate method the dataset of a collection of professional basketball player trajectories (Zhan et al. 2020) with the goal of recovering policies that can generate trajec- tories with diverse player- mevermt styles. The basketball trajectories real players in the NBA. We primarily focus on two movement styles: (1) The Destination, which the distance from the final position fixed destination the court (e.g. the basket), and (2) The Curvature, which measures the players propensity to change directions.
collected from tracking
Figure 4 and Figure 5) present the calibration results of different styles for the BC-PMI algorithm in the dimensions of Destination and Curvature. The environment a half-court basketball setting, where each player's trajectory can be categorized into different styles based on their movement destination and curvature. In Figure 4 Destination is close to the ball frame, Destination 2 is in the middle count, and Destination 3 far from the ball frame. The three destinations are separated by green lines in the figure. Similarly, three different movement curvature styles are illustrated in Figure The results indicate that the BC-PMI algorithm can effectively imitate policies of different movement styles from real human data.
D
D
b)
Destination
(b) Destination
(e) Destination
Figure 4: Visualization of different destination styles.

---

Published as a conference paper at ICLR 2025
D
(a) Curvature
(b) Curvature
(c) Curvature 3
Figure 5: Visualization of different curvature styles.
Table 4: Comparison of style calibration (%) across different styles in Basketball
Style
Class
BC
CBC
CTVAE
InfoGAIL
SORL
BC-PMI
Class
31.1
81.6
78.3
79.6
80.9
93.3
Destination
Class 2
42.7
84.4
81.2
75.5
76.1
92.7
Class
19.4
73.1
77.8
74.5
79.5
89.8
Class
17.1
67.3
58.4
60.9
51.6
79.4
Curvature
Class 2
30.6
66.2
58.7
57.2
49.4
80.9
Class 3
41.8
71.6
62.1
62.7
53.8
81.6
We compare style calibration for classes of Destination and 3 classes of Curvature Table Due to the lack of style label information, BC can only learn approximately average strategy, so should be evenly distributed among various styles. However, due to factors such as initialization position, the samples in each category do not strictly follow uniform distribution. Among the various methods with diversity mechanisms, the BC-PMI method, which focuses on state-action pairs highly relevant to the style, surpasses other baselines in terms of style calibration.
DISCUSSION
6.1 CONCLUSION
In this paper, investigate how to recover diverse policies from set of expert trajectories. We propose a new diverse policy recovering method by leveraging the relevance of the state-action pair with the trajectory styles. The highlight of our method lies our approach to the problem of policy diversity from different perspective, which involves the introduction of Pointwise Mutual Information model the relevance between each (s, a) pair and the style. By utilizing a unique and straightforward approach, we achieved results that surpassed previous state-of-the-ee methods.
6.2 LIMITATIONS AND FUTURE WORK
In our setting, our goal recover policies from diverse offline data, assuming that the data within the trajectories already meet the performance requirements. If the policy that generates the offline trajectories performs poorly, our method cannot further enhance the performance of the learned pol- icy. Hence, we propose exploring the combination of PMI and RL in future research. We believe that integrating PMI weights with RL methods can enhance diversity while simultaneously improving performance. Numerous studies have already demonstrated the effectiveness of promoting diversity in RL (Hong et al. 2018 Eysenbach et al. 2018 Holderrr--rrrrrr et al.) 2020 Peng et al., 2020 Tang al.;
incorporate aeeentnnglemet methods (Cao et al. 2022 Locatello et al. 2019) between the stylized code and the cate-actionn pair. For example, adapt Ray and MED to weight the BC loss by replacing MI with PMI.
Yao et al. 2024). Another potential direction
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Joshua Achiam, Harrison Edwards, Dario Amodei, and Pieter Abbeel. Variational option discovery algorithms. arXiv preprint arXiv: /807. /0299, 2018.
Saurabh Arora and Prashant Doshi. A survey of inverse reinforcement learning: Challenges, meth- ods and progress. Artificial Intelligence, 297: 2021.
Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm. Mutual information neural estimation. In International Conference on Machine Learning, pp. 531 PMLR, 2018.
Suneel Belkhale, Yuchen Cui, and Dorsa Sadigh. Data quality in imitation learning. Advances in Neural Information Processing Systems, 36, 2024.
Gerlof Bouma. Normalized (pointwise) mutual information in collocation extraction. Proceedings of GSCL, 30:31-40, 2009.
Pierre Brémaud. An introduction to probabilistic modeling. Springer Science & Business Media, 2012.
Víctor Campos. Alexander Trott, Caiming Xiong, Richard Socher, Xavier Giró-i Nieto, and Jordi Torres. Explore, discover and learn: Unsupervised discovery of state-covering skills. Interna- tional Conference on Machine Learning, pp. 1317-1327. PMLR, 2020.
Jinkun Cao, Ruiqian Nai, Qing Yang, Jialei Huang, and Yang Gao. An empirical study on dis- entanglement of negative-free contrastive learning. Advances in Neural Information Processing Systems, 15:1210-1222 2022.
Xi Chen, Yan Duan, Rein Houthooft, John Schulman, Ilya Sutskever, and Pieter Abbeel. Info- gan: Interpretable representation learning by information maximizing generative adversarial nets. Advances in Neural Information Processing Systems, 29, 2016.
Kenneth Church and Patrick Hanks. Word association norms, mutual information, and lexicography. Computational Linguistics, 16(1):23 (6222222- 1990.
Felipe Codevilla, Matthias Müller, Antonio López, Vladlen Koltun, and Alexey Dosovitskiy. End- to-end driving via conditional imitation learning. In 2018 IEEE international conference on robotics and automation (ICRA), pp. 4693-4700. IEEE, 2018.
Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time. iv. Communications Pure and Applied Mathematics, 36(2):183- 212, 1983.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. arXiv preprint arXiv:1802.06011 2018
Justin Fu, Katie Luo, and Sergey Levine. Learning robust rewards with adversarial inverse rein- forcement learning. arXiv preprint arXiv: 1248, 2017.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pp. 2052-2066 PMLR, 2019.
Karol Hausman, Yevgen Chebotar, Stefan Schaal, Gaurav Sukhatme, and Joseph J Lim. Multi-modal imitation learning from unstructured demonstrations using generative adversarial nets. Advances in neural information processing systems, 30, 2017.
Devon Hjelm, Alex Fedorov, Samuel Lavoie-Marchilde Karan Grewal, Phil Bachman, Adam Trischler, and Yoshua Bengio. Learning deep representations by mutual information estimation and maximization. arXiv preprint (XXi:18008006600 2018.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. Advances in Neural Information Processing Systems, 29, 2016.

---

Published as a conference paper at ICLR 2025
Zhang-Wor Hong, Tzu- Yun Shann, Shih- Yang Su, Yi-Hsiang Chang, Tsu-Jui Fu, and Chun-Yy Yi Lee. Diversity-drives exploration strategy for deep reinforcement learning. Advances in neural infor- mation processing systems, 31, 2018.
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Computing Surveys (CSUR), 50,2):133,,
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. arXiv preprint arXiv:21 10.06169, 2021.
Longyuan Li, Jian Yao, Li Wenliang. Tong He, Tianjun Xiao, Junchi Yan, David Wipf, and Zheng Zhang. Grin: Generative relation and intention network for multi-agent trajectory prediction. 2021.
Advances in Neural Information Processing Systems,
Yunzhu Li, Jiaming Song, and Stefano Ermon. Infogail: Interpretable imitation learning from visual demonstrations. Advances in Neural Information Processing Systems, 30, 2017.
Francesco Locatello, Stefan Bauer, Mario Lucie, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. International Conference on Machine Learning, pp. 4114- 4124. PMLR, 2019.
Christopher Manning and Hinrich Schutze. Foundations of statistical natural language processing. MIT press, 1999.
Yihuan Mao, Chengjie Wu, Xi Chen, Hao Hu, Jiang. Tianze Zhou, Tangjie Lv, Changjie Fan, Zhipeng Hu, Yi Wu, et al. Stylized offline reinforcement learning: Extracting diverse high-quality behaviors from heterogeneous datasets. In International Conference on Learning Representa- tions, 2023.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. International Conference on Machine Learning, pp. 1928-1937, PMLR, 2016.
Andrew Y Ng, Stuart Russell, et al. Algorithms for inverse reinforcement learning. In International Conference on Machine Learning, volume 1, pp. 2, 2000.
Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, Jan Peters, et al. An algorithmic perspective on imitation learning. Foundations and Trends@ Robotics, 7(1-2): 1-179, 2018.
Jack Parker-Holder, Aldo Pacchiano, Krzysztof M Choromanski, and Stephen J Roberts. Effective diversity in population based reinforcement learning. Advances in Neural Information Processing Systems, 33:1 8050-18062, 2020.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. wevintage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint 17X:0:1910.00177,, 2019.
Zhenghao Peng, Hao Sun, and Bolei Zhou. Non-local policy optimization via diversity-regullrrize collaborative exploration. arXiv preprint arXiv:2006.0778), 2020.
Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neu- ral Computation, 3(1):88-97, 1991.
Alexander J Ratner, Christopher M De Sa, Sen Wu, Daniel Selsam, and Christopher Ré. Data programming: Creating large training sets, quickly. Advances in Neural Information Processing Systems, 29, 2016.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. reduction of imitation learning and struc- tured prediction to no-regret online learning. In Proceedings of the Fourteenth International Con- ference on Artificial Intelligence and Statistics, pp. 635.. JMLR Workshop and Conference Proceedings, 2011.
12

---

Published as a conference paper at ICLR 2025
Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning & modes with one stone. Advances in neural information processing sys- tems, 35:22955- 2022.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-aware unsupervised discovery of skills. arXiv preprint urXiv:1907.000 2019.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Zhenggang Tang, Chao Yu, Boyuan Chen, Huazhe Xu, Xiaolong Wang, Fei Fang. Simon Shaolei Du, Yu Wang, and Yi Wu. Discovering diverse multi-agent strategic behavior via reward random- ization. In International Conference Learning Representations.
Ziyu Wang, Josh S Merel, Scott E Reed, Nando de Freitas, Gregory Wayne, and Nicolas Heess. Robust imitation of diverse behaviors. Advances in Neural Information Processing Systems, 30, 2017.
Shuang Wu, Jian Yao, Haobo Fu, Tian, Chao Qian, Yaodong Yang, Qiang Fu, and Yang Wei. Quality-similar diversity via population based reinforcement learning. In The Eleventh Interna- fional Conference on Learning Representations, 2023.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning.
arXiv preprint arXiv: 1rXhv:1911.11361 2019.
Haoran Xu, Xianyuan Zhan, Honglei Yin, and Huiling Qin. Discriminator- offline imitation learning from suboptimal demonstrations. In International Conference on Machine Learning, pp. 24725-24745 PMLR, 2022.
Georgios N Yannakakis and Julian Togelius. Artificial intelligence and games, volume 2. Springer, 2018.
Jian Yao, Weiming Liu, Haobo Fu, Yaodong Yang, Stephen McAleer, Qiang Fu, and Wei Yang. Pol- icy space diversity for non-transitive games. Advances in Neural Information Processing Systems, 36, 2024.
Eric Zhan, Albert Tseng. Yisong Yue, Adith Swaminathan, and Matthew Hausknecht. Learning cal- ibratable policies using programmatic style-consistency. In International Conference on Machine Learning, pp. 11001-11011 PMLR, 2020.
Ruohan Zhang, Zhuode Liu, Luxin Zhang, Jake A Whritner, Karl S Muller, Mary Hayhoe, and Dana H Ballard. Agil: Learning attention from human for visuomotor tasks. In Proceedings of the European Conference on Computer Vision, pp. 663-679, 2018.
Ruohan Zhang, Calen Walshe, Zhuode Liu, Lin Guan, Karl Muller, Jake Whritner, Luxin Zhang, Mary Hayhoe, and Dana Ballard. Atari-head: Atari human eye-tracking and demonstration dataset. In Proceedings of the AAAI conference on Artificial Intelligence, volume 34, pp. 6811- 6820, 2020.
Yuke Zhu, Ziyu Wang, Josh Merel, Andrei Rusu, Tom Erez, Serkan Cabi, Saran Tunyasuvunakool, János Kramár, Raia Hadsell, Nando de Freitas, et al. Reinforcement and imitation learning for diverse visuomotor skills. arXiv preprint 4rXhv:1802009564, 2018
13

---

Published as a conference paper at ICLR 2025
A THEORETICAL DERIVATION
A.1 PROOF OF PROPOSITION
Proof. (a) We first introduce the Gibbs' inequality, which was proved in (Brémaud] 2012))
Lemma (Gibbs' inequality). Let p(z), ((x) be the probability mass function, fixed pla)), that the object function -Ez-p(a) log (2)) = Σx -p(x) log q(z) takes the minimun when q(x) p(x).
By Lemma 1, we know
LBC = E |- log (a(a)))] (s,a)~D.
=
E
[-
log [a(ss))]
(16)
DnD (alsmp(als)
which takes the minimum when Ta(a)s) p(als).
When I(2;S, S, A) = 0, it means that the latent variable Z is independent of the state-action pair (S, A), and the PMI term in the BC-PMI objective becomes:
p(als.a) a) p(2) p(2) p(z)
1.
(17)
Hence,
LBC-PMI(0)
| log loggs
(e,0,))<D,
IR
s)
D
plg
D.
12
Σ
p(2) Σ
(18)
D>
where the third equation is because z is independent with (s, a). Since p(2) is non-negative, when En -p(als) ) log [a(a)s, 2) takes the minimum for all 2, the objective reach its the minimum. Using Lemma again, we can see when T (a|s, =) = p(a]s), ()) PMI(()) reach its minimum.
(b) when H A) 0, means that given state-action pair (s, a), the style variable Z can be determined with high certainty. In other words, for any (8, a, z) ~ De., we have:
p(z|s, a) = 1(2 = Zs.a),
(19)
where za.a is the unique style label corresponding to (s, a).
Assuming there are K styles in the Dataset. Let z, be the unique style label corresponding to trajectory т, ( be the subset of state-action pairs with style label and D!! means the whole
trajectory from ( Denote
(0) (
|- log To (aa))
Σ
Σ
- log no (a') i).
(20)
TeD! (st,"')ET
14

---

Published as a conference paper at ICLR 2025
We have
p(=18, log To(als,)))
(BC-PMI(0)
E
(s,
p(2)
1(2 2g,a)
E (s,a,;)~DD
| log (o
p(2)
k
1 (2 =)
Σ
log No (a'1s''
De
plz i)
TEDe t=1
(e' (e',a')Er
1
K
E
Σ
1(2r i)[- log TO (a'ls'''
Del
De i=l TED, (slaa))Er
D
K
Σ
- log TO (a)
D
==
TED
(d',a')E
K
(0) (0)
(21)
Note that by the definition, (0) is the behavior cloning loss on the subset of data with style label i. Hence, the last equation above means when H (Z],A,)) S) equals 0, the BC-PMI objection function can be viewed separating the (s,a) a) by the style and optimizing it respectively. This is equivalent to the clustering-basee behavior cloning objective, where each trajectory is first assigned to specific style cluster based on its style label, and then behavior cloning performed within each cluster. This result reveals the close connection between BC-PMI and clustering- based behavior cloning in the extreme case of zero style overlap.
B IMPLEMENTATION DETAILS
B. COMPUTATIONAL RESOURCE
All experiments in this paper are implemented with PyTorch and executed on NVIDIA Tesla T4 GPUs. All the runs in experiments use random seeds.
B.2 COMMON HYPERPARAMETERS
Table 5: Common hyperparameters setting.
Hyperparameter
Circle 2D
MsPacman
SpaceInvaders
Basketball
learning rate
0.001
0.001
0.001
0.0002
optimizer
Adam
Adam
Adam
Adam
epoch
10
30
30
30
batch size
128
512
512
128
hidden dim
32
64
64
128
B.3 MINE NETWORK STRUCTURE
Figure 6 shows the MINE network structure, which is used in Atari tasks.
B.4 STYLES IN SPACEINVADERS GAME
In the SpaceInvaders game, we employ two styles: the area style and the fire rate style. The former divides the map into three distinct areas, distinguishing the agents preferences for moving toward each area. The latter divides the agent's 's fire rate into three levels: (1) [o, 0.1); (2) 10.1, 0.3) and (3) [0.3, 1].
15

---

Published as a conference paper at ICLR 2025
MINE Output
FC 64
ELU
FC 64
Concat
ELU
Style Input
FC 64
Act Input
ELU
Conv. 4x4, stride
ELU
Conv. 8 8, stride
/255
Obs Input
84 x
Figure Architecture of the MINE networks used the Atari environment.
C COMPARISON OF WEIGHTS BETWEEN THE BC-PMI AND THE CBC METHOD
relevance with style
strong relevance with style
conditional state-action space
weak relevance with style
(a) Conditional BC
(b) BC using PMI weights
Figure 7: The weight distributions for (s, a) pairs different imitation learning paradigms are given follows: (a) In clustering condition based methods, the weights for (s, a) pairs belonging to the same style are set to I, while the weights for other pairs are set to O; (b) the proposed PMI weighting method, the weight distribution is approximated using Pointwise Mutual Information, which considers the relevance between each (s, a) pair and the style.
Lacking consideration to recognize the relevance of each state-action pair to the style is quite com- mon in various methods (e.g.. vanilla BC in (). The unsupervised clustering based or simply conditional methods can be regarded using indicator functions differentiate state-action pairs and assign corresponding style labels (Figure 7)). State-action pairs that belong to a specific style weight while pairs not belonging that style weight O, which can be denoted the following
16

---

Published as a conference paper at ICLR 2025
objective:
min
E
[ log 1, 1s(2(s,a))].
(22)
(s,a)~D.
where A(s, a) is a clustering function. To leverage the diversity information, our proposed method utilizes PMI to approximate the relevance between (s,a) a) pairs and styles, thereby obtaining more appropriate weights (Figure (p).
D THE LABEL FUNCTION
In our method, we focus on behavioral styles that can be represented by trajectory-level label function, denoted as J, which used to determine the style to which given trajectory belongs:
((T.) = [()])=(([
(2,)~ \(s,a)
> sc.
(23)
where À is a state-action-leve label function used to distinguish the behavioral styles, and e is style threshold. practice, we may only use a segment of the trajectory (+ to get the labels.
In this paper, following the common routine, approach this by specifying a task-relevant style embedding (represented by 2) in advance. For example, in atari, similar to Wu et al. (2023), we pre- define the styles such fire rate; in basketball, are commonly used as the style labels or metrics (Zhan et al. 2020 Li et al. 2021). A Python-style label function for fire rate is presented follows:
define the "destination" and "curvature," which
for traj in range (len (al1 (len(a]]traajs)
traj_name "().txt" format (all_trajs [traj]) with open (os, (os.path) join (data_dir,, traj_name), (r') as fl, open (os-path joonn dir, sraj_name]] 'w' ) #2:
lines f1. read .tp.aaidisslitlineesss
idx np. (1en (I. sub_idx [idxli: 1+sublen] for i in range (0, len (idx) -sublen, sublen)
max line = len (idx) sublen for si sub_idx: idx:
f_fire
fire_rate None count = 0
for sj in si:
data lines
[:8]
data [s.strip() for in data]
if len (data) >- and data[-3) !- (nallf;
count += 1
action (datal-311
if action
L_fire +-
fire_rate round (float fnre / count), 2)
if fire_rate < 0.1:
style =
elif fire_rate >= 0.1 and fire_rate 0.3:
style 2
elif fire_rate > 0.3: style =
for sj in si: data - lines [sj+s[sssppppcic.. data (s.strip() for in data] data. iata.insert(C (6, str (stylee)
[:8]
lines [s++] =
(data)
f2.write (n .ooin (lines [ max max_line]))
17

---

Published as a conference paper at ICLR 2025
A significant benefit is that labeling functions often simple scripts that can be quickly applied to the dataset, which is much cheaper than manual annotations and more reliable than unsupervised methods. In addition, classification ratios of each category label across the Atari dataset under each label function. The results in Table [6] show that the data differences for each category label are not significant, which is beneficial for the accuracy of the MINE network estimation. If the data for particular category label sparse, the estimated mutual information value may have a considerable error.
conducted a data analysis experiment the label function, comparing the
Table 6: The proportion of various styles of data labeled using the label function.
Env
Style
Category
Category
Category
Category
Pos
15.19%
34.67%
36.82%
13.32%
Alien
Range
37.08%
43.82%
19.10%
Ms
Pos
15.37%
33.76%
31.47%
19.40%
Pacman
Range
26.97%
43.82%
29.21% 1%
Space
Pos
32.35%
32.35%
35.29%
Invaders
Fire
18.52%
37.04%
44.44%
Our method focuses on the relevance between each (s, a) pair in a trajectory and the style. When the MINE function is well-fitted, samples that are completely irrelevant to the style will have PMI values close to zero, implying that these samples have minimal impact on the policy. However, noise introduced by the label function may affect the training of the MINE function. To analyze the impact of noise sample ratios on calibration, we use the position style in Space Invasers environment as an example. We define samples in a trajectory that fall outside the position range as noise samples. Ta- ble presents the proportion of noise samples for different styles and the corresponding calibration values. The results indicate that as the noise ratio increases, the improvement BC-PMI over CBC becomes more significant.
Table Comparison of calibration Under different noise ratios.
Pos Left Middle
Label Noise
CBC
BC-PMI
16.25%
0.75
0.84
31.23%
0.62
0.80
Right
19.28%
0.74
0.86
E ADDITIONAL RESULTS
E. THE NUMBER OF STYLES
We increased the fire rate to 6 categories in the Space Invaders environment to observe the impact of the increase in style categories on calibration. The results are shown in Table X Increasing the number of style categories may slightly decrease calibration, but compared to CBC, there is still noticeable improvement.
E.2 IMPLEMENTATION OF OTHER VARIANT METHODS
We implemented simple version of the GAIL algorithm combined with PMI. Specifically, first developed a Conditional GAIL (Ho & Ermon 2016) (CGAIL) variant, where joint information 2,) is used determine whether sample was generated by the actor. During interactions between the actor and the environment, the style label z is also provided as input when selecting actions. The PMI value is incorporated weight for the reward generated by the discriminator.
18

---

Published as a conference paper at ICLR 2025
Table 8: Comparison of calibration for more style categories.
Fire Rate
Proportion
CBC
BC-PMI
[0.0, 0.1)
18.52%
0.39+0.07
0.45
[0.1, 0.2)
22.17%
0.41
0.52
[0.2, 0.3)
14.85%
0.41#0.11
0.46+0.00
[0.3, 0.4)
17.42%
0.31#0.04
0.43
[0.4, 0.5)
15.89%
0.36+#0.05
0.39:00.09
[0.5, 1.0]
11.13%
0.46+40.07
0.53++0.05
When the reward is positive, samples with higher relevance between a) and receive a greater reward. The framework of the proposed method is illustrated in Figure 8
For the DWBC (Xu al. 2022 method, the original DWBC algorithm uses d(s,, a, log T) to deter- mine whether the current sample is generated by the expert policy. In our setting, we use d(s,, 2) to determine whether the current sample belongs style z. We compare the influence of the weights generated by the discriminator and the PMI weights policy diversity. Finally, the BC-Classifier method uses the softmax output of a classifier to estimate p(2|s, a) as the weight of the sample. The results shown in T
= - log(Dos, a,2) PMI
Actor
(s,a,2)
Env
Discriminator
(s',a',z))
Data
log(1 DD''', loog(((s, a,))
Figure 8: The Preliminary Framework of GAIL-PMI Method.
Table 9: Comparison of style calibration across different variant methods in Atari
Env
Style
CGAIL
CGAIL-PMI
BC-Classifier
DWBC
BC-PMI
Pos
0.31
0.39
0.47
0.44
0.61
Alien
Range
0.37++0.12
0.37+0.10
0.55=0.10
0.41+0008
0.55
MS
Pos
0.39+0.11
0.38+0.08
0.54=0.09
0.40+0.09
0.57+10.0
Pacman
Range
0.41+10.09
0.45+0.11
0.57-0.11
0.39
0.56:00.14
Space
Pos
0.48
0.60
0.75
0.64
0.83-10.06
Invaders
Fire
0.44+0.07
0.49-00.1
0.40=0.07
0.33++0.09
0.51++0.10
19


Reviews:
[{"review_id": "bmlIex7WtW", "reviewer": "Reviewer_6kbz", "summary": "The paper explores a methodology for deriving diverse policies from expert trajectory data. Rising from the traditional conditional behavior cloning (BC) algorithm, the authors  introduce an additional importance weight based on Pointwise Mutual Information (PMI), to assess the correspondence between state-action pairs and trajectory styles. The experimental results prove the proposed method can improve policy diversity with PMI weighting.", "questions": "1. Which behavior cloning (BC) algorithm is selected as the baseline method? Does the proposed PMI general well across different BC algorithms?\n2. From Table 1, the improvement of the proposed BC-PMI seems marginal compared with CBC.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The topic studied in this paper is quite interesting, with an objective to improve policy diversity.\n2. This article clearly defines the problem and provides a detailed exposition of the method.\n3. This article undertakes extensive comparative experiments across Circle 2D, Atari games and professional basketball player dataset, which valiate the effectiveness of the proposed algorithm compared with the baseline methods.", "weaknesses": "1. This paper aims to enhance policy diversity by incorporating a weighting method based on pointwise mutual information to the Conditional Behavioral Cloning framework. The key proposal of this paper is merely the introduction of a data weighting algorithm, which, in my view, does not represent adequate technical contribution.\n2. The manuscript does not include sufficient ablation experiments to validate the efficacy of the proposed data weighting method that utilizes pointwise mutual information. It remains unclear how BC-PMI compares to Conditional Behavioral Cloning that is conditioned on style, where each sample point is assigned equal weight. Additionally, its performance against other traditional data weighting methods has not been adequately compared.", "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": "8CXDVCDz9p", "reviewer": "Reviewer_ihL7", "summary": "In this work, the authors focus on developing a method for imitation learning better diversity of policy. It works under assumption that the expert trajectories in training are collected by stylized experts. The key is to using a pointwise mutual information-based weighting strategy to determine the policy importance. The state-action pair with higher posterior probability is given higher importance. The proposed method achieves good experiment results and also is demonstrated to be covered in the extreme cases, such as zero mutual information or no overlap among different policy styles in a single state-action pair. For the extreme cases, the theoretical evidence is provided to show that the proposed strategy degrades to usual strategy without failing.", "questions": "Please see my comments in previous sectors.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method is motivated by an intuitive insight and the failure of existing alternatives. The use of PMI and MINE are well connected for the purpose.\n2. The proposed method achieves good experimental results on the datasets of CIRCLE2D , Atari and Basketball Player dataset.", "weaknesses": "1. Given that stylized experts make unbalanced trajectory generation in training, it might be helpful to discuss the disentanglement between the stylized code and the state-action pair. Some previous works have built related benchmark for the purpose such as [1] and [2]. Though the benchmarks may not be originally designed for imitation learning, the corresponding metrics might be worth reference. Such evaluation can be adapted by measuring the mutual information or other metrics between the stylized code and the pair.\n\n2. It remains not unclear, at least not generally convincing across datasets, that PMI shows a significant advance than the usual MI for BC. More evidence or discussion about this can be helpful to enhance the significance of the proposed method as PMI is claimed as a main contribution in this paper. If replacing MI with PMI can result in generalizable performance boosting across datasets, the proposed method can be supported with more experimental significance. \n\nReference:\n\n[1] “**An Empirical Study on Disentanglement of Negative-free Contrastive Learning**”, NeurIPS 2022.\n\n[2] “Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations.”, ICML 2019", "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": "BGZ68kwklz", "reviewer": "Reviewer_fKFr", "summary": "The paper introduces Behavioral Cloning with Pointwise Mutual Information Weighting (BC-PMI), a new approach for recovering diverse policies in imitation learning. Here, the expert data consists of trajectories coming from different experts, each reflecting a unique approach or style within the task. The authors’ key idea is that, even within a single style, not all state-action pairs equally represent that style’s characteristics. To address this, BC-PMI weights state-action pairs based on their relevance to the style, calculated using pointwise mutual information between $(s,a)$ pairs and the style code $z$. These weighted pairs are then used to learn a style-conditioned policy $\\pi(a|s,z)$ using behavioral cloning. This selective weighting allows the model to prioritize pairs that are most representative of the target behavior, thereby enhancing its ability to learn style-conditioned policies that more accurately capture the style within expert trajectories. \n\nThe paper also provides some theoretical insights showing that BC-PMI smoothly interpolates between traditional behavioral cloning and clustering-based behavior cloning, depending on the mutual information between style and state-action pairs.\n\nThe authors evaluate BC-PMI in three settings: Circle 2D, a simple 2D setup; Atari Games (Alien, MsPacman, and Space Invaders); and the Professional Basketball Dataset, featuring NBA player movement data. Results show that BC-PMI achieves high style calibration across various metrics, outperforming baselines. However, it is worth noting that the baselines operate in an unsupervised manner and do not have access to the style labels available to BC-PMI.", "questions": "1. Why is MINE used to estimate the pointwise mutual information? This can be estimated by simply learning a classifier across different styles, since you are not interested in the $I(S,A;Z)$ rather $p(z|s,a)$. It would be valuable to compare the calibration performance of BC-PMI using a classifier-based approach against the current method with MINE’s statistic network, to assess if simpler alternatives might achieve similar results.\n2. What are the calibration metrics used in Tables 3/4? Is it DTW, or ED, or KL?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well written and clearly communicates its key ideas.\n2. The idea of selectively weighting state-action pairs based on their relevance to a given style, rather than treating entire trajectories uniformly, is both intuitive and impactful. Existing methods typically operate at the trajectory level, but this paper demonstrates that not all state-action pairs within a trajectory equally represent the intended style.", "weaknesses": "1. A key disadvantage of the method is that it requires style markers for each trajectory in the training dataset. These can be expensive to obtain. A brief note is made about using labeling functions to obtain these markers, but no experiments are included. This also makes the comparison with existing works unfair since they are unsupervised methods. To make the comparisons fair, the authors should include the supervision in other methods, e.g., they can train the posterior network in InfoGAIL using the style labels.\n\n2. The experimental section is relatively brief, lacking analysis on the method’s performance as the number of styles increases or when labels are noisy. The latter case is quite important given that labels generated by programmable labeling functions (as mentioned in L296) may be imperfect. For a comprehensive understanding of these factors, it would be beneficial to include experiments that measure calibration as the number of styles increases, as well as tests assessing robustness to label noise introduced by such labeling functions.", "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": "CZcsGl1BDy", "reviewer": "Reviewer_BUhV", "summary": "This paper investigates how to recover diverse policies from expert trajectories, proposing a new method that leverages the relevance of state-action pairs to trajectory styles. By introducing Pointwise Mutual Information to model this relationship, the method approaches the problem of policy diversity from a different perspective, ultimately achieving results that surpass previous state-of-the-art methods.", "questions": "Please refer to the weakness.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The motivation is simple but effective, and it will be highly beneficial for the efficient use of imitation learning data in subsequent robot tasks, contributing significantly to the development of the community.\n2. The paper is well-written and easy to follow.\n3. The experimental results are extensive and convincing.", "weaknesses": "1. I would like to know if there are other reweighting methods that can be used for comparison. If so, could some comparative experiments be included?\n2. Additionally, the current method evaluates the diversity of state-action pairs. As far as I understand, there are also studies specifically focused on state diversity and action diversity. I wonder if these could serve as a standard for comparison?\n\n[1] Suneel Belkhale, Yuchen Cui, Dorsa Sadigh. Data quality in imitation learning. NeurIPS 2023", "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)