Lipschitz-constrained Unsupervised Skill
Discovery

We consider a Markov decision process (MDP) ${\mathcal{M}}=({\mathcal{S}},{\mathcal{A}},p)$ without external rewards, where ${\mathcal{S}}$ is a (continuous) state space, ${\mathcal{A}}$ is an action space, and $p(s_{t+1}|s_{t},a_{t})$ is a stochastic transition dynamics function. We represent a skill with a latent variable $z\in{\mathcal{Z}}$ and a latent-conditioned policy $\pi(a|s,z)$. The skill latent space ${\mathcal{Z}}$ can be either discrete or continuous; we use $N$ to denote the number of skills in the discrete case, and $d$ to denote the dimensionality of the skill latent space in the continuous case. Given a skill $z$ and a skill policy $\pi(a|s,z)$, a trajectory $\tau=(s_{0},a_{0},\ldots,s_{T})$ is sampled with the following generative process: $p^{\pi}(\tau|z)=p(s_{0})\prod_{t=0}^{T-1}\pi(a_{t}|s_{t},z)p(s_{t+1}|s_{t},a_{t})$. We use uppercase letters to denote random variables, and $h(\cdot)$ and $I(\cdot;\cdot)$ to represent the differential entropy and the mutual information, respectively.

A number of previous methods maximize $I(Z;S)$ to learn diverse skills. One line of research employs the identity $I(Z;S)=h(Z)-h(Z|S)\geq\mathbb{E}_{z\sim p(z),s\sim p^{\pi}(s|z)}[\log q(z|s)]-\mathbb{E}_{z\sim p(z)}[\log p(z)]$ or its variants, where the skill discriminator $q(z|s)$ is a variational approximation of the posterior $p(z|s)$ (Barber & Agakov, 2003). VIC (Gregor et al., 2016) maximizes the MI between the last states and skills given the initial state. DIAYN (Eysenbach et al., 2019) optimizes the MI between individual states and skills. VALOR (Achiam et al., 2018) also takes a similar approach, but considers the whole trajectories instead of states. VISR (Hansen et al., 2020) models the variational posterior as the von-Mises Fisher distribution, which results in an inner-product reward form and hence enables combining with successor features (Barreto et al., 2017). HIDIO (Zhang et al., 2021) examines multiple variants of the MI identity, and jointly learns skills with a hierarchical controller that maximizes the task reward. Choi et al. (2021) point out the equivalency between the MI-based objective and goal-conditioned RL, and show that Spectral Normalization (Miyato et al., 2018) improves the quality of learned skills. DADS (Sharma et al., 2020) maximizes the opposite direction of the mutual information identity $I(Z;S)=h(S)-h(S|Z)$ with the skill dynamics model $q(s_{t+1}|s_{t},z)$, which allows zero-shot planning on downstream tasks. However, these methods share a limitation: they do not always prefer to reach distant states or to learn dynamic skills, as we can maximize $I(Z;S)$ even with the smallest state variations. One possible way to address this issue is to use heuristics such as feature engineering; for example, the $x$-$y$ prior (Sharma et al., 2020) enforces skills to be discriminated only by their $x$-$y$ coordinates so that the agent can discover locomotion skills.

On the other hand, a couple of methods overcome this limitation by integrating with exploration techniques. EDL (Campos Camúñez et al., 2020) first maximizes the state entropy $h(S)$ with SMM exploration (Lee et al., 2019), and then encodes the discovered states into skills via VAE (Kingma & Welling, 2014). APS (Liu & Abbeel, 2021a) combines VISR (Hansen et al., 2020) with APT (Liu & Abbeel, 2021b), an exploration method based on $k$-nearest neighbors. Yet, we empirically confirm that such a pure exploration signal is insufficient to make large and consistent transitions in states. IBOL (Kim et al., 2021) takes a hierarchical approach where it first pre-trains a low-level policy to make reaching remote states easier, and subsequently learns a high-level skill policy based on the information bottleneck framework (Tishby et al., 2000). While IBOL can discover skills reaching distant states in continuous control environments without locomotion priors (e.g., $x$-$y$ prior), it still has some limitations in that (1) IBOL cannot discover discrete skills, (2) it still capitalizes on input feature engineering in that they exclude the locomotion coordinates from the low-level policy, and (3) it consists of a two-level hierarchy with several additional hyperparameters, which make the implementation difficult. On the contrary, our proposed LSD can discover diverse skills in both discrete and continuous settings without using any feature engineering, and is easy to implement as it requires no additional hyperparameters or hierarchy. Table 1 overviews the comparison of properties between different skill discovery methods.

