SUSD: Structured Unsupervised Skill Discovery through State Factorization

TL;DR

This work tackles unsupervised skill discovery in reinforcement learning, addressing the limitation that mutual-information based methods yield static, transformation-invariant skills and distance-based methods often miss coverage of harder-to-control factors. The authors propose Structured Unsupervised Skill Discovery (SUSD), which factorizes the environment's state into independent factors, assigns per-factor skill variables, and introduces a curiosity-driven factor weighting mechanism to balance learning across factors. SUSD optimizes a factorized DSD objective via dual gradient descent and trains a low-level policy with SAC while a high-level controller composes skills for downstream tasks, yielding a factorized latent space suitable for hierarchical RL. Experimental results across factorized and unfactorized environments show that SUSD discovers richer, more diverse skills and achieves superior factor coverage and downstream performance compared to prevailing USD/DSD baselines, with publicly available code.

Abstract

Unsupervised Skill Discovery (USD) aims to autonomously learn a diverse set of skills without relying on extrinsic rewards. One of the most common USD approaches is to maximize the Mutual Information (MI) between skill latent variables and states. However, MI-based methods tend to favor simple, static skills due to their invariance properties, limiting the discovery of dynamic, task-relevant behaviors. Distance-Maximizing Skill Discovery (DSD) promotes more dynamic skills by leveraging state-space distances, yet still fall short in encouraging comprehensive skill sets that engage all controllable factors or entities in the environment. In this work, we introduce SUSD, a novel framework that harnesses the compositional structure of environments by factorizing the state space into independent components (e.g., objects or controllable entities). SUSD allocates distinct skill variables to different factors, enabling more fine-grained control on the skill discovery process. A dynamic model also tracks learning across factors, adaptively steering the agent's focus toward underexplored factors. This structured approach not only promotes the discovery of richer and more diverse skills, but also yields a factorized skill representation that enables fine-grained and disentangled control over individual entities which facilitates efficient training of compositional downstream tasks via Hierarchical Reinforcement Learning (HRL). Our experimental results across three environments, with factors ranging from 1 to 10, demonstrate that our method can discover diverse and complex skills without supervision, significantly outperforming existing unsupervised skill discovery methods in factorized and complex environments. Code is publicly available at: https://github.com/hadi-hosseini/SUSD.

Figures (13)

• Figure 1: Illustration of the SUSD Method.(a) In the skill learning stage, factorized embedding $\phi$ of the current and next state is passed through the corresponding mapping function $\phi(.)$ to obtain a skill latent embedding. Additional details about factorized embedding are shown in (b), where factor $s^i$, $i\in \{1,...,N\}$, is mapped to its embedding through the function $\phi_i$. These embeddings, together with the skill factor inputs, are used to compute the intrinsic reward $r^i$ of factor $i$. In the curiosity-based factor weighting module, a density model takes the full current state as input and estimates the probability of each next-state factor given the current state $-\log{q(s^i_{t+1}|s_t})$. These probabilities are then used as weights to scale the factor-wise intrinsic rewards, which are summed to form the final intrinsic reward for training the skill policy. (c) In the task learning stage, the learned skill policy is frozen as a low-level policy, while a high-level policy $\pi_{\text{high}}$ is trained to select a skill $z$ every $L$ steps by maximizing the task reward $r^{\text{task}}$.
• Figure 2: Left: The state $s$ is passed to the density model, which estimates the mean and variance of $q(s'|s)$. These statistics are then partitioned by factors to obtain ${q(s'_i|s_t)}_{i=1}^{2}$. Right: Point $x$ shows high probability in factor 1 but low probability in factor 2—a distinction that cannot be leveraged by the CSD method, which assigns a single weight to the entire state transition rather than to individual state factors.
• Figure 3: Training curves of SUSD and baseline methods on multiple downstream tasks in the Multi-Particle and Kitchen environments. Each plot shows the mean and standard deviation of returns over 3 random seeds.
• Figure 4: Comparison of state/factor coverage across different factorized environments.
• Figure 5: Benchmark Environments

Theorems & Definitions (1)

• Lemma 4.1