Offline Discovery of Interpretable Skills from Multi-Task Trajectories

TL;DR

This work tackles offline, multi-task hierarchical imitation learning by learning reusable skills from unannotated demonstrations. It introduces LOKI, a three-stage framework that combines Stage 1 macro segmentation via an alignment-enforced VQ-VAE (EVQ-VAE), Stage 2 micro segmentation with a sequential VAE and iterative clustering, and Stage 3 a hierarchical policy built on canonical skill sequences with a learned termination function $\beta$ and a conditioned low-level policy. Key technical contributions include a task-conditioned EVQ-VAE with a codebook divergence term and a KL-regularized micro-segmentation objective using a prior $p(z|s_i,c)$, plus a canonical skill sequencing strategy that aligns across tasks. Experiments on the D4RL Kitchen benchmark demonstrate high success rates, interpretability and compositionality of the discovered skills, and superior robustness over baselines, indicating strong practical potential for scalable offline skill discovery in robotics.

Abstract

Hierarchical Imitation Learning is a powerful paradigm for acquiring complex robot behaviors from demonstrations. A central challenge, however, lies in discovering reusable skills from long-horizon, multi-task offline data, especially when the data lacks explicit rewards or subtask annotations. In this work, we introduce LOKI, a three-stage end-to-end learning framework designed for offline skill discovery and hierarchical imitation. The framework commences with a two-stage, weakly supervised skill discovery process: Stage one performs coarse, task-aware macro-segmentation by employing an alignment-enforced Vector Quantized VAE guided by weak task labels. Stage two then refines these segments at a micro-level using a self-supervised sequential model, followed by an iterative clustering process to consolidate skill boundaries. The third stage then leverages these precise boundaries to construct a hierarchical policy within an option-based framework-complete with a learned termination condition beta for explicit skill switching. LOKI achieves high success rates on the challenging D4RL Kitchen benchmark and outperforms standard HIL baselines. Furthermore, we demonstrate that the discovered skills are semantically meaningful, aligning with human intuition, and exhibit compositionality by successfully sequencing them to solve a novel, unseen task.

Figures (7)

• Figure 1: Two-Stage Skill Discovery: From long-horizon multi-task trajectories to reusable independent skills. While traditional pure skill discovery through prior knowledge venkatraman2023reasoningdiffusion or information theory eysenbach2018diversity on this problem often struggles, our approach first processes a trajectory into task-related extrinsic General Skills(GS) and task-unrelated intrinsic GS. Subsequently, these are decomposed into granular, human-knowledge-aligned Independent Skills through weak-supervision alignment. This methodology enables robust skill discovery and facilitates flexible skill recomposition for novel tasks.
• Figure 2: LOKI Framework Diagram. The left panel illustrates the two-stage skill discovery and segmentation process: (light green) Macro-segmentation using the Enforced VQ-VAE, followed by (light blue) Self-Supervised Micro-segmentation with Clustering Iterative Refinement, and Task Skill Alignment, yielding Refined Independent Segments (IS). The right panel depicts the third stage of the framework: constructing executable hierarchical policies from the discovered skills, which includes learning a Termination Function ($\beta$), Low-Level Policy, and deriving a canonical skill sequence policy for task execution. Skill smooth execution is shown as a visual output.
• Figure 3: Diagram of the EVQ-VAE architecture that leverages weak task labels $c_i$ to guide an Enforced Quantization from embeddings $z_e(x)$ to codebook vector $z_q(x;c)$. The mechanism achieves an inherent clustering of task-relevant states and actions while maintaining a diverse representation for multi-task behaviors, thereby enabling granular trajectory decomposition. $M$ is the dimension of codebook vector.
• Figure 4: An overview of the different subtasks in the D4RL Kitchen dataset fu2020d4rl. The complete trajectory dataset includes seven subtasks: (a) Turn on Bottom Burner; (b) Open Hinge Cabinet; (c) Move Kettle; (d) Turn on Slide Cabinet; (e) Turn on Light Switch ; (f) Turn on Top Burner; (g) Open Microwave, and each trajectory contains four subtasks.
• Figure 5: Change point detection results from our EVQ-VAE on three complex tasks. The y-axis represents the log entropy, calculated from the L2 distance between a transition's embedding and each codebook vector. Dotted lines of different colors (blue, red, green) correspond to the log probability of the embedding belonging to each of the three primary codebook vectors. A low-entropy state, where one codebook vector is dominant (one dotted line drops significantly), indicates an extrinsic (task-specific) skill. A high-entropy state, where the lines are tangled and no single codebook is dominant, represents an intrinsic (task-agnostic) skill. The colored shaded regions denote the ground-truth duration of each subtask, while the solid black vertical lines are the change points detected by our method.