SMART: Self-supervised Multi-task pretrAining with contRol Transformers
Authors: Yanchao Sun
Shuang Ma
Ratnesh Madaan
Rogerio Bonatti
Furong Huang
Ashish Kapoor
Venue: ICLR 2023 notable top 25% (2023)
Keywords: pretrain
transformer
multi-task reinforcement learning
sequential decision making
self-supervised

Abstract: Self-supervised pretraining has been extensively studied in language and vision domains, where a unified model can be easily adapted to various downstream tasks by pretraining representations without explicit labels. When it comes to sequential decision-making tasks, however, it is difficult to properly design such a pretraining approach that can cope with both high-dimensional perceptual information and the complexity of sequential control over long interaction horizons. The challenge becomes combinatorially more complex if we want to pretrain representations amenable to a large variety of tasks. To tackle this problem, in this work, we formulate a general pretraining-finetuning pipeline for sequential decision making, under which we propose a generic pretraining framework \textit{Self-supervised Multi-task pretrAining with contRol Transformer (SMART)}. By systematically investigating pretraining regimes, we carefully design a Control Transformer (CT) coupled with a novel control-centric pretraining objective in a self-supervised manner. SMART encourages the representation to capture the common essential information relevant to short-term control and long-term control, which is transferrable across tasks. We show by extensive experiments in DeepMind Control Suite that SMART significantly improves the learning efficiency among seen and unseen downstream tasks and domains under different learning scenarios including Imitation Learning (IL) and Reinforcement Learning (RL). Benefiting from the proposed control-centric objective, SMART is resilient to distribution shift between pretraining and finetuning, and even works well with low-quality pretraining datasets that are randomly collected. The codebase, pretrained models and datasets are provided at https://github.com/microsoft/smart.

Published as a conference paper at ICLR 2023
SMART: SELF-SUPERVISED MULTI-TASK PRETRAIN- ING WITH CONTROL TRANSFORMERS
Yanchao Sun Shuang Ma¹, Ratnesh Madaan Rogerio Bonatti', Furong Huang and Ashish Kapoor
'University of Maryland, College Park. {yos, furongh} @umd. edu
Microsoft Redmond, WA. (shuama, ratnesh madaan, rbonatti, akapoor) wtp.saaaaaaooeeeootsottto..o wmwrrooottmmooooocccoommmm
ABSTRACT
Self-supervised pretraining has been extensively studied in language and vision domains, where a unified model can be easily adapted to various downstream tasks by pretraining representations without explicit labels. When it comes to sequential decision-making tasks, however, it is difficult properly design such a pretraining approach that can cope with both high-dimensional perceptual in- formation and the complexity of sequential control over long interaction hori- zons. The challenge becomes combinatorially more complex if we want to pre- train representations amenable in this work, we formulate a general pvetranning-finttniin tial decision making, under which we propose a generic pretraining framework Self-supervised Multi-task pretrAining with contRol Transformer (SMART). By systematically investigating pretraining regimes, we carefully design a Control Transformer (CT) coupled with novel control-centric pretraining objective in a self-supervised manner. SMART encourages the representation to capture the common essential information relevant to short-term control and long-term con- trol, which is transferrable across tasks. We show by extensive experiments DeepMind Control Suite that SMART significantly improves the learning effi- ciency among seen and unseen downstream tasks and domains under different learning scenarios including Imitation Learning (IL) and Reinforcement Learning (RL). Benefiting from the proposed control-centric objective, SMART is resilient to distribution shift between pretraining and finetuning, and even works well with low-quality pretraining datasets that are randomly collected. Our codebase, pre- trained models and datasets are provided https://giiuuh.ccomm/irrrssttttu
large variety of tasks. To tackle this problem,
pipeline for sequen-
INTRODUCTION
Self- Self-supervised pretraining has been successful in wide range of language and vision problems. Examples include BERT (Devlin et al. 2019),, GPT (Brown et al. 2020). MoCo (He et al.) 2020), and CLIP (Radford et al. 2021). These works demonstrate that single pretrained model can be easily finetuned to perform many downstream tasks, resulting in a simple, effective, and data- efficient paradigm. When it comes sequential decision making, however, is not clear yet whether the successes of pretraining approaches can be easily replicated.
There research efforts that investigate application of pretrained vision models facilitate control tasks (Parisi et al., 2022 Radosavovic et al.). However, there are challenges unique to sequential decision making and beyond the considerations of existing vision and language pretraining. We highlight these challenges below: (1) Data distribution shift: Training data for decision making tasks usually composed of trajectories generated under some specific behavior policies. As a result, data distributions during pretraining, downstream task finetuning and even during deployment can be drastically different, resulting in a suboptimal performance (Lee et al. 2021). (2) Large discrepancy between tasks: In contrast to language and vision where the underlying semantic information is often shared across tasks, decision making tasks span large variety of task- specific configurations, transition functions, rewards, as well as action and state spaces. Consequently, it is hard to obtain generic representation for multiple decision making tasks. (3) Long term reward maximization: The general goal of sequential decision making is learn policy that maximizes long-term reward.

---

