Diverse Policies Recovering via Pointwise Mutual Information Weighted Imitation Learning

## Abstract

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 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 recovery. 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 policies from expert data.

## 1 Introduction

Imitation Learning (IL) is about observing expert demonstrations in performing a task and learning to mimic those actions (Hussein et al., ; Osa et al., ). Vanilla behavioral cloning (BC) (Pomerleau, ) learns a 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, ) mitigates the issue via learning both a discriminator and a policy. Despite their wide applications, these methods lack mechanisms to 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., ; Wang et al., ; Zhan et al., ; Shafiullah et al., ; Mao et al., ), which can generally divided into two categories. In one category, the latent style $z$ of a trajectory is inferred in an unsupervised manner, for instance, by an expectation maximization (EM) procedure. In the other category, the style $z$ of a trajectory is determined by a user-specified function, for instance, a programmatic labeling function. No matter in which category, samples that are used to train a style-conditioned policy $\pi(a|s,z)$ 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 a basketball game, the diversity of shooting position behaviors is primarily related to the part leading up to the shot, while it is less relevant to other parts in the trajectory. In a nutshell, policies with different styles could have significantly different degrees of overlap in different areas of the state-action space.
Therefore, when learning a style-conditioned policy, the relevance between the state-action pairs in the trajectory and the behavioral style should be taken into consideration.

In this paper, we propose a 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 $z$ 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, ; Manning & Schutze, ; Bouma, ), in traditional conditional BC, to quantify the relevance of state-action pairs with the conditioned style. Intuitively, state-action pairs with a larger posterior of the corresponding style $p(z|s,a)$ are given a larger weight.
In practical implementation, we utilize Mutual Information Neural Estimation (MINE) (Belghazi et al., ) 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 no style diversity in expert data, our objective degenerates to vanilla BC, which views the data as generated from one policy. By contrast, when policies with different 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 (IL) methods are designed to mimic the behaviors of experts.
Behavioral Cloning (BC) (Pomerleau, ), a well-known IL 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 handle distributional shifts during evaluation (Ross et al., ; Fujimoto et al., ; Wu et al., ; Peng et al., ; Kostrikov et al., ).
Inverse Reinforcement Learning (IRL) (Ng et al., ; Arora & Doshi, ), another type of IL, learns a reward function that explains the expert behavior and then uses this reward function to guide the agent’s learning process.
Popular IRL approaches like Generative Adversarial Imitation Learning (GAIL) (Ho & Ermon, ) and Adversarial Inverse Reinforcement Learning (AIRL) (Fu et al., ) use adversarial training to learn a policy that is similar to the expert policy while being robust to distributional shift.
However, these methods are limited to imitating a 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 et al., ).
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 to cover all styles of the expert’s diverse behaviors Wang et al. ().

#### Policy Diversity

Diversity is crucial in imitation learning algorithms, especially in practical control tasks and multi-player games (Zhu et al., ), as diverse policies in control tasks can enhance the robustness of adapting to various environments. In contrast, AI with diverse policies can maximize the player’s experience and the ornamental value in games and competitions (Yannakakis & Togelius, ).
Some approaches utilize information-theoretic methods to address this issue and learn the behavioral styles of policies.
InfoGAIL (Li et al., ) and Intention-GAN (Hausman et al., ) augment the objective of GAIL with the mutual information between generated trajectories and the corresponding latent codes.
Wang et al. () use a variational autoencoder (VAE) module to encode expert trajectories into a continuous latent variable.
Eysenbach et al. () propose a general method called Diversity Is All You Need (DIAYN) for learning diverse skills without explicit reward functions.
DIAYN focuses on discovering skills in online reinforcement learning tasks (Campos et al., ; Sharma et al., ; Achiam et al., ), whereas our method specifically considers pure offline imitation learning scenarios.
More recently, Mao et al. () introduced Stylized Offline RL (SORL), which utilizes unsupervised learning methods to cluster styles and fits policies to each cluster separately, optimizing an equal number of policies while using the KL divergence as a constraint.
However, these methods focus on inferring the latent style or clustering the trajectories with different styles without considering the relevance between the state-action pairs in the trajectory and the behavioral style.
In contrast, we focus on 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 to focus on the samples that are highly relevant to the style.

## 3 Preliminary

### 3.1 Problem Setting and Vanilla Behavior Cloning

