\method: Efficient Policy Learning by Extracting Transferrable Robot Skills from Offline Data

Feature extraction. We leverage vision-language models (VLMs), trained to align large corpora of images with natural language descriptions [23, 47, 48, 49], to extract high-level features used to label skills. Although our approach does not require the use of language, we utilize VLMs because, as VLMs were trained to align images with language, VLM image embeddings represent a semantically aligned embedding space. However, one main issue precludes the naïve application of VLMs in robotics. In particular, VLMs do not inherently account for object variations or robot arm starting positions across images [50, 24, 51, 52]. But in robot manipulation, high-level behaviors should be characterized by changes in arm and object positions across a trajectory—picking up a cup should be considered the same skill regardless of if the cup is to the robot’s left or right. Our initial experiments of using the embeddings directly resulted in skills specific to one type of environment layout or object. Therefore, to capture high-level behaviors, we use trajectory-level embedding differences by taking the difference of each VLM image embedding with the first one in the trajectory:¹¹1To ensure that each timestep has an embedding, we assign embedding ${e}_{1}$ to be identical to ${e}_{2}$.

Skill label assignment. After creating embeddings ${e}_{t}$ for each image $s_{t}$, we assign skill labels in an unsupervised manner based on these features. Inspired by classical algorithms from speaker diarization, a long-studied problem in speech processing where the objective is to assign a “speaker label” to each speech timestep [53], we first perform unsupervised clustering with K-means on the entire dataset of embedding differences $e_{i}$ to assign per-timestep skill labels (the label is the cluster ID), then we smooth out the label assignments with a simple median filter run along the trajectory sequence to reduce the frequency of single or few-timestep label assignments. See Figure 3 for a visual demonstration of this process.

In summary, we first extract observation embedding difference features with a VLM and then perform unsupervised K-means clustering to obtain skill labels for each trajectory timestep. This forms the skill-labeled dataset $\mathcal{D}_{d}=\{\tau_{d}^{1},\tau_{d}^{2},...\}$, where each $\tau_{d}$ is a trajectory of sequential $(s,a)$ tuples that all belong to one skill $d$. Next, we perform skill learning on $\mathcal{D}_{d}$.

We aim to learn a discrete set of skills, parameterized by continuous arguments, similar to a functional API over skills (see Figure 2 middle). Therefore, we train a generative skill decoder $p_{a}(\bar{a}\mid z,d)$ to convert a discrete skill choice $d$ and a continuous argument for that skill, $z$, into an action sequence. As alluded to in Section 3, we build upon SPiRL by Pertsch et al. [6]. However, they train their decoder to decode fixed-length action sequences from a single continuous latent $z$. In contrast, we automatically extract a set of variable-length skill trajectories with labels denoted $d$ and parameterize each skill by a learned, continuous latent argument $z$.²²2To simplify notation, we use $z$ for both our method and SPiRL. However, it is important to note that $z$ uniquely determines the skill in SPiRL, while $z$ denotes a continuous latent argument in our method.

We train an autoregressive VAE [54] consisting of the following learned neural network components: a skill argument encoder $q(z\mid\bar{a},d)$ map** to a continuous latent $z$ conditioned on a discrete skill choice $d$ and an action sequence $\bar{a}$, and an autoregressive skill decoder $p_{a}(\bar{a}\mid z,d)$ conditioned on the latent $z$ and the discrete skill choice $d$. Because the action sequence $\bar{a}$ can be of various lengths, the decoder also learns to produce a continuous value $l$ at each autoregressive timestep representing the proportion of the skill completed at the current action. This variable is used during online RL to stop the execution of the skill when $l$ equals 1 (see Section B.1 for further details).

