Skill-Based Reinforcement Learning with Intrinsic Reward Matching

The skill learning literature has long sought to design agents that autonomously acquire structured behaviors in new environments [33, 32, 26]. Recent work in competence-based unsupervised RL proposes generic objectives encouraging the discovery of skills representing diverse and useful behaviors [7, 29, 19]. A skill is defined as a latent code vector $z\in\mathcal{Z}$ that indexes the conditional policy $\pi(a|s,z)$. In order to learn such a policy, this class of skill pretraining algorithms maximizes the mutual information between sampled skills and their resulting trajectories $\tau$ [13, 6, 28] :

Commonly, methods consider $\tau=s$ or $\tau=(s,s^{\prime})$, and we leave more details in Appendix 8.2. Since the mutual information $I(s;z)$ is intractable to calculate in practice, competence-based methods instead maximize a variational lower bound proposed in [2] which is parameterized by a learned neural network function $q_{\phi}(\tau,z)$ called a skill discriminator. This discriminator, along with other terms independent of $z$, parameterizes an intrinsic reward that the skill policy $\pi(\cdot|s,z)$ maximizes during pretraining. Given an unseen task specification, the agent needs to infer which skill will finetune to solve the task with the fewest samples. For more detailed explanations of the mutual information decompositions of various skill discovery algorithms, refer to Appendix 8.2.

Pretrained Multitask Reward Functions: We observe that the intrinsic reward function learned during skill pretraining can be viewed as a multitask reward function, where the continuous skill code $z$ determines the task. In other words, we have some function:

where $\text{VLB}\leq I(\tau,z)$ is the variational lower bound proposed in [2] ($\tau$ is a trajectory representation such as $(s,s^{\prime})$). Since skill discovery algorithms aim to maximize $I(\tau,z)$, we can view its parameterized lower bound VLB as a multitask reward function: scoring transitions based on their alignment with a skill code [19].

We can formalize a general notion of reward function similarity by equivalent-policy invariant comparison (EPIC) as established in [12]. EPIC defines a distance metric between two reward functions such that similar reward functions induce similar optimal policies. We consider the case of action-independent reward:

where $D_{\rho}(X,Y)=\sqrt{\frac{1-\rho(X,Y)}{2}}$ is the Pearson distance between two random variables $X$ and $Y$, $s_{P},s_{P}^{\prime}$ are samples from the Pearson distribution $D_{P}$, and $S_{C},S_{C}^{\prime}$ are batches sampled from the Canonical distribution $D_{C}$. We compute the Pearson distance over Pearson samples $s_{P},s_{P}^{\prime}$, with additional canonicalization with batches $S_{c},S_{c}^{\prime}$ to ensure invariance over constant shifts and scaling. The canonicalized reward function is defined as:

where $R:S\times S\rightarrow\mathop{\mathbb{R}}$ is a reward function and $\gamma$ is the discount factor [12]. The expectation is taken over the Canonical distribution $D_{C}$; for simplicity, we sample these batches $S_{C}$, $S_{C}^{\prime}\sim D_{C}$ ahead of time. The canonicalization ensures invariance to reward shaping such that rewards that have different shaping but induce similar optimal policies are close in distance. In practice, the final term can be omitted as the Pearson correlation is invariant to constant shifts and scaling.

The EPIC psuedometric has typically been used for benchmarking algorithms and evaluating learned reward functions [23, 21]. To the best of our knowledge, IRM is the first to use EPIC in an optimization objective, particularly for complex tasks and larger state dimensions.

A multitask reward function that can supervise the learning of diverse behaviors is useful in its own right. However, in the case of skill-based RL, we have additionally learned a corresponding $\pi(a|s,z)$. Therefore, for any “task" that can be specified by our intrinsic reward function, we already have an optimal policy, so long as we condition on the corresponding skill. If we have learned a sufficiently diverse library of skills, we might expect that some of our skills share behavioral similarity to the optimal policy for the downstream task. It thus also holds that the corresponding intrinsic reward for that skill is a semantically similar task specification to the downstream task.

Given this interpretation of intrinsic reward, we posit that the task of identifying which our pretrained skills to apply to a downstream task can be reframed as inferring which task in our multitask reward function is most similar to the downstream task. Moreover, we should hope to find the skill code $z$ that produces the reward function most semantically aligned with the downstream task reward.

