SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Algorithm 1: Ess-InfoGAIL
Input: Unlabeled and labeled demonstrations  $d^{E}(\mathbf{s},\mathbf{a})$ $d^{EL}(\mathbf{s},\mathbf{a},y)$  , initial parameters of the policy, value function, discriminator and encoders  $\theta_0,\beta_0,\phi_0,\psi_0$  , initial latent skill distribution  $p_0(\mathbf{c})$  , weighting coefficients  $\lambda_{1},\lambda_{2},\lambda_{3}$  Output: Learned policy  $\pi_{\theta^*}$    
1 for  $i = 0,1,2,\dots$  do   
2 Sample a batch of latent variables  $\epsilon \sim p(\epsilon),\mathbf{c}\sim p_i(\mathbf{c})$    
3 Interact with the environment using  $\pi_{\theta_i}$ $\epsilon$  and c, and obtain  $d^{\pi}(\mathbf{s},\mathbf{a})$    
4 Sample a batch of state-action pairs  $b^{\pi}\sim d^{\pi}$ $b^{E}\sim d^{E}$ $b^{EL}\sim d^{EL}$    
5 Update  $\phi_{i}$  to  $\phi_{i + 1}$  by ascending with gradients according to VGAIL   
6 Update  $\psi_{i}$  to  $\psi_{i + 1}$  by descending with gradients according to VEss-InfoGAIL   
7 Update  $\theta_{i}$ $p_i(\mathbf{c})$  to  $\theta_{i + 1}$ $p_{i + 1}(\mathbf{c})$  by taking a policy step using the PPO update rule according to VEss-InfoGAIL   
8 Update the value function using  $TD(\lambda)$    
9 end
In this section, we validate our method by conducting experiments in a variety of environments, including a simple 2D trajectory environment, as well as 4 challenging MuJoCo environments (Reacher, Pusher, Walker-2D, and Humanoid). First, we perform the quantitative analysis to demonstrate the superiority of our method in discovering disentangled behavior representations from imbalanced demonstrations with limited labeled data, as compared to baseline methods. Second, we analyze the amount of labeled data required for the model to effectively encode the semantic meaning of the labels. Finally, we analyze the effect of varying the degree of data imbalance and the number of behavior modes on the imitation of multi-modal behaviors.
Task setup (1) 2D trajectory environment (Fig. 1) is a plane where the agent can move freely at a constant velocity by selecting its direction  $\mathbf{p}_t$  at discrete time  $t$ . The agent's observations at time  $t$  consist of positions from  $t - 4$  to  $t$ . The expert demonstrations consist of 4 distinct modes, each produced by a stochastic expert policy that generates a circular trajectory. The setting is similar to the experiment in [13]. (2) Reacher environment (Fig. 3 Left) consists a 2-DoF robot arm. The goal is to move the robot's end effector close to a target position. The original environment is modified by setting 6 evenly distributed targets on a circle with a radius of 1.5, and the target position is removed from the observation space. The initial angles of the two joints are randomly sampled between  $[- \pi, \pi]$  and  $[- \pi/2, \pi/2]$ , respectively. (3) Pusher environment (Fig. 3 Right) consists a 4-DoF robot arm. The goal is to move a target cylinder to a goal position using the robot's end effector. Similar to Reacher, we use 6 evenly distributed goal positions and remove the goal position from the observation space. (4) Walker-2D environment (Fig. 4 Left) is a 6-DoF bipedal robot consisting of two legs
Figure 3: Left: Reacher with 6 targets: random initial state, reaching one target, reaching another target. Right: Pusher with 6 goal positions: random initial state, reaching one goal position, reaching another goal position.
Figure 4: Left: Walker-2D with 3 goals: running forwards, running backwards, balancing. Right: Humanoid with 3 goals: running forwards, running backwards, balancing.
and feet attached to a common base. The goal is to make coordinate both sets of legs and feet to move the robot in the right direction, including running forward, running backward and balancing. (5) Humanoid environment (Fig. 4 Right) is a high-dimensional robot with 17 degrees of freedom. The goal is similar to the Walker-2D, including running forward, running backward and balancing.
Multi-modal demonstrations preparation In the experiments, we first pre-train K expert policies, each corresponding to K different goals (or K behavior modes). Subsequently, we use these K expert policies to sample K sets of expert demonstrations. From each set of expert demonstrations, we extract a small portion and label them with the one-hot behavioral categories, while the remaining expert demonstrations are randomly sampled and mixed to create imbalanced unlabeled expert data. Moreover, in the real-world scenario, one can directly use the raw motion capture data (e.g., motion capture data of an animal over a day) without the need to train additional expert policies.
Baselines and evaluation metrics We design various baselines to demonstrate the efficiency of our method in learning multi-modal behaviors from imbalanced demonstrations. (1) GAIL: The original GAIL. (2) InfoGAIL: The original InfoGAIL with a fixed uniform categorical distribution. (3) ACGAIL: Reproduction of ACGAIL [15] using an auxiliary classifier with full labeled data. (4) Elastic-InfoGAIL: Modification of Elastic-InfoGAN [38] in imitation tasks with a learnable latent distribution and a contrastive loss. (5) Ess-InfoGAIL $\backslash$ GS: Ess-InfoGAIL without using the Gumbel-Softmax technique. (6) Ess-InfoGAIL\RIM: Ess-InfoGAIL without using the RIM technique. (7) Ess-InfoGAIL: Our final method.
To evaluate the behavior disentanglement quality of the above methods, we use a pre-trained behavior classifier and employ two commonly used metrics in clustering algorithms similar to [38]: (a) Average Entropy (ENT) evaluates two properties: i) whether the state-action pairs generated for a given latent variable belong to the same ground-truth class and ii) whether each ground-truth class is associated with a single unique latent variable. (b) Normalized Mutual Information (NMI) quantifies the correlation between two clusterings, with a value ranging from 0 to 1. A higher NMI value indicates a stronger correlation between the clusterings.
Additionally, we use the normalized average task reward to assess the quality of learned policies. Note that, we modify the original MuJoCo environment and compute the task reward differently for each behavior mode. The final task reward, derived by averaging across all modes, serves exclusively as an evaluation metric and does not affect the policy training process.
4.1 Comparison with Baseline Methods
We quantitatively analyze Ess-InfoGAIL's behavior disentanglement performance, as presented in Table 1. In complex tasks with numerous behavior modes, both GAIL and InfoGAIL often encounter mode collapse, where their policies manifest behavior corresponding to only a subset of the behavior modes due to a lack of prior behavior representation knowledge. For instance, in the Reacher environment, the robot arm reaches only a few target positions. Leveraging limited labeling
Table 1: Behavior disentanglement quality measured by NMI (↑) and ENT (↓). Where Ess-InfoGAIL\GS removes Gumbel-Softmax in Ess-InfoGAIL, and Ess-InfoGAIL\RIM removes RIM in Ess-InfoGAIL. For each experiment, we use 10 different random seeds and collect 50 episodes for each seed. The results demonstrate the efficiency of Ess-InfoGAIL in learning multi-modal behaviors compared to baseline methods. (See supplementary for error bars.)
2D trajectory	Reacher	Pusher	Walker-2D	Humanoid
Metrics	NMI	ENT	NMI	ENT	NMI	ENT	NMI	ENT	NMI	ENT
GAIL [8]	0.392	0.529	0.153	1.055	0.376	0.723	0.439	0.406	0.388	0.597
InfoGAIL [13]	0.742	0.371	0.301	1.113	0.604	0.551	0.657	0.284	0.550	0.487
ACGAIL [15]	0.783	0.324	0.537	0.781	0.754	0.409	0.658	0.340	0.544	0.478
Elastic-InfoGAIL	0.773	0.330	0.311	1.101	0.650	0.537	0.615	0.351	0.503	0.498
Ess-InfoGAIL\GS	0.892	0.157	0.607	0.661	0.857	0.194	0.740	0.262	0.638	0.360
Ess-InfoGAIL\RIM	0.893	0.159	0.575	0.725	0.875	0.188	0.715	0.274	0.642	0.347
Ess-InfoGAIL (Ours)	0.910	0.131	0.662	0.587	0.906	0.144	0.755	0.237	0.696	0.206
guidance, our approach, Ess-InfoGAIL, improves the average NMI by 0.436 and 0.215 and decreases the average ENT by 0.401 and 0.300 compared to GAIL and InfoGAIL, respectively, across all environments. Notably, Ess-InfoGAIL attains superior behavior disentanglement quality with fewer labels than ACGAIL, which utilizes full labeled data. This success might stem from imbalanced data hindering ACGAIL's auxiliary classifier learning, while our method balances the small number of labels by extracting equal-duration data for each behavior mode.
Furthermore, we adapt Elastic-InfoGAN to the imitation learning framework using Gumbel-Softmax and contrastive learning loss, augmenting negative samples with Gaussian noise. However, contrastive learning's efficacy is limited in visual-free imitation tasks, and it indirectly addresses behavior disentanglement, yielding uncertain performance. In contrast, Ess-InfoGAIL directly tackles mode collapse and achieves semantically interpretable behavior disentanglement via limited labels. It exhibits an average NMI improvement of 0.215 and an average ENT decrease of 0.302 over Elastic-InfoGAIL across all environments. We perform ablation experiments, confirming the enhancements brought by Gumbel-Softmax and RIM techniques under the semi-supervised framework and imbalanced data. Ess-InfoGAIL demonstrates an average NMI increase of 0.039 and 0.046 and an average ENT reduction of 0.066 and 0.078 compared to Ess-InfoGAIL\GS and Ess-InfoGAIL\RIM, respectively.
To assess the quality of learned policies, we analyze the normalized average task rewards within the Reacher environment, depicted in Fig. 5, demonstrating the superiority of our method over all baseline methods. Results for different environments can be found in the supplementary material. Due to mode collapse and data imbalance, some methods (e.g., GAIL and InfoGAIL) may end up learning only one or two behavior modes, leading to lower average task rewards compared to a random policy.
4.2 Amount of Labeled Data Requirements
In real-world scenarios, obtaining labeled behavior data (e.g., from motion capture or raw video) is difficult due to the interleaving of various behavior modes. Thus, minimizing labeling costs is
Figure 5: Normalized average task reward of each method during training (only for evaluation).
Figure 6: Classification accuracy using different amount of labeled data.
crucial. We assess the prediction accuracy of encoders trained with limited labeled data across varying proportions within all unlabeled data. Results are shown in Fig. 6 (using Reacher environment as an example). Proportions of labeled data are set at  $0.1\%$ ,  $0.5\%$ ,  $1\%$ ,  $10\%$ , and  $100\%$ . For Reacher,  $0.1\%$  data equates to 2 episodes (100 time steps) per mode. The encoder demonstrates nearly  $92\%$  prediction accuracy with just  $0.5\%$  labeled data. Walker-2D and Humanoid tasks, with higher state dimensions, require more labeled data for effective behavior mode separation. Default label ratios, if not specified, are: 2D-Trajectory:  $1\%$ , Reacher:  $0.5\%$ , Pusher:  $1\%$ , Walker-2D:  $1\%$ , Humanoid:  $2\%$ . Guided by limited labeled data, the agent adeptly disentangles specific styles from labeled data and more varied behavioral styles from unlabeled data.
Table 2: Degree of data imbalance.
InfoGAIL	Ess-InfoGAIL
Metrics	NMI	ENT	NMI	ENT
20	0.328	1.035	0.704	0.510
40	0.317	1.050	0.693	0.521
60	0.313	1.061	0.687	0.583
80	0.306	1.073	0.681	0.561
100	0.301	1.113	0.662	0.587
200	0.291	1.121	0.625	0.607
Table 3: Learning more behavior modes.
InfoGAIL	Ess-InfoGAIL
Metrics	NMI	ENT	NMI	ENT
2	0.329	1.031	0.711	0.504
4	0.319	1.073	0.703	0.522
6	0.301	1.113	0.662	0.587
8	0.308	1.087	0.651	0.619
10	0.303	1.076	0.617	0.657
12	0.298	1.115	0.613	0.653
4.3 Degree of Data Imbalance
The degree of data imbalance directly affects the learning quality of the policy, discriminator, and encoder. We set up 6 levels of data imbalance in the Reacher environment, ranging from 20 to 200. For example, at a data imbalance level of 20, the majority class has at most 20 times the number of samples as the minority class. The proportions of labeled data is set to  $0.5\%$ . At the maximum data imbalance level of 200, the minority class has only about 2 episodes of trajectories for each mode. The experimental results are shown in Table 2. It can be observed that as the degree of data imbalance increases, the NMI gradually decreases while the ENT gradually increases. However, even when the degree of the data imbalance is 200, the NMI of Ess-InfoGAIL only decreases by 0.079 compared to the case of a data imbalance of 20, and the performance is still better than InfoGAIL with a data imbalance of 20. Our proposed method demonstrates the ability to achieve effective disentanglement of behaviors, even in the presence of highly imbalanced data distributions.
4.4 Learning More Behavior Modes
The disentanglement task becomes increasingly challenging with an increase in the number of behavior modes. We set up 6 levels of behavior modes ranging from 2 to 12 in the Reacher environment, where the agent needs to learn disentangled behavior representations from these unstructured behavior data. We maintain a fixed degree of data imbalance at 100 and the proportions of labeled data is set to  $0.5\%$ . The experimental results are presented in Table 3. The results demonstrate that even with 12 behavior modes, our method consistently achieves higher NMI values and lower ENT values, compared to InfoGAIL with only 2 behavior modes. This affirms the controllability and scalability of our method in learning disentangled behavior representations while preserving category semantic information.
In this work, we proposed a novel semi-supervised imitation learning method, Ess-InfoGAIL, which utilizes a learnable latent distribution and an improved RIM to enable the agent to learn better disentangled multi-modal behaviors from unstructured expert demonstrations. We qualitatively and quantitatively validated the efficiency of our method, compared with other baseline methods. We also demonstrated that our method only requires a very small amount of labeled data to achieve good results, and further verified its controllability and scalability in scenarios with higher imbalance degrees and more behavior modes, which are crucial for multi-modal imitation tasks in the real world. Although there are some limitations, such as the need to provide a small amount of labeled data for each category and preset the number of modal categories, we believe that our work lays the foundation for imitation learning in large-scale unstructured demonstrations.
This work was supported in part by the National Natural Science Foundation of China (No. 62073160).
[1] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[2] Takahiro Miki, Joonho Lee, Jemin Hwangbo, Lorenz Wellhausen, Vladlen Koltun, and Marco Hutter. Learning robust perceptive locomotion for quadrupedal robots in the wild. Science Robotics, 7(62):eabk2822, 2022.
[3] Huiqiao Fu, Kaiqiang Tang, Peng Li, Wenqi Zhang, Xinpeng Wang, Guizhou Deng, Tao Wang, and Chunlin Chen. Deep reinforcement learning for multi-contact motion planning of hexapod robots. In International Joint Conferences on Artificial Intelligence, pages 2381-2388, 2021.
[4] Andrew Y Ng, Daishi Harada, and Stuart Russell. Policy invariance under reward transformations: Theory and application to reward shaping. In International Conference on Machine Learning, volume 3, pages 278-287, 1999.
[5] Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. Deep reinforcement learning that matters. In AAAI Conference on Artificial Intelligence, page 3207-3214, 2018.
[6] Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural computation, 3(1):88-97, 1991.
[7] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, and Anind K Dey. Maximum entropy inverse reinforcement learning. In AAAI Conference on Artificial Intelligence, volume 8, pages 1433-1438, 2008.
[8] Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in Neural Information Processing Systems, volume 29, pages 4565-4573, 2016.
[9] Jiaming Song, Hongyu Ren, Dorsa Sadigh, and Stefano Ermon. Multi-agent generative adversarial imitation learning. In Advances in Neural Information Processing Systems, volume 31, pages 7461-7472, 2018.
[10] Xue Bin Peng, Ze Ma, Pieter Abbeel, Sergey Levine, and Angjoo Kanazawa. Amp: Adversarial motion priors for stylized physics-based character control. ACM Transactions on Graphics (TOG), 40(4):1-20, 2021.
[11] Alejandro Escontrela, Xue Bin Peng, Wenhao Yu, Tingnan Zhang, Atil Iscen, Ken Goldberg, and Pieter Abbeel. Adversarial motion priors make good substitutes for complex reward functions. In IEEE/RSJ International Conference on Intelligent Robots and Systems, pages 25–32, 2022.
[12] Xue Bin Peng, Yunrong Guo, Lina Halper, Sergey Levine, and Sanja Fidler. Ase: Large-scale reusable adversarial skill embeddings for physically simulated characters. ACM Transactions On Graphics (TOG), 41(4):1-17, 2022.
[13] Yunzhu Li, Jiaming Song, and Stefano Ermon. Infogail: Interpretable imitation learning from visual demonstrations. In Advances in Neural Information Processing Systems, volume 30, pages 3812-3822, 2017.
[14] Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In International Conference on Machine Learning, pages 4114-4124, 2019.
[15] Jiahao Lin and Zongzhang Zhang. Acgail: Imitation learning about multiple intentions with auxiliary classifier gans. In 15th Pacific Rim International Conference on Artificial Intelligence, pages 321-334, 2018.
[32] Cong Fei, Bin Wang, Yuzheng Zhuang, Zongzhang Zhang, Jianye Hao, Hongbo Zhang, Xuewu Ji, and Wulong Liu. Triple-gail: a multi-modal imitation learning framework with generative adversarial nets. In International Joint Conference on Artificial Intelligence, pages 2929-2935, 2020.
[33] Guo Haixiang, Li Yijing, Jennifer Shang, Gu Mingyun, Huang Yuanyue, and Gong Bing. Learning from class-imbalanced data: Review of methods and applications. Expert systems with applications, 73:220-239, 2017.
[34] Xinting Hu, Yi Jiang, Kaihua Tang, Jingyuan Chen, Chunyan Miao, and Hanwang Zhang. Learning to segment the tail. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 14045-14054, 2020.
[35] Aaron Reite, Scott Kangas, Zackery Steck, Steven Goley, Jonathan Von Stroh, and Steven Forsyth. Unsupervised feature learning in remote sensing. In Applications of Machine Learning, volume 11139, pages 132-144, 2019.
[36] Mengye Ren, Tyler R Scott, Michael L Iuzzolino, Michael C Mozer, and Richard Zemel. Online unsupervised learning of visual representations and categories. arXiv preprint arXiv:2109.05675, 2021.
[37] Tanmoy Dam, Sreenatha G Anavatti, and Hussein A Abbass. Improving clustergan using self-augmented information maximization of disentangling latentspaces. arXiv preprint arXiv:2107.12706, 2021.
[38] Utkarsh Ojha, Krishna Kumar Singh, Cho-Jui Hsieh, and Yong Jae Lee. Elastic-infogan: Unsupervised disentangled representation learning in class-imbalanced data. In Advances in Neural Information Processing Systems, volume 33, pages 18063-18075, 2020.
[39] Liyiming Ke, Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, and Siddhartha Srinivasa. Imitation learning as f-divergence minimization. In Algorithmic Foundations of Robotics XIV, pages 313-329, 2021.
[40] Matt J Kusner and José Miguel Hernández-Lobato. Gans for sequences of discrete elements with the gumbel-softmax distribution. arXiv preprint arXiv:1611.04051, 2016.
[41] Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint arXiv:1812.05905, 2018.
[42] Attri Ghosal, Arunima Nandy, Amit Kumar Das, Saptarsi Goswami, and Mrityunjoy Panday. A short review on different clustering techniques and their applications. In Emerging Technology in Modelling and Graphics, pages 69–83, 2020.
[43] Yves Grandvalet and Yoshua Bengio. Semi-supervised learning by entropy minimization. In Advances in Neural Information Processing Systems, volume 17, pages 529-536, 2005.
[44] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
[45] John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438, 2015.