Hierarchical Empowerment:
Towards Tractable Empowerment-Based
Skill Learning

A core concept in our work is Goal-Conditioned MDPs (Kaelbling, 1993; Schaul et al., 2015), which represent a decision making setting in which an agent needs to learn a distribution of tasks. We define Goal-Conditioned MDPs by the tuple $\{\mathcal{S},p(s_{0}),\mathcal{G},p(g),\mathcal{A},T(s_{t+1}|s_{t},a_{t}),r(s_{t},g,a_{t})\}$. $\mathcal{S}$ is the space of states; $p(s_{0})$ is the distribution of starting states; $\mathcal{G}$ is the space of goals and is assumed to be a subset of $\mathcal{S}$; $p(g)$ is the distribution of goals for the agent to achieve; $\mathcal{A}$ is the action space; $T(s_{t+1}|s_{t},a_{t})$ is the transition dynamics of the environment; $r(s_{t},g,a_{t})$ is the reward function for a (state,goal,action) tuple. A policy $\pi_{\theta}(\cdot|s_{t},z)$ in a Goal-Conditioned MDP is a mapping from states and goals to a distribution over the action space.

The objective in a goal-conditioned MDP is to learn the parameters $\theta$ of the policy that maximize the sum of discounted rewards averaged over all goals and goal-conditioned trajectories:

in which $\tau$ is a trajectory of states and actions $(s_{0},a_{0},s_{1},a_{1},\dots)$. The goal-conditioned MDP objective can be optimized with Reinforcement Learning, such as policy gradient methods (Sutton et al., 1999a).

Central to our approach for skill learning is the observation that skills are an instance of noisy channels from Information Theory (Cover & Thomas, 2006). The proposed skill is the message input into the noisy channel. The state that results from executing the skill is the output of the channel. We will refer to this noisy channel as the skill channel. The skill channel is defined by the tuple $\{\mathcal{S},\mathcal{Z},p(s_{n}|s_{0},z)\}$, in which $\mathcal{S}$ is the set of states, $\mathcal{Z}$ is the set of skills that can be executed from a start state $s_{0}\in\mathcal{S}$, $n$ is the number of actions contained within a skill, and $p(s_{n}|s_{0},z)$ is the probability of a skill $z$ that started in state $s_{0}$ terminating in state $s_{n}$ after $n$ actions. That is, $p(s_{n}|s_{0},z)=\int_{a_{0},s_{1},a_{1},\dots,s_{n-1},a_{n-1}}\pi(a_{0}|s_{0},z)p(s_{1}|s_{0},a_{0})\pi(a_{1}|s_{1},z)\dots\pi(a_{n-1}|s_{n-1},z)p(s_{n}|s_{n-1},a_{n-1})$. The distribution $p(s_{n}|s_{0},z)$ thus depends on the skill-conditioned policy $\pi(a_{t}|s_{t},z)$ and the transition dynamics of the environment $p(s_{t+1}|s_{t},a_{t})$.

The mutual information of the skill channel, $I(Z;S_{n}|s_{0})$, describes the number of distinct skills that can be executed from state $s_{0}$ given the current skill distribution $p(z|s_{0})$ and the skill-conditioned policy $\pi(a_{t}|s_{t},z)$. Mathematically, the mutual information of the skill channel can be defined:

That is, the number of distinct skills that can be executed from state $s_{0}$ will be larger when (i) the entropy $H(Z|s_{0})$ (i.e., the size) of the skill space is larger and/or (ii) the conditional entropy $H(Z|s_{0},S_{n})$ is smaller (i.e., the skills are more differentiated as they target distinct regions of the state space).

Given our goal of equipping agents with as large a skillset as possible, the purpose of this paper is to develop an algorithm that can learn a distribution over skills $p_{\phi}(z|s_{0})$ and a skill-conditioned policy $\pi_{\theta}(s_{t},z)$ that maximize the mutual information of the skill channel:

The maximum mutual information (equivalently, the channel capacity) of the skill channel is known as Empowerment (Klyubin et al., 2005). Note that in the literature empowerment is often defined with respect to more simple skills, such as single primitive actions or open loop action sequences (Klyubin et al., 2005; Mohamed & Rezende, 2015), but we use the more general definition developed by Gregor et al. (2017) that permits learnable closed loop skills.

