Year: 2021

Abstract:
Pre-training deep neural network models using large unlabelled datasets followed by finetuning them on small task-specific datasets has emerged as a dominant paradigm in natural language processing (NLP) and computer vision (CV). Despite the widespread success, such a paradigm has remained atypical in reinforcement learning (RL). In this paper, we investigate how we can leverage large reward-free (i.e. task-agnostic) offline datasets of prior interactions to pre-train agents that can then be fine-tuned using a small rewardannotated dataset. To this end, we present Pre-trained Decision Transformer (PDT), a simple yet powerful algorithm for semi-supervised Offline RL. By masking reward tokens during pre-training, the transformer learns to autoregressivley predict actions based on previous state and action context and effectively extracts behaviors present in the dataset. During fine-tuning, rewards are un-masked and the agent learns the set of skills that should be invoked for the desired behavior as per the reward function. We demonstrate the efficacy of this simple and flexible approach on tasks from the D4RL benchmark with limited reward annotations.

text: Pre-training deep neural network models using large unlabelled datasets followed by finetuning them on small task-specific datasets has emerged as a dominant paradigm in natural language processing (NLP) and computer vision (CV). Despite the widespread success, such a paradigm has remained atypical in reinforcement learning (RL). In this paper, we investigate how we can leverage large reward-free (i.e. task-agnostic) offline datasets of prior interactions to pre-train agents that can then be fine-tuned using a small rewardannotated dataset. To this end, we present Pre-trained Decision Transformer (PDT), a simple yet powerful algorithm for semi-supervised Offline RL. By masking reward tokens during pre-training, the transformer learns to autoregressivley predict actions based on previous state and action context and effectively extracts behaviors present in the dataset. During fine-tuning, rewards are un-masked and the agent learns the set of skills that should be invoked for the desired behavior as per the reward function. We demonstrate the efficacy of this simple and flexible approach on tasks from the D4RL benchmark with limited reward annotations.

heading: 1 INTRODUCTION
text: Over the past decade, we have seen tremendous progress in artificial intelligence, due in large part to our ability to train deep neural network models using massive amounts of data. The initial success of training such deep models was achieved by using large labeled datasets (e.g. ImageNet (Deng et al., 2009)), which provided rich and direct supervision for various tasks (e.g. in CV and NLP). However, the need for human annotations limited the scalability of such an approach, especially for tasks requiring detailed annotations like instance segmentation or machine translation (Lin et al., 2014; Bojar et al., 2016). Self-supervision has emerged as a dominant paradigm to overcome this limitation, where deep models are pre-trained using large unlabeled datasets, followed by fine-tuning with a small labelled dataset, considerably reducing the annotation requirements (Chen et al., 2020; He et al., 2020; Grill et al., 2020).
Analogously, progress in RL has required human annotations in the form of detailed reward functions. While some exceptions exist where the environment computes the rewards (e.g. games like Atari and Go); in most real-world applications of RL, a human is required to design reward functions or annotate frames for rewards. Such well-engineered reward functions are crucial for an agent’s ability to learn successful policies, which is hard to scale. At the same time, collecting reward-free and task-agnostic datasets for RL is becoming increasingly easier through use of offline datasets (Fu et al., 2020), tele-operation and play datasets (Rajeswaran et al., 2018; Lynch et al., 2019), or reward-free exploration (Pathak et al., 2017; Eysenbach et al., 2018; Liu & Abbeel, 2021). In this work, we aim to design RL agents that can utilize large reward-free datasets for unsupervised pre-training as well as small reward-annotated datasets for supervised finetuning to solve a variety of downstream tasks.
To this end, we consider the problem of semi-supervised offline RL, which is conducive for investigating our research question. In this setting, the agent has access to a large reward-free (task-agnostic) offline dataset of environment interactions. Such datasets are easier to scale than reward-labeled data and can be re-used for several downstream tasks with different rewards, similar to the trend of pre-training on unlabeled data and adapting to supervised downstream tasks in CV (Dosovitskiy et al., 2020) and NLP (Devlin et al., 2018; Brown et al., 2020). The agent further has access to a small dataset containing reward annotations, which anchors the agent to learn the desired behavior or task.
For this setting, we introduce a new approach: Pre-trained Decision Transformers (PDT), an architectural extension to the recently introduced Decision Transformer (DT; Chen et al. (2021)), to make it amenable to semi-supervised (offline) RL. DT frames RL as a sequence modeling task using an autoregressive and causally masked transformer model, and was shown to be competitive with traditional RL algorithms based on temporal difference learning. As shown in Figure 1, our framework extends DT and consists of two steps: pre-training and fine-tuning. First, from a large and fixed offline data, we use the state and action context of trajectories to predict the next actions autoregressively with the DT architecture. To accommodate fine-tuning with reward signals, during this pre-training step, we mask the reward tokens by zeroing them out as well as modifying the causal mask of the self-attention layers to generate actions without reward tokens. This can enable the model to implicitly learn the skills present in the offline un-labeled dataset. Then in the fine-tuning step, we un-mask the reward tokens and fine-tune the transformer end-to-end, predicting both reward and action tokens. This allows the agent to use the information present in the small reward-annotated dataset to adapt its set of skills towards achieving high rewards in the downstream task.
We summarize our contributions as follows:
1. We propose a new semi-supervised offline RL setting to bring RL closer to the self-supervised learning paradigm in CV and NLP.
2. We introduce Pre-trained Decision Transformers (PDT) - an extension of the recently proposed Decision Transformers which masks the reward tokens during pre-training and can automatically generate reward-to-gos during evaluation.
3. We show how PDT can efficiently be fine-tuned for a range of diverse downstream tasks using the same pre-trained model in a few-shot manner, achieving leading aggregate performance on D4RL environments in the low-label regime.
4. We provide insight into how different parts of PDT contribute to the final performance.