Before deriving our objective for discovery of continuous skills, we review variational MI-based skill discovery algorithms and discuss why such methods might end up learning only simple and static skills. The MI objective $I(Z;S)$ with continuous skills can be written with the variational lower bound as follows (Eysenbach et al., 2019; Choi et al., 2021):

where we assume that a skill $z\in{\mathbb{R}}^{d}$ is sampled from a fixed prior distribution $p(z)$, and $q(z|s)$ is a variational approximation of $p(z|s)$ (Barber & Agakov, 2003; Mohamed & Rezende, 2015), parameterized as a normal distribution with unit variance, ${\mathcal{N}}(\mu(s),\mathbf{I})$ (Choi et al., 2021). Some other methods are based on a conditional form of mutual information (Gregor et al., 2016; Sharma et al., 2020); for instance, the objective of VIC (Gregor et al., 2016) can be written as

where we assume that $p(z|s_{0})=p(z)$ is a fixed prior distribution, and the posterior is chosen as $q(z|s_{0},s_{T})=\mathcal{N}(\mu(s_{0},s_{T}),\mathbf{I})$ in a continuous skill setting.

One issue with Equations 2 and 3 is that these objectives can be fully maximized even with small differences in states as long as different $z$’s correspond to even marginally different $s_{T}$’s, not necessarily encouraging more ‘interesting’ skills. This is especially problematic because discovering skills with such slight or less dynamic state variations is usually a ‘lower-hanging fruit’ than making dynamic and large differences in the state space (e.g., $\mu$ can simply map the angles of Ant’s joints to $z$). As a result, continuous DIAYN and DADS discover only posing skills on Ant (Figures 2(a) and 17) in the absence of any feature engineering or tricks to elicit more diverse behaviors. We refer to Appendix H for quantitative demonstrations of this phenomenon on MuJoCo environments.

Further decomposition of the MI objective. Before we address this limitation, we decompose the objective of VIC (Equation 3) to get further insights that will inspire our new objective in Section 3.2. Here, we model $\mu(s_{0},s_{T})$ with $\phi(s_{T})-\phi(s_{0})$ to focus on the relative differences between the initial and final states, where $\phi:{\mathcal{S}}\to{\mathcal{Z}}$ is a learnable state representation function that maps a state observation into a latent space, which will be the core component of our method. This choice makes the latent skill $z$ represent a direction or displacement in the latent space induced by $\phi$. Then, we can rewrite Equation 3 as follows:

where we use the fact that $\mathbb{E}_{z}[z^{\top}z]$ is a constant as $p(z)$ is a fixed distribution. This decomposition of the MI lower-bound objective provides an intuitive interpretation: the first inner-product term \raisebox{-1.0pt}{\small1}⃝ in Equation 5 encourages the direction vector $\phi(s_{T})-\phi(s_{0})$ to be aligned with $z$, while the second term \raisebox{-1.0pt}{\small2}⃝ regularizes the norm of the vector $\phi(s_{T})-\phi(s_{0})$.

The LSD objective. We now propose our new objective that is based on neither a skill discriminator (Equation 3) nor mutual information but a Lipschitz-constrained state representation function, in order to address the limitation that $I(Z;S)$ can be fully optimized with small state differences. Specifically, inspired by the decomposition in Equation 5, we suggest using the directional term \raisebox{-1.0pt}{\small1}⃝ as our objective. However, since this term alone could be trivially optimized by just increasing the value of $\phi(s_{T})$ to the infinity regardless of the final state $s_{T}$, we apply the 1-Lipschitz constraint on the state representation function $\phi$ so that maximizing our objective in the latent space can result in an increase in state differences. This leads to a constrained maximization objective as follows:

where $J^{\text{LSD}}$ is the objective of our proposed Lipschitz-constrained Skill Discovery (LSD).

