Semi-supervised Offline Reinforcement Learning with Pre-trained Decision Transformers

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 . 1 INTRODUCTION . 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 . 2 RELATED WORK . 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 can not 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 . 3 APPROACH . 3.1 PRELIMINARIES . 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 .

Summary: 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.