Skill Disentanglement for Imitation Learning from Suboptimal Demonstrations

/in /by

Publication Date: 8/10/2023

Event: 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining (KDD 2023)

Reference: pp. 3513-3524, 2023

Authors: Tianxiang Zhao, Pennsylvania State University; Wenchao Yu, NEC Laboratories America, Inc.; Suhang Wang, Pennsylvania State University; Lu Wang, East China Normal University; Xiang Zhang, Pennsylvania State University; Yuncong Chen, NEC Laboratories America, Inc.; Yanchi Liu, NEC Laboratories America, Inc.; Wei Cheng, NEC Laboratories America, Inc.; Haifeng Chen, NEC Laboratories America, Inc.

Abstract: Imitation learning has achieved great success in many sequential decision-making tasks, in which a neural agent is learned by imitating collected human demonstrations. However, existing algorithms typically require a large number of high-quality demonstrations that are difficult and expensive to collect. Usually, a trade-off needs to be made between demonstration quality and quantity in practice. Targeting this problem, in this work we consider the imitation of sub-optimal demonstrations, with both a small clean demonstration set and a large noisy set. Some pioneering works have been proposed, but they suffer from many limitations, e.g., assuming a demonstration to be of the same optimality throughout time steps and failing to provide any interpretation w.r.t knowledge learned from the noisy set. Addressing these problems, we propose method by evaluating and imitating at the sub-demonstration level, encoding action primitives of varying quality into different skills. Concretely, SDIL consists of a high-level controller to discover skills and a skill-conditioned module to capture action-taking policies and is trained following a two-phase pipeline by first discovering skills with all demonstrations and then adapting the controller to only the clean set. A mutual-information-based regularization and a dynamic sub-demonstration optimality estimator are designed to promote disentanglement in the skill space. Extensive experiments are conducted over two gym environments and a real-world healthcare dataset to demonstrate the superiority of SDIL in learning from sub-optimal demonstrations and its improved interpretability by examining learned skills.

Publication Link: https://dl.acm.org/doi/10.1145/3580305.3599506