Group-Invariant Unsupervised Skill Discovery: Symmetry-aware Skill Representations for Generalizable Behavior

TL;DR

This work tackles unsupervised skill discovery in reinforcement learning by exploiting geometric symmetries through Group-Invariant Skill Discovery (GISD). It proves that in group-symmetric environments, the Wasserstein dependency measure $I_{\mathcal{W}}$ has a globally optimal solution consisting of an equivariant policy and a group-invariant scoring function, motivating a Group-Invariant Wasserstein Dependency Measure ($I_{\mathcal{W}}^{G}$). The authors parameterize the invariant scoring function in group Fourier space, yielding an intrinsic reward $r(s,z,s') = \langle \phi_F(s')-\phi_F(s), z \rangle$ that is symmetry-consistent and conducive to generalization under group transformations. Experiments on state-based and pixel-based locomotion tasks show that GISD achieves broader state-space coverage and faster, more reliable downstream task learning than a strong baseline, validating the practical impact of symmetry-aware skill discovery.

Abstract

Unsupervised skill discovery aims to acquire behavior primitives that improve exploration and accelerate downstream task learning. However, existing approaches often ignore the geometric symmetries of physical environments, leading to redundant behaviors and sample inefficiency. To address this, we introduce Group-Invariant Skill Discovery (GISD), a framework that explicitly embeds group structure into the skill discovery objective. Our approach is grounded in a theoretical guarantee: we prove that in group-symmetric environments, the standard Wasserstein dependency measure admits a globally optimal solution comprised of an equivariant policy and a group-invariant scoring function. Motivated by this, we formulate the Group-Invariant Wasserstein dependency measure, which restricts the optimization to this symmetry-aware subspace without loss of optimality. Practically, we parameterize the scoring function using a group Fourier representation and define the intrinsic reward via the alignment of equivariant latent features, ensuring that the discovered skills generalize systematically under group transformations. Experiments on state-based and pixel-based locomotion benchmarks demonstrate that GISD achieves broader state-space coverage and improved efficiency in downstream task learning compared to a strong baseline.

Figures (6)

• Figure 1: Overview of Group-Invariant Skill Discovery (GISD). Our method learns an equivariant mapping $\phi_F$ from the state space $\mathcal{S}$ to the Group Fourier space$Z$. By aligning state transitions with latent skill vectors $z$ in the group Fourier domain, the discovered skills inherently respect the underlying geometric symmetry. This structure enables generalization: a policy trained for a specific trajectory $\tau$ automatically generalizes to any group-transformed trajectory $g\tau$ by simply shifting the skill vector along its group orbit to $gz$.
• Figure 2: Benchmark Environments. We evaluate GISD in state-based Ant and pixel-based Quadruped locomotion. In the pixel-based setting, we redesign the floor so that color varies radially from the origin, making the observations consistent with the horizontal flip symmetry.
• Figure 3: State-space coverage during skill discovery.(a) State-based Ant (4 seeds). (b) Pixel-based Quadruped (5 seeds). Shaded regions show standard error. GISD achieves higher coverage and better sample efficiency than METRA in both environments.
• Figure 4: Visualization of discovered skills. We plot trajectories for 48 randomly sampled skills. Colors indicate different latent skills. GISD discovers symmetry-consistent skills (Ant: $C_4$ rotations; Quadruped: horizontal flips), leading to substantially broader and more uniform state-space coverage than METRA.
• Figure 5: Downstream task performance.(a) Average return in the state-based Ant environment (4 seeds). (b) Average return in the pixel-based Quadruped environment (3 seeds). Shaded regions show standard error. GISD consistently achieves higher performance and better sample efficiency than METRA.

Theorems & Definitions (18)

• Theorem 1: Existence of Equivariant Optima
• Definition 1: Group-Invariant WDM
• Proposition 1
• Definition 2: Group-averaged Scoring Function
• Proposition 2: Properties of the group-averaged scoring function
• Remark 1: Exactness for Equivariant Policies
• Definition 3: Fixed-interval High-level Transition
• Theorem 2: Group-invariant Fixed-interval Semi-MDP
• Theorem \ref{thm:eq-opt}: Existence of Equivariant Optima
• proof