With this formalism, we can formulate the task inference problem as performing the following optimization:

in order to find $z^{*}$ most aligned with the task reward. Moreover, Equation 5 performs a minimization of a novel loss we name the EPIC loss with respect to the skill parameter $z$. By EPIC’s equivalence class invariance, we know that if the EPIC loss is small for some $z^{*}$, and $\pi(a|s,z^{*})$ is near optimal for $R^{\textit{int}}(\tau,z^{*})$, then $\pi(a|s,z^{*})$ approaches the optimal policy for the task as specified by $R^{ext}$. Notably, we require access to the task reward function $R_{ext}$ to compute the EPIC loss. Leveraging a known task reward function is a divergence from conventional skill selection methods.

Computing $R^{int}$ during reward matching During pretraining, for some methods such as [19, 29], we require negative samples in order to compute the variational objective in Equation 2 and avoid a degenerate optimization where all embedded trajectories have high similarity with all skills. However, during selection when skills are fixed, the negative sampling component amounts to a reward offset which does not impact the task semantics. Furthermore, since we may not in general have access to a large amount of negative samples on a given downstream task, we choose to simplify the objective to the following:

where $q_{\phi}$ is the skill discriminator. This parameterization of the intrinsic reward preserves the alignment semantics of VLB without the normalization by negative samples. For more details regarding the discriminator parameterization of the intrinsic reward for [19, 29] refer to Appendix 8.3 and Appendix 8.4.

We make a number of sample-based approximations of various unknown quantities in order to concretize the continuous optimization Equation 5 as a tractable loss minimization problem.

Canonical State Distribution Approximation: In order to canonicalize our reward functions, we estimate the expectation over the state and next state distributions with a sample-based average over 1024 samples. These distributions can be entirely arbitrary, though using heavily out-of-distribution samples with respect to pretraining can weaken the accuracy of the approximation. We choose to instantiate a uniform distribution bounded by known workspace constraints for both of these distributions.

Sampling Distribution for Pearson Correlation: We find that generating samples uniformly roughly within the environment workspace bounds, just as with the reward canonicalization, often leads to strong approximations. Furthermore, as both sample generation and relatively inexpensive function evaluation are independent of the online-finetuning phase, we can perform the full skill optimization as a self-contained preprocess to downstream policy adaptation without any environment samples. Rough knowledge of workspace bounds represents some amount of prior environment knowledge. We leave more general options such as sampling from a learned generative model over trajectories encountered during pretraining or sampling from saved pretraining data to future work. We ablate various sampling distribution choices in Table 7 and present the full algorithm in detail in Algorithm 1.

Many realistic downstream tasks derive additional complexity from temporally extended planning horizons. In contrast to hierarchical reinforcement learning (HRL) approaches, which aim to stitch together pretrained skills at the policy level with a higher-level manager policy, we can extend the task matching framework of IRM to efficiently solve the problem of skill sequencing, entirely doing away with the manager policy. Consider the long-horizon setting where we have a sequence of reward functions over task horizon $H$. Central to the finetuning problem is determining over what time intervals should potentially different pretrained skills be selected. In this work, we predetermine a fixed skill horizon $\lfloor H/N\rfloor$, where $N$ is the number of rewards. This skill horizon could also be specified as a parameter and learned from the task reward signal.

Next, in order to perform skill selection over each time interval, we perform the IRM algorithm in parallel or sequentially for each reward. We note the key assumption that IRM requires access to the reward functions for each of the subtasks. For example, for a sequential goal reaching task, we divide the episode into $N$ segments for each of the $N$ goals and corresponding rewards. We perform the IRM skill selection algorithm for each reward to select the optimal skill over each interval. After selecting the skills, we freeze our selections and finetune the skill policies jointly.

Reach Target We first evaluate IRM on the Reach Target task, where the Fetch robot is rewarded for reaching a target position. IRM outperforms or performs similarly to environment-rollout methods while requiring no environment samples to perform skill selection. As shown in Table 1, the random skill policy performs particularly poorly and with very high variance relative to the IRM and environment-rollout based methods. Moreover, appropriate skill selection is required for strong zero-shot performance as certain skills obtain much higher rewards than others. Reward relabelling fails to consistently select optimal skills across the benchmark, likely because its options are limited to the finite set of skills sampled during pretraining. IRM by contrast leverages continuous optimization in the infinite skill space to find the best skill for the task. Figure 3 shows the finetuning performance of the methods on the downstream task reward.