The LSD objective encourages the agent to prefer skills with larger traveled distances, unlike previous MI-based methods, as follows. First, in order to maximize the inner product in Equation 6, the length of $\phi(s_{T})-\phi(s_{0})$ should be increased. It then makes its upper bound $\|s_{T}-s_{0}\|$ increase as well due to the 1-Lipschitz constraint (i.e., $\|\phi(s_{T})-\phi(s_{0})\|\leq\|s_{T}-s_{0}\|$). As a result, it leads to learning more dynamic skills in terms of state differences.

Note that LSD’s objective differs from VIC’s in an important way. Equation 3 tries to equate the value of $z$ and $\mu(s_{0},s_{T})$ (i.e., it tries to recover $z$ from its skill discriminator), while the objective of LSD (Equation 6) only requires the directions of $z$ and $\phi(s_{T})-\phi(s_{0})$ to be aligned.

We also note that our purpose of enforcing the Lipschitz constraint is very different from its common usages in machine learning. Many works have adopted the Lipschitz continuity to regularize functions for better generalization (Neyshabur et al., 2018; Sokolic et al., 2017), interpretability (Tsipras et al., 2018) or stability (Choi et al., 2021). On the other hand, we employ it to ensure that maximization of the reward entails increased state variations. The Lipschitz constant $1$ is chosen empirically as it can be easily implemented using Spectral Normalization (Miyato et al., 2018).

Per-step transition reward. By eliminating the second term in Equation 5, we can further decompose the objective using a telescoping sum as

This formulation enables the optimization of $J^{\text{LSD}}$ with respect to the policy $\pi$ (i.e., reinforcement learning steps) with a per-step transition reward given as:

Compared to the per-trajectory reward in Equation 3 (Gregor et al., 2016), this can be optimized more easily and stably with off-the-shelf RL algorithms such as SAC (Haarnoja et al., 2018a).

Connections to previous methods. The per-step reward function of Equation 8 is closely related to continuous DIAYN (Eysenbach et al., 2019; Choi et al., 2021) and VISR (Hansen et al., 2020):

where $\tilde{z}$ and $\smash{\tilde{\phi}(s_{t})}$ in VISR are the normalized vectors of unit length (i.e., $z/\|z\|$ and $\phi(s_{t})/\|\phi(s_{t})\|$), respectively. We assume that $q^{\text{DIAYN}}$ is parameterized as a normal distribution with unit variance, and $q^{\text{VISR}}$ as a von-Mises Fisher (vMF) distribution with a scale parameter of $1$ (Hansen et al., 2020).

While it appears that there are some similarities among Equation 9–(11), LSD’s reward function is fundamentally different from the others in that it optimizes neither a log-probability nor mutual information, and thus only LSD seeks for distant states. For instance, VISR (Equation 10) has the most similar form as it also uses an inner product, but $\phi(s_{t+1})-\phi(s_{t})$ in $r_{z}^{\text{LSD}}$ does not need to be a unit vector unlike VISR optimizing the vMF distribution. Instead, LSD increases the difference of $\phi$ to encourage the agent to reach more distant states. In addition, while Choi et al. (2021) also use Spectral Normalization, our objective differs from theirs as we do not optimize $I(Z;S)$. We present further discussion and an empirical comparison of reward functions in Appendix D.

Implementation. In order to maximize the LSD objective (Equation 8), we alternately train $\pi$ with SAC (Haarnoja et al., 2018a) and $\phi$ with stochastic gradient descent. We provide the full procedure for LSD in Appendix C. Note that LSD has the same components and hyperparameters as DIAYN, and is thus easy to implement, especially as opposed to IBOL (Kim et al., 2021), a state-of-the-art skill discovery method on MuJoCo environments that requires a two-level hierarchy with many moving components and hyperparameters.

Another advantage of the LSD objective $J^{\text{LSD}}$ (Equation 6) is that the learned state representation $\phi(s)$ allows solving goal-following downstream tasks in a zero-shot manner (i.e., without any further training or complex planning), as $z$ is aligned with the direction in the representation space. Although it is also possible for DIAYN-like methods to reach a single goal from the initial state in a zero-shot manner (Choi et al., 2021), LSD is able to reach a goal from an arbitrary state or follow multiple goals thanks to the directional alignment without any additional modifications to the method. Specifically, if we want to make a transition from the current state $s\in\mathcal{S}$ to a target state $g\in\mathcal{S}$, we can simply repeat selecting the following $z$ until reaching the goal:

and executing the latent-conditioned policy $\pi(a|s,z)$ to choose an action. Here, $\alpha$ is a hyperparameter that controls the norm of $z$, and we find that $\alpha\approx\mathbb{E}_{z\sim p(z)}[\|z\|]$ empirically works the best (e.g., $\alpha=\smash{2^{-\frac{1}{2}}\Gamma(1/2)}\approx 1.25$ for $d=2$). As will be shown in Section 4.3, this zero-shot scheme provides a convenient way to immediately achieve strong performance on many goal-following downstream tasks. Note that, in such tasks, this scheme is more efficient than the zero-shot planning of DADS with its learned skill dynamics model (Sharma et al., 2020); LSD does not need to learn models or require any planning steps in the representation space.

Continuous LSD can be extended to discovery of discrete skills. One might be tempted to use the one-hot encoding for $z$ with the same objective as continuous LSD (Equation 6), as in prior methods (Eysenbach et al., 2019; Sharma et al., 2020). For discrete LSD, however, we cannot simply use the standard one-hot encoding. This is because an encoding that has a non-zero mean could make all the skills collapse into a single behavior in LSD. For example, without loss of generality, suppose the first dimension of the mean of the $N$ encoding vectors is $c>0$. Then, if the agent finds a final state $s_{T}$ that makes $\|s_{T}-s_{0}\|$ fairly large, it can simply learn the skill policy to reach $s_{T}$ regardless of $z$ so that the agent can always receive a reward of $c\cdot\|s_{T}-s_{0}\|$ on average, by setting, e.g., $\phi(s_{T})=[\|s_{T}-s_{0}\|,0,\ldots,0]^{\top}$ and $\phi(s_{0})=[0,0,\ldots,0]$. In other words, the agent can easily exploit the reward function without learning any diverse set of skills.

To resolve this issue, we propose using a zero-centered one-hot vectors as follows:

where $N$ is the number of skills, $[\cdot]_{i}$ denotes the $i$-th element of the vector, and $\mathrm{unif}\{\ldots\}$ denotes the uniform distribution over a set. Plugging into Equation 8, the reward for the $k$-th skill becomes

This formulation provides an intuitive interpretation: it enforces the $k$-th element of $\phi(s)$ to be the only indicator for the $k$-th skill in a contrastive manner. We note that Equation 14 consistently pushes the difference in $\phi$ to be as large as possible, unlike prior approaches using the softmax function (Eysenbach et al., 2019; Sharma et al., 2020). Thanks to the Lipschitz constraint on $\phi(s)$, this makes the agent learn diverse and dynamic behaviors as in the continuous case.

Visualization of skills. We train LSD and baselines on the Ant environment to learn two-dimensional continuous skills ($d=2$). Figure 2(a) visualizes the learned skills as trajectories of the Ant agent on the $x$-$y$ plane. LSD discovers skills that move far from the initial location in almost all possible directions, while the other methods except IBOL fail to discover such locomotion primitives without feature engineering (i.e., $x$-$y$ prior) even with an increased skill dimensionality ($d=5$). Instead, they simply learn to take static postures rather than to move; such a phenomenon is also reported in Gu et al. (2021). This is because their MI objectives do not particularly induce the agent to increase state variations. On the other hand, LSD discovers skills reaching even farther than those of the baselines using feature engineering (IBOL-O, DADS-XYO and DIAYN-XYO).

Figure 2(b) demonstrates the results on the Humanoid environment. As in Ant, LSD learns diverse skills walking or running consistently in various directions, while skills discovered by other methods are limited in terms of the state space coverage or the variety of directions. We provide the videos of skills discovered by LSD at https://shpark.me/projects/lsd/.