heading: 2 RELATED WORK
text: Behavioral Cloning (BC). When a reward-free but expert-quality dataset is available, imitation learning is often considered as a simple and yet powerful baseline for learning effective policies Osa et al. (2018); Rahmatizadeh et al. (2018). Our work differs in the motivation to utilize sub-optimal reward-free experience in the pre-training phase of the agent.
Offline RL. Offline RL is a paradigm for RL where agents are trained on offline datasets rather than directly interacting with the environment. The main issue in utilizing existing value-based off-policy RL algorithms in the offline setting is that these methods often erroneously produce optimistic value function estimates that can encourage the policy to pick actions that are out-of-distribution (OOD) compared to the offline dataset, resulting in failure. One way to mitigate this problem is to constraint the learned policy to be "close" to the data generating policy (also known as behavior policy) , by means of KL-divergence (Jaques et al., 2019; Wu et al., 2019a; Siegel et al., 2020; Peng et al., 2019a), Wasserstein distance (Wu et al., 2019a), or MMD (Kumar et al., 2019), and then using the sampled actions from this constrained policy in Bellman backup or applying a value penalty based on these divergence measures. Another way is to be conservative on value estimation by adding penalty terms that enforces the models to learn a lower bound on the value functions (Kumar et al., 2020; Fujimoto et al., 2019). All these methods however require the offline datasets to be reward-annotated, and rely crucially on rewards for learning.
Another line of work extends Model Based Reinforcement Learning (MBRL) methods to offline settings. Similar to model-free methods, existing MBRL methods cannot be readily used in offline settings due to the distribution shift problem and require some safeguard to mitigate exploiting model inaccuracies. To mitigate this issue, Kidambi et al. (2020) and Yu et al. (2020) both incorporate pessimism in learning the dynamics and estimating uncertainties in reward learning. While the dynamics models themselves can be learned without reward functions, planning or policy learning using the learned model again relies crucially on having access to a reward function or a large reward-annotated dataset.
Skill Extraction with Behavioral Priors. Methods that leverage behavioral priors utilize reward-free environment interactions to learn the set of task-agnostic skills via either maximizing likelihood estimates (Pertsch et al., 2020; Ajay et al., 2020; Singh et al., 2020) or maximizing mutual information (Eysenbach et al., 2018; Sharma et al., 2019; Campos et al., 2020). Behavioral priors learned through maximum likelihood latent variable models have been used for structured exploration in RL (Singh et al., 2020), to solve complex long-horizon tasks from sparse rewards (Pertsch et al., 2020), and regularize offline RL policies (Ajay et al., 2020; Wu et al., 2019b; Peng et al., 2019b; Nair et al., 2020). However, these skill extraction algorithms have high architectural complexity, requiring a generative model to learn a latent skill space and an RL agent to learn policies over skills as its action space. In contrast, PDT has one simple architecture for both skill extraction by modeling state-action sequences and control by conditioning on the desired reward.

heading: 3 APPROACH

