SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Huiqiao Fu $^{1}$ , Kaiqiang Tang $^{1}$ , Yuanyang Lu $^{1}$ , Yiming Qi $^{1}$ , Guizhou Deng $^{1}$ , Flood Sung $^{2}$ , Chunlin Chen $^{1*}$
$^{1}$ Nanjing University, China,  $^{2}$ Moonshot AI, China
{hqfu, kqtang, yylu, ymqi}  $@$  smail.nju.edu.cn,
{guizhoudeng, floodsung}  $@$  gmail.com, clchen@nju.edu.cn
Imitation learning aims to reproduce expert behaviors without relying on an explicit reward signal. However, real-world demonstrations often present challenges, such as multi-modal, data imbalance, and expensive labeling processes. In this work, we propose a novel semi-supervised imitation learning architecture that learns disentangled behavior representations from imbalanced demonstrations using limited labeled data. Specifically, our method consists of three key components. First, we adapt the concept of semi-supervised generative adversarial networks to the imitation learning context. Second, we employ a learnable latent distribution to align the generated and expert data distributions. Finally, we utilize a regularized information maximization approach in conjunction with an approximate label prior to further improve the semi-supervised learning performance. Experimental results demonstrate the efficiency of our method in learning multi-modal behaviors from imbalanced demonstrations compared to baseline methods.
The recent success of Reinforcement Learning (RL) has shown great potential in solving complex decision-making tasks [1, 2, 3]. A key prerequisite for RL is to design a reward function, which can be challenging in many real-world applications [4, 5]. By contrast, Imitation Learning (IL) methods, such as Behavior Cloning (BC) [6], Inverse Reinforcement Learning (IRL) [7] and Generative Adversarial Imitation Learning (GAIL) [8], can reproduce expert behaviors without relying on an explicit reward signal. Among them, GAIL performs a more effective learning process by learning a policy directly from demonstrations, and has been widely applied in various fields [9, 10, 11, 12].
Real-world demonstrations are typically coupled by multi-modal behaviors. To disentangle interpretable and meaningful behavior representations, InfoGAIL [13] utilizes the concept of maximizing mutual information between discrete or continuous latent variables to generate corresponding samples in an unsupervised manner. However, unsupervised learning of desired decodable factors from unstructured data can be challenging, particularly in the case of behavior imitation, where commonly used methods such as contrastive learning often underperform. Recent researches have uncovered that unsupervised learning of disentangled representations is inherently difficult without the presence of inductive biases on both the models and data [14]. Additionally, real-world demonstrations often consist of imbalanced behavior modes. For example, walking constitutes the majority of human behavior, whereas jumping only represents a small fraction. It is difficult for InfoGAIL to learn disentangled behavior representations from imbalanced demonstrations, as shown in Fig. 1. One intuitive solution is to introduce supervised learning with fully labeled data [15]. However, manually segmenting and labeling from the raw demonstrations can be a costly endeavor.
(a) Expert
(b) GAIL
Figure 1: 2D trajectory environment. An agent can move freely within a plane, seeking to mimic 4 expert trajectories, each represented by a different color to indicate a distinct mode of expert behavior. (a) The imbalanced expert demonstrations, where the red and green trajectories are predominant, while the blue and purple trajectories only represent a small portion. (b) Trajectories generated by GAIL. (c) Trajectories generated by InfoGAIL. (d) Trajectories generated by the proposed EssInfoGAIL. Our method is able to learn high quality disentangled behavior modes from the imbalanced expert demonstrations.
(c) InfoGAIL
(d) Ess-InfoGAIL
To tackle these challenges, we introduce a novel semi-supervised imitation learning architecture that can acquire disentangled behavior representations from imbalanced demonstrations, guided by only a limited amount of labeled data (as little as  $0.5\%$  to a maximum of  $2\%$  of the dataset). Our model, Elastic semi-supervised InfoGAIL (Ess-InfoGAIL), makes three key improvements to the related InfoGAIL. First, we draw inspiration from ss-InfoGAN [16] and focus on learning meaningful and controllable behavior representations by maximizing two mutual information terms: i) The mutual information between a latent skill variable, drawn from a categorical distribution, and the labeled expert state-action transitions, ii) and the mutual information between a latent shifting variable, drawn from a continuous uniform distribution, and the state-action transitions produced by a policy. The former term enforces that the latent skill variable corresponds to label categories in a semi-supervised manner, while the latter allows for style shifting within a given skill. Second, we tackle the issue of imbalanced data and align the distribution of generated transitions with that of expert demonstrations by utilizing a latent skill variable drawn from a differentiable Gumbel-Softmax distribution [17], instead of using a latent skill variable drawn from a fixed categorical distribution. Third, to make use of the intrinsic information in the unlabeled imbalanced data and improve the efficiency of the semi-supervised learning, we leverage the Regularized Information Maximization (RIM) [18] with a label prior that approximates the learned latent skill distribution.
In the experiment, we first validated our method in a simple 2D trajectory environment to obtain a more intuitive visualization, as shown in Fig. 1, and subsequently tested our method in 4 challenging MuJoCo environments (Reacher, Pusher, Walker-2D, and Humanoid). The experimental results illustrate the efficiency of our method in disentangling multi-modal behaviors from imbalanced demonstrations compared to baseline methods. To the best of our knowledge, our work is the first to tackle the problem of learning disentangled behavior representations from imbalanced demonstrations, with the aid of only a limited amount of labeled data. The code is available at https://github.com/tRNAoO/Ess-InfoGAIL. In particular, the contributions of this work are threefold:
1. We extend the InfoGAIL with a semi-supervised imitation learning architecture, resulting in Ess-InfoGAIL, which enables the learning of desired disentangled behavior representations.
2. We utilize a learnable latent distribution to align the generated and expert distributions in the imbalanced data setting, and leverage RIM with an approximate label prior to improve the semi-supervised learning.
3. We demonstrate the efficiency of our method in 4 challenging MuJoCo environments, showcasing that even in scenarios with highly imbalanced demonstrations, the policy can still reproduce the desired behavior modes.
Generative adversarial imitation learning GAIL adapts techniques developed for Generative Adversarial Networks (GANs) [19], and has demonstrated promising results in challenging imitation
learning tasks, such as robot control [10], imitation from raw visual data [20], multi-agent gaming [9], text generation [21], and driver behavior imitation [22]. Subsequent research has proposed various improvements to GAIL [23, 24], resulting in more efficient and reliable imitation results. However, discovering rich disentangled factors from an unstructured dataset using GAIL can be challenging due to the absence of inductive biases on the data or model. One promising approach to alleviate this issue is to introduce a latent variable model and use the technique of maximizing the mutual information [25] between the latent variable and the data, which can enable effective disentanglement of behaviors.
Learning multi-modal behaviors Tasks can often be accomplished in multiple ways, and the expert may exhibit different modes of behavior [26, 27]. Therefore, it is crucial to extend GAIL to learn and replicate these multi-modal behaviors. One intuitive approach is to train separate discriminators for each labeled mode [28]. Or one can use an auxiliary classifier allowing label-conditional imitation learning about multiple intentions [15]. These supervised approaches have demonstrated their effectiveness and reliability in various applications, from controlling wheeled robots [28] to task allocation in taxi companies [29]. However, they are constrained by their reliance on labeled data. To overcome this limitation, recent research has focused on unsupervised methods that typically leverage mutual information between latent variables and the generated data [30, 12]. For example, InfoGAIL maximizes this mutual information to learn from visual demonstrations [13], while Burn-InfoGAIL maximizes it from the perspective of Bayesian inference to draw modal variables from burn-in demonstrations [31]. However, unsupervised methods often struggle to distinguish latent factors without considering semantic information or task context [32], and proved unreliable in theory [14]. Therefore, we propose a novel semi-supervised approach that leverages the advantages of both supervised and unsupervised methods to learn semantically interpretable behavior representations with a small portion of labeled data.
Learning from imbalanced data Real-world data often exhibits imbalanced distributions, which can impede learning by causing the model to become biased towards the dominant category and potentially overlook the minority category. In a supervised learning setting, the problem of imbalanced data can be mitigated through various techniques such as data re-sampling and category re-weighting [33, 34]. These methods are useful when the category distributions are known beforehand. Several methods have been proposed for the unsupervised setting to deal with the imbalanced data with unknown category distributions [35, 36, 37]. A recent approach called Elastic-InfoGAN [38] has been proposed to address the issue of distribution alignment in GANs by utilizing a learnable latent distribution to handle imbalanced data. Unlike these methods, our work focuses on multi-modal semi-supervised imitation learning from raw imbalanced data without relying on prior knowledge of category distributions.
In this part, we first provide a brief background on RL, GAIL, and InfoGAIL, then provide a detailed description of how to extend InfoGAIL to a semi-supervised architecture and address the disentangled behavior learning using both the imbalanced unlabeled expert demonstrations and limited labeled expert demonstrations.
Reinforcement learning We consider a standard Markov Decision Process (MDP) setting [1] represented by a tuple  $\langle S, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma \rangle$ , where  $S$  is the state space,  $\mathcal{A}$  is the action space,  $\mathcal{P}(\mathbf{s}_{t+1} | \mathbf{s}_t, \mathbf{a}_t)$  is the transition probability of state  $\mathbf{s}_{t+1}$  at time step  $t+1$  given state  $\mathbf{s}_t$  and action  $\mathbf{a}_t$  at time step  $t$ ,  $\mathcal{R}(\mathbf{s}, \mathbf{a})$  is the reward function, and  $\gamma \in (0, 1)$  is the discount factor. In an RL process, an agent interacts with the environment and learns a policy  $\pi(\mathbf{a} | \mathbf{s})$  to maximize the expected discounted return  $\mathbb{E}_{p(\tau | \pi)} \left[ \sum_{t=0}^{T-1} \gamma^t \mathcal{R}(\mathbf{s}_t, \mathbf{a}_t) \right]$ , where  $p(\tau | \pi) = p(\mathbf{s}_0) \prod_{t=0}^{T-1} p(\mathbf{s}_{t+1} | \mathbf{s}_t, \mathbf{a}_t) \pi(\mathbf{a}_t | \mathbf{s}_t)$  represents the likelihood of a trajectory  $\tau = \{\mathbf{s}_0, \mathbf{a}_0, r_0, \mathbf{s}_1, \dots, \mathbf{s}_{T-1}, \mathbf{a}_{T-1}, r_{T-1}, \mathbf{s}_T\}$  under  $\pi$ . The reward function serves as an interface through which users can specify the task that an agent needs to perform. However, creating effective reward functions that elicit the desired behaviors from an agent can often be a laborious and challenging process.
GAIL and InfoGAIL As one of the apprenticeship learning method, GAIL adapts techniques developed for GANs to directly learn a stylized policy through expert demonstrations without the need for handcrafted rewards. The formal GAIL objective is denoted as
$$
\min  _ {\pi} \max  _ {D} V _ {\text {G A I L}} (\pi , D) = \mathbb {E} _ {d ^ {\pi} (\mathbf {s}, \mathbf {a})} [ \log D (\mathbf {s}, \mathbf {a}) ] + \mathbb {E} _ {d ^ {E} (\mathbf {s}, \mathbf {a})} [ \log (1 - D (\mathbf {s}, \mathbf {a})) ] - H (\pi), \tag {1}
$$
where  $D$  is a discriminator which tries to distinguish state-action pairs  $(\mathbf{s},\mathbf{a})$  drawn from the generated state-action distribution  $d^{\pi}(\mathbf{s},\mathbf{a})$ , induced by  $\pi$ , and the state-action distribution of the expert demonstration  $d^{E}(\mathbf{s},\mathbf{a})$ .  $H(\pi)$  is the causal entropy of the policy  $\pi$ . The GAIL objective is optimized alternately by taking a gradient step to increase Equation 1 with respect to  $D$  and an RL step to decrease Equation 1 with respect to  $\pi$ . The Jensen-Shannon divergence between  $d^{\pi}(\mathbf{s},\mathbf{a})$  and  $d^{E}(\mathbf{s},\mathbf{a})$  is eventually minimized [39].
GAIL may encounter difficulties in learning multi-modal behaviors from unstructured demonstrations in the absence of meaningful representations. To overcome this limitation, InfoGAIL introduces an additional regularization term to maximize the mutual information  $I(\mathbf{c};\mathbf{s},\mathbf{a})$  between the latent variable  $\mathbf{c}$  and state-action transitions  $(\mathbf{s},\mathbf{a})$ , and the objective is then approximated by a variational lower bound [25]:
$$
L _ {\mathrm {I}} (\pi , Q) = \mathbb {E} _ {p (\mathbf {c}), d ^ {\pi} (\mathbf {s}, \mathbf {a} | \mathbf {c})} [ \log Q (\mathbf {c} | \mathbf {s}, \mathbf {a}) ] + H (\mathbf {c}) \leq I (\mathbf {c}; \mathbf {s}, \mathbf {a}), \tag {2}
$$
where  $Q(\mathbf{c}|\mathbf{s},\mathbf{a})$  is an approximation of the true posterior  $P(\mathbf{c}|\mathbf{s},\mathbf{a})$ . The lower bound is tight if  $Q = P$ . This regularization allows InfoGAIL infer the interpretable latent structure in an unsupervised manner, and the final objective becomes:
$$
\min  _ {\pi , Q} \max  _ {D} V _ {\text {I n f o G A I L}} (\pi , D, Q) = V _ {\text {G A I L}} (\pi , D) - \lambda_ {1} L _ {\mathrm {I}} (\pi , Q), \tag {3}
$$
where  $\lambda_{1}$  is a weighting coefficient.
3.2 Semi-supervised Imitation Learning from Imbalanced Demonstrations
Figure 2: Schematic overview of the Ess-InfoGAIL network architecture. Here,  $\mathbf{c}$  represents a latent skill variable sampled from the differentiable Gumbel-Softmax distribution,  $\epsilon$  denotes a latent shifting variable sampled from a continuous uniform distribution, and  $y$  signifies the category label.  $D_{\mathrm{body}}$  serves as a shared backbone network for extracting features from state-action pairs. Parameters  $\theta$ ,  $\dot{\phi}$ , and  $\psi$  correspond to the policy  $\pi_{\theta}$ , the discriminator  $D_{\phi}$ , and the encoders  $Q_{\psi}^{\epsilon}$ ,  $Q_{\psi}^{\mathbf{c}}$ , respectively.
As InfoGAIL is learned in an unsupervised manner, obtaining the desired disentangled behavior representations can prove to be challenging. The difficulty is compounded when the expert demonstrations are imbalanced, as shown in Fig. 1. To address this, we propose Elastic semi-supervised InfoGAIL (Ess-InfoGAIL) with three improvements to InfoGAIL: i) A semi-supervised learning architecture, ii) a learnable latent skill distribution and iii) RIM with an approximate label prior. The network architecture of Ess-InfoGAIL is shown in Fig. 2.
Semi-supervised InfoGAIL To utilize both labeled and unlabeled data, we decompose the latent variable into a semi-supervised part  $\mathbf{c}$  and an unsupervised part  $\epsilon$ , inspired by ss-InfoGAN [16]. Unlike ss-InfoGAN, here we focus on imitation learning under sequential decision-making tasks. Additionally, we use  $\mathbf{c}$  as a latent skill variable which is sampled from a categorical distribution,
allowing for encoding the same information as the label  $y$ , and use  $\epsilon$  as a latent shifting variable which is sampled from a continuous uniform distribution, allowing for style shifting within a given skill.
Then we seek to maximize two mutual information terms  $I(\epsilon; \mathbf{s}, \mathbf{a})$  and  $I(\mathbf{c}; \mathbf{s}, \mathbf{a})$ , where the first term can be approximated using Equation 2 with a posterior approximation  $Q_{\psi}^{\epsilon}(\epsilon | \mathbf{s}, \mathbf{a})$ , and the second term can be approximated using the following variational lower bounds:
$$
L _ {\mathrm {I S}} ^ {1} \left(Q _ {\psi} ^ {1}\right) = \mathbb {E} _ {p (y), d ^ {E L} (\mathbf {s}, \mathbf {a} | y)} \left[ \log Q _ {\psi} ^ {1} (y | \mathbf {s}, \mathbf {a}) \right] + H (y) \leq I (y; \mathbf {s}, \mathbf {a}), \tag {4}
$$
$$
L _ {\mathrm {I S}} ^ {2} (\pi , Q _ {\psi} ^ {2}) = \mathbb {E} _ {p (\mathbf {c}), d ^ {\pi} (\mathbf {s}, \mathbf {a} | \mathbf {c})} [ \log Q _ {\psi} ^ {2} (\mathbf {c} | \mathbf {s}, \mathbf {a}) ] + H (\mathbf {c}) \leq I (\mathbf {c}; \mathbf {s}, \mathbf {a}), \tag {5}
$$
where, Equation 4 is a supervised term that uses the state-action distribution of the labeled expert demonstration  $d^{EL}(\mathbf{s},\mathbf{a},y)$  to predict  $y$ , and Equation 5 is an unsupervised term that utilizes the state-action distribution produced by a policy  $\pi_{\theta}$  to learn the inherent semantic meaning of  $y$ . The technique of lower bounding mutual information is known as variational information maximization [25]. With  $Q_{\psi}^{1} = Q_{\psi}^{2} = Q_{\psi}^{\mathbf{c}}$  we obtain the semi-supervised regularization term:
$$
L _ {\mathrm {I S}} \left(\pi_ {\theta}, Q _ {\psi} ^ {\mathbf {c}}\right) = L _ {\mathrm {I S}} ^ {1} \left(Q _ {\psi} ^ {\mathbf {c}}\right) + L _ {\mathrm {I S}} ^ {2} \left(\pi_ {\theta}, Q _ {\psi} ^ {\mathbf {c}}\right), \tag {6}
$$
where,  $L_{\mathrm{IS}}^{1}$  is optimized with respect to  $Q_{\psi}^{\mathbf{c}}$ , and  $L_{\mathrm{IS}}^{2}$  is optimized with respect to  $\pi_{\theta}$ . Specifically,  $L_{\mathrm{IS}}$  is trained as a classification process with the cross-entropy function, and  $L_{\mathrm{I}}$  in Equation 2 is trained as a regression process with the mean absolute error.
Learnable latent skill distribution The majority of GAN and GAIL algorithms assume that the latent variable is drawn from a uniform distribution with fixed parameters. However, in the presence of imbalanced data, this assumption may lead to suboptimal optimization results. To align the state-action distribution produced by a policy  $\pi_{\theta}$  with the state-action distribution of the imbalanced expert demonstration, we utilize a latent skill variable  $\mathbf{c}$  drawn from a differentiable Gumbel-Softmax distribution  $p(\mathbf{c})$  [17], similar to Elastic-InfoGAN [38]. The Gumbel-Softmax distribution is a continuous distribution over the simplex that can approximate samples from a categorical distribution in a differentiable way, and this trick has been widely used in many areas [40, 41].
Let  $p_1, p_2, \dots, p_K$  be the class probabilities of a categorical distribution, the sampling of a K-dimensional latent skill variable  $\mathbf{c}$  can be done as
$$
c _ {i} = \frac {\exp \left(\left(\log \left(p _ {i}\right) + g _ {i}\right) / \tau\right)}{\sum_ {j = 1} ^ {K} \exp \left(\left(\log \left(p _ {j}\right) + g _ {j}\right) / \tau\right)} \quad \text {f o r} i = 1, \dots , K, \tag {7}
$$
where  $g_{1} \cdots g_{K}$  are i.i.d samples drawn from Gumbel(0, 1). The coefficient  $\tau$  is a temperature parameter that determines the degree of smoothing applied to the probability distribution. A smaller value of  $\tau$  leads to a distribution that is closer to the one-hot encoding, while a larger value of  $\tau$  results in a distribution that is closer to the categorical distribution.
Since the objective of GAIL is to minimize the Jensen-Shannon divergence between  $d^{\pi}(\mathbf{s},\mathbf{a})$  and  $d^{E}(\mathbf{s},\mathbf{a})$ , we can eventually estimate the label distribution of the unlabeled expert demonstrations by passing the gradient back to the latent skill distribution  $p(\mathbf{c})$  using Gumbel-Softmax reparameterization trick.
RIM with an approximate label prior Discriminative clustering techniques can automatically identify the boundaries or distinctions between categories in unlabeled data [42]. To leverage the intrinsic information in the unlabeled imbalanced demonstrations and enhance the efficiency of the semi-supervised learning process in Equation 6, we employ Regularized Information Maximization (RIM) [18] with a label prior that approximates the learned latent skill distribution:
$$
L _ {\mathrm {R I M}} \left(Q _ {\psi} ^ {\mathbf {c}}, p (\mathbf {c})\right) = \underbrace {- \frac {1}{N} \sum_ {i} H \left(Q _ {\psi} ^ {\mathbf {c}} (\hat {y} | \mathbf {s} _ {i} , \mathbf {a} _ {i})\right)} _ {1} \underbrace {- D _ {K L} \left(Q _ {\psi} ^ {\mathbf {c}} (\hat {y}) \| p (\mathbf {c})\right)} _ {2} \underbrace {- R (\psi)} _ {3}, \tag {8}
$$
where the first term is the cluster assumption [43], corresponding to the idea that datapoints should be classified with large margin. The second term is performed to avoid degenerate solutions by minimizing the KL divergence between the empirical label distribution  $Q_{\psi}^{\mathbf{c}}(\hat{y}) \approx \frac{1}{N}\sum_{i}Q_{\psi}^{\mathbf{c}}(\hat{y} | \mathbf{s}_i, \mathbf{a}_i)$  and the learned latent skill distribution  $p(\mathbf{c})$ , while the last term is a parameter regularization (e.g., L2 regularization) performed to avoid complex solutions.
3.3 Overall Algorithm
With the above three improvements, the objective function of Ess-InfoGAIL eventually becomes:
$$
\begin{array}{l} \min  _ {Q \in \operatorname {G a r} _ {p} (\varepsilon)} \max  _ {D} V _ {\text {E s s - I n f o G A I L}} \left(\pi_ {\theta}, D _ {\phi}, Q _ {\psi} ^ {\epsilon}, Q _ {\psi} ^ {\mathbf {c}}, p (\mathbf {c})\right) \\ \begin{array}{r l} \pi_ {\theta}, Q _ {\psi} ^ {\epsilon}, Q _ {\psi} ^ {-}, p (\mathbf {c}) & D _ {\phi} \\ & = V _ {\text {I n f o G A I L}} \left(\pi_ {\theta}, D _ {\phi}, Q _ {\psi} ^ {\epsilon}, p (\mathbf {c})\right) - \lambda_ {2} L _ {\mathrm {I S}} \left(\pi_ {\theta}, Q _ {\psi} ^ {\mathbf {c}}, p (\mathbf {c})\right) - \lambda_ {3} L _ {\mathrm {R I M}} \left(Q _ {\psi} ^ {\mathbf {c}}, p (\mathbf {c})\right), \end{array} \tag {9} \\ \end{array}
$$
where  $\lambda_{2}$  and  $\lambda_{3}$  are two weighting coefficients.  $\theta, \phi, \psi$  are parameters for the policy  $\pi_{\theta}$ , the discriminator  $D_{\phi}$ , and the encoders  $Q_{\psi}^{\epsilon}$  and  $Q_{\psi}^{\mathbf{c}}$ , respectively. We optimize  $D_{\phi}, Q_{\psi}^{\epsilon}$  and  $Q_{\psi}^{\mathbf{c}}$  with the stochastic gradient method.  $\pi_{\theta}$  is optimized using Proximal Policy Optimization (PPO) [44], with advantages computed using  $GAE(\lambda)$  [45], and the value function is updated using  $TD(\lambda)$  [1]. The latent skill variable  $\mathbf{c}$ , concatenated with the state, serves as an input to the policy, where  $\mathbf{c}$  is sampled from the learnable latent skill distribution  $p(\mathbf{c})$  which is updated along with the policy. The schematic overview of the Ess-InfoGAIL network architecture is shown in Fig. 2, and the pseudo code is shown in Algorithm 1. In the experiments, we show that the proposed Ess-InfoGAIL can learn disentangled behavior representations efficiently from imbalanced demonstrations using limited labeled data.