Quantitative evaluation. For numerical comparison, we measure the state space coverage (Kim et al., 2021) of each skill discovery method on Ant and Humanoid. The state space coverage is measured by the number of occupied $1\times 1$ bins on the $x$-$y$ plane from $200$ randomly sampled trajectories, averaged over eight runs. For EDL, we report the state space coverage of its SMM exploration phase (Lee et al., 2019). Figure 3 shows that on both environments, LSD outperforms all the baselines even including those with feature engineering.

We train discrete LSD on Ant, HalfCheetah and Humanoid with $N=6$, $8$, $16$, where $N$ is the number of skills. While continuous LSD mainly discovers locomotion skills, we observe that discrete LSD learns more diverse skills thanks to its contrastive scheme (Figure 4, Appendix J). On Ant, discrete LSD discovers a skill set consisting of five locomotion skills, six rotation skills, three posing skills and two flipping skills. On HalfCheetah, the agent learns to run forward and backward in multiple postures, to roll forward and backward, and to take different poses. Finally, Humanoid learns to run or move in multiple directions and speeds with unique gaits. We highly recommend the reader to check the videos available on our project page. We refer to Section G.1 for quantitative evaluations. To the best of our knowledge, LSD is the only method that can discover such diverse and dynamic behaviors (i.e., having large traveled distances in many of the state dimensions) within a single set of skills in the absence of feature engineering.

As done in Eysenbach et al. (2019); Sharma et al. (2020); Kim et al. (2021), we make comparisons on downstream goal-following tasks to assess how well skills learned by LSD can be employed for solving tasks in a hierarchical manner, where we evaluate our approach not only on AntGoal and AntMultiGoals (Eysenbach et al., 2019; Sharma et al., 2020; Kim et al., 2021) but also on the more challenging tasks: HumanoidGoal and HumanoidMultiGoals. In the ‘-Goal’ tasks, the agent should reach a uniformly sampled random goal on the $x$-$y$ plane, while in the ‘-MultiGoals’ tasks, the agent should navigate through multiple randomly sampled goals in order. The agent is rewarded only when it reaches a goal. We refer to Appendix I for the full details of the tasks.

We first train each skill discovery method with continuous skills (without rewards), and then train a hierarchical meta-controller on top of the learned skill policy (kept frozen) with the task reward. The meta-controller observes the target goal concatenated to the state observation. At every $K$-th step, the controller selects a skill $z\in[-2,2]^{d}$ to be performed for the next $K$ steps, and the chosen skill is executed by the skill policy for the $K$ steps. We also examine zero-shot skill selection of LSD, denoted ‘LSD (Zero-shot)’, where the agent chooses $z$ at every step according to Equation 12.

Results. Figure 5 shows the performance of each algorithm evaluated on the four downstream tasks. LSD demonstrates the highest reward in all of the environments, outperforming even the baselines with feature engineering. On top of that, in some environments such as AntMultiGoals, LSD’s zero-shot skill selection performs the best, while still exhibiting strong performance on the other tasks. From these results, we show that LSD is capable of solving downstream goal-following tasks very efficiently with no further training or complex planning procedures.

In order to demonstrate that LSD can also discover useful skills in environments other than locomotion tasks, we make another evaluation on three Fetch robotic manipulation environments (Plappert et al., 2018). We compare LSD with other skill discovery methods combined with MUSIC-u (Zhao et al., 2021), an intrinsic reward that maximizes the mutual information $I(S^{a};S^{s})$ between the agent state $S^{a}$ and the surrounding state $S^{s}$. For a fair comparison with MUSIC, we make use of the same prior knowledge for skill discovery methods including LSD to make them focus only on the surrounding states (Zhao et al., 2021), which correspond to the target object in our experiments.

Results. Figure 6(a) visualizes the target object’s trajectories of 2-D continuous skills learned by each algorithm. They suggest that LSD can control the target object in the most diverse directions. Notably, in FetchPickAndPlace, LSD learns to pick the target object in multiple directions without any task reward or intrinsic motivation like MUSIC-u. Figures 6(b) and 6(c) present quantitative evaluations on the Fetch environments. As in Ant and Humanoid, LSD exhibits the best state space coverage (measured with $0.1\times 0.1$-sized bins) in all the environments. Also, LSD and LSD (Zero-shot) outperform the baselines by large margins on the three downstream tasks.