Push Block to Goal Next, we evaluate IRM on a more complex manipulation task involving pushing a block to a goal position. We report the zero-shot IRM skill selection performance in Table 1 and finetuning performance in Figure 3. This more complex task similarly benefits from bootstrapping the appropriate pretrained skill policy as evidenced by the performance gap of the selection based methods over random skill selection. We remark that even for more complex manipulation tasks, IRM is robust in consistently guiding optimal skill selection without requiring any interaction with the environment. Although Env Rollout CEM is the strongest interaction baseline in terms of zero-shot reward, it far exceeds a reasonable computational budget of interactions entirely on skill selection. For illustrative purposes, we display the plot although heavily shifted.

Franka Kitchen Manipulation Lastly, we evaluate IRM on 4 manipulation tasks in the Franka Kitchen benchmark [10]. We make no changes to the default environment rewards for each task, which do not include any shaping for reaching to the correct object, and thus present a harder exploration problem. We similarly observe favorable zero-shot performance of the IRM method, in the case of hinge cabinet and slide cabinet, solving the task zero-shot without any finetuning (Table 3, Figure 3). This zero-shot adaptation behavior presents a compelling means for agents to immediately deploy optimal policies in highly sample-sensitive applications.

Long-Horizon Manipulation Building on the results in Section 4.1, we demonstrate that IRM fully generalizes to solving long-horizon tasks in the setting of tabletop manipulation. During the unsupervised pretraining phase, skill discovery methods can acquire useful skills such as directional block pushing or pushing the block to certain spatial locations. We show that IRM can intelligently select a sequence of such skills to finetune via reward matching, avoiding learning a hierarchical manager policy that finetunes at the policy level.

For the Fetch Reach Sequential tasks, we consider an extended horizon where the agent is tasked with reaching a sequence of goals in a particular order. For the Fetch barrier tasks, we consider the environments such as depicted in Figure 2, where the agent must navigate around a barrier introduced during the finetuning phase in order to reach the goal. The tunnel environments consists of a goal enclosed in a tunnel; the agent must navigate around and into the tunnel to complete the task. We compare IRM methods to an environment rollout baseline (Env Seq) and a hierarchical RL baseline (HRL). The ‘IRM Seq’ methods select skills based on each defined sub-task’s reward function according to the IRM optimization scheme. ‘Env Seq’ chooses the best combination of skills based on extrinsic reward from rollouts. ‘HRL’ is initialized with random skills and simultaneously optimizes a manager policy over skills and the skill policies themselves. In both settings, IRM methods outperform the environment rollout method and in some cases the HRL method in identifying (Table 2) and finetuning skills (Figure 3). Implementation details are provided in Appendix 8.9.

Matching Metric Ablations We validate the importance of employing the EPIC pseudometric for formulating the matching loss by ablating its contribution against more naive selections in Table 5. L1 and L2 losses are common metrics in supervised regression problems but are poor choices for comparing task similarity with rewards. Moreover, rewards can have arbitrary differences in scaling and shaping that L1 and L2 are not invariant to. To strengthen these comparisons, we include a learned reward scaling parameter for L1 and L2 and similarly observe that EPIC is a superior matching metric.

Skill Discovery Algorithm Ablations IRM is fully general to any mutual information maximization based, RL pretraining algorithm as shown in Table 4. We validate on the Fetch Reach task that IRM CEM and IRM Rand convincingly outperform all episode rollout baselines in zero-shot episode reward.

Markov Decision Process: The goal of reinforcement learning is to maximize cumulative reward in an uncertain environment it interacts with. The problem can be modelled as a Markov Decision Process (MDP) defined by $(\mathcal{S},\mathcal{A},\mathcal{P},r,\gamma)$, where $\mathcal{S}$ is the set of states, $\mathcal{A}$ is the set of actions, $\mathcal{P}$ is the transition probability distribution, $r$ is the reward function and $\gamma$ is the discount factor.