Recall that SPiRL also trains a skill prior network $p_{z}(z\mid s)$ that predicts which $z$ is useful for an observation $s$; this prior is used to guide a high-level policy toward selecting reasonable $z$ while performing RL. In contrast with SPiRL where $z$ uniquely represents a skill, we train two prior networks, one to guide the selection of the skill $d$, $p_{d}(d\mid s)$, and one to guide the selection of its argument $z$ given $d$, $p_{z}(z\mid s,d)$. These are trained with the observation from the first timestep of the sampled trajectory, $s_{1}$, to be able to guide a skill-based policy during online RL in choosing $d$ and $z$. Our full objective for training this VAE is to maximize the following:

where the stop-gradient $\mathbf{sg}(\cdot)$ prevents prior losses from influencing the encoder. The first two terms are the $\beta$-VAE objective [55]; the last two train priors to predict the correct skill $d$ and continuous argument $z$ given $s_{1}$.

Additional fine-tuning. On extremely challenging transfer scenarios, demonstrations may still be needed to warm-start reinforcement learning [56]. \method can also flexibly be applied to this setting by using the same K-means clustering model from Section 4.1, which was trained to cluster $\mathcal{D}_{d}$, to assign skill labels to an additional, smaller demonstration dataset. After pre-training on $\mathcal{D}_{d}$, we then fine-tune the entire model on that labeled demonstration dataset before performing RL.

Finally, we describe how we perform RL for new tasks by training a skill-based policy to select skills and their arguments to solve new tasks. See Figure 2, right, for an overview of online RL.

Policy parameterization. After pre-training the decoder $p_{a}(\bar{a}\mid z,d)$, we treat it as a frozen lower-level policy that a learned skill-based policy can use to interact with a new task. Specifically, we train a skill-based policy $\pi(d,z\mid s)$ to output a $(d,z)$ tuple representing a discrete skill choice and its continuous argument. We parameterize this policy as a product of two policies: $\pi(d,z\mid s)=\pi_{d}(d\mid s)\pi_{z}(z\mid s,d)$ so that each component of $\pi(d,z\mid s)$ can be regularized with our pre-trained priors $p_{d}(d\mid s)$ and $p_{z}(z\mid s,d)$. Intuitively, this parameterization separates decision-making into what skill to use and how to use it. The complete factorization of the skill-based policy follows:

Policy learning. We can train the skill-based policy with online data collection using any entropy-regularized RL algorithm, such as SAC [57], where we regularize against the skill priors instead of against a max-entropy uniform prior. Because we have factorized $\pi(d,z\mid s)$ into two separate policies, we can easily regularize each with the priors trained in Section 4.2. The training objective for the policy with SAC is to maximize over $\pi_{d},\pi_{z}$:

where $\alpha_{z}$ and $\alpha_{d}$ control the prior regularization weights. The critic objective is also correspondingly modified (see Appendix Algorithm 4).

In summary, \method first extracts a set of discrete skills from offline image-action data (Section 4.1), then trains an action decoder to take low-level actions in the environment conditioned on a discrete skill and continuous latent (Section 4.2), and finally performs prior-guided reinforcement learning over these skills online in the target environment to learn new tasks (Section 4.3). See Algorithm 1 (appendix) for the pseudocode and Section B.1 for additional implementation details.

We evaluate \method on two long-horizon, continuous-control, robotic manipulation domains: Franka Kitchen [58] and LIBERO [59]. All environments use image observations and sparse rewards. For both Franka Kitchen and LIBERO, our method \method uses the R3M VLM [47] and K-means with $K=8$ for offline skill extraction (Section 4.1). We list specific details below; see Section B.3 for more.

Franka Kitchen: This environment, originally from Gupta et al. [22], Fu et al. [58] contains a Franka Panda arm operating in a kitchen environment. Similarly to Pertsch et al. [6], we test transfer learning of a sequence of 4 subtasks never performed in sequence in the dataset. Agents are given a reward of 1 for completing each subtask.