One major difficulty with calculating empowerment (i.e., maximizing the mutual information of the skill channel) is that computing the posterior distribution $p(z|s_{0},s_{n})$, located within the conditional entropy term $H(Z|s_{0},S_{n})$, is often intractable. Analytically calculating the posterior distribution would require integrating over the intermediate actions and states $a_{0},s_{1},a_{1}\dots,s_{n-1},a_{n-1}$. Inspired by Mohamed & Rezende (2015), many empowerment-based skill-learning approaches (Gregor et al., 2017; Eysenbach et al., 2019; Zhang et al., 2021; Strouse et al., 2021; Warde-Farley et al., 2018; Baumli et al., 2020; Achiam et al., 2018; Hansen et al., 2020) attempt to overcome the issue by optimizing a variational lower bound on the mutual information of the skill channel. In this lower bound, a learned variational distribution $q_{\psi}(z|s_{0},s_{n})$ parameterized by $\psi$ is used in place of the problematic posterior $p(z|s_{0},s_{n})$. (See Barber & Agakov (2003) for original proof of the lower bound.) The variational lower bound is accordingly defined

in which $\theta$ represents the parameters of the skill-conditioned policy $\pi_{\theta}(a_{t}|s_{t},z)$.

The common approach in empowerment-based skill learning is to train the learned variational distribution $q_{\psi}$ with maximum likelihood learning so that the variational distribution approaches the true posterior. The skill distribution $p(z)$ is typically held fixed so $H(Z|s_{0})$ is constant and can be ignored. The last term in the variational lower bound, as observed by (Gregor et al., 2017), is a Reinforcement Learning problem in which the reward for a (state,skill,action) tuple is $r(s_{t},z,a_{t})=0$ for the first $n-1$ actions and then $r(s_{n-1},z,a_{n-1})=\mathbb{E}_{s_{n}\sim p(s_{n}|s_{n-1},a_{n})}[\log q_{\psi}(z|s_{0},s_{n}))]$ for the final action.

Unfortunately, this reward function does not encourage skills to differentiate and target distinct regions of the state space (i.e., increase mutual information) when the skills overlap and visit similar states. For instance, at the start of training when skills tend to visit the same states in the immediate vicinity of the start state, the reward $q_{\psi}(z|s_{0},s_{n})$ will typically be small for most $(s_{0},z,s_{n})$ tuples. The undifferentiated reward will in turn produce little change in the skill-conditioned policy, which then triggers little change in the learned variational distribution and reward $q_{\psi}$, making it difficult to maximize mutual information.

Choi et al. (2021) observed that goal-conditioned RL, in the form of the objective below, is also a variational lower bound on mutual information.

That is, $q(\cdot|s_{0},s_{n})$ is a diagonal Gaussian distribution with mean $s_{n}$ and a fixed standard deviation $s_{0}$ set by the designer. The expression in 3 is a goal-conditioned RL objective as there is an outer expectation over a fixed distribution of goal states $p(g|s_{0})$ and an inner expectation with respect to the distribution of trajectories produced by a goal state. The reward function is 0 for all time steps except for the last when it is the $\log$ probability of the skill $g$ sampled from a Gaussian distribution centered on the skill-terminating state $s_{n}$ with fixed standard deviation $\sigma_{0}$. Like any goal-conditioned RL reward function, this reward encourages goal-conditioned policy $\pi_{\theta}(s_{t},g)$ to target the conditioned goal as the reward is higher when the goal $g$ is near the skill-terminating state $s_{n}$.

The goal-conditioned objective in Equation 3, which reduces to $\mathbb{E}_{g\sim p(g|s_{0}),s_{n}\sim p(s_{n}|s_{0},g)}[\log q(g|s_{0},s_{n})]$, is also a variational lower bound on the mutual information of the skill channel. In this case, the variational distribution is a fixed variance Gaussian distribution. Unlike the typical variational lower bound in empowerment-based skill learning, goal-conditioned RL uses a reward function that explicitly encourages skills to target specific states, even when there is significant overlap among the skills in regards to the states visited.

