Learning Transferable Skills using Goal-Conditioned Bisimulation

Unsupervised skill discovery has emerged as a promising approach for leveraging reward-free datasets to pretrain general-purpose policies. However, current skill discovery methods either require access to expert data or exhibit limited generalization, failing to transfer effectively to previously unseen layouts. A key challenge is to learn representations that capture the temporal structure of the environment while remaining robust to variations across layouts. To address this issue, we present an objective for learning action-aware temporal representations that satisfy the functional equivariance property while preserving the local temporal structure of the environment. Building upon this embedding, we further propose unsupervised skill discovery using bisimulation, which learns transferable skills by conditioning the behavior of skills exclusively on the subset of state features that directly affect their execution. This enforces invariant behavior across different layouts, enabling skills to transfer effectively to other configurations. Finally, through comprehensive empirical evaluations, we show that skills learned in a given environment can be effectively applied to solve downstream tasks in various environment layouts, demonstrating strong out-of-distribution generalization.

Given an offline, unlabeled dataset of trajectories collected by an unknown behavior policy in a single layout, our goal is to learn a skill-conditioned policy that discovers a diverse set of behaviors capable of spanning the state space. These skills should not only provide state-space coverage within the training layout, but also transfer effectively to previously unseen layouts. We then use the learned skills to solve downstream tasks in novel layouts, either using test-time planning or by training a high-level controller online to compose and select the skills.

We first train a state-goal pair encoder psi using goal-conditioned bisimulation, which maps state-goal pairs with similar underlying action sequences to nearby representations. We then learn an action-aware temporal representation (ATR) phi that captures local temporal dependencies by enforcing triangle inequality while preserving the functional equivariance induced by psi. Next, using the similarity between local structures in the phi representation, we train a state representation beta via bisimulation, such that dynamically similar states are mapped to nearby latent vectors. Finally, we train a skill-conditioned policy using HILP's directional state-spanning rewards, conditioning the skills on the learned beta representations. By conditioning on only the dynamically relevant information of states, USB skills learn monotonic behaviors across dynamically similar states. This enables the resulting skills to transfer to previously unseen layouts.

We visualize the temporal embeddings learned by HILP and ATR on a sample maze layout across two environments. ATR’s functional equivariance preserves the underlying structure of the environment and aligns similar action primitives along parallel directions. In contrast, HILP distorts this structure by encoding global temporal distances, resulting in behaviors that vary continuously across the layout.