LIBERO: LIBERO [59] consists of a Franka Panda arm interacting with many objects and drawers. We test transfer to four task suites, LIBERO-{Object, Spatial, Goal, 10} consisting of 10 unseen environments/tasks each, spanning various transfer scenarios (40 total tasks). LIBERO tasks are language conditioned (e.g., “turn on the stove and put the moka pot on it”); for pre-training and RL, we condition all methods on the language instruction. Due to LIBERO’s difficulty [60], for all pre-trained methods, we first fine-tune to a provided additional target task dataset with 50 demos per task before performing RL. During RL, we fine-tune on all tasks within each suite simultaneously. To the best of our knowledge, we are the first to report successful RL results on LIBERO tasks.

Baselines and Comparisons. We compare against: (1) an oracle (RAPS [9]), which is given ground truth discrete skills with continuous input arguments designed by humans specifically for the Franka Kitchen environment; (2) methods that pre-train with the same data—namely SPiRL [6] which extracts sequences of fixed-length random action trajectories as skills, and BC, behavior cloning using the same offline data but no temporally extended skills; and (3) SAC [57], i.e., RL without any offline data. See Appendix B for further implementation details. Unless otherwise stated, all reported experimental results are means and standard deviations over 5 seeds.

We first test \method’s ability to discover meaningful, well-aligned skills during skill extraction. In Figure 4, we plot K-means ($K=8$) skill assignments in Franka Kitchen. We project VLM embedding differences down to 2-D with PCA for visualization. These skill assignments demonstrate that unsupervised clustering of VLM embedding differences can create distinctly separable clustering assignments. For example, skill 4 (Figure 4, top left) demonstrates a cabinet opening behavior. See additional visualizations for all environments in Section C.1. We also analyze quantitative clustering statistics in Section C.2. Next, let’s see how these skills help with learning new tasks.

We investigate the ability of all methods to transfer to new tasks in Figure 5. In Kitchen, \method matches the oracle performance while being 10x more sample-efficient than SPiRL, with SPiRL needing 3M timesteps to reach the same performance of \method at ~300k. While SPiRL performed well in Franka Kitchen in their paper using ground truth environment states instead of RGB images, it struggles in our much more challenging image-based experiments. In all LIBERO task suites, \method performs best; it outperforms SPiRL the most in LIBERO-10, the task suite with the longest-horizon tasks. Meanwhile, SAC and BC perform poorly, indicating our tasks are difficult to solve with just standard RL or offline data without skills.

This improvement of our method over SPiRL comes from the semantically aligned, discrete skill-space that \method learns instead of SPiRL’s continuous, randomly assigned skill-space. For example, to open a drawer, \method can learn to select a single discrete drawer-opening skill with minor argument modifications when its gripper is near that drawer. With SPiRL, the robot must memorize the continuous skill representing each type of drawer-opening behavior for each way to open a drawer; these continuous skills must also be distinguished from others for completely different behaviors. \method also enables easier exploration later in the task; the policy can easily try the same discrete skill if the robot hand is near a different drawer that needs to be opened. For further analysis, see Appendix E. Next, we perform an ablation study on \method components.

VLMs. We first ablate the use of VLMs from selecting features for clustering. Therefore, we compare against Action, where skill labels are generated by clustering robot action differences. We also compare against State where skills are labeled by clustering ground truth state differences (e.g., robot joints, states of all objects). State represents an oracle scenario as ground truth states of all relevant objects are difficult to obtain in the real world. We plot results in Franka Kitchen in Figure 6. \method with VLM-extracted skills performs best, as both ground truth state and raw environment action differences can be difficult to directly obtain high-level, semantically meaningful skills from.

Number of Clusters. Finally, we ablate the number of K-means clusters. Intuitively, too few or too many clusters can make downstream RL more difficult as it trades off the ease of the RL policy selecting the correct discrete skill and the difficulty of deciding the correct continuous argument for that skill. In Figure 7, we plot average returns at 1M timesteps of \method in Kitchen with $K=3,5,8,15$. Final returns are relatively constant, with performance drop** only at $K=15$. This indicates that \method is robust to the number of discrete skills unsupervisedly discovered.