One issue with using goal-conditioned RL as a variational lower bound on mutual information is the hand-crafted distribution of skills $p(g|s_{0})$. In addition to requiring domain expertise, if this is significantly smaller or larger than the region of states that can be reached in $n$ actions by the goal-conditioned policy, then maximizing the goal-conditioned RL objective would only serve as a loose lower bound on empowerment because the objective would either learn fewer skills than it is capable of or too many redundant skills. Another issue with the objective is that it is not a practical objective for learning lengthy goal-conditioned policies given that RL has consistently struggled at performing credit assignment over long horizons using temporal difference learning (Levy et al., 2019; Nachum et al., 2018; McClinton et al., 2021).

The Goal-Conditioned Empowerment objective is defined

$s_{0}$ is the starting state for which empowerment is to be computed. $I^{GCE}(Z;S_{n}|s_{0})$ is a variational lower bound on the mutual information of the skill channel. $\pi_{\theta}(s_{t},z)$ is the goal-conditioned policy parameterized by $\theta$. We will assume $\pi_{\theta}$ is deterministic. $p_{\phi}(z|s_{0})$ is the learned distribution of continuous goal states (i.e., skills) parameterized by $\phi$. The mutual information variational lower bound within Goal-Conditioned Empowerment is defined

Next, we will analyze the variational lower bound objective in equation 5 with respect to the goal-space policy and goal-conditioned policy separately. An immediate challenge to optimizing Equation 5 with respect to the goal space parameters $\phi$ is that $\phi$ occurs in the expectation distribution. The reparameterization trick (Kingma & Welling, 2014) provides a solution to this issue by transforming the original expectation into an expectation with respect to exogenous noise. In order to use the reparameterization trick, the distribution of goal states needs to be a location-scale probability distribution (e.g., the uniform or the normal distributions). We will assume $p_{\phi}(z|s_{0})$ is a uniform distribution for the remainder of the paper. Applying the reparameterization trick to the expectation term in Equation 5 and inserting a fixed variance Gaussian variational distribution similar to the one used by (Choi et al., 2021), the variational lower bound objective becomes

in which $d$ is the dimensionality of the goal space and $z=g(\epsilon,\mu_{\phi}(s_{0}))$ is reparameterized to be a function of the exogenous unit uniform random variable $\epsilon$ and $\mu_{\phi}(s_{0})$, which outputs a vector containing the location-scale parameters of the uniform goal-space distribution. Specifically, for each dimension of the goal space, $\mu_{\phi}(s_{0})$ will output (i) the location of the center of the uniform distribution and (ii) the length of the half width of the uniform distribution. Note also that in our implementation, instead of sampling the skill $z$ from the fixed variational distribution (i.e., $\mathcal{N}(z;s_{n},\sigma_{0})$), we sample the sum $h(s_{0})+z$, in which $h(\cdot)$ is the function that maps the state space to the skill (i.e. goal) space. The change to sampling the sum $h(s_{0})+z$ will encourage the goal $z$ to reflect the desired change from the starting state $s_{0}$. The mutual information objective in Equation 6 can be further simplified to

Equation 7 has the same form as a maximum entropy bandit problem. The entropy term $H_{\phi}(Z|s_{0})$ encourages the goal space policy $\mu_{\phi}$ to output larger goal spaces. The reward function $\mathbb{E}_{\epsilon\sim\mathcal{U}^{d}}[R(s_{0},\epsilon,\mu_{\phi}(s_{0}))]$ in the maximum entropy bandit problem encourages the goal space to contain achievable goals. This is true because $R(s_{0},\epsilon,\mu_{\phi}(s_{0}))$ measures how effective the goal-conditioned policy is at achieving the particular goal $z$, which per the reparameterization trick is $g(\epsilon,\mu_{\phi}(s_{0}))$. The outer expectation $\mathbb{E}_{\epsilon\sim\mathcal{U}^{d}}[\cdot]$ then averages $R(s_{0},\epsilon,\mu_{\phi}(s_{0}))$ with respect to all goals in the learned goal space. Thus, goal spaces that contain more achievable goals will produce higher rewards.

The goal space policy $\mu_{\phi}$ can be optimized with Soft Actor-Critic (Haarnoja et al., 2018), in which a critic $R_{\omega}(s_{0},\epsilon,a)$ is used to approximate $R(s_{0},\epsilon,a)$ for which $a$ represents actions by the goal space policy. In our implementation, we fold in the entropy term into the reward, and estimate the expanded reward with a critic. We then use Stochastic Value Gradient (Heess et al., 2015) to optimize. Section F of the Appendix provides the objective for training the goal-space critic. The gradient for the goal-space policy $\mu_{\phi}$ can be determined by passing the gradient through the critic using the chain rule:

With respect to the goal-conditioned policy $\pi_{\theta}$, the mutual information objective reduces to a goal-conditioned RL problem because the entropy term can be ignored. We optimize this objective with respect to the deterministic goal-conditioned policy with Deterministic Policy Gradient (Silver et al., 2014).

In our implementation, we optimize the Goal-Conditioned Empowerment mutual information objective by alternating between updates to the goal-conditioned and goal-space policies—similar to an automated curriculum of goal-conditioned RL problems. In each step of the curriculum, the goal-conditioned policy is trained to reach goals within and nearby the current goal space. The goal space is then updated to reflect the current reach of the goal-conditioned policy. Detailed algorithms for how the goal-conditioned actor-critic and the goal-space actor-critic are updated are provided in section D of the Appendix.

To scale Goal-Conditioned Empowerment to longer horizons we borrow ideas from Goal-Conditioned Hierarchical Reinforcement Learning (GCHRL), sometimes referred to as Feudal RL (Dayan & Hinton, 1992; Levy et al., 2017; 2019; Nachum et al., 2018; Zhang et al., 2021). GCHRL approaches have shown that temporally extended policies can be learned by nesting multiple short sequence policies. In the Hierarchical Empowerment framework, this approach is implemented as follows. First, the designer implements $k$ levels of Goal-Conditioned Empowerment, in which $k$ is a hyperparameter set by the designer. That is, the agent will learn $k$ goal-conditioned and goal space policies $(\pi_{\theta_{0}},\mu_{\phi_{0}},\dots,\pi_{\theta_{k-1}},\mu_{\phi_{k-1}})$. Second, for all levels above the bottom level, the action space is set to the learned goal space from the level below. (The action space for the bottom level remains the primitive action space.) For instance, for an agent with $k=2$ levels of skills, the action space for the top level goal-conditioned policy $\pi_{\theta_{1}}$ at state $s$ will be set to the goal space output by $\mu_{\phi_{0}}(s)$. Further, the transition function at level $i>0$, $T_{i}(s_{t+1}|s_{t},a_{i_{t}})$, fully executes subgoal action $a_{i}$ proposed by level $i$. For instance, for a $k=2$ level agent, sampling the transition distribution at the higher level $T_{1}(s_{t+1}|s_{t},a_{1_{t}})$ for subgoal action $a_{1_{t}}$ involves passing subgoal $a_{1_{t}}$ to level 0 as the level 0’s next goal. The level 0 goal-conditioned policy then has at most $n$ primitive actions to achieve goal $a_{1_{t}}$. Thus for a 2-level agent, the top level goal-space $\mu_{\phi_{1}}$ and goal-conditioned policies $\pi_{\theta_{1}}$ will thus seek to learn the largest space of states reachable in $n$ subgoals, in which each subgoal can contain at most $n$ primitive actions. With this nested structure inspired by GCHRL, the time horizon of the computed empowerment can grow exponentially with the number of levels $k$, and no goal-conditioned policy is required to learn a long sequence of actions, making optimizing Goal-Conditioned Empowerment with RL more tractable. Figure 2 visualizes the proposed hierarchical architecture in one of our domains. Figure E in the Appendix provides some additional detail on the hierarchical architecture.

The Hierarchical Empowerment framework has two significant limitations as currently constructed. The first is that the framework can only learn long horizon skills in domains in which the state space has large regions of reachable states that can be covered by $d$-dimensional goal space boxes. This limitation results from the assumption that the learned goal space takes the form of a uniform distribution across a $d$-dimensional subset of the state space. In the framework, the uniform goal space distribution will only expand to include achievable states. If regions adjacent to the goal space are not reachable, the goal space will not expand. Consequently, if the state space contains large regions of unreachable states (e.g., a pixel space or a maze), the framework will not be able to learn a meaningful space of skills.