heading: 3.1 PRELIMINARIES
text: Offline Reinforcement Learning. We consider the standard RL setting of a Markov Decision Process (MDP) defined by the tupleM = (S,A,R, T ). The MDP tuple consists of state s ∈ S, actions a ∈ A, rewards r = R(s, a), and transition dynamics T (s′|s, a). A policy in RL aims to pick actions that will maximize the expected total return of the MDP E [∑T t=0R(st, at) ] . One episode can be represented as a trajectory in the
MDP consisting of states, actions and rewards τ = (s1, a1, r1, s2, a2, r2, ..., sT , aT , rT ).
Offline RL is a setting of RL that seeks to find the optimal policy from a fixed offline dataset of interactions. D = {si, ai, ri}Ni=1. As a result, this setting is more difficult than the standard RL setting due to lack of access to the environment, and therefore is prone to distribution shift.
Decision Transformers. The usage of transformers (Vaswani et al., 2017) to tackle reinforcement learning problems has been a topic of recent study in Chen et al. (2021) and Janner et al. (2021). In Decision Transformer, rather than learning explicit policies or value functions, trajectories are modeled as sequences of state, actions, and reward-to-go. A reward-to-go (RTG) is defined as the sum of future rewards R̂t = ∑T t′=t rt′ . Then the problem setting becomes a sequence modeling problem, where actions are generated conditioned on the past context, current state, and the desired reward-to-go. With a context of length H at current time step t (i.e. C(H)t ), the policy becomes:
p(at|st, R̂t−1, at−1, st−1, ...R̂t−H , at−H , st−H︸ ︷︷ ︸ C
(H) t
)
The architecture of Decision Transformer first applies a linear token embedding to states, actions, and RTGs, and then adds a learned positional embbedding (shared across tokens within the same time-step) to differentiate between similar tokens at different positions along the sequence. It uses a GPT (Radford et al., 2019) architecture to learn a policy that auto-regressively generates actions based on prior context. The loss during training DT is a mean-squared error loss between predicted actions and ground truth actions.

heading: 3.2 OFFLINE RL WITH PRE-TRAINED DECISION TRANSFORMERS
text: Our approach is composed of two stages: Pre-training and Fine-tuning. During pre-training the model learns to extract the common temporally extended sequences of actions from the pre-training dataset. Our hope is that by fine-tuning this model on the downstream reward-annotated dataset it will be able to quickly learn task-specific knowledge including a reward prediction model and a reward-conditioned policy. We now elaborate each part in more details:

heading: 3.2.1 PRE-TRAINING DT
text: During pre-training, a sequence of trajectory is modeled as:
[s1], [a1], [RTG]︸ ︷︷ ︸ Pos.1 , [s2], [a2], [RTG]︸ ︷︷ ︸ Pos.2 , ..., [sT ], [aT ], [RTG]︸ ︷︷ ︸ Pos.T
Where [RTG] is the placeholder token for reward-to-go which is hard-coded to zero during pre-training. Note that the same positional encoding is added to state, RTG, and action of the same time-step. We also modify the causal mask in the self-attention layer to disregard any attention paid to the [RTG] tokens. This is crucial for the model to be reward-agnostric, and therefore, allowing for a quicker adaptation to downstream RTGs fed in at fine-tuning time. Pre-training is done by computing the maximum-likelihood estimates (MLE) of the neural network parameters.
LPT (θ) := E(st,at,Ct)∼DPT [ − log pθ ( at|st, C(H)t )] (1)

heading: 3.2.2 FINE-TUNING PDT
text: Similar to the fine-tuning strategy of Devlin et al. (2018), we simply fine-tune all PDT parameters on the downstream reward-annotated data. Since we have access to reward signals, we can now include them in the input sequence:
[s1], [a1], [R̂1]︸ ︷︷ ︸ Pos.1 , [s2], [a2], [R̂2]︸ ︷︷ ︸ Pos.2 , ..., [sT ], [aT ], [R̂T ]︸ ︷︷ ︸ Pos.T
During fine-tuning, instead of just predicting actions conditioned on the current state and context, i.e. pθ ( at|st, C(H)t ) , we also predict the RTGs conditioned on the current state, action, and context at every step,
i.e. pθ ( R̂t|st, at, C(H)t ) . We found that empirically, this helps the model infer the RTG distribution which
can help with the quality of the fine-tuned policy.
Overall, the training process in this step is similar to the training in the Decision Transformer, except for the addition of reward-prediction loss term. The fine-tuning loss is modified to contain both the MLE term for actions and the prediction loss for rewards, as described in Eq. 2. Note that the historical context C(H)t also contain the now unmasked reward-to-go tokens.
LFT (θ) = − E(st,at,Ct,R̂t)∼DFT [ log pθ ( at|st, C(H)t ) + log pθ ( R̂t|st, at, C(H)t )] (2)
Return Prediction. One limitation of the reward-to-go-conditioned structure of Decision Transformer is that each environment and each task requires a hand-picked initial reward-to-go to feed in. This can be un-intuitive and difficult: picking the best initial value can require prior knowledge about the training dataset and what rewards are feasibly achievable. Additionally, a fixed initial reward-to-go does not work well in environments that have random resets, such as D4RL’s (Fu et al., 2020) pointmaze environment, since the distances to goal are different based on the randomized initial locations. In such environments, the best feasible reward-to-go depends on the initial state, and cannot be a fixed value. To overcome this while still learning a reward-to-go-conditioned actor, we learn to predict the best feasible reward-to-go based on the context and state. The model is trained to predict a probability distribution over the possible RTG tokens from the training dataset. At evaluation time, we pick the next RTG by sampling N RTGs (e.g. N = 250) from the outputted probability distribution and picking the one with the highest value.
R̂t = max(R̂ 1 t , R̂ 2 t . . . R̂ N t ), R̂ i t ∼ pθ(R̂t|st, at, C (H) t )∀i ∈ (1, . . . N) (3)
One detail necessary for return prediction in continuous-reward environments (e.g. locomation mujoco tasks) is that we discretize the reward-to-gos into bins in order to learn a probability distribution over discerete values of RTGs rather than predicting scalar real-value returns.