We consider the standard Markov Decision Process (MDP) (Sutton & Barto, ) as the mathematical framework for modeling sequential decision-making problems, which is defined by a tuple $\left\langle\mathcal{S},\mathcal{A},P,r,d_{0},T\right\rangle$, where $\mathcal{S}$ is a finite set of states, $\mathcal{A}$ is a finite set of actions, $P:\mathcal{S}\times\mathcal{A}\times\mathcal{S}\to\mathbb{R}$ is the transition probability function, $r:\mathcal{S}\to\mathbb{R}$ is the reward function, $d_{0}$ is the initial distribution, and $T$ is an episode horizon.
A policy $\pi:\mathcal{S}\times\mathcal{A}\to[0,1]$ maps from state to distribution over actions.
Let $d^{\pi}_{t}$ and $d^{\pi}=\frac{1}{T}\sum_{t=1}^{T}d^{\pi}_{t}$ denote the distribution over states at time step $t$ and the average distribution over $T$ time steps induced by $\pi$, respectively.
The vanilla BC loss function is:

$$ $\mathcal{L}_{\mathrm{BC}}=\mathop{\mathbb{E}}\limits_{(s,a)\sim\mathcal{D}_{e} }\big{[}-\log\pi(a|s)\big{]}.$ (1) $$

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 $\mathcal{D}_{e}$, which consist of many diverse trajectories $\tau$, are collected by stylized expert policies denoted as $\{\pi_{e}^{(1)},\pi_{e}^{(2)},\dots,\pi_{e}^{(K)}\}$.
Let $z\in\mathcal{Z}$ denotes the variable indicating which stylized policy $\tau$ belongs to, and $p(\tau|z=i)$ denotes the probability of $\tau$ sampled under policy $\pi_{e}^{(i)}$.
Our objective is to learn a conditioned policy $\pi(a|s,z)$ such that trajectories generated by $\pi(a|s,z)$ closely match the demonstrations in $\mathcal{D}_{e}$ that exhibit the corresponding style $z$.

### 3.2 Mutual Information Neural Estimation

Mutual Information Neural Estimation (MINE) is a powerful technique for estimating the mutual information between two random variables using neural networks (Belghazi et al., ).
It has been widely used in various domains, including representation learning (Hjelm et al., ), generative modeling (Chen et al., ), and imitation learning (Eysenbach et al., ).
The key idea behind MINE is to formulate the estimation of mutual information as an optimization problem.
Given two random variables $X$ and $Y$, the mutual information $I(X;Y)$ can be expressed as the Kullback-Leibler (KL) divergence between the joint distribution $P_{XY}$ and the product of marginal distributions $P_{X}\otimes P_{Y}$:

$$ $I(X;Y)=D_{KL}(P_{XY}||P_{X}\otimes P_{Y}).$ (2) $$

MINE approximates this KL divergence using a lower bound based on the Donsker-Varadhan representation (Donsker & Varadhan, ):

$$ $I(X;Y)\geq\sup_{T_{\theta}\in\mathcal{F}}\mathbb{E}_{P_{XY}}[T_{\theta}]-\log( \mathbb{E}_{P_{X}\otimes P_{Y}}[e^{T_{\theta}}]),$ (3) $$

where $\mathcal{F}$ is a class of functions $T:\mathcal{X}\times\mathcal{Y}\to\mathbb{R}$.
In MINE, this function class is parameterized by a neural network $T_{\theta}$, which takes as input samples from the joint distribution $P_{XY}$ and the product of marginal distributions $P_{X}\otimes P_{Y}$.
The network is trained to maximize the lower bound, equivalent to estimating the mutual information.

## 4 Behavioral Cloning with Pointwise Mutual Information Weighting

As discussed in the introduction, in many real-world applications, the impact of different state-action pairs on the style can vary greatly, and often only a part of the trajectory is highly relevant to the style. Hence, when training a style-conditioned policy, we 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 on the style, thereby achieving a more generalized imitation objective:

$$ $\min_{\pi}\mathop{\mathbb{E}}\limits_{(s,a,z)\sim\mathcal{D}_{e}}\left[-\log \pi(a|s,z)\cdot\sigma(s,a,z)\right],$ (4) $$

where $\sigma(s,a,z)$ is a weighting function. Intuitively, state-action pairs that are more exclusive to a 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, ; Manning & Schutze, ; Bouma, ) to quantify the contribution of $(s,a)$ when learning a style-conditioned policy:

$$ $\mathcal{P}(z;s,a)=\log\frac{p(s,a,z)}{p(s,a)p(z)}=\log\frac{p(z|s,a)}{p(z)}.$ (5) $$

We utilize the exponential of $\mathcal{P}(z;s,a)$ 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 $z$ given $(s,a)$ is larger than the prior, i.e., $p(z|s,a)>p(z)$, this means there is a high relevance between $(s,a)$ and style $z$, and we should give a larger weight.
If $(s,a)$ and style $z$ are nearly independent, then $p(z|s,a)\approx p(z)$.
When we have $p(z|s,a)<p(z)$, this means (s,a) is more likely generated by other styles, and we should give a lower weight here. Accordingly, Eq.(4) turns into the following form:

$$ $\mathcal{L}_{\mathrm{BC-PMI}}(\theta)=\mathop{\mathbb{E}}\limits_{(s,a,z)\sim \mathcal{D}_{e}}\left[-\log\pi_{\theta}(a|s,z)\cdot e^{\mathcal{P}(z;s,a)} \right]=\mathop{\mathbb{E}}\limits_{(s,a,z)\sim\mathcal{D}_{e}}\left[-\log\pi_ {\theta}(a|s,z)\cdot\frac{p(z|s,a)}{p(z)}\right]$ (6) $$

### 4.1 Theoretical analysis

We give some theoretical justifications for our BC-MI objective in Eq.(6). Considering an extreme case where the mutual information between $(s,a)$ and style $z$ approaches zero, which means policies with different styles have nearly the same state-action distribution.
In this case, the BC-PMI objective 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 overlap, we find that the BC-PMI objective degenerates into behavior cloning on each style.
Formally, we have the following proposition:

###### Proposition 1 .

(a). When the mutual information $I(Z;S,A)$ equals to 0 0, 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:

$$ $\operatorname*{arg\,min}_{\theta}\mathcal{L}_{\text{BC-PMI}}(\theta)= \operatorname*{arg\,min}_{\theta}\mathcal{L}_{\text{BC}}(\theta).$ (7) $$

(b). When the conditional entropy $H(Z|S,A)$ equals 0 0, it indicates that there is a significant distinction in the trajectory style corresponding to all the state-action pairs. In this case, the BC-PMI objective degenerates to the behavior cloning on each style:

$$ $\mathcal{L}_{\text{BC-PMI}}(\theta)=\sum_{i=1}^{K}\mathcal{L}_{\mathrm{BC}}^{( i)}(\theta),$ (8) $$

where $\mathcal{L}_{\mathrm{BC}}^{(i)}(\theta)$ is the behavior cloning loss on the subset of data with style label $i$.

Proof. Refer to Appendix A.1.

Equation 8 means when $H(Z|S,A)$ equals to 0 0, the BC-PMI objection function can be viewed as separating the $(s,a)$ by the style $z$ and optimize it respectively.
This is associated with the clustering-based behavior cloning objective, where each trajectory is first assigned to a specific style cluster based on its style label, and then behavior cloning is performed within each cluster.

The above analysis provides insights into the mechanisms of BC-PMI. In practice, the mutual information $I(Z;S,A)$ often lies between the two extremes of 0 0 and $H(Z)$.
By adjusting the weight of a state-action sample in learning a 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.

###### Proposition 1 .

### 4.2 Practical Implementation

To practically estimate the PMI values in Eq.(6), we employ the Mutual Information Neural Estimation (MINE) (Belghazi et al., ) method. MINE is a neural network-based approach that can effectively estimate mutual information between high-dimensional random variables. By leveraging the Donsker-Varadhan representation (Donsker & Varadhan, ) of the Kullback-Leibler (KL) divergence, MINE allows for the estimation of mutual information using neural networks.

Let $T_{\phi}:\mathcal{S}\times\mathcal{A}\times\mathcal{Z}\to\mathbb{R}$ be a neural network parameterized by $\phi$. The mutual information between $(s,a)$ and $z$ can be estimated as:

$$ $\mathcal{I}(s,a;z)\geq\mathop{\mathbb{E}}\limits_{(s,a,z)\sim\mathcal{D}_{e}} \left[T_{\phi}(s,a,z)\right]-\log\left[\mathop{\mathbb{E}}\limits_{(s,a)\sim \mathcal{D}_{e},\bar{z}\sim p(z)}\left[e^{T_{\phi}(s,a,\bar{z})}\right]\right].$ (9) $$

The neural network $T_{\phi}$ is trained to maximize the lower bound in Eq.(9), which is equivalent to minimizing the KL divergence between the joint distribution $p(s,a,z)$ and the product of marginals $p(s,a)p(z)$. According to the proof of Theorem 1 in MINE (Belghazi et al., ), the optimal solution for $T_{\phi}$ is:

$$ $T_{\phi}^{*}(s,a,z)=\log\frac{p(s,a,z)}{p(s,a)p(z)}=\log\frac{p(z|s,a)}{p(z)},$ (10) $$

which is exactly the PMI we aim to estimate.

In practice, we can train the MINE network $T_{\phi}$ using samples from the expert demonstrations $\mathcal{D}_{e}$ and the style distribution $p(z)$. The training objective for $T_{\phi}$ is:

$$ $\max_{\phi}\mathop{\mathbb{E}}\limits_{(s,a,z)\sim\mathcal{D}_{e}}\left[T_{ \phi}(s,a,z)\right]-\log\left[\mathop{\mathbb{E}}\limits_{(s,a)\sim\mathcal{D} _{e},\bar{z}\sim p(z)}\left[e^{T_{\phi}(s,a,\bar{z})}\right]\right].$ (11) $$

By optimizing Eq.(11), we obtain an approximation of the PMI values, which can be used as the weights in the behavioral cloning objective in Eq.(6).
The weights can be denoted as:

$$ $\sigma(s,a,z)=\exp[T_{\phi}^{*}(s,a,z)].$ (12) $$

Furthermore, in order to reduce the variance of the gradient in optimizing Eq.(6), we can subtract an optimal baseline from the weight, like what has been done in A3C (Mnih et al., ). This turns Eq.(6) to the following objective:

$$ $\min_{\pi}\mathop{\mathbb{E}}\limits_{(s,a,z)\sim\mathcal{D}_{e}}\Big{[}-\log \pi(a|s,z)\cdot\big{[}{\color[rgb]{0.25,0.41,0.88}\exp(T_{\phi}^{*}(s,a,z)){ \color[rgb]{0.9,0.17,0.31}-\widetilde{b}}}\big{]}\Big{]},$ (13) $$

where
$\widetilde{b}=\mathop{\mathbb{E}}\limits_{(s,a,z)\sim\mathcal{D}_{e}}\Big{[}
\exp(T_{\phi}^{*}(s,a,z))\Big{]}$.
In practice, we can estimate $\widetilde{b}$ using the moving average.
The pseudo-code for the BC-PMI algorithm is shown in Algorithm 1.

## 5 Experiments

In this experimental section, we aim to address the following questions:

- Q1.
Can our method recover diverse and controllable policies from diverse style trajectories?
- 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?

Figure: Algorithm 1 Behavioral Cloning with Pointwise Mutual Information Weighting (BC-PMI)

### 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 on manual annotations and uncontrollable when using unsupervised approaches.
Instead, the programmable labeling functions (Ratner et al., ; Zhan et al., ) can be used to automatically generate style labels.

We compare PMI-BC to 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., ); (4) InfoGAIL (Li et al., ), which infers the latent style of trajectories by maximizing the mutual information between the latent codes and trajectories; (5) SORL (Mao et al., )(^*^**For the SORL algorithm, we only used the EM clustering part.), which use Expectation-Maximization (EM) algorithm to classify the trajectories from the heterogeneous dataset into clusters where each represents a distinct and dominant motion style.
Specifically, we first introduce a toy example called Circle 2D to provide a simple visualization and analysis.
Then in Atari games, detailed analysis 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 a 2D plane where an agent can freely move at a constant velocity by selecting its direction, denoted as $p_{t}$, at time step $t$.
For the agent, the observation at time step $t$ includes the state information from time step $t-4$ to $t$.
The offline expert trajectories consist of four different styles, each generated by a random expert policy.
The expert policy generates trajectories that resemble circular patterns after a 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.

Figure: (a) Expert(without noise)
Refer to caption: extracted/5944743/Figure/circle_exp.png

This task is a 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, $(s,a)$ that moves according to the curvature of a specific style after 75 steps will be assigned a higher PMI, and the linear motion part or the movement curvature which are not related to the style will be assigned a lower PMI.
The visualized results are shown in Figure 1.
By introducing PMI, the calibration of the BC policies has been further improved.
Table 1 is a comparison of numerical results.
The metrics used are as follows: (1)DTW, which is Dynamic Time Warping, a dynamic programming algorithm for calculating the similarity of sequences of different lengths; (2)ED, which is Euclidean Distance; (3)KL, which is used to calculate the state-action distribution difference of the trajectory.

