Learning Task Agnostic Skills with Data-driven Guidance

The overall goal of skill discovery is to find a policy $\pi_{\vartheta}$ capable of carrying out different tasks that are learned without extrinsic supervision for each type of behavior. We consider policies of the form $\pi_{\vartheta}(a|s,z)$ that specify different distributions over actions depending on which skill $z$ they are conditioned on. Although this general framework does not constrain how $z$ should be represented, we define it as a discrete variable since it has been empirically shown to perform better than continuous alternatives (Eysenbach et al., 2018).

We follow the framework proposed by the "diversity is all you need" (DIAYN) algorithm (Eysenbach et al., 2018), in which skills are learned by defining an intrinsic reward that promotes diversity. Intuitively, each skill should make the agent visit a unique section of the state space. This can be expressed as maximising the mutual information $\mathcal{I}(S;Z)$ of the state visitation distributions for different skills (Salge et al., 2014). To ensure that the visited areas of the state space are spaced sufficiently far apart, we use a soft policy that maximises the entropy of the action distribution. Formally, we maximize the following objective function:

The first term $\mathcal{H}[A|S,Z]$ means that the policy should act as randomly as possible and can be optimized by maximizing the policy’s entropy. The second term $-\mathcal{H}[Z|S]$ dictates that each visited state should (ideally) identify the current skill. The third term is the entropy of the skill distribution, which can be maximized by deliberately sampling skills from a uniform distribution during training. Unfortunately, $-\mathcal{H}[Z|S]$ requires knowledge about $p(z|s)$, which is not readily available. Consequently, we approximate the true distribution by training a classifier $q_{\phi}(z|s)$, leading to a lower bound:

The lower bound follows from the non-negative property of the Kullback-Leibler divergence $D_{KL}(p||q)=\mathbb{E}_{p}[\log p(z|s)-\log q(z|s)]\geq 0$, which can be rearranged to $\mathbb{E}_{p}[\log p(z|s)]\geq\mathbb{E}_{p}[\log q(z|s)]$ (Agakov, 2004).

The classifier $q_{\phi}$ is fitted throughout training with maximum likelihood estimation over the sampled states and active skills. This leads to a scenario where the policy is rolled out for a (uniformly) sampled skill, and the classifier is trained to detect the skill based on the states that were visited. The policy is given a reward proportional to how well the classifier could detect the skill in each state. In the end, this should make the policy favor visiting disjoint sets of states for each skill, leading to a cooperative game between $q_{\phi}$ and $\pi_{\vartheta}$.

A major challenge that arises when maximizing the objective in Equation 2, particularly in applications with high-dimensional spaces, is that it becomes trivial for each skill to find a sub-region of the state space where it is easy to be recognised by $q_{\phi}$. In preliminary experiments, we observed that the existing methods discovered behaviours that covered small parts of the state space. For the HalfCheetah environment (Brockman et al., 2016) this resulted in many skills generating different types of static poses (see Figure 1) and not many skills exhibiting "interesting" behaviour such as locomotion.

Optimising for $\mathcal{H}[A|S,Z]$ should mitigate this issue to some extent. Increasing the policy’s entropy incentivises the skills to progressively visit regions of the state space that are so far apart that not even highly stochastic actions will cause them to overlap accidentally. However, it has been shown that mutual information based algorithms have difficulties spreading out to novel states due to low values of $\log q_{\phi}(z|s)$ for out-of-sample states (Campos et al., 2020).

We consider linear projections of continuous factored state representations on the form $f_{\chi}:\mathbb{R}^{|S|}\xrightarrow[]{}\mathbb{R}^{|E|}$ with $e=f_{\chi}(s)=\chi s$, $\chi\in\mathbb{R}^{|E|\times|S|}$, and $|E|<|S|$. In principle, the idea should apply to more complex mappings, such as a multi-layer perceptron. However, we want to limit the scope of skill discovery to a hyperplane within the original state space.

For the same reason, we also omit any non-linearities in the encoder. Squeezing the output through a Sigmoidal function would limit discriminability at the (potentially interesting) extremes of the encoding. Similarly, a ReLU function would effectively eliminate all exploration along the negative direction of $f_{\chi_{i}}$. In summary, the objective of the DIAYN skill classifier becomes:

We learn the parameters $\chi$ for the projection through an auxiliary discriminative objective. Specifically, a binary classifier $h_{\psi}:\mathbb{R^{|E|}}\xrightarrow[]{}\{0,1\}$ is trained to predict whether an (encoded) state was sampled from the marginal state visitation distribution of a random policy $\pi_{rand}$ or from the distribution of a reference (expert) policy $\pi^{*}$. Let $x\sim\mathcal{D}$ denote whether a state $s$ was visited by the reference policy or not (in dataset $\mathcal{D}$), then the parameters of $f_{\chi}$ are obtained through joint pretraining with $h_{\psi}$ by maximizing the log likelihood over $\mathcal{D}$:

where the dataset $\mathcal{D}$ is collected prior to training the main RL algorithm. The first half (random samples) are collected by rolling out $\pi_{rand}$ whereas the second half (reference samples) are collected by rolling out $\pi^{*}$. After the objective in Equation 4 is optimized, the discriminator $h_{\psi}$ is discarded and the projection encoding $f_{\chi}(s)$ is extracted to be used for the objective in Equation 3. Analogous to autoencoders (Hinton & Salakhutdinov, 2006), the idea is that the embeddings produced by $f_{\chi}(s)$ should now contain a more compact representation of the state space without collapsing the dimensions that make "interesting" behaviour stand out.

While the use of a reference data changes our approach from a strictly unsupervised skill discovery algorithm, the discriminative objective in equation 4 resembles the objectives used in adversarial inverse reinforcement learning (e.g. (Fu et al., 2018)). However, it differs in that it makes no attempts at matching the behaviour of a reference policy as it is used only as a prior for simplifying the state space. This approach could also be used with samples from several different reference policies with substantially different marginal state distributions. As long as their variation can be explained sufficiently without full use of the entire state space, a projection should simplify skill discovery.

For learning diverse skills, we use DIAYN as a basis framework. DIAYN uses the Soft Actor-Critic (SAC) algorithm (Haarnoja et al., 2018) that is optimized using policy gradient style updates in contrast to the reparameterized version (DDPG style updates (Lillicrap et al., )). They also use a Squashed Gaussian Mixture Model to represent the policy $a\sim\pi_{\vartheta}=\tanh GMM(\mu_{\vartheta}(s),\sigma_{\vartheta}(s))$. The learning objective is to maximize the mutual information between the state and skill $I(S;Z)$. This objective is optimized by replacing the task rewards with a pseudo-reward

where $q_{\phi}$ is trained to discriminate between skills and p(z) is the fixed uniform prior over skills (Eysenbach et al., 2018). A skill is sampled from $z\sim p(z)$ and used throughout a full episode.

In contrast to DIAYN, we use two Q-functions $Q^{1}_{\theta}(s,a)$ & $Q^{2}_{\theta}(s,a)$ where both Q-functions attempt to predict the same quantity. This allows us to sample differentiable actions and climb the gradient of the minimum of the two Q-functions (DDPG-style update (Lillicrap et al., )), giving us this objective:

Like in DIAYN, we also use a Squashed Gaussian Mixture Model to promote diverse behaviour.