Unsupervised Skill Discovery: In competence-based unsupervised RL the aim is to learn skills that generate diverse and useful behaviors [7]. The broad aim is to learn policies that are skill-conditioned and generalizable. Formally, we also learn skills $z\in\mathcal{Z}$ and take actions according to $a\sim\pi(\cdot|s,z)$. As an illustrative example, applying this formalism to the Mujoco Walker domain, we might hope to find a skill-conditioned policy and skills $z_{\mathrm{walk}},z_{\mathrm{run}}$ such that $\pi(\cdot|s,z_{\mathrm{walk}})$ makes the agent walk, while $\pi(\cdot|s,z_{\mathrm{run}})$ makes it run. Further, if we allow for continuous skills, we can also imagine being able to use the policy to “jog" at different speeds by interpolation the $z_{\mathrm{walk}}$ and $z_{\mathrm{run}}$ skills. That is, taking $z_{\mathrm{jog}}^{\alpha}=\alpha\cdot z_{\mathrm{walk}}+(1-\alpha)\cdot z_{\mathrm{run}}$ should, intuitively, yield a policy $\pi(\cdot|s,z_{\mathrm{jog}}^{\alpha})$ that makes the agent jog at speed dictated by the parameter $\alpha$.

Finetuning Pretrained Skills: With a skill-conditioned policy $\pi(\cdot|s,z)$, an agent needs to infer which skill to index for a downstream task (e.g. identifying if it needs to use $z_{\mathrm{walk}}$ or $z_{\mathrm{run}}$) during finetuning. This is a relatively under-explored area, with the most universal approach being a coarse, discretized grid search. Least squares regression has also been investigated in the context of successor features [22].

Competence-based skill discovery algorithms aim to maximize the mutual information between trajectories and skills:

Since the mutual information $I(s;z)$ is intractable to calculate, in practice, competence-based methods maximize a variational lower bound. Many mutual information maximization algorithms, such as Variational Intrinsic Control [13] and Diversity is All You Need [6], use the estimate $I(\tau;z)=\mathcal{H}(z)-\mathcal{H}(z|\tau)$. Other competence-based methods, such as Dynamics-Aware Unsupervised Discovery of Skills [28], Active Pretraining with Successor Features [22], and Contrastive Intrinsic Control (CIC) [19], maximize a lower bound for $\mathcal{H}(\tau)-\mathcal{H}(\tau|z)$.

While the decompositions of the mutual information objective are equivalent, algorithms make different design choices regarding how to approximate entropy, represent trajectories, and embed skills. These choices affect the distillation of skills: for instance, without explicit maximization of $\mathcal{H}(\tau)$ in the decomposition of mutual information, behavioral diversity may not be guaranteed when the state space is much larger than the skill space [19].

Contrastive Intrinsic Control (CIC) [19] is a state of the art algorithm for competence-based skill discovery. CIC maximizes a lower bound for $I(\tau;z)=\mathcal{H}(\tau)-\mathcal{H}(\tau|z)$ through a particle estimator for $\mathcal{H}(\tau)$ and a contrastive loss from Contrastive Predictive Coding (CPC) [34] for $\mathcal{H}(\tau|z)$. The lower bound for $I(\tau;z)$ is:

where $\mathcal{H}_{\text{particle}}(\tau)\propto\sum_{i=1}^{n}\log||h_{i}-h_{i}^{*}||$, $h_{i}^{*}$ is the k-Nearest Neighbors embedding, $N_{k}$ is the number of k-NNs used to approximate entropy, and $N-1$ is the number of negative samples.

We additionally use Dynamics-Aware Unsupervised Discovery of Skills (DADS) [29] for skill discovery, as it is one of the few skill discovery algorithms to successfully scale up to continuous skills. DADS maximizes a lower bound for $I(\tau;z)=\mathcal{H}(\tau)-\mathcal{H}(\tau|z)$ through learning skill-conditioned transition distributions. The lower bound for $I(\tau;z)$ is:

For our experiments, we reimplement the on-policy DADS algorithm in PyTorch. We follow the default hyperparameters and train for 20 million environment steps, per [29].

For our Fetch Reaching environment, we use the Gym Robotics Fetch environment [5]. We set the time limit to 200. For the fetch push environment, we partition the continuous action space into 4 actions, which involve pushing the block forward, backward, left, and right. We set the time limit to 10 for skill learning.