**Table 1: Comparison of style calibration across different metrics in Circle 2D.**
| Style | Metrics | BC | CBC | BC-PMI |
| --- | --- | --- | --- | --- |
| Class 1 (Red) | DTW | 88.590 $\pm$ 9.329 | 8.130 $\pm$ 0.732 | 7.511 $\pm$ 0.545 |
| ED | 133.100 $\pm$ 20.210 | 22.039 $\pm$ 1.383 | 21.591 $\pm$ 1.091 |  |
| KL | 8.037 $\pm$ 0.870 | 0.044 $\pm$ 0.006 | 0.037 $\pm$ 0.004 |  |
| Class 2 (Green) | DTW | 45.729 $\pm$ 5.293 | 7.695 $\pm$ 0.814 | 7.631 $\pm$ 0.821 |
| ED | 88.180 $\pm$ 7.541 | 35.355 $\pm$ 2.691 | 35.777 $\pm$ 2.912 |  |
| KL | 1.111 $\pm$0.072 | 0.043 $\pm$ 0.003 | 0.037 $\pm$ 0.003 |  |
| Class 3 (Orange) | DTW | 77.619 $\pm$ 7.812 | 18.972 $\pm$ 2.323 | 11.576 $\pm$ 1.788 |
| ED | 130.839 $\pm$ 10.322 | 107.833 $\pm$ 8.213 | 74.267 $\pm$ 8.002 |  |
| KL | 16.839 $\pm$ 1.022 | 0.271 $\pm$ 0.031 | 0.135 $\pm$ 0.014 |  |
| Class 4 (Blue) | DTW | 69.527 $\pm$ 10.924 | 7.755 $\pm$ 1.832 | 7.577 $\pm$ 1.051 |
| ED | 97.066 $\pm$ 14.239 | 27.230 $\pm$ 2.110 | 26.749 $\pm$ 2.347 |  |
| KL | 36.238 $\pm$ 2.981 | 0.401 $\pm$ 0.061 | 0.219 $\pm$ 0.019 |  |

### 5.3 Atari Games

In this experiment, we concentrate on three widely recognized Atari games: Alien, MsPacman and SpaceInvaders.
The datasets utilized in this study are sourced from Atari-Head (Zhang et al., ; ), an extensive collection of human game-play data.
Atari-Head are meticulously recorded in a semi-frame-by-frame manner, ensuring high data quality and granularity, which facilitates in-depth analysis and robust evaluation of our proposed method.

This experiment consists of three parts.
Firstly, we demonstrate the convergence of the lower bound of mutual information in Eq.(9), which indicates the relevance between state-action pairs and styles.
Secondly, we provide interpretability of PMI weights by assessing the extent to which they appropriately 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 the Alien and MsPacman game, we 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 on the map as a Gaussian distribution, differentiating the variance of the distribution.
A higher variance indicates a tendency for the agent to move across areas, while a lower variance indicates a 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.

Figure: Figure 2: MI between state-action pairs and styles. FR: Fire Rate style; AR: Movement Area style; RG: Movement Range style.
Refer to caption: x1.png

The results in Figure 2 indicate a high relevance between the agent’s state-action distribution and the style.
Moreover, as the relevance increases, the converged MI values also increase.
The MsPacmanAR style divides the map into four areas: top-left, top-right, bottom-left, and bottom-right.
The agent’s trajectory exhibits clear distinctions under this style, which is reflected in the larger converged MI values, while the MsPacmanRG style represents the style based on the range of movement, which is less distinct compared to the movement area style.

Figure: Figure 3: The PMI weight values related to style for each frame along a trajectory, with the corresponding style being the bottom-left area style (green color) in the game frame. The agent in the game frame is indicated by a white circle, and the corresponding actions are indicated by white arrows. No arrow indicates that the action is NOOP.
Refer to caption: x2.png

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 a 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 to a higher relevance with the style and higher PMI values. Conversely, in the third and 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 a lower relevance with the style.

To verify whether different styles of policies affect the cumulative reward, we conducted experiments in three different Atari environments, as shown in Table 2.
The results indicate that most styles, such as those related to position and range of movement, do not impact the cumulative reward 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 a higher firing frequency, significantly increases the agent’s cumulative reward, while a lower firing frequency reduces the cumulative reward.

**Table 2: 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 |
| --- | --- | --- | --- | --- | --- | --- |
| Alien | Pos | 460.5$\pm$173.4 | 416.0$\pm$160.8 | 432.0$\pm$174.7 | 500.3$\pm$211.7 | 496.1$\pm$191.5 |
| Range | 460.5$\pm$173.4 | 440.0$\pm$217.0 | 501.2$\pm$247.9 | 423.2$\pm$240.9 | - |  |
| Ms | Pos | 701.5$\pm$182.4 | 665.4$\pm$216.3 | 803.5$\pm$211.2 | 674.3$\pm$146.5 | 666.2$\pm$172.2 |
| Pacman | Range | 701.5$\pm$182.4 | 636.1$\pm$265.1 | 759.5$\pm$280.6 | 738.3$\pm$260.8 | - |
| Space | Pos | 208.2$\pm$73.0 | 178.9$\pm$54.5 | 187.0$\pm$54.9 | 200.2$\pm$56.4 | - |
| Invaders | Fire | 208.2$\pm$73.0 | 94.8$\pm$47.2 | 207.9$\pm$70.3 | 270.6$\pm$86.7 | - |

Finally, we compared the calibration of several different style policies across these Atari environments, as 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 is 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.