Published as a conference paper at ICLR 2023
Thus, good representation for downstream policy learning should capture information relevant for both immediate and long-term planning, which is usually hard in tasks with long horizons, partial observability and continuous control. (4) Lack of supervision and high-quality data: Success in representation learning often depends on the availability of high quality expert demonstrations
and ground-truth rewards (Lee al. [022] Stooke et al. (02T). However, for most real-world sequential decision making tasks, high-quality data and/or supervisory signals are either non non-existent or prohibitively expensive to obtain.
Under these challenges, we strive for pretrained representations for control tasks that are (1) Versatile so as to handle wide variety of downstream control tasks and variable downstream learning methods such as imitation and reinforcement learning (IL., RL) etc.
(2) Generalizable to unseen tasks and domains spanning multiple rewards and agent dynamics, and
(3) Resilient and robust to varying-quality pretraining data without supervision.
We propose a general pretraining framework named Self-supervised Multi-task pretrAining with contRol Transformer (SMART), which aims to satisfy the above listed properties. We introduce Control Transformer (CT) which models state-action interactions from high-dimensional observa- tions through causal attention mechanism. Different from the recent transformer- -based models for sequential decision making/Chen Chen et al. (202T) which directly learn reward- bbaedd policies, CT is de- signed to learn reward-agnoste representations, which enables it as a unified model to fit different learning methods (e.g. II. and RI.) and various tasks. Built upon CT, we propose a control-centric pretraining objective that consists of three terms: forward dynamics prediction, inverse dynamics prediction and random masked hindsight control. These terms focus on policy-independent transi- tion probabilities, and encourage CT to capture dynamics information of both short-term and long- term temporal granularities. In contrast with prior pretrained vision models (Oord et al.) 20188] Parisi et al., 2022) that primarily focus on learning object-centric semantics, SMART captures the essential coltrol-relevant information which is empirically shown to be more suitable for interactive decision making. SMART produces superior performance than training from scratch and state-of- the-art (SOTA) pretraining approaches on a large variety of tasks under both II. and RI. Our main contributions are summarized as follows:
1. We propose SMART, generic pretraining framework for multi-task sequential decision making-
2. We introduce the Control Transformer model and a control-centric pretraining objective to learn representation from offline interaction data, capturing both perceptual and dynamics information with multiple temporal granularities.
3. We conduct extensive experiments on DeepMind Control Suite TTasa et al. 2018). By evalu- ating SMART on large variety of tasks under both IL and RL regimes, SMART demonstrates its versatile usages for downstream applications. When adapting to unseen tasks and unseen do- mains, SMART shows superior generalizability SMART can even produce compelling results when pretrained on low-quality data that is randomly collected, validating its resilience property.
2 RELATED WORKS
Offline Pretraining of Representation for Control. Many recent works investigate pretraining rep- resentations and finetuning policies for the same task. Yang & Nachum (2021) investigate several pretraining objectives MuloCo with vector state inputs. They find that many existing repre- sentation learning objectives fail to improve the downstream task, while contrastive self-prediction obtains the best results among all tested methods. [Schwarzer et al. (2021) pretrain a convolutional encoder with a combination of several self-supervised objectives, achieving superior performance on the Atari 100K. However, these works just demonstrated the single-task pretraining scenario, it is not clear yet whether the methods can be extended to multi-task control. Stooke et al. (2021) propose ATC, a contrastive learning method with temporal augmentation. By pretraining an encoder on ex- pert demonstrations from one or multiple tasks, ATC outperforms prior unsupervised representation learning methods in downstream online RL tasks, even in tasks unseen during pretraining.
Pretrained Visual Representations for Control Tasks. Recent studies reveals that visual represen- tations pretrained on control-free datasets can be transferred to control tasks. Shah & Kumar (2021) show that that ResNet encoder pretrained on ImageNet is effective for learning manipulation tasks. Some recent papers also show that encoders pretrained with control-free datasets can generalize well to RL settings (Nar et al. 2022 Seo et al., 2022 Parisi et al. 2022). However, the generalizability of the visual encoder can be task-dependent. Kadavath et al. 20221) 1) point out that ResNet pretrained
on ImageNet does not help in DMC (IICCTTun (Tunyasuvunakkoll et al. 2020 environments.
2

---

Published as a conference paper at ICLR 2023
Pretraining Phase
Finetuning Phase
Pretraining Objectives
Opeenal FLY
Learned
Eood
representations
Transforor block
Control Transformer
Control Transformer
Takan umboddings
EXIT
Tokenbers
Ippul rawsinahs
Figure 1: Architecture of Control Transformer. In the pretraining phase, we use the control-centric objective introduced in Section 5.2 to train representation over multiple tasks; in the finetuning phase where a specific task is given, learn policy based on the pretrained representation (pretrained weights are shown in grey blocks). The construction of the policy head can vary for different downstream datasets or learning methods.
Unsupervised RL. Unsupervised RL (URL) focuses on learning exploration policies (Pathak et al,
2017 Liu & Abbeel )021), goal-conditioned policies (Andrychowicz et al., 2017 Mendonca al. 2021 or diverse skills (Eysenbach et al.) 2019) in a task without external rewards, and finetuning the
policy later when reward is specified. The unsupervised learning phase of URL is usually interactive and prolonged. Our goal, in contrast, is to train representations of states and actions from fixed offline datasets, with focus capturing essential and important information from raw inputs.
Sequential Decision Making with Transformers. There is a growing body of work that uses Trans- former Vaswani et al. 2017) architectures to model and learn sequential decision making problems. Chen et al. (2021) propose Decision Transformer (DT) for offline RL, which takes a sequence of returns, observations and actions, and outputs action predictions. Trajectory Transformer (Janner et al. (02T) also models the trajectory as sequence of states, actions and rewards, while dis- cretizing each dimension of state/actions. [Bonatti et al. (2022) propose a pretraining scheme for state-action representations in navigation scenarios using a causal transformer, which can then be finetuned with imitation learning towards different tasks for the same robot. Furuta et al. (2022) propose Generalized DT that unifies a family of algorithms for future information matching with transformers. iZheng et al. (2022) extend DT to online learning by blending offline pretraining and online finetuning. Transformers
also be used world models for model-based RI. (Chen al., 2022 Micheli et al.) 2022). Recent studies show that wransformer-based models can be scaled up with diverse multi-task datasets to produce generalist agents (Reed al. 2022 Lee et al. 2022). Our proposed method has a similar structure that regards RL trajectories as sequential inputs. How- ever, differently from most existing transformer models that learn policy from returns, our SMART focuses on learning control-relevan representations with self-supervised pretraining.
3 PRELIMINARIES
Partially Observable Markov Decision Process. We model control tasks and environments as a partially observable Markov decision process (POMDP) M = (S, A, O, P, R. E), which is a generalization of Markov decision process (MDP). Here, S is the underlying state space, A is the action space, O is the observation space, P is the transition kernal, R is the reward function, and E is the observation emission function with E(oss) being the probability of observing o given state s. In practice, the observation space can be high dimensional. For example, for a mobile robot navigating with camera sensors, the images are observations and its odometry (location, orientation, and associated velocities) and the ground- truth obstacle locations form the underlying state.
Learning Agent and Controlling Policy. At every step t, the agent receives an observation or based on the underlying state 8t (hidden from the agent), takes action az, and obtains reward ra and the environment proceeds to the next state 8141. Given history of observation-action pairs
of length L and the current observation, h (ot L. at-L, L+-1+1, 0t-L+1, ,o),, the agent exccutes action ag according to policy т: a, = π (h;). The agent's goal is to learn an optimal policy n² that maximizes the agent's is cumulative reward Ep 1
Reinforcement Learning (RL) and Imitation Learning (IL). RL (Sutton & Barto 2018 is the process where an agent seeks to maximize its policy returns. RL agents learn in online man- ner by interacting with the environment, or learn from offline data with pre-collected interactions. Differently from supervised learning, the training data for RL stems from policy-dependent inter-
3

---

Published as a conference paper at ICLR 2023
actions with the environment, rendering non-i..d. training regime. Another effective method for learning control policies is IL (Hussein et al., 2017; Osa et al. 2018), where an agent obtains su- pervision from expert demonstrations. For both RL and IL, there could be discrepancy between training and deployment data distributions due to the different distribution of states induced by the imperfect policy and environment uncertainty.
4 PROBLEM SETUP: PRETRAINING AND FINETUNING PIPELINE
Multi-Task Control with Shared Representation. We consider a of multiple tasks T with the same dimensionality in observation space. In this work we select TI from different environment in DeepMind Control Suite (DMC) TTssa al. 2018), in which the agent observes RGB image of the current state. Tasks in T can have entirely different state spaces S, different action spaces A and different environment dynamics (P, R, E). We also define the concept of domain to differentiate tasks that have different state/action spaces. For example, in DMC, "hopper" and "walker" belong to different domains because they posses distinct action spaces, while "walker-walk' and "walker-run" are different tasks within the same domain. In this paper we use the term mulli-task to refer tasks spanning potentially multiple domains.
Table 1: A comparison between pretraining and finetuning.
Pretraining- Saainn-FFinetung Finetuning Pipeline. Al- though pretraining methods are widely
Pretraining phase
Finetuning phase
Learn generic representation
Learn policy
applied in many areas, it is not yet clear
Offline
Offiine or online
what role pretraining should play in se-
Multiple tasks
One task, seen
quential decision making tasks, espe-
Reward or expert demon-
Has reward supervision or expert demonstration
cially when considering the multi-task setup. In this work, follow ideas es-
stration may be absent
More samples
Fewer samples
tablished in vision and language community to explicitly define our pretraining and finetuning pipeline, which we summarize in [Tble Specifically, during the pretraining phase we train repre- sentations with a possibly large offline dataset collected from set of training tasks Tpre = {M: Then, given specific downstream task M which may or may not be contained in Teee we attach simple policy head on top of the pretrained representation" and train with IL or with RL. The cen- tral tenet of pretraining is to learn generic representations which allow downstream task finetuming to be simple, effective and efficient, even under low-data regimes.
This cestraining-fintuning pipeline is general extension of many prior settings of pretraining for decision making. For example, al. (202T) pretrain an encoder on one or multiple tasks, then learn RL policy in downstream tasks. Their pretraining dataset is composed of expert demon- stration, and the finetuning process focuses on online RL. In addition, [Schwarzer et al. (2021) learn representations with offline datasets, but perform pretraining and finetuning within the same task.
5 OUR PRETRAINING MODEL AND APPROACH
We propose Self-supervised Multi-task pretrAining with contRol Transformer (SMART), general pretraining approach for multi-task sequential decision making. We first give an overview of the proposed Control Transformer (CT) and illustrate how it fits in our pretraining-finttuni pipeline in Section 5.1
Then, we introduce our control-centric pretraining objective in Section 5.2
5.1 APPROACH OVERVIEW AND MODEL ARCHITECTURE
Model Architecture of Control Transformer. Inspired by the recent success of transformer models in sequential modeling (Chen et al. 202T Janner et al. (02T), we propose a Control Transformer (CT). The input to the model is control sequence of length 2L composed of observations and ac- tions: (O,, at Ot+1, 0.+1. ... Or+LL L. Different from Decision Transformer (DT) (Chen et al.) 2021 purposefully do not include the reward signal in the control sequence keep our repre- sentations awward-ggosstic, as explained the end of Section Each element of the sequence is embedded into a d-dimensional token, with modality-specificc tokenizer jointly trained with Trans- former blocks. We also learn an additional positional embedding and sum it with each token. The outputs of CT correspond token embeddings representing each observation and action, and represented by ((o:) and Φ(a;), respectively Figure depicts the CT architecture.
'The pretrained encoder can he either frozen or finetuned with the policy, depending on the task. "The implementation details are provided in Appendix A.3
4

---

Published as a conference paper at ICLR 2023
Forward Dynamics Prediction
Inverse Dynamics Prediction
Random Masked Hiiddiigtt Control
A+
ETTTT
Masked sel-attention
Masked setf- eell-ttention
Seltaattetiion
TBB
MonswALLY apdita
Figure 2: The three terms of our proposed pretraining objective. The red shaded areas denote the attention span, while the grey regions masked.
Pretraining of SMART. We generate an offline dataset for pretraining which contains control tra- jectories generated by some behavior policies for set of diverse tasks THe spanning multiple do- mains. During pretraining, we append several prediction heads to the transformer output, and train the entire model by minimizing the control-centric pretraining objective introduced in Section [5.2 These prediction heads are used to learn desired representations, will be dropped in finetuning.
Downstream Finetuning of SMART. As discussed in Section 4 the pretrained representation can be used to learn policies for different tasks. To do so, we append a policy head π to the observa- tion representation, such that π ((((:)) predicts the proper action for observation oG. We can train the policy head using both IL and RL. For our IL experiments we use behavior cloning with ex- pert demonstrations, where we feed Φ(oz) into a policy head to get action predictions. For RL we use return-to-oo (RTGG)-conditooeed policy with trajectories that contain reward values. We feed ((o,) along with an RTG embedding to get the policy head's action predictions. Online RL with transformer-based models is still a novel field and is not the focus of this work. But we show advan- tages of our pretraining method for online RL finetuning, following the same settings as online DT presented by Zheng et al. (022). See Appendix B.4 for more discussion and results.
Comparison with Prior Decision-making Transformers. Recent works leverage transformer ar- chitectures for modeling sequential decision making problems (Chen et al. [202] Janner al. [021] Lee et al. 2022) as summarized in Section 2 Most of these models use reward information in the input sequence, as their goal is to directly learn policy for specific task. In contrast, our goal is to pretrain representations for various downstream tasks, and thus our CT uses reward-free control sequence as the model input; rewards are only used for downstream task when needed. There are several benefits of making representation agnostic to reward during pretraining. (1) A pretrained model requiring reward input does not flexibly fit some downstream learning scenarios such as be- havior cloning, while CT can be unified model for various learning methods. A user easily learn a policy under different learning methods (II. or RL) without modifying transformer blocks. (2) Reward distribution be significantly different when the task or policy changes, making reward-dependent representation less resilient to distribution shift. We show in Section 6.3 that utilizing reward during pretraining may hurt the overall downstream performance.
5.2 CONTROL-CENTRIC SELF-SUPERVISED PRETRAINING OBJECTIVES
Our pretraining objective employs three terms: forward dynamics prediction, inverse dynamics pre- diction, and random masked hindsight control. The first two terms focus on local and short-term dynamics, while the third term is designed capture more global and long-term temporal depen- dence. As motivated in Section 5.1 these terms are based on control sequences and are reward-free, such that they can be used for multiple tasks. Figure 2 illustrates each objective.
I. Forward Dynamics Prediction. For each observation-action pair (0),0t) in control sequence, we aim to predict the next immediate latent state. Let g be the observation tokenizer being trained. We maintain momentum encoder I the exponential moving average of g, generate the target latent state from the next observation 0++1, i.c., $ := g(or+1). ). The idea of momentum encoder is
widely used when the target value is not fixed (He et al. 2020 Mnih et al.) 2015). in order to make training more stable. Then, we train a linear head to predict $1+1 based on (o(o:), p(a.)) This forward prediction captures the local transition information in the embedding space.
II. Inverse Dynamics Prediction. For each consecutive observation pair (or, 1)+),, we learn to recover the action that leads o 0++1 Note that in causal transformer, at is visible the model
5

---

Published as a conference paper at ICLR 2023
when generating representation of ot+1 which can lead to trivial solutions, modify the original causal mask and mask out a from the attention of 01+1- Therefore, the observation representation is forced to contain information for relevant actions and transitions.
Both forward and inverse predictions focus on local dynamics induced by the transition kernel P. However, fitting local dynamics only may result in collapsed representation, i.e., the model learns the same representation for two semantically different observations (Rakelly et al.) (2021). To perform well in downstream tasks, long-term temporal dependence should also be captured in the represen- tation. We achieve this by a novel random masked hindsight control term in pretraining.
III. Random Masked Hindsight Control. Given a control sequence h (0t, Ot,""' L, a++ +.), we randomly mask & actions and k' observations, and recover the masked actions based on the re- maining incomplete sequence. This idea of masked token prediction is related to BERT (Devlin et al. 2019) for language modeling, but note that we only predict the masked actions for the purpose of control. More rationale behind the selection of l and k'' is explained in Appendix A. This objective is akin to asking the question "what actions should I take to generate such trajectory?" Therefore, we replace the causal attention mask with a non-causal one, to temporarily allow the model "see the future", as shown in Figure righht). As a result, we encourage the model to learn controllable representations and global temporal relations, and to attend to the most essential repre- sentations for multi-step control. The idea of our random masked hindsight control is also related to the multi-step inverse prediction proposed by a concurrent work (Lamb et al. 2022), which predicts a, given 8t and 81+1 for random interger and theoretically shows the effectiveness of this method in discovering controllable states. Our random masked hindsight control is different as it predicts multiple actions altogether from randomly masked sequences with a transformer model, which can efficiently learn the control information in large-scale tasks, and avoid ambiguity caused by different paths between states. An empirical comparison between our formulation and the multi-step inverse prediction (Lamb et al. 2022 is provided in Appendix B.5.2
Finally, our pretraining objective is the summation of the above three terms with equal weights, which in experiments renders good performance. The mathematical formulations and implementa- tion details of the objective are explained in Appendix A.
6 EXPERIMENTS
We provide empirical results to demonstrate the effectiveness of our proposed pretraining method, while aiming to answer the following questions: (1) Can SMART effectively improve learning effi- ciency and performance in variety of downstream tasks under different learning methods? (2) How well can SMART generalize to out-of-distribution tasks? (3) Is SMART resistant to low-quality pre- training data? (4) How does SMART compare state-of-the-art pretraining techniques? (5) How do different pretraining objectives affect the downstream performance?
6.1 EXPERIMENTAL SETUP
We evaluate SMART the DeepMind Control (DMC) suite (Tasaa al. 2018), which contains series of continuous control tasks with RGB image observations. There are multiple domains (phys- ical models with different state and action spaces) and multiple tasks (associated with particular MDP) within each domain, which creates diverse scenarios for evaluating pretrained representa- tions. Our experiments use 10 different tasks spanning over 6 domains. In pretraining, we an offline dataset collected over tasks, while the other 5 tasks (with unseen domains) are held out to test the generalizability of SMART. A full list of tested domains and tasks is in Appendix A.2
Pretraining Tasks and Datasets. We pretrain SMART on 5 tasks: cartpole-swingup hopper-hop.
cheetah-run, walker-stand and walker-run. For each task, we adopt/train behavior policies to collect the following two types of offline datasets. (Details of dataset collection are in Appendix (A.2)
Trajectories with random environment interactions, with 400K timesteps per task. Expl cratory: Trajectories generated in the exploratory stage of multiple RI. agents with dif- ferent random seeds, with 400K timesteps per task.
Random:
Downstream Tasks, Learning Methods and Datasets. We evaluate the pretrained models in the 5 seen tasks and another 5 unseen tasks: cartpole-balance hopper-stand, walker-walk, pendulum- swingup and finger-spin (the last two tasks are from unseen domains with different state-action spaces). We consider two learning methods: return-ta coum-to-goo to-go-conditioned policy learning (RTG) and behavior cloning (BC). For RTG, use the Sampled Replay ay dataset containing randomly sam-

---

Published as a conference paper at ICLR 2023
pled trajectories from the full replay buffer collected by learning agent. For BC, the Expert trajectories with the highest returns from the full replay buffer of learning agents. The downstream dataset for every task only has 100K timesteps, making it challenging to learn from scratch.
Implementation Details, Our implementation of CT is based on a GPT model (Radford et al. 2018 with 8 layers and 8 attention heads. We context length L 30 and embedding size d == 256. As explained in Appendix [.1] & and k' are linearly increased from to L and L/2, respectively. The observation tokenizer i8 a standard 3-layer CNN. Action tokenizer, return tokenizer, and all single-layer linear prediction heads. We found that freezing the pretrained weights in downstream tasks works well in relatively simple environments, but fails in harder ones. Therefore, we finetune the entire model including transformer blocks for all downstream tasks. Since actions in different domains have different dimensions and physical meanings, we project the raw actions into a larger common action space to train the action tokenizer. When there is novel downstream task with a different action space, we simply re-initialize the action tokenizer and finetune it. Please see Appendix A.3 for more implementation and hyperparameter details.
Baselines. We compare SMART with the following transformer- pretraining baselines:
Scratch trains a policy with randomly initialized CT representation weights.
ACL (Yang & Nachum 202T) is modified BERT (Devlin et al. 2019) that randomly masks and predicts tokens with contrastive loss, pretrained on the same dataset as SMART.
DT (Chen et al. 2021) pretrained on the same dataset as ours but uses extra reward supervision.
CT-single is variant of SMART, which pretrains CT with a single-task dataset containing trajectories from the downstream environment.
For fair comparisons, we use the same network architecture for the baseline models (except for DT where we keep their original network structure with RTG as transformer inputs) and train them with the same configurations. We also compare SMART with other state-of-the-art pretraining works, such as CPC (Oord et al. 2018) and ATC (Stooke et al. (02T). using the same cetrraining-finetuiinn pipeline. However, these approaches are built upon ResNet backbones, direct comparison of Transformer against a ResNet could be not straightforward. Hence we refer readers to Appendix (B.3 [.. for more discussions and results.
To evaluate the quality of pretrained representations we report the average cumulative reward obtained in downstream tasks after finetuning. We deploy the trained policies in each environment for 50 episodes and report average returns. For evaluation of the RTG-conditioned policies, we use the expert score of each task as the initial RTG, done in DT (Chen et al. 2021 Detailed evaluation settings as described in Appendix A.4
Evaluation Metrics.
6.2 EXPERIMENTAL RESULTS
We first evaluate the versatility of SMART: 1) whether a single pretrained model can be finetuned with different downstream learning methods (i.e. RTG and BC); and 2) whether the pretrained model can adapt towards various downstream tasks. Figure 3 compares the the reward curve of SMART with Scratch and CT-Single, where models are pretrained with Exploratory datase 3 To avoid overlapping of curves, comparison with ACL and DT is shown and discussed later in Fig- u It can be seen that pretrained CT from both single-task dataset (CT-single singlel) and multi-task dataset (SMART) can achieve much better results than training from scratch. In general, under both RTG and BC finetuning, pretrained models have a warm start, a faster convergence rate, and a rel- atively better asymptotic performance in a variety of downstream tasks, In most cases, pretraining CT from multi-task dataset (SMART) yields better results than pretraining with only in-task data (CT-s per-singl- in) although it harder to accommodate multiple different tasks with the same model capacity, which suggests that SMART can extract common knowledge from diverse tasks.
Next, we show the generalizability of SMART. Figure shows the performance of SMART pre- trained on Explor lopttrry _ compared to Scratch and CT-single e on 5 unseen tasks. We can see that SMART is able to generalize to unseen tasks and even unseen domains, whose dis- tributions have larger discrepancy as compared to the pretraining dataset. Surprisingly, SMART achieves better performance than CT-single in most tasks, even though CT CT-single has already seen the downstream environments. This suggests that good generalization ability can be obtained from learning underlying information which might be shared among multiple tasks and domains,
Models pretrained with Ra ndom dataset show similar results, as can he seen in Appendix B.1

---

Published as a conference paper at ICLR 2023
Carae folk
Cherai
Hoper
Wilce bond Leantyy
Pablo
Cach
RTG
Cherch Clae
Miloe folovio
Maater
SHAET
RTG BC
Figure Downstream learning rewards SMART (red) compared with pretraining CT with single-task data (blue) and training from scratch (gray). ). Results are averaged over 3 random seeds.
Sepper Stad Soe
Vutur Voco Strites Palin Loaning
Fause Spe suadt
Buyac Bobrne Dhie
Nabor Fut, Albrrs Usen
Preddle
Pager
MAdI
BC
Figure 4: Downstream learning rewards in unseen tasks and domains of SMART (red) compared with pre-
training CT with single-task data (blue) and training from scratch (gray). Results are averaged over 3 seeds.
spanning diverse set of distributions. Appendix B.J provides additional results in other challeng- ing tasks which have larger discrepancy with pretramiig tasks, where SMART still performs well.
Next we evaluate the resilience of SMART by comparing with all aforementioned baselines, as visu- alized in Figure We aggregate the results in all tasks by averaging the normalized reward (dividing raw scores by expert scores) in both RTG and BC settings. When using the for pretraining. SMART outperforms ACL, and comparable to DT which has extra information of reward. When pretrained with the Random dataset, SMART is significantly better than DT and ACL, while ACL fails to outperform training from scratch. This result show that SMART is robust to low-quality data compared to other baseline methods.
dataset
KIGG-condiioned Policy Leurning
Dehavinr Cloning
Streatch Explecatory
Normalized Roward
Exploratary Random
Nemilize Reward
1:
0.6
11
0.4
0.2
0.2
Scratch
CT-Single
SMAKT
Scrathh
CT-SSigle
SMART
Figure 5: Downstream learning rewards (normalized by expert score) of all methods using Exploraton ory Random dataset. The gap between each pair of green and red bars corresponds to the resilience of each method to pretraining data quality, and our SMART shows the best resilience among all basclines.
6.3 ABLATION AND DISCUSSION
Ablation of Pretraining Objectives. As analyzed in Section 5.2 forward prediction and inverse prediction aim to learn the information of short-term control, while the random masked hindsight control (in short, Mask-Ctl) learns the long-temn control information. To show the effectiveness of combining these two kinds of information, we conduct ablation study over these three terms. Figure D demonstrates their relative improvements wrt Scratch defined in Appendix A.4 We
8

---

Published as a conference paper at ICLR 2023
Pooriiiing with Explocacory Data (RTG Dowastrean)
Preniéeing wit Sancos Dax (RDO Desmsteoo))
0.4
Nomaboofleuat
0.4
644
Ong Muc-C1
Alv-Mab-CC
SAAARI
Only M644-31
w...mmm
SMAKI
Figure 6: Ablation study on proposed pretraining objective. Rewards are averaged over tasks. Both long-term control information (Mask-Ctl) and short-term control information (Forward and Inverse) important.
Chenal Bac
Aides
(b) Overall performance v.S. model capacity. Figure 7: Ablation study of the effect of reward in pretraining and comparison various model capacity.
(a) Reward- reward-free pretraining.
that only learning long-term control information (Only Mask-Ctl) only learning short-term control information (w/o Mask-Ctl) renders much lower performance than the original SMART that leverages both types of information. addition, we find that both the performance slightly drops either forward prediction or inverse prediction is removed, as the combination of these two terms leads more stable and comprehensive representation for short-term control. This ablation study verifies the effectiveness of each component in the proposed pretraining objective. More results in other settings are provided in Appendix [B..T]
Reward v.s. Reward-baser Pretraining, As discussed in Section 5.1 including reward in- formation pretraining objectives is not necessarily helpful. We study the effects of rewards in pretraining by adding two reward-based objectives in the pretraining phase: immediate reward pre- diction and RTG-based action prediction. We evaluate this reward-based variant of SMART using exploratory dataset and RTGG-conditioned downstream learning. Note that the RTG-based action prediction used both in pretraining and finetuning of the reward-hased variant, providing strong supervision for the pretrained model. However, (surprisingly) observe that such supervision does not improve the downstream performance in many tasks, as shown in Figure 7a A potential rea- son is that reward-based objectives are more fragile distribution shifts, which also explains the non-ideal performance of DT pretrained from random data.
Discussion on Model Capacity. In large-scale training problems, performance usually benefits from larger model capacity (Kaplan et al. (220).. We investigate if this also applies to sequential decision making tasks by varying the embedding size (width) and the number of layers (depth) in CT. The aggregated results averaged over all tasks are show in Figure From the comparison, we can see that in general, increasing the model depth leads to better performance. However, when embedding size gets too large, the performance further drops, as a large representation space might allow for irrelevant information. Per-task comparison is provided in Appendix B.2
Other Potential Variants of SMART. Following the basic idea of SMART, there could be many other variants and extensions on the design of both architecture and pretraining objective (e.g., does an additional contrastive loss help?). We investigate several variants of SMART and discuss the results in Appendix [B.5.2 summary, adding more components to SMART could further improve its performance in certain scenarios, although such improvement may be limited and task dependeend. Exploring more extensions of SMART in various settings is an interesting future work.
7 CONCLUSION
This paper studies how to pretrain a versatile, generalizable and resilient representation model for multi-task sequential decision making. We propose a self-supervised and control-centric objective that encourages the transformer-based model to capture control-relevant representation. Empirical results in multiple domains and tasks demonstrate the effectiveness of the proposed method, as well its robustness to distribution shift and low-quality data. Future work includes strengthen the attention mechanism on both spatial observation space and temporal otate-observation interactions, well investigating its potential generalization in wider range of application scenarios.

---

Published as a conference paper at ICLR 2023
8 ETHICS STATEMENT
SMART a method that learns pretrained representations for controls tasks involving sequencial decision-making. As such, our models will reproduce patterns and biases found in its pretraining and finetuning datasets.
REPRODUCIBILITY STATEMENT
The codebase, pretrained models and datasets are provided at https://gthubb.ccmm//mrrroottttt
REFERENCES
Marcin Andrychowicz. Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, OpenAl Picter Abbeel, and Wojciech Zaremba. Hindsight experience re- play. Advances in neural information processing systems, 30, 2017.
Andrea Banino, Adria Puigdomenech Badia, Jacob C Walker, Tim Scholtes, Jovana Mitrovic, and Charles Blundell. CoBERL: Contrastive BERT for reinforcement learning. In International Con- ference on Learning Representations, 2022. URL. https::/openteviiet net. / forum?
Id=sRZ3GGmeqq:
Rogerio Bonatti, Sai Vemprala, Shuang Ma, Felipe Frujeri, Shuhang Chen, and Ashish Kapoor. Pact: Perception-action causal transformer for autoregressive robotics pre-training. arXiv preprint arXiv:2209.11133 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot leamners Advances in neural information processing systems, 33:1877-1901, 2020. 2020.
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint K2V:22202.09481, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 15084-15097. Cur- ran Associates, Inc., 2021. URL / wwwpwinnnnnnnnuurneeeeeeeerrrree nwwrips. wwwp:rripscccpaapeee cctppc/eeppcppppp///
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Compulational Linguisties: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pp. 4171- 4186. Association for Computational Linguistics, 2019. doi: 10.18653//1/111--44222 URL https::/doi.. /tttps://ddl.oor/1.... 1ttp:/00..115////1111/33
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity all you need: Learning skills without reward function. In International Conference on Learning Representa- fions, 2019. URL //tpttte////eeeeeeeeeee
http:///ppeeevvveeeee hettetewwweeeeeeeeeeeeeeeeeeeeee feeet//formm?i:ssf d= JRQFm
Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. In International Conference on Learning Representations, 2022.
URL htttps:://ppereviiiii net www.:neetoourrn forum id= ixVodll
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. PMLR, 2018.
Kaiming He, Haoqi Fan, Yuxin WU, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and patlern recognition, pp. 9729-9738, 2020.
10

---

Published as a conference paper at ICLR 2023
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: A survey of learning methods. ACM Comput. Surv. 50(2), apr 2017. ISSN 0360-0300. doi: 10.1 10.1145/3054912. htthtts:: htt.://..oorrr1111113333135000
Michael Janner. Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34, 2021.
Saurav Kadavath, Samuel Paradis, and Brian Yao. Pretraining & reinforcement learning: Sharpening the axe before cutting the tree. arXiv preprint arXiv:2110.02497, 2021.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprini arXiv:2001.0836/, 2020.
Alex Lamb, Riashat Islam, Yonathan Efroni, Aniket Didolkar, Dipendra Misra, Dylan Foster, Lekan Molu, Rajan Chari, Akshay Krishnamurthy, and John Langford. Guaranteed discovery of con- trollable latent states with multi-step inverse models. arXiv preprint arXiv:2207.08223) 2022.
Kuang-Hus Huee Lee, Ofir Nachum, Mengjiao Yang, Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, et al. Multi-game decision transform- ers. arXiv preprini arXiv: 15241, 2022.
Seunghyun Lee, Younggyo Seo, Kimin Lee, Pieter Abbeel, and Jinwoo Shin. Ofline-to-online reinforcement learning via balanced replay and pessimistic q-ensemble. In 5th Amnual Conference on Robot Learning, 2021.
Hao Liu and Pieter Abbeel. Behavior from the void: Unsupervised active pre-training. Advances in Neural Information Processing Systems, 34: 18449-184777 2021.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization, International Confer-
ence Learning Representations, 2019. URL http::://peeneereee..eet///rrrre
Bkg6RiCqY7
Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak. Discover- ing and achieving goals via world models. Advances in Neural Information Processing Systems, 34, 2021.
Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers sample efficient world models. arXiv preprint arXiv:2209.008: 00588, 2022.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforeement learning. nature, 518(7540):529-5333 533, 2015.
Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta. R3m: univer- sal visual representation for robot manipulation. arXiv preprint arXiv:2203. 2022.
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predic- tive coding, arXiv preprint arXiv:
2018
Takayuki Osa, Joni Pajarinen, Gerhard Neumann, J Andrew Bagnell, Pieter Abbeel, Jan Peters, et al. An algorithmic perspective on imitation learning. Foundations and Trends@ in Robotics, 7(1-2): 1-179, 2018.
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effectiveness of pre-trained vision models for control. arXiv preprint arXiv:2203.03588, 2022.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, pp. 2778-2787. PMLR, 2017.
Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al. Improving language under- standing by generative pre-training. 2018.
11

---

Published as a conference paper at ICLR 2023
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning, pp. 8748-8763. PMLR, 2021.
Ilija Radosavovic, Tete Xiao, Stephen James, Pieter Abbeel, Jitendra Malik, and Trevor Darrell. Real world robot learning with masked visual pre-training. In 6th Annual Conference on Robot
Learning.
Kate Rakelly, Abhishek Gupta, Carlos Florensa, and Sergey Levine. Which mutual- information representation learning objectives are sufficient for control? In M. Ran- zato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp. 26345-26333 Curran Asso- httptttss//ppoceeeette ciates, Inc., 2021. 2t0e///5555077770000004422623633 3544545588688885554477775888
neur i ps. cc./paperr/221//11ee. pdf
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06/73. 2022.
Max Schwarzer, Nitarshan Rajkumar, Michael Noukhovitch, Ankesh Anand, Laurent Charlin, R Devon Hjelm, Philip Bachman, and Aaron Courville. Pretraining representations for data- efficient reinforcement learning. In M. Ranzato, A. Beygelzimer. Y. Dauphin, P.S. Liang, and
J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, volume 34, pp- 12686- 12699. Curran Associates, Inc., 2021. URL https::/prrccedddntttt ngs neurips. pdf
cc/
/ttp://111ee6eba24644664 f 1eeaeeleeeeaaaaeeeaaa-e----
Younggyo Seo, Kimin Lee, Stephen James, and Pieter Abbeel. Reinforcement learning with action- free pre-training from videos. arXiv preprint arXiv:2203. 33880, 2022.
Rutav Shah and Vikash Kumar. Rrl: Resnet as representation for reinforcement learning. In Inter- national Conference on Machine Learning. PMLR, 2021.
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning, pp. 9870-9879. PMLR, 2021.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, al. Deepmind control suite. arXiv preprint arXiv:1801.00690 2018
Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa. dm control: Software and tasks for continuous control. Software Impacts, 6:100022, 2020. ISSN 2665-9638. doi: wciencedi com/
https://doi.org/00..00 wttp:i.a.aiaaaa...0000....0. http:://bbiooool/l.l.bu..mm.mopp URL
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz. Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- fion processing systems, 30, 2017.
Mengjiao Yang and Ofir Nachum. Representation matters: offline pretraining for sequential decision making. In International Conference on Machine Learning, pp. 11784-11794, PMLR, 2021.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer: In Kamalika Chaudhuri, Stefanie Jegelka, Le Song. Csaba Szepesvari. Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning, volume 162 of Pro- ceedings of Machine Learning Research, pp. 27042-27059. PMLR, 17-23 Jul 2022. hhtps://proceee.e
URL.
mlr. mtr.presrsrrrrrslm.mm / v1 62/ 2h.ng2cc.html
12

---

Published as a conference paper at ICLR 2023
Appendix
A IMPLEMENTATION DETAILS
A.1 IMPLEMENTATION OF PRETRAINING OBJECTIVES
I. Forward Dynamics Prediction.
For each observation-action pair (ot, (e)) in a control sequence, the forward prediction loss is con- structed as follows. Let go be the observation tokenizer being trained where 0 denotes the parameter- ization. We maintain momentum encoder ge whose parameters updated by 0 r(0) + (1-7)8 With the next observation 0s+1, we have $141 := (G(ō)(0(+1)), where SG refers to stop gradient. Then the forward prediction loss is defined as:
Lud == MSE (frma(0(0))) 8
where find is linear prediction head.
II. Inverse Dynamics Prediction. For each consecutive observation pair (or, 01+1), inverse prediction tries to recover the action that leads Og to Ot+1, which gives the loss function
Limy == MSE (Jim( ((o++)));
where fin is a linear prediction head.
III. Random Masked Hindsight Control.
In our implementation, we use predefined mask-token [-1,- -11 1] to replace the original
tokens. Let h denote the masked control sequence, and 0 81, 82, 8k S L be the selected indices for masked actions. Then the loss function can be defined as:
Limask-inv :-
Σ i=1 (ssi)
if 8; < L where ((s,) MSE (fmask-ine (/mast-ime((b/m (Mee aetse).
if 8; = L
where faskk inv is linear prediction head, and dM is the transformer model without causal mask. Note that we do not predict 0++ 1. as it is infeasible to recover it without future observations.
Schedule of Masking Size. Theoretically, it is possible recover full action sequence for given observation sequence, which implies that & = a reasonable setup. But in environments with complex dynamies, directly recovering all actions is hard in the start of training. Hence, we adjust the difficulty of random masked hindsight control in a curriculum way, by gradually increasing the value of in the following schedule:
current epoch + 1 +1)
k: max 1,int
L
(4)
totaleppohhs
On the other hand, if we predict actions from all observations, is possible that the model mainly relies local dynamics, i.e., predict a; mainly based 0,- 1, Ot-2, 0r-3, conflicting with our desire of learning long-term dependence. Therefore, we also mask a subset of observations with k'. In the extreme case k' L, the objective becomes similar to goal-conditioned modeling with o, being the start and 0t+L is the goal. However, in a control environment, there usually exist multiple paths from or to 0t+ f. making the action prediction ambiguous. [Schwareer al. (202I) uses con-ccondtionedd objective which finds the shortest path. However, extra value fitting and planning are required, which may lead to higher cost and compounding errors. Therefore, we intentionally make k smaller than the context length (half of L), such that the model is able to predict the masked actions based on revealed subsequence observations and actions. The schedule of is as below.
k' max (1,int
current.cpoch + +1) ))
(5)
total_ppochs
The overall pretraining objective is
min
Lswa + Limy + Lemask inv-
(6)
O Ponddafs fine Iandiam
13

---

Published as a conference paper at ICLR 2023
A.2 ENVIRONMENT AND DATASET.
lists all domains and tasks from DMC used in our experiments, and their relations are further depicted in Figure
Phase
Domain cartpole
Task swingup hopper hop
Expert Score by SAC
875
200
Pretraining & Finetuning
cheetah run
850
walker stand
980
walker run
700
cartpole
balance
1000
hopper
stand
900
Finetuning only
walker walk
950
pendulum
swingup
1000
finger
spin
Table 2: A list of domains and tasks used in pretraining and finetuning. The first 5 tasks are used for pretraining. In the finetuning phase, the pretrained model learn policies in all 10 tasks, including the last 5 tasks that are unseen during pretraining.
Usedd Protzwining
Only in
Cartpole
Chestah
Finge
Figure 8: Graphical relations of all tasks involved.
To generate our datasets for pretraining and finetuning, we first train 5 agents (corresponding to different random seeds) for each task with ground truth vector states using SAC (Haarnoja et al.) 2018) for 1M steps, and collect the full replay buffer with corresponding RGB images (3x84x84) rendered by the physical simulator. Then, we divide the replay buffers and create the following datasets of different qualities.
Random: Randomly generated interaction trajectories. This dataset has 400K timesteps per task. Exploratory The first 80K timesteps of each SAC learner, corresponding the exploratory stage. For all agents, this leads to a cumulative of 400K timesteps per task.
Note that different tasks have different difficulties for the SAC agent to converge. For example, in cartpole, the agent converges with much less samples than in walker. Therefore, we slightly adjust the proportion of data from different tasks in multi-task pretraining to avoid overfitting simple tasks and underfitting to harder tasks. In pretraining, we use 40K timesteps from each cartpole behavior agent (200K by 5 agents), and 90K timesteps from each walker behavior agent (450K by agents), and 80K for all other agents (400K by 5 agents).
Datasets for downstream learning:
Sampled Replay (for RTG): We randomly sample 10% trajectories from the full replay buffer of SAC agent, resulting in dataset of size 100K per task, with diverse return distribution. Expert (for BC): We select 10% trajectories with the highest returns from the full replay buffer of SAC agent, resulting in an expert dataset of size 100K per task.
A.3 MODEL AND HYPERPARAMETERS
Following the implementation of Decision Transformer (Chen et al. 202I), our transformer back- bone based on the minGPT implementation https:/gitthbbcoo/maaaattttttttt the default
14

---

Published as a conference paper at ICLR 2023
AdamW optimizer (Loshchilov & Hutter 2019). For both pretraining and finetuning, the learning rate is set to be x 10 and batch size 256. For learning rate, linear warmup and cosine decay are used. A context length 30 is used in all tasks for both training and execution. We tested different context lengths in preliminary experiments, and shorter context length (5/10/20) does not work as well 30 when training from scratch. We use 8 attention heads, 8 layer attention blocks, and an embedding size 256 in all experiments, for both our model and baselines. There are 10.8 M train- able parameters in the model. We test SMART with varying layer numbers and embedding sizes in Appendix B.2
For the execution of learned RTG-conditioned policies, we set the expected RTG as the expert scores in Table 2 Tuning the RTG setting may further increase the results. But since our focus is to show the effectiveness of pretraining. we did not explore other possibilities.
All models are trained for 10 epochs in pretraining, and 20 epochs for each downstream task. The performance of the best checkpoint is reported, as detailed the next section.
A.4 EVALUATION METRICS.
For both BC and RTG downstream learning, we report the average cumulative reward of the learned policy by interacting with the environment for 50 episodes. In Figure 3 we calculate the normalized reward based on expert scores.
In Figure O and Figure 7a we report the relative improvement of each ablated method calculated by the following formula.
relative improvement rewhod - scrathhhrwaaaacrarraaaa
scratch reward
where the scratch_rewar rewwad is the best reward of training from scratch using the same learning config- urations.
B ADDITIONAL EXPERIMENT RESULTS
B.1 MORE RESULTS OF DOWNSTREAM LEARNING
Heppr
RIG
Koocs
1040
Copplc Seaugrc Petorte Cher
Checki Kateoo Chee
Hapas Kitassa Cine
Nobor Unforen
Pallo Aabrn
baps
MACI
MAVG
S0AT
BC
Figure Downstream learning rewards of SMART (red) compared with pretraining CT with single-task data (blue) and training from scratch (gray), using the Random pretraining dataset. Results are averaged over random seeds.
Curves with the Random Dataset. The results in Figure 3 and Figure 4 are generated with the Exploratory pretraining dataset. Now, we show the performance of models pretrained using the Random dataset in tasks and unseen tasks in Figure and Figure T0 respectively.
Although the single-task pretrained model consistently outperforms training from scratch when pre- trained with the Exploraton ory dataset, it sometime gets worst-than-scratch downstream perfor- mance when pretrained with the Random dataset. Therefore, is challenging to overcome the distri- bution shift problem. In contrast, SMART still achieves much better performance than training from
15

---

Published as a conference paper at ICLR 2023
Coryle Sabess fale
130
RTG
Tacc
Toch
104n
Eopprr Siatte Dire
Soloo Fut Mirre Ohey
Ploge Age Bitanie Cisee
Page Salic
Scructy
MALII
SHAR
SHART
Downstream learning rewards in unseen tasks and domains of SMART (red) compared with pre- training CT with single-task data (blue) and training from scratch (gray), using the Random pretraining dataset. Results are averaged over random seeds.
Figure 10:
scratch in all tested tasks, which verifies the resilience of SMART due to multi-task self-supervised pretraining.
Generalizability of SMART: Tasks That Have Larger Discrepancy with Prefraining Tasks. In addition to the generalizability test show in Section 6.2 evaluate the performance of the pretrained SMART model in other more challenging domains/tasks from DMC, that sig- nificantly different from the pretraining tasks. These additional domain-tasks are: ball-in-cup-catch, finger- turn-hard, fish-swim, swimmer-swimmert and swimmer- swimmer15. 5. Note that these agents have significantly dif- ferent appearance and moving patterns compared to pre- training tasks, as visualized in Figure
Figure 11: Discrepancy between pretrain- ing domains and selected downstream do- mains: (left) Walker domain. (right) Swim- mer domain (6 and 15 links)
The results are shown in Figure 12 and Figure 13 where we can see that the pretrained model can still work in most cases, even under such a large task discrepancy. Note that here CT-Si Siglle pretrained with data from exactly the downstream task, where SMART has never seen a sample from the downstream tasks and is pretrained on significantly different domains. Therefore, it unsurpris- ing that CT FT-SS Single is generally better than SMART in this setting. However, it is interesting to see that SMART is comparable with or even better than CT Sinningle in some tasks, suggesting the strong generalizability of SMART.
On the other hand, one can imagine that is unavoidable that the performance of a pretrained model will decrease the discrepancy between pretraining tasks and downstream tasks increases. Therefore, we stress the importance of using diverse multi-task data for pretraining in practice.
Toge That loved Lrendey
And
KIG
]
Spack
Cock Aitatet
Tagy You Aibeken
Celenoon Solercet
apl
Auaoch
BC
Figure Downstream learning rewards of SMART (red) in challenging tasks that have larger discrepancy with pretraining tasks, using the Exploratory pretraining dataset. Results are from random seed.
16

---

Published as a conference paper at ICLR 2023
Tagg Testial boad Leuiin
RTG
The
BC
GovRT
Figure Downstream learning rewards of SMART (red) challenging tasks have larger discrepancy with pretraining tasks, using the Random pretraining dataset. Results are from random seed.
B.2 MODEL CAPACITY TEST
Figure 7b shows the results of varying model capacities averaged over tasks. To demonstrate the influence model capacity on different tasks, we provide the per-task comparison in Figure T4
Relutive Improvement
Embedding Size=128
Embedding Size=256
Embedding Size=512
cartpole
cheetah run
hopper hop
walker stand
walker run
cartpole balance
hopper stand
walker
swingup
walk
Relative Improvement
Layer=6
Layer=8
Layer=10
cartpole
cheetah run
hopper hop
walker
walker
carppole balance
hopper
walker
swingup
stand
walk
Figure 14: Comparison of varying model capacities (emhedding size and layer number) in different tasks in terms of relative improvement wrt training from scratch.
Averge Reluve inpovenerr wnt Scaza
Further Decreasing Model Capacity. In Figure 7b we can see that 6-layer transformer backbone is slightly better than the 8- layer transformer. It may suggest that the current network we use overly large for the tasks. To whether it is the case, fur-
0.8
ther decrease the number of transformer layers to 4 and 2. The results depicted in Figure [5] where we can see that 4-layer a 2-layer transformer is clearly worse than the 6-layer and 8-layer transformers. Therefore, the current architecture not too heavy for the DMC tasks. The exact numbers in Figure [5 are different from the corresponding ones in Figure 7b because Figure ab gates results from 3 random seeds but Figure 15 is from random seed.
0.05
Number
Tranfommer Bllcks
Figure 15: More model capacity tests. Results are from one ran- dom seed and are averaged over tasks.
17

---

Published as a conference paper at ICLR 2023
B.3 COMPARISON WITH PRETRAINED RESNET MODELS
Our empirical evaluation is done for multiple transformer- based pretraining approaches. In parallel to them, there are some existing pretraining paradigms that use large-capacity models as ResNet instead of transformers the backbones. Although it is hard directly compare the performance of totally different model architectures, we still provide the results of ResNet pretraining, to better posit this work in literature and verify the significance of our results.
Baselines and Implementation Details. We use the following two state-of-the-art pretraining ap- proaches.
CPC (Oord et al., 2018) self-supervised representation learning approach with contrastive predictive coding. It has been demonstrated success in many vision applications. When leveraging it in sequential decision making, state representations can be pretrained by decoupled from policy learning.
ATC (Stooke et al. 2021 is another contrastive learning approach target on decision making tasks. It also decouples representation learning from policy learning. Different from CPC which performs InfoNCE (Oord et al. (2018) loss in a predictive manner, ATC propose an Augmented Temporal Contrast directly compute InfoNCE loss among temporally augmented clips.
For both CPC and ATC, we use 3D-RReNet18 as the encoder backbone. During pretraining. a 2-layer 2D-ConwNet is used as the prediction head for both CPC and ATC. For CPC, a single-layer Con- vGRU is used as the aggregation network. Note that, during finetuning, both prediction heads and aggregation network are dropped. Only the pretrained encoder (3D-ResNetl 8) is used to produce the pretrained representations. During finetuning, we simply attach a single linear layer as the action prediction head on top of the pretrained encoder, which is under the same setting with ours.
For a fair comparison, we pretrain and finetune both CPC and ATC with a context length of 30 such that all comparing models are seeing the same time horizon. Follow the widely used training protocol in decision making tasks, we also leverage a frame stacking with stack size as 3 when training both of them. We keep the other hyperparameters the same with default setting.
Results. We compare our SMART with CPC and ATC in RTG and BC downstream tasks, with Exploratory and Random dataset. The results are shown in Figure 16 and Figure 17 respec- tively. Although training a transformer and training a ResNet model from scratch usually produces different rewards in the same task, the results show that pretrained models are able to their cor- responding train-from-scrate baselines. When using RTG downstream learning, we can see that SMART outperforms CPC and ATC in all downstream tasks. When using BC as downstream learn- ing method, the transformer backbone fails to get a high score and so does SMART pretrained mod- els. But remaining tasks, SMART is still significantly better than CPC and ATC. This suggests the advantages of our pretraining framework SMART and the proposed model Control Transformer.
Transformer from Scratch ResNet from Scratch
CPC (ResNet)
ATC (ResNet)
Ours (Transformer
Sepa
Polae 2014
60-
RTG
BC
Figure 16: Downstream learning rewards of SMART (red) compared with CPC (darkblue) and ATC (darkor- ange), using the Expl oratory pretraining datasct. CPC and ATC are based on ResNet, whose performance of training from scratch is shown darkgreen in comparison to training transformer from scratch (grey).
18

---

Published as a conference paper at ICLR 2023
Transformer from Scratch
ResNet from Scratch
CPC (ResNet)
ATC (ResNet)
Ours (Transformer)
Oveek
Pellerion 2012
Saive Kolk orTs
1pach
Comple Riy
Cherc Roc
Send
Vilr Aah
BCC
Figure 17: Comparison with ResNet-based pretrained models using the Random pretraining datasct.
B.4 PRETRAIN CT FOR ONLINE FINETUNING
Similar to most transformer-bast based decision making models, we consider policy learning in the form of IL and offline RL. For online RL, special care is needed due to the environment uncertainty and the need exploration. A recent work by Zheng et al. (2022) proposes online decision transformer (ODT), which first pretrains the model with offline trajectories and then finetune the model in an online manner. Zheng et al. 2022) demonstrate the effectiveness of ODT in a series of MuJoCo tasks with groundtruth state being observations. It is shown that the offline pretraining phase is crucial for online fintuning with transformer-basee models. However, the pretraining phase of ODT is DT with reward supervision. With the self-supervised control-centrie pretraining objective (Trwd, Trmes Lmaas--in) proposed in this work, we would like to ask the following questions. (1) Can we replace the supervised DT pretraining objective with our self-supervised pretraining objective that does not require reward supervision? (2) Can we improve ODT by combining our objectives with it?
We follow the open-soureed implementation of ODT and evaluate our proposed objective using the default model and hyperparameter settings of ODT. We first replace the DT pretraining loss with our self-supervised losses. The results shown Figure T8 suggests that even without any reward information, pretraining ODT with our self-supervised losses can achieve comparable performance with pretraining with reward supervisions. Although our pretrained model takes more steps to warm up due to the lack of supervised pretraining, quickly converges to similar results ODT in online finetuning. Therefore, in practical applications where only reward-free pretraining trajectories are available, DT pretraining is infeasible while our pretraining be used without sacrificing the finetuning performance.
Sald Foi pec (raral
3.00
Rowad
2.00
1000
1,00-
20
Sepvvied Promlt
Sel Sctaggeniol Pern Hem
5d Sibogerbodd.ooo tomm
12
120
100
125
Dpach
Speck
Epecc
Figure 18: Comparison of our self- self-supervised pretraining objective and the supervised pretraining objective of ODT in hopper and walker2d environments with medium and random pretraining dataset.
To answer the second question, we combine our pretraining objective with ODT. There are two potential ways of such combination: (1) only combine in pretraining, and (2) add our self-supervised losses during both pretraining and finetuning (the losses serve auxiliary tasks in online finetuning). We evaluate both versions in our experiments, as shown in Figure 19 The results suggest that including pretraining objective improve the performance of ODT, especially when the data quality is low pretraining, it suffers from high variance in downstream learning. To address this issue, we find that
dataset is used). However, when our objective is only used during
19

---

Published as a conference paper at ICLR 2023
adding objective auxiliary loss in finetuning effectively improve the result and decrease the variance.
Kibat puaic
OCT
CDT (anilaaal
7000
EDT
60
OETT
2081+
OET
1.00
OET Dan (oddlo))
CGT 0an (edd
13
60
20
90
1peck
Figure 19: Combining our pretraining objectives with ODT produces better results, especially when the pre- training dataset has low quality. Results are averaged over 5 random seeds.
It is also interesting to see that the online learning performance of ODT depends the quality of pretraining dataset. When the random dataset used, ODT converges to suboptimal policy in finetuning, although it could have explored better solutions. Adding our self-supervised pretraining objective improves the performance, but it still cannot match the performance of models pretrained in higher-quality dataset. We hypothesize that this is due very large distribution shift between the random dataset and the ideal dataset for online learning. That is, the model pretrained with random data overfits to the random behavior and fails to further explore and exploit. How make a transformer-based model adapt to an online environment with low-quality pretraining data is still an open problem that aim to study in our future work.
B.5 ADDITIONAL RESULTS OF ABLATION STUDY
B.5.1 ABLATION OF PRETRAINING OBJECTIVE
In addition to Figure 6 which uses RTG the downstream learning objective, we also provide the ablation results of BC downstream learning in Figure Similar the results of RTG, removing any term from the control-centric pretraining objective usually causes performance drop, which verifies the effectiveness of our objective design. One exception that removing the inverse prediction gives better result for BC with exploratory pretraining data. Combined with results in other cases, can find that inverse prediction is not as helpful as forward prediction and random masked hindsight control average. But in most cases, especially when pretraining data is low-quality (Random), including the inverse prediction still improves the performance.
Implementation Details. For the ablation study, we pretrain the model with the same hyperparam- eter settings on the same datasets using different versions of (ablated) pretraining objectives. The cartpole-balance hopper-hop, hopper-stand,
results are averaged over tasks: portole-swinggeee cheetah-run, walker-stand walker-run, and walker-walk, spanning both seen and unseen tasks. Due the large number of experiments (4 variants 2 learning scenarios downstream tasks 128 experiments), the results in Figure 6 and Figure 20 are from one random seed. Although randomness exists, we believe that the conclusion is meaningful as the results are averaged over multiple tasks and multiple learning scenarios.
dataset selections
Expaarar
Nemalord
0.4
0.5
0.2
Figure 20: Ablation study our pretraining objective with BC (behavior cloning) being downstream learn- ing algorithm. Both long-term control information (Mask Ctl) and short-term control information (Forward and Inverse) are important.
B.5.2 VARIANTS OF PRETRAINING OBIECTIVE
We also investigate some other possible variants of the pretraining objective.
(1) Multi-step Inverse (inspired by Lamb et al. (202))): masking all states/actions but 8; and 8
20

---

Published as a conference paper at ICLR 2023
and predicting dg-
(2) Max Fixed Mask: using a fixed mask size l = L and k' = L/2 instead of gradually increasing k and K'
(3) + Masked State Prediction: in addition to the original objective, this variant also predicts the masked state token in the third term. (The state token target is generated from the momentum encoder, as in forward prediction.)
+ Contrastive Loss; in addition the original objective, this variant also employs an additional contrastive loss, which takes the state/action token and their representations for the same timestep as positive pairs. while regarding representations from other timesteps as negative samples (inspired by Banino et al. (2022)).
The implementation and comparison of these variants follow the same setup with the ablation study in Appendix [B.5.]] Note that (3) and (4) are both adding new losses to the original method, which renders higher computational cost.
The results of comparison with these variants are shown in Figure 21 with RTG being the down- stream learning objective. The results show that Multi-step Inverse and Max Fixed Mask are in general worse than SMART, suggesting the effectiveness of our design. Interestingly, optimizing extra state prediction loss or a contrastive loss gives better performance when the pretraining data is Random, while hurts the performance when the pretraining data is Expleratory We hypothe- size that this because such two additional losses emphasize more on capturing visual information from pixel observations, which more important for learning and transferring representation from the Random dataset. In other words, a better visual representation is more crucial under a larger distribution shift (on behavior policies) that happened when a Random dataset is used for pretrain- ing. On the other hand, when the pretraining data is of higher quality, adding these two components downgrades the original SMART. potential reason that with smaller distribution shift, the control-centric objective can already capture sufficient information from data, while adding new losses may make the optimization harder and less stable.
Premsising Exptoratory
Promuning Sandos (KTG (lwnnmnoo)))
Nomatoou Nosed
Noved
0.75
0.8
0.03
0.5
Nax Makk Makal Phal Deettatcc
SMART
Naklo Paal Maroed Phal
Gettzath
SMART
Figure 21: Comparison between SMART and its variants. Adding masked state prediction and contrastive loss SMART could increase the performance, but improvement not across all scenarios, with a cost more computations.
In summary, although none of the tested variants can outperform SMART in all scenarios, some of them can render better results in certain cases by adding extra pretraining losses to the original objective, at a cost of more computations and possibly more challenging optimization. Therefore, whether to use these variants in practice depends the specific use case and computation resources. deeper and more thorough understanding of the relationships among these pretraining losses de- sired for better application of pretraining methods, which would be our future work. We emphasize that multi-task pretraining for control tasks a relatively new area. This paper takes one step fur- ther towards general large-scale pretraining models for sequential decision making. Exploring more possibilities of SMART would be an interesting and important future direction to further improve its performance variety of real-world tasks.
21


Reviews:
[{"review_id": "yQotcA095N", "reviewer": "Reviewer_Jx8K", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors consider the problem of unsupervised pretraining representation for  high-dimensional sequential control in various downstream tasks. The authors propose a framework as a pretraining then finetuning pipeline for sequential decision making, it consists of a Control Transformer (CT) which is coupled with a novel control-centric pretraining objective in a self-supervised manner. The authors show that the proposed method, SMART, significantly improves the learning efficiency in both seen and unseen downstream tasks and domains. ", "summary_of_the_review": "This paper tackles the problem of unsupervised pretraining for sequential decision making. The proposed algorithm is simple and effective, with extensive experiments to show its effectiveness. There are some minor limitations of this paper that I would like the authors to address.  ", "strength_and_weaknesses": "Strength\n- This work proposes a simple and effective method, for unsupervised pretraining and adaptation to downstream imitation learning and RL tasks. \n- The experiments are extensive and show the strength of SMART, and I appreciate the authors conduct pretraining experiments on random and exploratory datasets since otherwise simple behavior cloning would be sufficient.  \n\nWeakness\n- How does SMART perform in goal conditioned tasks and model-based tasks? Showing the results could help have a better understanding of what SMART captures from pretraining, especially since it’s trained with forward, inverse dynamic prediction and random hindsight prediction. \n- The Figure 11 scalability experiment is interesting because it shows after a certain amount of model capacity, there is no further relative improvement. I wonder if the authors tried smaller model sizes since the current smallest model is arguably quite large for simpler tasks. \n\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "mSQ3c37y7ow", "reviewer": "Reviewer_JbXg", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "In this paper, authors present a method for pretraining a generalisable and resilient SSL model with control transformer for multi-task sequential decision making. They propose a self-supervised and control-centric objective that encourages the transformer-based model to capture control-relevant representation. The evaluation is performed on multiple domains and tasks. The results show model’s effectiveness and robustness to distribution shift. \n\n", "summary_of_the_review": "The proposed SMART method for training SSL model with control transformer for multi-task sequential decision making is shown as effective and robust. This paper clearly presents the proposed approach and conducts extensive experiments on multiple tasks and comparison with multiple baselines. \n\nI think authors should consider improving section 6.3 with more discussion on the results of ablation experiments.\n\nAdditonally, I find this paper has limited contribution with control transformers. \n\n\n\n", "strength_and_weaknesses": "Strength: This paper proposes an effective approach for pertaining SSL model for multi-task sequential decision making. This paper clearly presents the proposed approach, with carefully designed extensive experiments on DeepMind Control Suite and comparison with a set of baselines.\n\nWeakness: Overall, I find this paper has limited contribution. Also, the section 6.3 (Ablation and Discussion) is brief towards discussion and ablation studies can be extended.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "M39Gzjie53x", "reviewer": "Reviewer_CXdh", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper presented a new self-supervised pre-training method for sequential decision-making tasks. Specifically, a new general pre-training-finetuning pipeline named SMART (Self-supervised multi-task pertaining with control transformer) was proposed with a Control Transformer (CT) and a control-centric pre-training objective for the model supervision. Experiments in DeepMind Control Suite show the effectiveness of the proposed method over several tasks. The main contributions of this paper are the proposed new self-supervised pre-training pipeline with the corresponding objective and also the carefully designed experimental analysis.", "summary_of_the_review": "The studied problem is clearly stated with convincing motivations, the paper is well-written and the idea is mostly clearly presented. The experiments also well validate the claims with shown effectiveness of the proposed method over other solutions. Although there are a few concerns as mentioned in the Weaknesses section above, they could be addressed in the rebuttal phase, and overall this paper would be a good contribution to the community. As a result, I would like to recommend 'accept' for this paper.\n\n-----------------After rebuttal---------------------\n\nThanks to the authors' responses, which well addressed my concerns. The additional included experiments and the clarification also make the paper stronger.", "strength_and_weaknesses": "**Strengths**\n\n\\+ The paper well summarised the challenges unique to sequential decision-making scenarios and accordingly proposed a solution to address the issues.\n\n\\+ The idea of the proposed control transformer (CT) is interesting and well-motivated, especially the purposefully designed exclusion of the reward signal for pre-training.\n\n\\+ The experiments were well designed to validate the proposed method from several different perspectives in response to the questions mentioned at the beginning of Sec. 6.\n\n\\+ The proposed method was shown to have better resilience when facing pretraining data with different qualities, as shown in Fig. 5. \n\n\\+ The paper is well-written and easy to follow.\n\n\n**Weaknesses**\n\n\\- The novelty of the proposed random masked hindsight control is a bit limited. Although the authors claimed the difference to BERT-like models, the general idea of mask-and-predict is the same and the only difference is the items to mask. This idea is also similar to those visual pre-training models like ViT etc.\n\n\\- It is unclear why only masking the actions while leaving out the observations. The authors claimed the reason to be \"force the model to learn global temporal relations\", but there is no clear evidence to support this.\nWhat if including the observations as well, or partially? or what if similar to the random scheme as proposed, randomly predict the observations as well?\n\n\\- From the description, the overall objective used all three mentioned terms as mentioned in Sec. 5.2 / Fig. 2. The description also seems to suggest the effectiveness of the 3rd term of random masked hindsight control, but it is unclear how would the model performs if only using this term.\n\n\\- For the experiment of ‘versatility’ as shown in Fig. 3, there are some cases where the CT-single performs on par with or even better than the proposed SMART, i.e. single-task performs better. It would be better if some explanation or insights could have been provided.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "-L7gu-8xmR", "reviewer": "Reviewer_EgV4", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The manuscript proposes a new method which learns essential representations in a pretraining phase, including both long-term and short-term dependency information, in a sequential decision making process using reinforcement learning. For this, authors have proposed a novel control-centric objective which contains three components: forward dynamic prediction, inverse dynamics prediction, and random-masked hindsight control. First two components exploit local dependency information and the third component learns long-term dependency. The proposed method is first pretrained on multi-task dataset to learn reward agnostic representations which are utilized on downstream tasks. The method is evaluated on DeepMind Control (DMC) suite, demonstrating superior performance compared to scratch, single task pretrained, and ACL method reported by Yang et.al. 2021.", "summary_of_the_review": "The proposed method is innovative and effective as demonstrated by their empirical study. The presentation of the paper is clear and code will be available to the research community to facilitate the reproduction of reported results and applications on other learning tasks. ", "strength_and_weaknesses": "Strengths:\n+ The proposed method learns reward-agnostic representations in pre-training phase to improve generalizability\n+ The proposed method effectively captures long-term dependency information using random-masked hindsight control mechanism.\n+ Ablation studies demonstrates effectiveness of each component of control-centric objective function.\n+ The paper is organized well and explained every aspect in clear terms.\n \nComments/Questions:\n- Figure 1 finetuning phase shows that R_t is shown as input at all different time points: t, t+1 and t+L. Are there typos? Should it be R_{t+1} at t+1 and R_{t+L} at t+L time stamps?\n- Proposed method was evaluated on DeepMind Control (DMC) suite tasks cheetah: run, cart-pole: swingup, cart-pole: balance, hopper: hop, hopper: stand, walker: stand, walker: run, walker: walk, pendulum: swingup, finger: spin. This dataset also contains other tasks such as ball-in-cup: catch, reacher: hard, finger: turn, manipulator: bring ball, swimmer: 6 links, swimmer: 15 links, fish: swim, cheetah: random action, hopper: random actions, walker: random actions, humanoid: random actions, humanoid: stand, humanoid: walk, humanoid: run, CMU motion capture. It would be interesting to see how the proposed method would perform on these slightly more challenging tasks. \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: notable-top-25%