heading: 3.2.3 EVALUATION
text: To rollout a trajectory using the model trained from PDT, we model it as a sequence of state-action-rewards. We alternate between two prediction phases for each time-step in the environment: (1) Using the context ( st−H , at−H , R̂t−H , ...st ) , predict the return-to-go R̂t (see: Eq. 3). (2) Using the updated context(
st−H , at−H , R̂t−H , ...st, R̂t )
predict the action at. In the first H steps of the trajectory when there is not enough context, we use zero padding.

heading: 4 EXPERIMENTS
text: In the experiments we are interested in answering the following questions: (i) Can PDT be efficiently finetuned on downstream offline datasets where reward-labeled data is limited? (ii) How does PDT compare to
TD learning offline RL algorithms, despite not fitting any explicit value function? (iii) What are the most important design choices for adapting DTs to the semi-supervised offline RL setting?

heading: 4.1 ENVIRONMENTS AND DATASETS
text: We evaluate the performance of PDT on three simulated mujoco (HalfCheetah, Walker2d, and Hopper) and one robotic manipulation environment (FrankaKitchen) from the D4RL benchmark.

heading: 4.1.1 DATASETS
text: Mujoco. For the mujoco locomotion datasets, we consider three simulated locomotion agents and three datasets per agents (medium, medium-replay, and medium-expert), summing to a total of nine datasets. The datasets are collected in the following ways:
• The medium dataset is collected with partially trained SAC agent. • The medium-replay dataset is the entire replay buffer of a SAC agent throughout training. • The medium-expert is a mix between trajectories from the medium dataset and an expert policy.
FrankaKitchen. This environment consists of a 7-DOF kitchen robotic arm that is tasked to manipulate 4 objects in a particular order. We use kicthen-partial-v0 which is collected with tele-operation and includes manipulating all objects, present in the environment, in different orders. However, each trajectory in the dataset is reward-labeled according to the order of the target task.
For pre-trained methods, we use the entire dataset in the D4RL benchmark to learn the behavioral prior. For fune-tuning we evaluate on different number of trajectories sampled from the dataset. In our experiments, full refers to using the entire reward-labeled dataset, 25 refers to using 25 randomly sampled trajectories, and 10 refers to using 10 randomly sampled trajectories. By doing so, we can investigate whether PDT can efficiently fine-tune the downstream offline tasks with small reward-labeled offline datasets?

heading: 4.1.2 EVALUATION PROTOCOL
text: For evaluating each of the results, we calculate each algorithm’s performance by taking the average returns of the last 5 epochs for training on the downstream dataset. For each epoch, we evaluate 10 trajectory rollouts in the environment and report the average return based on the downstream task reward function.
We use the following approaches for comparison:
We train BC on the entire reward-free dataset; CQL, and DT only on the reward-annotated dataset, and CQL+BC on both reward-free and reward-annotated dataset. CQL+BC, similar to PDT, leverages both the reward-free and reward-annotated datasets.
Results. Figure 2 shows the aggregate of the normalized returns across all tasks as a function of the number of reward-labeled fine-tuning trajectories. As we can see, compared to DT and CQL, the algorithms that leverage reward-free datasets (i.e. PDT and CQL+BC) perform significantly better in low data regimes. It is also noted that PDT’s performance is competitive with CQL+BC, but the simplicity and flexibility of its architecture makes PDT a suitable candidate for wide-spread adoption in semi-supervised offline RL settings.
To illustrate the granular results on each environment, we have included Table 1. The scores are normalized to the maximal return of the trajectories in each dataset. As we can see, even though both DT and CQL mostly achieve good results in full data regimes, their performance significantly drops in settings where labeled-data is limited. In the kitchen environment in particular, CQL+BC completely collapses as the desired multi-modal behavior of the agent cannot be captured by a simple BC regularization, while PDT can recover the task-specific behavior in low data regimes.