**Table 3: Comparison of style calibration across different styles in Atari**
| Env | Style | BC | CBC | CTVAE | InfoGAIL | SORL | BC-PMI |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Alien | Pos | 0.23$\pm$0.11 | 0.43$\pm$0.04 | 0.52$\pm$0.04 | 0.34$\pm$0.11 | 0.31$\pm$0.07 | 0.61$\pm$0.08 |
| Range | 0.31$\pm$0.09 | 0.51$\pm$0.13 | 0.57$\pm$0.09 | 0.31$\pm$0.11 | 0.27$\pm$0.07 | 0.55$\pm$0.12 |  |
| MS | Pos | 0.17$\pm$0.05 | 0.48$\pm$0.07 | 0.55$\pm$0.04 | 0.33$\pm$0.05 | 0.27$\pm$0.08 | 0.57$\pm$0.07 |
| Pacman | Range | 0.33$\pm$0.09 | 0.53$\pm$0.11 | 0.61$\pm$0.05 | 0.29$\pm$0.09 | 0.30$\pm$0.12 | 0.56$\pm$0.14 |
| Space | Pos | 0.36$\pm$0.06 | 0.71$\pm$0.08 | 0.66$\pm$0.10 | 0.41$\pm$0.06 | 0.37$\pm$0.04 | 0.83$\pm$0.06 |
| Invaders | Fire | 0.27$\pm$0.07 | 0.41$\pm$0.08 | 0.47$\pm$0.06 | 0.36$\pm$0.08 | 0.34$\pm$0.09 | 0.51$\pm$0.10 |

### 5.4 Professional Basketball Player Dataset

In this experiment, we validate our method on the dataset of a collection of professional basketball player trajectories (Zhan et al., ) with the goal of recovering policies that can generate trajectories with diverse player-movement styles.
The basketball trajectories are collected from tracking real players in the NBA.
We primarily focus on two movement styles: (1) The Destination, which is the distance from the final position to a fixed destination on the
court (e.g. the basket), and (2) The Curvature, which measures
the player’s propensity to change directions.

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 is 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 1 is close to the ball frame, Destination 2 is in the middle count, and Destination 3 is 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 5.
The results indicate that the BC-PMI algorithm can effectively imitate policies of different movement styles from real human data.

Figure: Figure 4: Visualization of different destination styles.
Refer to caption: x3.png

Figure: Figure 5: Visualization of different curvature styles.
Refer to caption: x4.png

**Table 4: Comparison of style calibration (%) across different styles in Basketball**
| Style | Class | BC | CBC | CTVAE | InfoGAIL | SORL | BC-PMI |
| --- | --- | --- | --- | --- | --- | --- | --- |
| Destination | Class 1 | 31.1 | 81.6 | 78.3 | 79.6 | 80.9 | 93.3 |
| Class 2 | 42.7 | 84.4 | 81.2 | 75.5 | 76.1 | 92.7 |  |
| Class 3 | 19.4 | 73.1 | 77.8 | 74.5 | 79.5 | 89.8 |  |
| Curvature | Class 1 | 17.1 | 67.3 | 58.4 | 60.9 | 51.6 | 79.4 |
| 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 3 classes of Destination and 3 classes of Curvature in Table 4.
Due to the lack of style label information, BC can only learn an approximately average strategy, so it should be evenly distributed among various styles.
However, due to factors such as initialization position, the samples in each category do not strictly follow a 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.

## 6 Discussion

### 6.1 Conclusion

In this paper, we investigate how to recover diverse policies from a 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 in our approach to the problem of policy diversity from a different perspective, which involves the introduction of Pointwise Mutual Information to 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-art methods.

### 6.2 Limitations and Future Work

In our setting, our goal is to 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 policy.
Hence, we consider the combination of PMI and RL as future research, in which we believe the PMI weight can be associated with the offline RL method to promote diversity while improving performance.

## Appendix A Theoretical Derivation

### A.1 Proof of Proposition 1

###### Proof.

(a) We first introduce the Gibbs’ inequality, which was proved in (Brémaud, ).

###### Lemma 1 (Gibbs’ inequality) .

Let $p(x)$, $q(x)$ be the probability mass function, fixed $p(x)$, that the object function $-\mathbb{E}_{x\sim p(x)}\log q(x)=\sum_{x}-p(x)\log q(x)$ takes the minimum when $q(x)=p(x)$.

By Lemma 1, we know

$$ $\mathcal{L}_{BC}(\theta)=\mathop{\mathbb{E}}\limits_{(s,a)\sim\mathcal{D}_{e}} [-\log\pi_{\theta}(a|s)]=\mathop{\mathbb{E}}\limits_{s\sim\mathcal{D}_{e}} \left[\mathop{\mathbb{E}}\limits_{a|s\sim p(a|s)}[-\log\pi_{\theta}(a|s)] \right],$ (16) $$