The second limitation is that a model of the transition dynamics is required. The most important need for the model is to simulate the temporally extended subgoal actions so that (state, subgoal action, next state, reward) tuples can be obtained to train the higher level goal-conditioned policies with RL. However, the existing empowerment-based skill-learning approaches that learn the variational distribution $q_{\psi}(z|s_{0},s_{n})$ also require a model in large domains to obtain unbiased gradients of the maximum likelihood objective. The maximum likelihood objective requires sampling large batches of skills $z$ from large sets of skill starting states $s_{0}$ and then executing the current skill-conditioned policy to obtain $s_{n}$. In large domains, a model of the transition dynamics is needed for this to be feasible.

Although there are some emerging benchmarks for unsupervised skill learning (Laskin et al., 2021; Gu et al., 2021; Kim et al., 2021), we cannot evaluate Hierarchical Empowerment in these benchmarks due to the limitations of the framework. The benchmarks either (i) do not provide access to a model of the transition dynamics or (ii) do not contain state spaces with large contiguous reachable regions that can support the growth of $d$-dimensional goal space boxes. For instance, some frameworks involve tasks like flipping or running in which the agent needs to achieve particular combinations of positions and velocities. In these settings, there may be positions outside the current goal space the agent can achieve, but if the agent cannot achieve those positions jointly with the velocities in the current goal space, the goal space may not expand.

Instead, we implement our own domains in the popular Ant and Half Cheetah environments. We implement the typical open field Ant and Half Cheetah domains involving no barriers, but also environments with barriers. All our domains were implemented using Brax (Freeman et al., 2021b), which provides access to a model of the transition dynamics. In all domains, the learned goal space is limited to the position of the agent’s torso ((x,y) positions in Ant and (x) in Half Cheetah). When limited to just the torso position, all the domains support large learned goal spaces. Also, please note the default half cheetah gait in Brax is different than the default gait in MuJoCo (Todorov et al., 2012). We provide more detail in section J of the Appendix.

Each experiment involves two phases of training. In the first phase, there is no external reward, and the agent simply optimizes its mutual information objective to learn skills. In the second phase, the agent learns a policy over the skills to solve a downstream task. In our experiments, the downstream task is to achieve a goal state uniformly sampled from a hand-designed continuous space of goals. Additional details on how the first and second phases are implemented are provided in section G of the Appendix.

We present our results in a few ways. Figure 3 provides charts for the phase two performances for most of the experiments. Results of the three versus four level Ant experiment is provided in section A of the Appendix. We also provide a video presentation showing sample episodes of the trained agents at the following url: https://www.youtube.com/watch?v=OOjtC30VjPQ. In addition, we provide before and after images of the learned goal spaces in section C of the Appendix. Section K of the Appendix discusses some of the ablation experiments that were run to examine the limitations of the framework.

The experimental results support both contributions of the framework. Goal-Conditioned Empowerment substantially outperforms DIAYN, HIDIO, and DADS in the Ant Field – Small and Half Cheetah – Small domains. None of the baselines were able to solve either of the tasks. As we describe in section I of the Appendix, in our implementation of DIAYN we tracked the scale of the learned variational distribution $q_{\phi}(z|s_{0},s_{t})$. As we hypothesized, the standard deviation of the variational distribution did not decrease over time, indicating that the the skills were not specializing and targeting specific states. On the other hand, per the performance charts and video, Goal-Conditioned Empowerment was able to learn skills to achieve goals in the phase two goal space.

The experimental results also support our claim that hierarchy is helpful for computing long-term empowerment. In both the two versus three level agents (i.e., one skill level vs. two skill level agents) and in the three versus four level comparison, the agents with the additional level outperformed, often significantly, the agent with fewer levels. In addition, the surface area of the skills learned by the four level Hierarchical Empowerment agents is significantly larger than has been reported in prior empowerment-based skill-learning papers. In our largest Ant domain solved by the four level agent, the phase 2 goal space with size 800x800 is over four orders of magnitude larger than the surface area reported in the DIAYN Ant task (4x4), and over 2 orders of magnitude larger than the area (30x30) partially achieved by Dynamics-Aware Unsupervised Discovery of Skills (DADS) (Sharma et al., 2020). Section B of the Appendix provides a visual of the differences in surface area coverage among Hierarchical Empowerment, DADS, and DIAYN. In terms of wall clock time, the longest goals in the 800x800 domain required over 13 minutes of movement by a reasonably efficient policy.