heading: 4.2 ABLATION STUDIES
text: In this section we study different components of the PDT algorithm to provide insight on the contribution of each part. In particular, we are interested in performing the following ablations:
The effect of reward prediction during fine-tuning. During the fine-tuning step of our framework we not only predict the next actions conditioned on the current state, reward-to-go, and context, but we also predict the distribution of possible remaining reward-to-gos. To this end, we study the effect of this component by comparing to a case where it is not used. We run an experiment with the maze2d-large-v1 D4RL in which the agent needs to know how far it is from the goal location to be able to navigate to it properly.
In this experiment, we use our model to output the best possible return-to-go given different starting locations in the maze. In Figure 3 (right) a darker square indicates a bin that represents a better return-to-go. This means that if the point mass starts in that state, it will be able to quickly reach the goal square of (9, 1). As expected, states closer to (9, 1) have a better reward-to-go.
The effect of masking reward token attention during pre-training. During pre-training we mask the reward attention to force the model to not attend to reward tokens, and to auto-regressively model the next actions conditioned on the current state, and context without rewards. Here, we compare the post fine-tuning performance of PDT with a version that does not apply attention masks during pre-training. We show three experiments on the medium-expert suite in mujoco, since the effect is seen most clearly in this dataset type. Including masking during pre-training, improves results regardless of how many trajectories the fine-tuning dataset includes.
Comparison to other choices of fine-tuning. PDT’s strategy for fine-tuning is the most straight-forward instantiation of this idea which is fine-tuning all pre-trained parameters. Here, we compare this particular design choice to other common fine-tuning patterns seen in transformer literature. In particular we compare,
fine-tuning with frozen self-attention layer and fine-tuning with both frozen self-attention and feed-forward layers. The later choice was proposed recently on Lu et al. (2021).
Empirically, we see that fine-tuning all of the parameters, on average, outperforms fine-tuning only a subset of the parameters, although the difference overall is not extreme. Interestingly, in the very-low data domain of 10 trajectories, it does appear that only freezing the attention parameters does lead to the best performance, but for 25 trajectories and the full trajectories, fine-tuning all parameters does best.

heading: 5 CONCLUSION
text: We presented PDT, a simple and effective architecture amenable to semi-supervised offline RL settings where the architecture is pre-trained on reward-free and task-agnostic datasets and then fine-tuned to a task-specific downstream dataset. Our technical contribution is extending Decision Transformers to this setting by masking reward tokens during pre-training and predicting remaining reward-to-gos during fine-tuning. We show how PDT can outperform existing offline RL methods in low data regimes. The simple training procedure of PDT makes it a great candidate for wide-spread adoption in semi-supervised offline RL settings.

Review summaries:
### What is the Problem / Question? Pre-training has not been thoroughly explored within RL, so it is unknown how effectively task-free unlabeled datasets can be used for PT in an RL context. The authors present a simple strategy for such PT for RL that shows benefits with limited reward annotations.  ### Why is it impactful? RL is a famously expensive discipline of ML which nonetheless has huge application areas and potential in theoretial and deployment contexts. Being able to translate the improvements offerred by pre-training to the RL context would help empower this already important technique significantly.  ### Why is it hard? Why have previous approaches failed? There are several previous approaches to this problem, including some the authors readily identify in their study, such as the CQL+BC method which uses unlabeled data to perform semi-supervised offline RL by adapting policy prior data via unlabeled data.   ### How do they solve it? The authors propose using the (already published) decision transformer model, but pre-training that model first on a dataset with only action sequences, without any reward information.  ### How do they validate their solution? The authors compare their proposed framework against existing models, including both the traditional DT, a BC baseline, and the CQL and CQL+BC methods (the latter of which is a semi-supervised approach). They find slight improvements over existing methods in the extreme few-shot settings (10 observed trajectories at FT time), with further reduced gains at larger #s of observed trajectories, culminating in worse performance than traditional methods at full data scale. 

This paper presents Pre-trained Decision Transformer (PDT) for semi-supervised offline reinforcement learning. PDT first pre-train a decision transformer model on the trajectory dataset without rewards, and then fine-tune on a smaller dataset with reward annotations. Empirically, PDT achieves compatible performance with offline reinforcement learning baselines.