SUSD: Structured Unsupervised Skill Discovery through State Factorization

AI summary

Overview

• Research area: Unsupervised Reinforcement Learning, specifically Unsupervised Skill Discovery (USD) and reward-free pretraining.
• Technical level: Advanced. The paper builds on Wasserstein dependency measures, factored Markov Decision Processes, dual gradient descent with Lagrange multipliers, and Soft Actor-Critic.
• Scope: The paper proposes SUSD, a skill-discovery method that factorizes the state space into controllable components and allocates separate skill variables to each, with a curiosity-based weighting mechanism that shifts learning emphasis toward harder-to-control factors.

What This Paper Is About

Unsupervised Skill Discovery lets an agent learn a repertoire of behaviors without any reward function, so those skills can later be reused for downstream tasks. Existing distance-maximizing methods (DSD) learn dynamic behaviors but tend to concentrate on the easiest-to-control parts of the environment, leaving whole entities or factors underrepresented. SUSD addresses this by treating the environment as a factored MDP, assigning a distinct sub-vector of the skill latent to each state factor, and adaptively up-weighting the factors the agent is currently failing to control.

Key Contributions

1. A factorized skill-learning formulation for DSD. The state space is decomposed into N factors, the mapping function φ is replaced by N separate networks φ_i(s^i), and each skill vector z is split into N factors of dimension D. The distance constraint in the original DSD objective is replaced by a sum-over-factors constraint.
2. A curiosity-driven factor weighting mechanism. A conditional density model q_θ(s′|s) = N(μ_θ(s), Σ_θ(s)) is fitted on (s, s′) tuples, and the Gaussian marginals are partitioned per factor to produce factor-wise weights sqrt(−log q_θ(s^i_{t+1}|s_t)). This is finer-grained than CSD, which assigns one weight to an entire state transition.
3. A reformulation of the DSD objective (Lemma 4.1) that moves the distance term from a constraint into the objective as a coefficient, adapted from prior work. The resulting intrinsic reward is the sum over factors of the weighted latent displacement along the skill direction.
4. Empirical evaluation against five baselines (LSD, CSD, METRA, DIAYN, DUSDi) on factorized and unfactorized environments, plus state-coverage, factor-decoding, and accidental-task-completion analyses.

Main Findings

• Downstream task performance: In the Multi-Particle and Kitchen environments, SUSD generally outperforms all compared methods "by a good margin," and a larger gap is observed in Kitchen, which is described as more complex with a high-dimensional observation space. Plots show mean and standard deviation over 3 random seeds (Figure 3).
• Unfactorized environments: On Ant and HalfCheetah, zero-shot goal-reaching is evaluated with 20K steps of reward accumulation across eight seeds. The paper states the results are shown in Figure 6 but does not report numerical outcomes in the text.
• State coverage across factors: With a randomly resampled skill every 200 steps and 20K rollout steps in Multi-Particle, SUSD achieves substantially better coverage than baselines, particularly DUSDi. Coordinates are rounded to two decimal places to count distinct states (e.g., (0.27392337, −0.46042657) to (0.27, −0.46)). The paper reports both Worst Agent State Coverage (minimum across the 10 agent factors) and Average Agent State Coverage, and states the same trend appears in 2D-Gunner.
• Factor decoding (Table 1, factor-wise MSE): Multi-Particle — SUSD 0.060, METRA 0.147, CSD 0.313, LSD 0.308. Kitchen — SUSD 0.014, METRA 0.028, CSD 0.049, LSD 0.038. 2D-Gunner — SUSD 0.080, METRA 0.186, CSD 0.404, LSD 0.224.
• Accidental task completion during skill learning (Table 2, Kitchen, 10K steps with skill switches every 50 steps, 8 independent runs): BiP — SUSD 39.875 ± 18.452, all four listed baselines 0.0 ± 0.0. MiP — SUSD 58.875 ± 25.784, DUSDi 2.5 ± 1.14, CSD/METRA/LSD 0.0 ± 0.0. PoS — SUSD 20.5 ± 17.965, DUSDi 1.275 ± 0.954, others 0.0 ± 0.0. PoT — SUSD 13.75 ± 6.923, others 0.0 ± 0.0.
• Disentangled control for downstream HRL: Because skills are factorized, the learned low-level policy can be frozen and a high-level policy can select a skill every L steps to maximize task reward.

Methodology in Plain English

The agent assumes the environment hands it a structured state vector that can be cut into meaningful pieces — one per object or controllable entity. Instead of one neural network mapping the whole state to a latent space, SUSD trains one small network per piece. Each piece of the skill vector is then responsible for moving its corresponding piece of the state, and the original DSD constraint (limiting how far the latent can move between adjacent states) is applied jointly across all pieces. Alongside this, a probabilistic model predicts the next state from the current one. When a transition is surprising under that model, it is treated as a sign that the corresponding factor is hard to control; the surprise score for each factor is square-rooted and used to scale that factor's intrinsic reward. The result is that factors the agent has already mastered contribute little to the learning signal while neglected factors receive more attention. Training follows a dual gradient descent scheme with a Lagrange multiplier λ, and the skill policy itself is learned with Soft Actor-Critic using the summed weighted reward as the intrinsic objective. On top of this, prior work's approach of folding the distance constraint into the objective as a coefficient is used, which requires a slack variable ε to keep the multiplier's gradient from always being nonnegative.

Why This Matters

Impact on research. The paper argues that factorization, previously explored for mutual-information-based skill discovery (DUSDi, SkiLD), had not been exploited within the DSD family, which dominates on dynamic behaviors. SUSD combines the two lines and, unlike DUSDi, does not impose an independence assumption between factors.

Real-world applications (as framed by the environments used):

• Robotic manipulation and cookery-style multi-step tasks (the Kitchen environment with an arm, butter, meatball, pot, stove, button, and target).
• Long-horizon control where a hierarchical controller chains reusable low-level skills to solve compositional goals.
• Multi-agent coordination, as in Multi-Particle, where a centralized controller drives 10 heterogeneous point-mass agents each tied to a specific station.
• Reward-free pretraining of general-purpose RL agents that can later be adapted to new tasks without hand-designed reward functions.

Industry relevance. The core motivation is that reward engineering does not scale across tasks. A method that learns a broad, disentangled skill library without supervision reduces the cost of deploying reinforcement learning in settings with sparse rewards or many objects, such as warehouse robotics and robot manipulation.

Future Directions

1. Pixel-based observations. SUSD assumes direct access to the underlying state vector. The paper notes that with only visual input, an encoder plus representation learning would be needed to extract disentangled factors — an untested extension.
2. Robustness to incorrect factorization. Appendix M analyzes the effects of distorted inductive biases in factorization, implying that the method's dependence on a correct state decomposition is an open concern.
3. Discrete skill spaces. The paper focuses on continuous skills but states the method is readily applicable to discrete skill spaces, with details deferred to Appendix F.
4. Scaling and unfactorized settings. The evaluation covers factors ranging from 1 to 10 and shows a large advantage in complex factorized environments, but the reported zero-shot goal-reaching results on Ant and HalfCheetah are not quantified in the text, leaving the method's standing in unfactorized settings an explicit open question (Q2).

Target Audience

Reinforcement learning researchers working on unsupervised skill discovery, hierarchical RL, and intrinsic motivation; roboticists interested in reward-free pretraining for manipulation and multi-entity control; and graduate students who already understand MDPs, mutual information objectives, and actor-critic methods. The paper requires prior familiarity with DSD baselines such as LSD, CSD, and METRA to follow the reformulation in Lemma 4.1 and the dual objectives.

Authors’ 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.