We evaluate sequential skill selection on Fetch Reach and Fetch Push. For the Fetch Push agent, we have a Barrier and Tunnel environment, each with 3 waypoints, depicted in Figure 2. We fix a time horizon of 30 pushes per waypoint. For Fetch Barrier, the easy task consists of the first two waypoints (L shape), and the hard task consists of all three (U shape). For Fetch Tunnel, the easy task consists of the last two (into the tunnel), and the hard task consists of all three waypoints (side of the tunnel, and then into the tunnel). For Fetch Reach, we consider 2 waypoints and a time horizon of 25 for each waypoint.

Our plane environment is a 2D world with observations in [-128, 128] x [-128, 128] and continuous actions in [-10, 10] x [-10, 10].

The Franka Kitchen Environment consists of a 9 dof Franka Arm in a Kitchen Environment with a variety of rigid and articulated objects. The benchmark specifies several goals with rewards for reaching those goals. The 7d action space includes the end-effector 3d translation, 3d rotation, and a gripper open/close dimension. The state includes the joint positions corresponding to each dof of the arm in addition to the joint positions of articulated bodies and 7-dimensional position and orientation information for rigid objects in the scene. The Kettle task involves moving the kettle to a target burner location on the stove. The Hinge and Slide Cabinet tasks involve opening cabinet drawers with revolute and prismatic joints respectively. The Light Switch task involves flipping the switch for the stove light. We use an episode length of 1000 for solving the tasks.

For the Fetch Reach environment, we use a skill dimension of 8 and a discriminator MLP hidden dimension of 64. We use an alpha value of 0 for entropy weighting. For the Fetch Push environment, we use a skill dimension of 16 and a discriminator MLP hidden dimension of 16. We use an alpha value of 0 for entropy weighting. For the Franka Kitchen environment we use an alpha value of 0.1 for entropy weighting along with a skill dimension of 2 and an MLP hidden dimension of 64. For all environments, we use a replay buffer size of 100k.

IRM CEM and Env Rollout CEM are trained for 5 iterations with 1000 samples at each iteration and 100 elites selected each iteration. Env Rollout CEM consumes the entire downstream finetuning budget on just skill selection. For illustrative purposes, we start its plot at 50k steps to show that finetuning still occurs, however, sample-inefficiency suffers due to excessive rollouts for skill selection. This problem only worsens for long time horizons. IRM Gradient Descent is trained for 5000 steps with a learning rate of 5e-3 and initialized at the skill vector of all 0.5s. IRM Random selects 100 random skills. Env Rollout trials 10 random skills for a fully episode. Grid Search coarsely trials 10 skills from the skill of all 0s to the skill of all 1s as in [20].

The planar goal reaching task consists of a simple 2D plane with a point with a 2D Cartesian state space that can displace in the x and y coordinates with a 2D action space. Skills learned tend to span the 2D space reaching to diverse locations distributed broadly across the environment. We show some sample zero-shot skill selection results over three different skill dimensions in Figure 6.

For sequential skill selection, we compare IRM Sequential and Environment Sequential skill selection. IRM Sequential consists of an iterative process. The first skill is chosen entirely free of environment samples, exactly identical to the single-skill tasks. Once the first skill is chosen, we roll out a trajectory with the skills we have chosen so far and use the latter half of the trajectory as the Pearson samples for our EPIC loss. We use Gaussian noise with variance 1 for our Canonical samples as described in Appendix 8.12.2. At each step of the skill selection process, we use the corresponding IRM optimization methods.

For our Environment Sequential skill selection method, we select skills iteratively as well. For each waypoint or subtask, we randomly sample $N$ skills and commit to the best, where $N=10/\text{n\_subtasks}$.

In order to validate the benefits of IRM’s offline skill selection, we compare against a baseline that leverages a conventional hierarchical RL algorithm to solve long-horizon, sequential tasks. We instantiate a TD3 manager agent that outputs into a skill action space from state input at a temporally abstract timescale. As in the IRM setup, this timescale is fixed to align with the changes in reward to encourage the manager to change its skill prediction according to the change in the reward semantics. The manager’s is then inputted to the low-level pretrained skill policy which is rolled out over many steps with the skill fixed. Both the manager policy and the low-level policy weights are updated during finetuning. The manager agent is randomly initialized such that its initial skill prediction is random.

We thank NVIDIA for providing compute resources through the NVIDIA Academic DGX Grant. We used a cluster of 8 NVIDIA A100s for experimentation.