which takes the minimum when $\pi_{\theta}(a|s)=p(a|s)$.

When $I(Z;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:

$$ $\frac{p(z|s,a)}{p(z)}=\frac{p(z)}{p(z)}=1.$ (17) $$

Hence,

$$ $\displaystyle\mathcal{L}_{\mathrm{BC-PMI}}(\theta)$ $\displaystyle=\mathop{\mathbb{E}}\limits_{(s,a,z)\sim\mathcal{D}_{e}}[-\log\pi _{\theta}(a|s,z)]$ $\displaystyle=\mathop{\mathbb{E}}\limits_{s\sim\mathcal{D}_{e}}\left[\sum_{a} \sum_{z}-p(a|s)p(z|a,s)\log\pi_{\theta}(a|s,z)\right]$ $\displaystyle=\mathop{\mathbb{E}}\limits_{s\sim\mathcal{D}_{e}}\left[\sum_{z} \sum_{a}-p(a|s)p(z)\log\pi_{\theta}(a|s,z)\right]$ $\displaystyle=\mathop{\mathbb{E}}\limits_{s\sim\mathcal{D}_{e}}\left[\sum_{z} \left[p(z)(\sum_{a}-p(a|s)\log\pi_{\theta}(a|s,z))\right]\right],$ (18) $$

where the third equation is because $z$ is independent with $(s,a)$.
Since $p(z)$ is non-negative, when $\sum_{a}-p(a|s)\log\pi_{\theta}(a|s,z)$ takes the minimum for all $z$, the objective reach its the minimum. Using Lemma 1 again, we can see when $\pi_{\theta}(a|s,z)=p(a|s)$, $\mathcal{L}_{\mathrm{BC-PMI}}(\theta)$ reach its minimum.

(b) when $H(Z|S,A)=0$, it means that given a state-action pair $(s,a)$, the style variable $Z$ can be determined with high certainty.
In other words, for any $(s,a,z)\sim\mathcal{D}_{e}$, we have:

$$ $p(z|s,a)=\mathbbm{1}(z=z_{s,a}),$ (19) $$

where $z_{s,a}$ is the unique style label corresponding to $(s,a)$.

Assuming there are $K$ styles in the Dataset. Let $z_{\tau}$ be the unique style label corresponding to trajectory $\tau$, $D_{e}^{(i)}$ be the subset of state-action pairs with style label $i$ and $\tau\in D_{e}^{(i)}$ means the whole trajectory from $D_{e}^{(i)}$. Denote

$$ $\mathcal{L}_{\mathrm{BC}}^{(i)}(\theta)=\mathop{\mathbb{E}}\limits_{(s^{\prime },a^{\prime})\sim D_{e}^{(i)}}[-\log\pi_{\theta}(a^{\prime}|s^{\prime},i)]= \frac{1}{|D_{e}^{(i)}|}\sum_{\tau\in D_{e}^{(i)}}\sum_{(s^{\prime},a^{\prime}) \in\tau}-\log\pi_{\theta}(a^{\prime}|s^{\prime},i).$ (20) $$

We have

$$ $\displaystyle\mathcal{L}_{\mathrm{BC-PMI}}(\theta)$ $\displaystyle=\mathop{\mathbb{E}}\limits_{(s,a,z)\sim\mathcal{D}_{e}}\left[- \log\pi_{\theta}(a|s,z)\cdot\frac{p(z|s,a)}{p(z)}\right]$ $\displaystyle=\mathop{\mathbb{E}}\limits_{(s,a,z)\sim D_{e}}[-\log\pi_{\theta} (a|s,z)\frac{\mathbbm{1}(z=z_{s,a})}{p(z)}]$ $\displaystyle=\frac{1}{|D_{e}|}\sum_{\tau\in D_{e}}\sum_{i=1}^{K}\mathbbm{1}(z _{\tau}=i)\sum_{(s^{\prime},a^{\prime})\in\tau}[-\log\pi_{\theta}(a^{\prime}|s ^{\prime},i)\cdot\frac{1}{p(z=i)}]$ $\displaystyle=\frac{1}{|D_{e}|}\sum_{i=1}^{K}\sum_{\tau\in D_{e}}\sum_{(s^{ \prime},a^{\prime})\in\tau}\mathbbm{1}(z_{\tau}=i)[-\log\pi_{\theta}(a^{\prime }|s^{\prime},i)\cdot\frac{|D_{e}|}{|D_{e}^{(i)}|}]$ $\displaystyle=\sum_{i=1}^{K}\sum_{\tau\in D_{e}^{(i)}}\frac{1}{|D_{e}^{(i)}|} \sum_{(s^{\prime},a^{\prime})\in\tau}-\log\pi_{\theta}(a^{\prime}|s^{\prime},i)$ $\displaystyle=\sum_{i=1}^{K}\mathcal{L}_{\mathrm{BC}}^{(i)}(\theta)$ (21) $$

Note that by the definition, $\mathcal{L}_{\mathrm{BC}}^{(i)}(\theta)$ 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 0, the BC-PMI objection function can be viewed as separating the $(s,a)$ by the style $z$ and optimizing it respectively. This is equivalent to the clustering-based behavior cloning objective, where each trajectory is first assigned to a specific style cluster based on its style label, and then behavior cloning is 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.
∎

###### Proof.

###### Lemma 1 (Gibbs’ inequality) .

## Appendix B Implementation Details

### B.1 Computational Resource

All experiments in this paper are implemented with PyTorch and executed on NVIDIA Tesla T4 GPUs.
All the runs in experiments use 5 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.

Figure: Figure 6: Architecture of the MINE networks used in the Atari environment.
Refer to caption: x5.png

### 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 agent’s preferences for moving toward each area.
The latter divides the agent’s fire rate into three levels: (1) [0, 0.1); (2) [0.1, 0.3) and (3) [0.3, 1].

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1. Which behavior cloning (BC) algorithm is selected as the baseline method? Does the proposed PMI general well across different BC algorithms?
2. From Table 1, the improvement of the proposed BC-PMI seems marginal compared with CBC.
soundness: 3
strengths: 1. The topic studied in this paper is quite interesting, with an objective to improve policy diversity.
2. This article clearly defines the problem and provides a detailed exposition of the method.
3. 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.
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.
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.
2. 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.

---

confidence: 2
contribution: 3
initial_rating: 6
presentation: 3
questions: Please see my comments in previous sectors.
soundness: 2
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.
2. The proposed method achieves good experimental results on the datasets of CIRCLE2D , Atari and Basketball Player dataset.
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.
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.

2. 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. 

Reference:

[1] “**An Empirical Study on Disentanglement of Negative-free Contrastive Learning**”, NeurIPS 2022.

[2] “Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations.”, ICML 2019

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.
2. What are the calibration metrics used in Tables 3/4? Is it DTW, or ED, or KL?
soundness: 3
strengths: 1. The paper is well written and clearly communicates its key ideas.
2. 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.
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. 

The 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.

The 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.
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.

2. 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.

---

confidence: 4
contribution: 4
initial_rating: 8
presentation: 3
questions: Please refer to the weakness.
soundness: 3
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.
2. The paper is well-written and easy to follow.
3. The experimental results are extensive and convincing.
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.
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?
2. 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?

[1] Suneel Belkhale, Yuchen Cui, Dorsa Sadigh. Data quality in imitation learning. NeurIPS 2023

Meta-review: Summary: This paper proposes Behavioral Cloning with Pointwise Mutual Information Weighting (BC-PMI), a method that enhances policy diversity by weighting state-action pairs based on their relevance to trajectory styles using Pointwise Mutual Information (PMI). Experimental results across three settings, Circle 2D, Atari Games, and the professional basketball datasets, show that BC-PMI outperforms baselines in terms of style calibration, while also providing insights into the smooth interpolation between traditional behavioral cloning and clustering-based behavior cloning.


Strengths and Weaknesses: 

Generally, the reviewers find the topic compelling, the problem clearly defined, and the method well-explained, with an intuitive motivation that addresses the shortcomings of existing alternatives and strong results across multiple datasets, making the paper a significant contribution to improving the efficiency of imitation learning.

While the paper presents a promising approach, several weaknesses were identified by the reviewers. The original manuscript lacks several key ablation experiments and comparisons with other reweighting methods or style-conditioned BC, as well as evaluations on state and action diversity, which are necessary to fully establish the superiority of the proposed method. There were also concerns regarding the fairness of the comparisons due to differing types of supervision. Additionally, the experimental section is brief and does not analyze the method's performance as the number of styles increases or in the presence of noisy labels, which is particularly important given the reliance on potentially imperfect labels. 

During the discussion phase, the authors' responses, along with the additional experimental results, effectively address the reviewers' concerns.

All reviewers recommend acceptance, and ACs support their recommendation. The authors should address the main points raised in the reviews, particularly the additional analyses and experiments, such as validating the effectiveness and applicability of PMI on a broader range of advanced BC policies and conducting more comprehensive comparisons with other data weighting methods, as promised in the rebuttal, when preparing the camera-ready version.

Decision: Accept (Poster)