S KILL D ECISION T RANSFORMER
A BSTRACT
Recent work has shown that Large Language Models (LLMs) can be incredibly effective for offline reinforcement learning (RL) by representing the traditional
RL problem as a sequence modelling problem (Chen et al., 2021; Janner et al., 2021).
However many of these methods only optimize for high returns, and may not extract much information from a diverse dataset of trajectories.
Generalized
Decision Transformers (GDTs) (Furuta et al., 2021) have shown that by utilizing future trajectory information, in the form of information statistics, can help extract more information from offline trajectory data.
Building upon this, we propose
Skill Decision Transformer (Skill DT).
Skill DT draws inspiration from hindsight relabelling (Andrychowicz et al., 2017) and skill discovery methods to discover a diverse set of primitive behaviors , or skills.
We show that Skill DT can not only perform offline state-marginal matching (SMM), but can discovery descriptive behaviors that can be easily sampled.
Furthermore, we show that through purely reward-free optimization, Skill DT is still competitive with supervised offline RL approaches on the D4RL benchmark.
I NTRODUCTION
Reinforcement Learning (RL) has been incredibly effective in a variety of online scenarios such as games and continuous control environments (Li, 2017).
However, they generally suffer from sample inefficiency, where millions of interactions with an environment are required.
In addition, efficient exploration is needed to avoid local minimas (Pathak et al., 2017; Campos et al., 2020).
Because of these limitations, there is interest in methods that can learn diverse and useful primitives without supervision, enabling better exploration and re-usability of learned skills (Eysenbach et al., 2018; Strouse et al., 2021; Campos et al., 2020).
However, these online skill discovery methods still require interactions with an environment, where access may be limited.
This requirement has sparked interest in Offline RL, where a dataset of trajectories is provided.
Some of these datasets (Fu et al., 2020) are composed of large and diverse trajectories of varying performance, making it non trivial to actually make proper use of these datasets; simply applying behavioral cloning (BC) leads to sub-optimal performance.
Recently, approaches such as the Decision Transformer (DT) (Chen et al., 2021) and the Trajectory Transformer (TT) (Janner et al., 2021), utilize Transformer architectures (Vaswani et al., 2017) to achieve high performance on Offline RL benchmarks.
Furuta et al.
(2021) showed that these methods are effectively doing hindsight information matching (HIM), where the policies are trained to estimate a trajectory that matches given target statistics of future information.
The work also generalizes DT as an information-statistic conditioned policy, Generalized Decision Transformer (GDT).
This results in policies with different capabilities, such as supervised learning and State Marginal Matching (SMM) (Lee et al., 2019), just by simply varying different information statistics.
In the work presented here, we take inspiration from the previously mentioned skill discovery methods and introduce Skill Decision Transformers (Skill DT), a special case of GDT, where we wish to condition action predictions on skill embeddings and also future skill distributions.
We show that Skill DT is not only able to discovery a number of discrete behaviors, but it is also able to effectively match target trajectory distributions.
Furthermore, we empirically show that through pure unsupervised skill discovery, Skill DT is actually able to discover high performing behaviors that
Our method predicts actions, conditioned by previous states, skills, and distributions of future skills.
Empirically, we show that Skill DT can not only perform SMM on target trajectories, but can also match or achieve higher performance on D4RL benchmarks (Fu et al., 2020) compared to other state-of-the-art offline RL approaches.
Skill DT also has the added benefit of using discrete skills, which are useful for easily sampling diverse behaviors.
R ELATED W ORK 2. S KILL D ISCOVERY
Many skill methods attempt to learn a latent skill conditioned policy π ( a | s, z ) , where state s ∼ p ( s ) and skill z ∼ Z , that maximizes mutual information between S and Z (Gregor et al., 2016;
Sharma et al., 2019; Eysenbach et al., 2018).
Another way of learning meaningful skills is through variational inference, where z is learned via a reconstruction loss (Campos et al., 2020).
Explore,
Discover and Learn (EDL) (Campos et al., 2020) is an approach, which discovers a discrete set of skills by encoding states via a VQ-VAE: p ( z | s ) , and reconstructing them: p ( s | z ) .
We use a similar approach, but instead of reconstructing states, we utilize offline trajectories and optimize action reconstruction directly ( p ( a | s, z ) ).
Since our policy is autoregressive, our skill encoding actually takes into account temporal information, leading to more descriptive skill embeddings.
Offline
Primitive Discovery for Accelerating Offline Reinforcement Learning (OPAL) (Ajay et al., 2020), also discovers offline skills temporally, but instead uses a continuous distribution of skills.
These continuous skills are then sample by a hierarchical policy that is optimized by task rewards.
Because our approach is completely supervised, we wish to easily sample skills.
To simplify this, we opt to use a discrete distribution of skills.
This makes it trivial to query the highest performing behaviors, accomplished by just iterating through the discrete skills.
2. S TATE M ARGINAL M ATCHING
State marginal matching (SMM) (Lee et al., 2019) involves finding policies that minimize the distance between the marginal state distribution that the policy represents p π ( s ) , and a target distribution p ∗ ( s ) .
These objectives have an advantage over traditional RL objectives in that they do not require any rewards and are guided towards exploration (Campos et al., 2020).
CDT has shown impressive SMM capabilities by utilizing binned target state distributions to condition actions in order to match the given target state distributions.
However, using CDT in a real environment is difficult because target distributions must be provided, while Skill DT learns discrete skills that can be sampled easily.
Also, CDT requires a low dimensional state space, while Skill DT in theory can work on any type of input as long as it can be encoded effectively into a vector.
P RELIMINARIES
In this work, we consider learning in environments modelled as Markov decision processes (MDPs), which can be described using varibles ( S, A, P, R ) , where S represents the state space, A represents the action space, and P ( s t +1 | s t , a t ) represents state transition dynamics of the environment.
3. G ENERALIZED D ECISION T RANSFORMER
The Decision Transformer (DT) (Chen et al., 2021) represents RL as a sequence modelling problem and uses a GPT architecture Alec Radford & Sutskever (2018) to predict actions autoregressively.
Specifically, DT takes in a sequence of RTGs, states, and actions, where R t = (cid:80) Tt r t , and trajectory τ =
( R 0 , s 0 , a 0 , ..., R | τ | , s | τ | , a | τ | ) .
DT uses K previous tokens to predict a t with a deterministic policy which is optimized by a mean squared error loss between target and predicted actions.
For evaluation, a target return ˆ R target is provided and DT attempts to achieve the targeted return in the actual environment.
Furuta et al.
(2021) introduced a generalized version of DT, Generalized
Decision Transformer (GDT).
GDT provides a simple interface for representing a variety of different objectives, configurable by different information statistics (for consistency, we represent variations of GDT with π gdt ):
Generalized Decision Transformer (GDT): π
[Equation]
) where τ = trajectory, I ϕ = information statistics
Decision Transformer (DT):
[Equation]
s [Equation] , a t ) , where I ϕdt ( τ t ) =
(cid:80)
Tt
[Equation]
∗ r Categorical Decision Transformer (CDT): π gdtcdt ( a t | I ϕcdt ( τ 0 ) , s 0 , a 0 , ..., I ϕcdt ( τ t ) , s t , a t ) , where I ϕcdt ( τ t ) = histogram ( s t , ..., s T )
CDT is the most similar to Skill DT – CDT captures future trajectory information using future state distributions, represented as histograms for each state dimension, essentially binning and counting the bin ids for each state dimension.
Skill DT instead utilizes learned skill embeddings to generate future skill distributions, represented as histograms of full embeddings.
In addition, Skill DT also makes use of the representation learnt by the skill embedding by also using it in tandem with the skill distributions.
S KILL D ECISION T RANSFORMER 4. F ORMULATION
Our Skill DT architecture is very similar to the original Decision Transformer presented in Chen et al.
While the classic DT uses summed future returns to condition trajectories, we instead make use of learned skill embeddings and future skill distributions , represented as a histogram of skill embedding indices, similar to the way Categorical Decision Transformer (CDT) (Furuta et al., 2021) utilizes future state counts.
Formally, we wish to learn a policy:
[Equation]
where K is the context length, and θ are the learnable parameters of the model.
States are encoded as skill embeddings ˆ z t , which are then quantized using a learned codebook of embeddings z = argmin n ||
ˆ
z − z n || 22 .
The future skill distributions are represented as the normalized histogram of summed future one hot encoded skill indices: Z t = (cid:80)
Tt one hot ( z t ) .
Connecting this to GDT, our policy can be viewed as:
[Equation] ϕskill [Equation] ϕskill
[Equation] t [Equation] ( z t , ..., z T ) , z t ).
4.1. H INDSIGHT S KILL R E - LABELLING
Hindsight experience replay (HER) is a method that has been effective in improving sampleefficiency of goal-oriented agents (Andrychowicz et al., 2017; Rauber et al., 2017).
The core concept revolves around goal relabelling , where trajectory goals are replaced by achieved goals vs. inteded goals.
This concept of re-labelling information has been utilized in a number of works (Ghosh et al., 2019; Zheng et al., 2022; Faccio et al., 2022), to iteratively learn an condition predictions on target statistics.
Bi-Directional Decision Transformer (BDT) (Furuta et al., 2021), utilizes an anticausal transformer to encode trajectory information, and passes it into a causal transformer action predictor.
At every training iteration, BDT re-labels trajectory information with the anti-causal transformer.
Similarly, Skill DT re-labels future skill distributions at every training iteration.
Because the skill encoder is being updated consistently and skill representations change during training, the re-labelling of skill distributions is required to ensure stability in action predictions.
4. A RCHITECTURE
VQ-VAE Skill Encoder .
Many previous works have represented discrete skills as categorical variables, sampled from a categorical distribution prior (Strouse et al., 2021;
Eysenbach et al., 2018).
VQ-VAEs (van den Oord et al., 2017) have shown impressive capabilities with discrete variational inference in the space of computer vision (Razavi et al., 2019; Esser et al., 2020), planning (Ozair et al., 2021), and online skill discovery (Campos et al., 2020).
Because of this, we use a VQ-VAE to quantize encoded states into a set of continuous skill embeddings.
To get skill embeddings from actual trajectory data, we feed states into a simple two layer MLP and output the nearest embedding in the VQ-VAE codebook for each state.
To ensure that our MLP output is close to the nearest embedding in the codebook, we minimize the loss:
V QLOSS ( z, ˆ z ) =
MSE ( z, ˆ z ) , where ˆ z is the output of the MLP encoder and z is the nearest embedding in the VQ-VAE codebook.
Optimizing this loss minimizes the distance of our skill encoder with their corresponding nearest embeddings.
This is analagous to clustering, where we are trying to minimize the distance between datapoints and their actual cluster centers.
In practice to stabilize training, we optimize the VQ-VAE using an exponential moving average, as detailed in Lai et al.
Causal Transformer .
The Causal Transformer portion of Skill DT shares a similar architecture to that of the original DT (Chen et al., 2021), utilizing a GPT (Alec Radford & Sutskever, 2018) model.
It takes in input the last K states s t − K : t , skill encodings z t − K : t , and future skill embedding distributions Z t − K : t .
As mentioned above, the future skill embedding distributions are calculated by generating a histogram of skill indices from timestep t : T , and normalizing them so that they add up to 1.
For states and skill embedding distributions, we use learned linear layers to create token embeddings.
To capture temporal information, we also learn a timestep embedding that is added to each token.
Note that we don’t tokenize our skill embeddings because we want to ensure that we don’t lose important skill embedding information.
It’s important to note that even though we don’t add timestep embeddings to the skill embeddings, they still capture temporal behavior because the attention mechanism (Vaswani et al., 2017) of the causal transformer attends the embeddings to temporally conditioned states and skill embedding distributions.
The VQ-VAE and Causal Transformer components are shown visually in Fig.
4. T RAINING P ROCEDURE
Training Skill DT is very similar to how other variants of GDT are trained (CDT, BDT, DT, etc.).
First, before every training iteration we re-label skill distributions for every trajectory using our VQVAE encoder.
Afterwards, we sample minibatches of sequence length K , where timesteps are sampled uniformly.
Specifically, at every training iteration, we sample τ = ( s t , ...
s t
+ K , a t , ...a t + K ) , where t is sampled uniformly for each trajectory in the batch.
The sampled states, ( s t , ...s t + K ) , are encoded into skill embeddings using the VQVAE encoder.
We then pass in the states, encoded skills, and skill distributions into the causal transformer to output actions, which are used to compute a reconstruction loss, either MSE or cross entropy for discrete actions.
Like the original DT (Chen et al., 2021), we also did not find it useful to predict states or skill distributions, but it could be useful for actively predicting skill distributions without having to actually provide states to encode.
This is a topic we hope to explore more in the future.
The simplified training procedure is shown in
Algorithm 1.
Algorithm Initialize offline dataset D , Causal Transformer f θ , VQVAE Encoder e ϕ , context length K , num updates per iteration J for training iterations i = 1 ...
N do
Sample timesteps uniformly: t ∈ 1 , ...max len
Label dataset trajectories with skill distributions Z τ t = (cid:80)
Tt one hot ( z t ) for all t, ..
| τ |
Sample batch of trajectory states: τ = ( s t , ...s t
+ K , a t , ...a t + K ) for j = 1...
J do ˆ z τ t : t + K
=
( e ϕ ( s t ) , ...
e ϕ ( s t ))
Encode skills z τ t : t + K = quantize (ˆ z τ t : t + K ) Quantize skills with VQVAE
ˆ a τ t : t + K
=
f
θ
( Z τ t , z τ t , s t , ..., Z τ t + K , z τ t + K , s t + K )
L θ,ϕ = 1 K (cid:80) t + K t ( a t − ˆ a t ) 2 + V QLOSS ϕ
( z τ t : t + K , ˆ z τ t : t + K ) backprop
L θ,ϕ w.r.t θ, ϕ end for end for
E XPERIMENTS
We empirically want to answer the following questions:
• Can Skill DT achieve near or competitive performance, using only unsupervised trajectory information, compared to supervised offline RL approaches?
• How well can Skill DT reconstruct target trajectories and perform SMM in a zero shot manner?
• How diverse and descriptive are the skills that Skill DT discovers?
5. T ASKS AND D ATASETS
For evaluating the performance of Skill DT, we use tasks and datasets from the D4RL benchmark (Fu et al., 2020).
D4RL has been used as a standard for evaluating many offline RL methods (Kumar et al., 2020; Chen et al., 2021; Kostrikov et al., 2021; Zheng et al., 2022).
We evaluate our methods on mujoco gym continuous control tasks, as well as two antmaze tasks.
A timeplapse of a behavior learned can be seen in
E VALUATING S UPERVISED R ETURN
Can Skill DT achieve near or competitive performance, using only trajectory information, compared to supervised offline RL approaches?
Other offline skill discovery optimize hierarchical policies via supervised RL, utilizing the learned primitives to maximize rewards of downstream tasks (Ajay et al., 2020).
However, because we are interested in evaluating Skill DT without rewards, we have to rely on learning enough skills such that high performing trajectories are represented.
To evaluate this in practice, we run rollouts for each unique skill and take the maximum reward achieved.
Surprisingly, through just pure skill discovery, we are able to achieve competitive results on Mujoco continuous control environments compared to state-of-the-art offline reinforcement learning algorithms (Kumar et al., 2020; Kostrikov et al., 2021; Chen et al., 2021).
As we can see in our results in Table 1, Skill DT does noticeably worse in the ”-replay” environments and dataset.
We hypothesize that Skill DT performs worse in these environment because the dataset follows a policy during training, meaning that the trajectories
Mujoco Mean Results
[Table]
[Figure]
behaviors may be inconsistent throughout the dataset.
We think that with additional return context or online play, Skill DT may be able to perform better in these environments, and we hope to explore this as future work.
Skill DT, like the original Decision Transformer (Chen et al., 2021), also struggles on harder exploration problems like the antmaze-medium environments.
Even though
Skill DT performs marginally better than DT, there is still a lot of room for improvement in future work.
D ISCUSSION 6. SMM WITH LEARNED SKILLS
How well can Skill DT reconstruct target trajectories and perform SMM in a zero shot manner?
Ideally, if Skill DT is effective at SMM, it should be able to use its learned diverse skill embeddings to embed and reconstruct a target trajectory in an actual environment.
That is, given a target trajectory, Skill DT should be able to rollout a similar trajectory after encoding it in skill space.
The
[Figure] [Figure]
process is fairly simple and detailed in Algorithm 2.
In addition to state trajectories, learned skill distributions of the reconstructed trajectory and the target trajectory should also be close.
We investigate this by looking into target trajectories from antmaze-umaze-v2, antmaze-umaze-diverse-v2, and ant-v2.
For the ant-v2 environment, we are able to match the target skill distribution almost identically in a real rollout (Fig. 2).
For a more challenging example, we handpicked a trajectory from antmaze-umaze-diverse that is unique in that it has a loop.
Even though the trajectory is unqique,
Skill DT is still able to roughly recreate it in a zero shot manner (Fig. 4).
[Figure] [Figure]
6. S KILL D IVERSITY AND D ESCRIPTIVENESS
How diverse and descriptive are the skills that Skill DT discovers? .
In order to evaluate Skill DT as a skill discovery method, we must show that behaviors are not only diverse but are descriptive, or more intuitively, distinguishable .
We are able to visualize the diversity of learned behaviors by plotting eat trajectory generated by a skill on both antmaze-umaze environments (Fig. 3).
We also can observe from Figs.
6 and 5 that Skill DT learns behaviors that describe trajectories.
In addition, we show that through pure reward-free optimization, we are able to achieve competitive results on a number of offline RL benchmarks.
6. L IMITATIONS AND F UTURE WORK
Our approach is powerful because it is unsupervised, but it is also limited because of it.
Because we do not have access to rewards, we rely on pure offline diversity to ensure that high performing trajectories are learned and encoded into skills that can be sampled.
Skill DT could benefit from borrowing concepts from hierarchical skill discovery (Ajay et al., 2020) to re-use learned skills on downstream tasks by using an additional return-conditioned model.
We believe that this could help Skill DT solve some of the harder antmaze-medium environments.
In addition, it would be interesting to explore an online component to the training procedure, similar to the work in Zheng et al.
This would enable Skill DT to continuously learn and fine tune skills.
C ONCLUSION
We proposed Skill DT, a variant of Generalized DT, to explore the capabilities of offline skill discovery with sequence modelling.
We showed that a combination of LLMs, hindsight-relabelling can be incredibly useful for extracting information from diverse offline trajectories..
On standard offline RL environments, we showed that Skill DT is capable of learning a rich set of behaviors and can perform zero-shot SMM through state-encoded skill embeddings.
Skill DT can further be improved by adding an online component, a hierarchical component that utilizes returns, and improved exploration.
A CKNOWLEDGEMENTS
Removed for blind review.
R EPRODUCIBILITY S TATEMENT
The code to to reproduce every experiment in this paper is available at [removed].
A A PPENDIX A. F EW SHOT TARGET TRAJECTORY RECONSTRUCTION
Skill DT, like other skill discovery methods, can use its state-to-skill encoder to guide its actions towards a particular goal.
In this case, we are interested in recreating target trajectories as close as possible.
The detailed algorithm for few shot skill reconstruction:
Initialize target trajectory τ , encoder π
Extract states from target
▷ rollout in real
[Figure]