Figure 2 illustrates the training process of the proposed expert-guided skill discovery. First, we train the encoder $f_{\chi}$ jointly with the auxiliary classifier $h_{\psi}$ using the external dataset $\mathcal{D}$. Secondly, we train the agent using an offline policy algorithm (SAC), in which the agent samples a skill $z\sim p(z)$, and then interacts with the environment by taking action $a_{t}$ according the skill-conditioned policy $\pi_{\vartheta}(a_{t}|s_{t},z)$. The environment, then, transits to a new state according to the transition probability $s_{t+1}\sim p(s_{t+1}|s_{t},a_{t})$. We add this transition $(s_{z},z,a_{t},s_{t+1})$ to the replay buffer $\mathcal{B}$. Simultaneously, the policy is updated by sampling a mini-batch from the replay buffer $\mathcal{M}\sim\mathcal{B}$, then encoding the next states $e_{t+1}=f_{\chi}(s_{t+1})$ and passing them through the discriminator $q_{\phi}(z|e_{t+1}$ to get the intrinsic reward. This reward is used by the Q-functions $Q_{\theta}(s_{t},a_{t},z)$ to minimize the soft Bellman residual and update the policy. A pseudocode for the proposed approach can be found in the supplementary material.

As an illustrative example, we begin by testing the algorithm on a simple 2D point-maze problem. The term maze is used very generously here, as the environment consists of an open plane in $\mathbb{R}^{2}$ enclosed by walls that restrict the agent to $\frac{1}{7}<x<\frac{6}{7}$ and $\frac{1}{7}<y<\frac{6}{7}$. At initialization, the agent is dropped down at $x,y=\frac{7+\epsilon}{14},\epsilon\sim U(-1,1)$ and incentivized to move towards the lower right by a reward proportional to a gaussian kernel centered at ($\frac{9}{14}$, $\frac{3}{14}$). The agent is free to move by $\pm\frac{1}{70}$ in both the x and y direction.

We train a SAC agent against the extrinsic environment reward to convergence and set its final policy as the reference policy $\pi*$. We then sample 10 trajectories of length 100 (green dots in Figure 3) with $\pi*$, as well as 10 trajectories of length 100 (blue dots in Figure 4) from a uniform random policy $\pi_{rand}$. The resulting dataset $\mathcal{D}$ consists of 2000 samples and is used to train $h_{\psi}(x|e)$ until it can distinguish states from $\pi*$ and $\pi_{rand}$ with around 98% accuracy. For this experiment, we project down from 2D to 1D, making $\chi\in\mathbb{R}^{1\times 2}$. The resulting projection axis is visualized as a red line in Figure 3 and is the only thing exported to the next stage of the algorithm.

We then train two versions of the DIAYN algorithm; a baseline using the states as-is in the classifier ($q_{\phi}(z|s)$) and our proposed method using the state projections ($q_{\phi}(z|f_{\chi}(s))$). All other hyperparameters are held equal in the two experiments, and the algorithms are trained for 400,000 environment interactions, each attempting to learn 10 distinct skills.

Figure 4 visualizes the results of the baseline and the state projection to the left and right, respectively. The top row shows the states that were visited for five rollouts of each skill. As expected, the baseline skills spread out in all directions (albeit slightly more so towards the left) and converge on locations that are easy to distinguish with a 2D state representation. In contrast, the skills generated in the left plot form lines along the state projection axis. Their wide lateral spread follows from the (unbounded) entropy maximization objective. Besides, any movement perpendicular to the projection axis will not affect the 1D vector passed to the classifier.

Next, we evaluate the algorithm on three continuous control problems from the OpenAI gym suite. (Brockman et al., 2016). The environments include HalfCheetah ($S=\mathbb{R}^{17}$, $A=\mathbb{R}^{6}$), Hopper ($S=\mathbb{R}^{11}$, $A=\mathbb{R}^{3}$) and Ant ($S=\mathbb{R}^{27}$, $A=\mathbb{R}^{8}$). We choose these environments because they involve substantially different locomotion methods. Additionally, they have observation spaces with different dimensionality, which enable us to better investigate the projection impact.

We do one baseline run without any state projection for all three problems, one with a projection down to $\mathbb{R}^{5}$, and one with a projection down to $\mathbb{R}^{3}$. We use our base SAC implementation to obtain reference policies $\pi^{*}$ and sample 10 trajectories of length 1000 with fairly high returns (Ant: $5063.9\pm 469.8$, Cheetah: $10656.4\pm 673.5$, Hopper: $3348.0\pm 316.0$). The DIAYN algorithm is otherwise identical to (Eysenbach et al., 2018) in terms of hyperparameters; $Q$, $\pi$, $q_{\phi}$, and $h_{\psi}$ use MLP architectures with 2 hidden layers of width 300, the entropy bonus weight $\alpha$ is set to 0.1, and the number of skills is set to 50. We limit each skill-discovery run to 2.5 million environment interactions but repeat each experiment 5 times with different random seeds (including training of SAC agents for $\pi^{*}$).

For quantitative evaluation, we look at the displacement along the target locomotion axis for the extrinsic objective. In our approach, we would expect to observe skills that cover this axis well, i.e., skills that run forward and backward at different speeds. To test this, we roll out each skill deterministically²²2Deterministic sampling from our GMM-based policy implies taking the mean of the component with the highest mixture probability., record its movement over 1000 time steps (or until it reaches a terminal state) and observe the inter-skill spread. A similar assessment is possible by only looking at the environment’s rewards. However, the environment reward also includes terms for energy expenditure, staying alive (for Ant/Hopper), and collisions (Ant), which would obscure the results. Figure 6 shows the displacement distribution of the 50 skills across all runs. The same information is summarized numerically in Table 1.

For a qualitative evaluation, we have also composed a video with every skill across all runs³³3Video of skills: https://www.youtube.com/watch?v=Xx7RVNmv1tY.