Masked Autoencoding for Scalable and Generalizable Decision Making
Authors: Fangchen Liu
Hao Liu
Aditya Grover
Pieter Abbeel
Venue: NeurIPS 2022 Accept (2022)
Keywords: reinforcement learning
unsupervised pretraining
zero-shot
transfer learning

Abstract: We are interested in learning scalable agents for reinforcement learning that can learn from large-scale, diverse sequential data similar to current large vision and language models. To this end, this paper presents masked decision prediction (MaskDP), a simple and scalable self-supervised pretraining method for reinforcement learning (RL) and behavioral cloning (BC). In our MaskDP approach, we employ a masked autoencoder (MAE) to state-action trajectories, wherein we randomly mask state and action tokens and reconstruct the missing data. By doing so, the model is required to infer masked out states and actions and extract information about dynamics. We find that masking different proportions of the input sequence significantly helps with learning a better model that generalizes well to multiple downstream tasks. In our empirical study we ﬁnd that a MaskDP model gains the capability of zero-shot transfer to new BC tasks, such as single and multiple goal reaching, and it can zero-shot infer skills from a few example transitions. In addition, MaskDP transfers well to offline RL and shows promising scaling behavior w.r.t. to model size. It is amenable to data efficient finetuning, achieving competitive results with prior methods based on autoregressive pretraining.

Masked Autoencoding for Scalable and Generalizable Decision Making
Fangchen Liu!' Hao Liu!'
Aditya Grover?
Pieter Abbeel'
Berkeley Al Research, UC Berkeley UCLA
Equal contribution
(fangchen_ liu, hao ww..iul@berkkeee..oo edu
Abstract
We are interested in learning scalable agents for reinforcement learning that can learn from large-scale, diverse sequential data similar to current large vision and language models. To this end, this paper presents masked decision prediction (MaskDP), simple and scalable self-supervised pretraining method for reinforce- ment learning (RL) and behavioral cloning (BC). In our MaskDP approach, we employ a masked autoencoder (MAE) to aate-action trajectories, wherein we ran- domly mask state and action tokens and reconstruct the missing data. By doing so, the model is required to infer masked-out states and actions and extract information about dynamics. We find that masking different proportions of the input sequence significantly helps with learning a better model that generalizes well to multiple downstream tasks. In our empirical study, we find that MaskDP model gains the capability of zero-shot transfer to new BC tasks, such as single and multiple goal reaching, and it can zero-shot infer skills from a few example transitions. In addition, MaskDP transfers well to offline RL and shows promising scaling behavior w.r.t. model size. It is amenable to data-efficient finetuning, achieving competitive results with prior methods based on autoregressive pretraining
Introduction
Self-supervised pretraining has made tremendous successes for unsupervised representation learning in natural language processing (NLP) and vision 12
y 3 4. These methods work by predicting a removed portion of the data, which is often referred to masked token prediction. By varying the masking patterns and architectures, different methods have been developed for NLP and vision, e.g.,
Transformer 3 GPT 4 BERT (}) and MAE [12]. These methods are simple to implement and scalable to large Internet-scale datasets and deep neural networks, leading to excellent flexibility and generalization for downstream tasks
In this work, we explore the generality of masked token prediction for generalizable and flexible reinforcement learning (RL). Prior work has explored sequence modeling for sequential decision making in the context of offline RI. e.g. decision transformer (DT) 3 and trajectory transformer (TT) 113) and black-box optimization e-g- transformer neural processes (TNP) [T These methods are based on autoregressive next token prediction, similar to GPT While promising, these works do not leverage diverse unlabeled data for generalization across various downstream tasks. In addition,
DT [6] needs reward-labeled high quality datasets, while TT [K] requires discretizing states and actions, further limiting its applicability. The flexibility of applying arbitrary masks for executing various task specifications in RL significantly lags behind NLP and vision.
Thee implementation of MaskDP is available wttttpss//ghhaai.ccooo//aaaaaaau
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

Pre-Tree Training
Downstream Tasks
Masked Decision Prediction
Goal Resching
Transorner
+L
+
FanL
Ridinectional Decader
Skill
Rdleectional Tratsfeersst
rast
Fast
BAA
Ridirectional Rncader
-..
4
4
Offinne
Figure Illustration of MaskDP. During pretraining stage, perform the masked token prediction task. And after pretraining, the model can be deployed various downstream tasks using different mask patterns.
We propose Masked Decision Prediction (MaskDP), a pretraining method learn generalizable mod- els that achieve data-efficient adaptation to various downstream tasks. MaskDP is a self-superrvied pretraining method that can leverage unlabeled diverse data. With MaskDP pretraining, the model can generalize well to both goal reaching and offline RL, two distinctive and popular RL paradigms.
Our first key observation is that masked token prediction with random masking similar to BERT ( and MAE 12 provides a general and flexible way for learning from unsupervised data. Unlike autoregressive action prediction used in prior works, random masking is strictly more general and requires the model to infer masked out states and actions, and thus leads to a single model that can reason about both the forward and inverse dynamics from each sample.
Our second key observation is that since states and actions are highly correlated temporally, trajecto- ries have significantly lower information density, i.e. it is easier to predict action or state based on nearby states and actions. Consequently, a high mask ratio (95%) necessary to make reconstruction task meaningful. Unlike in MAE [2] and BERT [9] where the goal is learning representations, we want directly apply MaskDP to various downstream tasks, and different mask ratios induce different pre-train and downstream gaps. For example, consider the goal-reaching task within certain time limit. Given current state, future goal and mask tokens between them, the model should be able to inpaint intermediate actions the goal-reaching plan. The mask ratio varies from short-term plans to long-term plans. Therefore, we combine multiple different mask ratios (e.g. 15%, 35%, 75%, and 95%), and mask portion of data using randomly sampled mask ratio. Our experiments show that doing so is crucial to achieving high performance. We show that self-supervised pretrained MaskDP achieves high performance in challenging multiple goals reaching setting, outperforming strong baselines in a zero-shot manner.
We highlight our key results here:
Single goal reaching: MaskDP achieves performance that exceeds or matches both training from scratch task-specific methods and other pretraining based methods.
Sequential multiple goal reaching: MaskDP can reach sequence of goals effectively, even without closed-loop execution, while outperforming iterative baselines significantly.
Offline RL: MaskDP achieves competitive results as specialized approaches. Notably, we demonstrate that non-autoregressive architecture works well for offline RI. tasks.
2
Related work
Masked modeling in language and vision. Large-scale language models are highly successful g
Later pretraining on a large amount of data, these pretrained representations generalize well to
various downstream tasks. Taking inspiration from the success in NLP, Transformer
[ based

---

methods have been proposed to model images 7 10 3 12
iGPT V operates on sequences of
pixels and predicts unknown pixels. BEiT [3] proposes to predict discrete tokens B2 22
MAE [2]
proposes to randomly mask patches of the input image and reconstruct the missing pixels. Since we apply random mask across states and actions, our work is also related to prior work on masked prediction across multiple input modalities [see e.g. 34
Masked modeling in RL Masking trajectories for pretraining has been studied in RL 21
23
39 MVP (35] studies transferring pretrained visual representations to RL tasks. MWM 25 studies
masked prediction over convolutional features, and learn latent dynamics model. Modelling inverse
dynamics has also been studied for robot learning from demonstrations and sim-to-real transfer S
30
TT
[K]
studies autoregressive next token prediction for model-based RL applications. DT
39
study masking autoregressive next token prediction conditioned on return. TNP 19] and BONET study autoregressive masking for sequential decision making for black-box optimization. ICM and SPR ( study predicting masked state and action in a transition tuple for exploration. Different from these works, MaskDP randomly masks a portion of trajectories and generalizes prior masking strategies such as inverse dynamics. In addition, MaskDP generalizes well to downstream tasks while prior work is task-specific. Concurrent our work, Uni[MASK] [5] proposes using the bidirectional transformer to predict masked states and actions and demonstrates that the resulting model performs well on large variety of tasks. The main difference between our works is that Uni[MASK] is more interested in comparing the performance between training regimes with different masking schemes, while our work focuses on solving a range of downstream tasks with minimal task-specific designs during pretraining.
14
Unsupervised pretraining in RL Our work falls under the category of self-supervvsed pretraining in RL. Self-supervised discovery of a set of aank-sgn behaviors by means of seeking to maximize an intrinsic reward has been explored intrinsic motivation 22 often with the goal of encouraging
exploration (27 20 APT [8] studies nonparametric entropy maximization for pretraining and is extended to learning skills 17 Proto-RI. 36 further improves pretraining by representation
learning. CIC [T combines contrastive learning with skill discover and improves results on URL benchmarks [13]. APV (26) shows successful transfer of pretrained representation across domains. Many of these methods are used to pretrain agents that are later adapted to specific reinforcement learning tasks. Using offline data for pretraining agents has also been explored in prior work [24 28. 38
SGI 24 proposes combining self- predictiee representation [ and inverse dynamics prediction. ATC
2 studies contrastive pretraining on trajectories and shows transferring
the representations to downstream tasks.
3
Method
The key idea in MaskDP is to mask and reconstruct state-action sequences during the pretraining stage. Post pretraining. MaskDP can be zero-shot deployed finetuned for various downstream tasks. The paradigm of the model for pretraining and finetuning is summarized in Figure !
3.1 MaskDP Pretraining
Random masking. For sequences with low information density, high masking ratio is typically applied to eliminate information redundancy and make the task sufficiently difficult to avoid trivial interpolation from visible neighbor tokens. However, unlike vision and language, where the goal is to learn good representations; we also consider directly deploying this model by leveraging its inpainting capability for various downstream tasks. For example, we can give the model a goal at timestep T and mask all the future inputs, the model can generate intermediate actions by inpainting the mask tokens. The mask ratio varies from goal to goal, depending the gap between training and deployment, we keep a set of mask ratios (i.e. 15% 35%, 50%, 75%, and 95%), and the data randomly masked with ratio sampled from this set. We find that masking multiple proportions of the input yields a meaningful self-supervisory task.
the time budget. To reduce
We apply random masking on state tokens and action tokens independently. By doing so, the model implicitly learning both forward and inverse dynamics. This also provides more flexibility as we can provide state action-level inputs but not transition-level.
3

---

Architecture Our encoder is a Transformer [33 but applied only on visible, unmasked states and actions, similar to MAE [2]. The states and actions are first embedded by separated linear layers, positional embeddings are then added, and lastly, the embeddings are processed by series of self-attentional blocks. The decoder operates on the full set of encoded visible state and action tokens and mask tokens. Each mask token is a shared, learned vector that indicates the presence of a missing token to be predicted. Similar to the encoder. the masked whole sequence will pass through separated linear projections added with positional embedding prior to being passed to the decoder. Both the encoder and decoder are bidirectional.
Prediction target Our MaskDP reconstructs the input by predicting the whole action and state sequences. The last layer of the decoder consists of two MLPs to decode states and actions separately. The loss function computes the mean squared error (MSE) between the reconstructed whole sequence and original inputs. Different from other masked prediction variants [12] 9. we found mask loss is not useful in our setting, as our goal obtain an scalable decision making model but not only for representation learning.
3.2 MaskDP Downstream Tasks
MaskDP for goal reaching We consider the problem of reaching one goal or multiple goals from given state. The model has to generate sequence of actions to reach goals within a certain amount of steps. MaskDP denoising pretraining objective fits the goal reaching scenario well as the model must learn to inpaint masked actions based on remaining states. In this task, the MaskDP encoder input is a concatenation of initial state and goals, and the decoder imput is a concatenation of initial state embedding, masked token sequence, and goal embeddings. Note that the number of masked tokens determines the number of timesteps the model expected reach the given goals. The model then generates a state-action sequence, where we can directly execute the whole action sequence (namely "open-loop"), or only execute the first action and forward the model again with the obtained new observation (namely "closed-loop").
MaskDP for skill prompting Skill prompting requires the model to generate a trajectory condi- tioned a given context. For example, consider a walker agent: if we prompt it with few state-action pairs of walkinng/ruming/ttannee Accordingly, append the observed initial state-action sequences with masked tokens for the future. The model can be rolled out once to generate the whole future sequence, or queried repeatedly to refill the masked tokens at each time step. Similar to goal-reaching task, we refer to these strategies as "open-loop" and "closed-loop" respectively.
it should continue to generate a trajectory in the same skill pattern.
MaskDP for offline RL In offline RL, the objective is to learn one model for maximizing the return for a task specified by a reward function. This is different from our self-supervised pretraining target, so extra finetuning is needed. We adopt standard actor-critic framework similar to TD3 [I] by adding a critic head and actor head, where the actor takes a state sequence input, and the critic takes the actton sequence as input. Both are mask-free. To match the setting in RL, we change the bidirectional attention mask in the transformer to causal attention mask. More details about RL finetuning can be found in section 4.2.3
4 Experiments
In our experiments, we evaluate transfer learning in downstream tasks using MaskDP. Section introduces the environments, pretraining, and the baselines compared in experiments. Section summarizes the results of MaskDP on goal reaching, skill prompting, and offline RL. Through further analysis in Section we present an ablation study on various design choices of our model.
4.1 Experiments Setup
Environments: domains vs. tasks We adopt the environment setup used in EXoRL 377 based on DeepMind control suite (2 where domain describes the type of agent (e.g. Walker) but tasks are specified by rewards ge.g., Walker walk, Walker run). We use domains (Walker, Cheetah and Quadruped) with 7 tasks in total. More details about the environments can be found in the Appendix.

---

Pretraining datasets Real-world pretraining data generally varies greatly in quality. To mimic this, we construct two different pretraining datasets to approximate different data quality scenarios.
Near-expert: For every task, we train a TD3 agent [1]] for 1M steps and freeze its parameters. We rollout the policy with Gaussian random noise and collect 4M experience on each task.
Mixed: This dataset consists of diverse data collected from various agents, including 2K near- expert trajectories for each task. Similar to ExoRI. 37. we collect 10M exploratory trajectories using intrinsic reward from Proto-RI. [66] for each domain. We also to use TD3 [[]] agent to maximize the sum of extrinsic reward and the Proto-RI. intrinsic reward, and store its 2M experience on each task.
For more details about the above datasets and more ablations the dataset quality, please refer to Section and Section B. respectively. We perform both single-task and multi-task pretraining using the above datasets. The former leverages task-specific data while the latter utilizes data from all tasks within the same domain. We pretrain agents for 400K gradient steps. Specifically, for the model pretrained on the near-expert dataset, we perform zero-sho eevaluattion of goal reaching and skill prompting, and finetuning for offline RL; for model trained on the mixed dataset, we provide the finetuning results in Section 4.3 and Section A.
Baselines
GPT. We train an autoregressive model similar to GPT which takes the past states and actions as input to predict the next state or action.
Goal-GPT. We specifically modify GPT to Goal-GPT to evaluate its performance on goal reaching tasks. The model takes current goal and observations as input, and predicts the action to reach this goal. The model is trained using behaviour cloning loss as [
Goal-MLP. Standard behavior cloning method that conditions on the goal. The major difference between this and Goal-GPT is here we do not use the causal Transformer architecture make the history visible.
By default, MaskDP uses a 3-layer encoder and 2-layer decoder, and the baselines based on GPT use 5 attention layers. MaskDP and all the above models are comparable with similar architecture design and size, and share the same training MaskDP and the above baselines can be found in Section
Details about the architecture and training of
4.2
Main Results
4.2.1
Goal Reaching
We consider both single and multiple goal-reaching settings. The agent is required to reach one or multiple goals from given state, which are all sampled from the same trajectory to guarantee reachabilty within a reasonable time budget. During evaluation, the agent rolls out to reach the given goal(s) within time budget. The evaluation dataset is also collected by the same RL agent in 3 environments with different seeds, which is unseen during pretraining. The detailed settings are:
single-goal reaching: For every trajectory in the validation set. randomly sample a start state and a future state in E [15, 20) steps as the goal. All the methods are evaluated on the same set of 300 state-gowe pairs with given budget of T + 3. We set the agent to the start state and report the 1.2 distance between the goal and the closest rollout state within this budget.
multi-goal reaching: For every trajectory in the validation set, we randomly sample a start state and 5 goal states at random future timesteps from [12, 60). We evaluate the same set of 100 state-goal sequences and add additional 5 timestep budgets for all the goals. Similar to single-gon goal reaching, We report the L2 distance between every goal and the closest rollout state before running out of its corresponding budget.
We show the zero-shot performance of MaskDP and baselines pretrained with the near-expert data (both in single-task and multi-task settings). We report L2 distance averaged over the states and goals
We directly evaluate the model unseen state-gor goal pairs in the validation dataset
5

---

MaskDP (clesed-locp)
Avcrug goa disanc
MaskDP (open-loop)
Goal MLP
Geal GPT
suathupor_yol yolk
walker
Walker_wak Jak
walker stand
cheebar_
cheefai, soodwwaddd
Figure 2: Single task pretraining followed by single goal reaching downstream task. MaskDP with closed-loop execution achieves the hest performance all the tasks, and get the most significant improvements in the Quadruped domain, which is higher dimensional.
Oou Distance
CeN COE Gan #3 Goe Goal
Gca Gml #5 Codl #3 Gol #4 Codl
Coo Goll 92 Gen #3 Goll ed Cod
Gol Goal #2 See #3 Gol m4 See
(closedilonp)
MaskDP (open-leop)
Goal GPT
Goai
Oma Ocal Godl Cheetah Backward
Quadruped
Quadruped
Figure 3: Single task pretraining followed by multiple goals reaching downstream task. MaskDP achieves significant improvement on all the tasks with better flcxibility in sequential goal reaching.
sampled based the above rules. Tables of numbers and standard derivations be found in Section
Figure show the results of reaching single goal. The y-axis is the L2 distance (the lower the better). We observe that both MaskDP (open-loop) and MaskDP (closed-loop) outperform Goal-GPT and Goal-MLP. Despite Goal-GPT being natural formulation for goal reaching. MaskDP reaches lower distance to the goal. We attribute the effectiveness to learning better understanding of the forward and inverse dynamics implicitly. We also observe that the advantages of MaskDP more significant in higher dimensional environments, such as Quadruped.
For the more challenging multi-goal reaching task, MaskDP has significant advantage in flexibility: we can just provide the goals at specific time budgets with interleaved masks and get an executable plan; however. for Goal-MLP and Goal-GPT, future multiple goals. As shown in Figure 3 MaskDP outperforms both goal-GPT and BC by large margin. In Figure 14] showed that having "foresight" about future goals can help the agent to generate a better plan.
have change goals certain timesteps to fulfill
We can get similar conclusions from the multi-task pretrained models in Figure and Figure 5 where our method consistently works well on all domains, with the most visible advantage multi-goal reaching setup.
4.2.2 Skill Prompting
We are interested in the learned behavior of pretrained models. We use prompting, which has become popular analyzing models ever since GPT 4 To do so, we give the agent a short state-action segment randomly cropped from expert trajectory, set the agent to the last state of the segment, and let the model continue to generate consecutive behaviors. We evaluate the quality of the generated sequence by comparing its obtained rewards with the rollout of skilled expert.

---

MaskDP (clased llcpp)
MaskDP (oper-loop)
Goal MLP
Goal GPT
Quedruced
quedruped walk
walter
Walker
Halker stand
cheetah
chostar bachward
Figure 4: Multiple tasks pretraining followed by single goal reaching downstream task, where MaskDP with closed-loop execution works the best, especially in the Quadruped domain.
Coal Dieancc
Coa Goa
COM #0 GON GoM
Goa Gol #5 CON #3 Gel #4 Cod
Goa Coll #2 #3 Coal
Gol Goal #2 See #1 Gall See
(closed-loop)
MaskDP (open-leop
Goal GPT
Goal MLP
Oca Ocelm2 Cheerah
Backward
Quadruged
Quetruped
Figure 5: Multiple task pretraining followed by multiple goals reaching downstream task.
output
prompt
Figure 6: Qualitative results on for skill prompting in the Walker domain. Given 5 initial states, the model learns to forecast future trajectories as in the expert-level behaviour.
To be specific, we prompt the mulli-task model trained with expert data and sample a 5-timestep
state-action segment from T E [100, 900). where the agent can be walking/running at low or high speeds. We prompt the model with this short segment and let the model generate rollouts for 20, 40, 60 timesteps. We provide both qualitative results and quantitative results in Figure and Figure b respectively. We can see in Figure 7 both our method and GPT can match the expert return. It shows that our method can perform well as autoregressive model in generation task.
4.2.3 Offline Reinforcement Learning
337 does, where the overall return of the data is quite low and thus the BC-based method cannot work well. ExoRL 37 simply shows that an offline TD3 agent works the best on diverse low-return offline data.
Evaluation We provide 2M buffer of the data collected by Proto-RL 36 as in ExoRL
We can modify MaskDP to this setting by adding additional actor and critic heads on top of the pretrained encoder, and performing RI training. We evaluate the efficiency the pretrained model by its return after certain TD gradient steps. The results are shown in Figure averaged over
7

---

Bun
Walker
MaskDP (open-loop)
MaskDP (closed-loop)
GPT
Prompt (oracle)
Cheetan Backward
Quadruped Run
Quadruped Wak
Figure 7: Quantitative results on learned behaviors using prompt. Both MaskDP and GPT can match or even slightly surpass the expert-level performance (right grey bar) in trajectory forecasting.
450
1000
1023
800
Task rewerd
209
730
200
MaskDP
MaskDP from scratch
200
GPT
50
steps (Walker Run)
steps
steps
Figure 8: Offline RL results on Walker domain. The result of MaskDP matches the GPT-style pretraining performance, and are all comparable to the SoTA in ExoRL 37
seeds. We observe MaskDP capable of adapting to downstream tasks quickly, outperforms training from scratch, and achieves similar results the GPT baseline. Note that in this setting, need to replace the bidirectional attention mask with the causal attention mask, so there is larger gap between pretraining and downstream tasks finetuning compared with GPT, which is trained with causal masking. Note that MaskDP from scratch almost the same as GPT from scratch (both with causal masking). From Figure & both MaskDP and GPT can match the best result ExoRL 37 from their offline TD3 agent, where BC-based method cannot successfully solve this task.
4.3 Analysis
Model scalability. also pretrain our agent and baselines on the diverse "mixed" dataset. We com- pare the smaller version of MaskDP and Goal-GPT on
We
38
MarkDP snall insp)
MastOP (dsseslespp
MaskDP arg lapan loop)
GPT email
the Quadruped goal reaching problem as shown in Figure g
16
CPT large
50%
5%
The x-axis is the finetuning gradient steps on the expert dataset, and y-axis is the L2 distance to the goal (the lower
Figure 9: Model scalability two tasks. X-axis represents number of gradient steps. With MaskDP pretraining, larger models outperform smaller models.
Both use attention layers.

---

30
26
26
24
MaskDP (open-loop)
MaskDP (closed-loop)
Conl
20
20
MaskDP (open-loop) default
MaskDP (closed-loop) default
14
14
0.15
0:50
0.95
0.15
0.50
0.95
Mask ratio
Mask ratio
Figure 10: Mask ratio ablation. We compare multiple ratio pretrained model with models trained with fixed ratios, where our masking strategy can achieve much better performance.
the better). We found for both zero-shot evaluation and finetuning, our model's performance improves when model size is enlarged, whereas for Goal-GPP the performance gain is not obvious.
Mask ratio ablation. Figure TO shows the influence of the masking ratio. With fixed mask ratio, we observe that an extremely high mask ratio (95%) generally does not work well and the typical mask ratio (15%) used in BERT seems to perform much worse than others. middle mask ratio 50% performs reasonably well, despite still being surprisingly high, similiar to the observations in MAE. However, our mixed mask ratio strategy strictly outperforms the above options.
Predicting unmasked tokens ablation. We also compare the model trained with mask loss vs. total loss. As shown in Fig- ure IT] empirically we do find mask loss has more advantages than total loss, even on the relatively clean expert dataset, it converges slower than using total loss. For the results on diverse data, please refer to Section
CoN dstance
Mask (closed loop)
(open-loop)
Tocal (clased loop)
Total lopen loop)
Figure 11: Masked loss and total loss ablation 20k pretrain- ing gradient steps. The model trained with total loss converges faster than one trained with masked loss.
5 Conclusion
This paper presents masked decision prediction (MaskDP), a simple and scalable self-supervised method for reinforcement learning (RL) inspired by current large language and vision models. MaskDP capable of learning scalable and generalizable agents for reinforcement learning that can learn from diverse-quality data sources and infer tasks in goal-reaching and skill-execution settings. Through our empirical study we find MaskDP models outperform past work in zero-shot goal reaching and transfer well to downstream RL tasks, performing competitively with prior pretraining and training from scratch methods.
Limitations and Future Work Computer vision and NLP domains have shown that the true promise of masking architectures lies with their ability to ingest diverse, fully unsupervised data. We study how MaskDP performs when trained without access to expert-level data, and evaluated on unseen proprioceptive states. In the future, we can extend our method to pixel inputs, and pretrain the model adapt far different downstream tasks.
The architecture used MaskDP closely resembles model-based method, as states are predicted sequentially from actions. In this paper, simplest and fastest approach. However, it is straightforward to extend MaskDP to plan through our learned model and compare against related baselines.
the predicted next actions directly this is the
Societal Impact This is an algorithm for training agents the style of recent large-scale CV and NLP models. While we do not anticipate particular social risks from our method, as algorithms become capable of ingesting large-scale, in-the-we wild data it important to ensure the dataset does not reinforce undesirable biases promote harmful behaviors.

---

References
[1] J.-B. Alayrac, J. Donahue, P. Luc, A. Miech, I. Barr, Y. Hasson, K. Lenc, A. Mensch, K. Millican, M. Reynolds, R. Ring, E. Rutherford, S. Cabi, T. Han, Z. Gong, S. Samangooei, M. Monteiro, J. Menick, S. Borgeaud, A. Brock, A. Nematzadeh, S. Sharifzadeh, M. Binkowski, R. Barreira,
O. Vinyals, A. Zisserman, and K. Simonyan. Flamingo: visual language model for few-shot learning. arXiv preprint arXiv: Arxiv-2204. 2022.
[2] G. Baldassarre and M. Mirolli. Intrinsically motivated learning in natural and artificial systems. Springer. 2013.
[3] H. Bao, L. Dong, and F. Wei. Beit: Bert pre-training of image transformers. arXiv preprini arXiv:2106. 08.054 2021.
[4] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, Shyam,
G. Sastry, A. Askell, et al. Language models are few-shot learners. Advances in neural
information processing systems, 33:1877-1901, 2020.
[5] M. Carroll, J. Lin, O. Paradise, R. Georgescu, M. Sun, D. Bignell, S. Milani, K. Hofmann, M. Hausknecht, A. Dragan, and S. Devlin. Uni[mask]: Unified inference in sequential decision problems. Advances in neural information processing systems, 2022.
[6] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and
I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. arXiv preprint arXiv: Arxiv-2106. 01345, 2021.
[7] M. Chen, A. Radford, R. Child, J. Wu, H. Jun, D. Luan, and I. Sutskever. Generative pretraining from pixels. International Conference on Machine Learning, pages 1691-1703, PMLR, 2020.
[8] P. Christiano, Z. Shah, I. Mordatch, J. Schneider, T. Blackwell, J. Tobin, P. Abbeel, and W. Zaremba. Transfer from simulation to real world through learning deep inverse dynamics model. arXiv preprini arXiv: 1610.0361 18, 2016.
[9] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 1810004805, 2018.
[10] A. Dosovitskiy, L. Beyer, Kolesnikov, D. Weissenborn, Zhai, Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 1 16x words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1929, 2020.
[11] Fujimoto, H. van Hoof, and Meger. Addressing function approximation error in actor-critic methods. arXiv preprint arXiv: Arxiv-1802.09477. 2018.
[12] K. He, X. Chen, S. Xie, Y. Li, P. Dollár, and R. Girshick. Masked autoencoders are scalable vision learners. arXiv preprint arXiv:2 1 2021.
[13] M. Janner, Q. Li, and S. Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34, 2021.
[14] S. Krishnamoorthy, S. M. Mashkaria, and A. Grover. Generative pretraining for black-box optimization. arXiv preprint arXiv:2206. 10786, 2022.
[15] M. Laskin, D. Yarats, Liu, Lee, A. Zhan, K. Lu, C. Cang. Pinto, and P. Abbeel. Urlb: Unsupervised reinforcement learning benchmark. arXiv preprint arXiv; Arxiv-2110.15191. 2021.
[16] M. Laskin, Liu, X. B. Peng, D. Yarats, A. Rajeswaran, and Abbeel. Cic: Contrastive intrinsic control for unsupervised skill discovery. arXiv preprini arXiv:2202.00166 2022.
[17] H. Liu and P. Abbeel. Aps: Active pretraining with successor features. In International Conference on Muchine Leurning, pages 6736-6747. PMLR, 2021.
[18) H. Liu and P. Abbeel. Behavior from the void: Unsupervised active pre-training. Advances in Neural Information Processing Systems, 34, 2021.
[19] T. Nguyen and A. Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. arXiv preprint arXiv:2207.04/73 2022.
[20] P.Y Oudeyer and F. Kaplan. What intrinsic motivation? a typology of computational approaches. Frontiers in neurorobotics, 1:6, 2009.
10

---

[21] D. Pathak, P. Agrawal, A. A. Efros. and T. Darrell. Curiosity-drive driven exploration by self- supervised prediction. In International conference on machine learning, pages 2778-2787. PMLR, 2017.
[22] A. Ramesh, M. Pavlov, G. Goh, S. Gray, C. Voss, A. Radford, M. Chen, and I. Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pages 8821-8831. PMLR, 2021.
[23] M. Schwarzer, A. Anand, R. Goel, R. D. Hjelm, A. Courville, and P. Bachman, Data-efficient reinforcement learning with seef-prredcctve representations, arXiv preprint arXiv:2007.05929, 2020.
[24] M. Schwarzer, N. Rajkumar, M. Noukhovitch, A. Anand, Charlin, R. D. Hjelm, P. Bachman, and A. C. Courville. Pretraining representations for data-efficient reinforcement learning. Advances in Neural Information Processing Systems, 34, 2021.
[25] Y. Seo, D. Hafner, H. Liu, F. Liu, S. James, K. Lee, and Abbeel. Masked world models for visual control. arXiv preprint arXiv:2206. 24444:22006.1244,, 2022.
[26] Y. Seo, K. Lee, S. James, and P. Abbeel. Reinforcement learning with action-free pre-training 22 2022.
from videos. arXiv preprint
[27] Ö. Şimşek and A. G. Barto. An intrinsic reward mechanism for efficient exploration. In Proceedings of the 23rd international conference on Machine learning, pages 833-840, 2006.
[28] A. Stooke, K. Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. In International Conference on Machine Learning, pages 9870-9879.
PMLR, 2021.
[29] Y. Tassa, Y. Doron, A. Muldal, T. Erez, Y. Li, D. de Las Casas, D. Budden, A. Abdolmaleki,
J. Merel, A. Lefrancq. T. Lillicrap, and M. Riedmiller. Deepmind control suite. arXiv preprint arXiv: (anp-1801.00690 2018.
[30] F. Torabi, G. Warnell, and P. Stone. Behavioral cloning from observation. arXiv preprint 1rXiv:1805.01955 2018.
[31] H. Touvron, M. Cord, M. Douze, Massa, A. Sablayrolles, and H. Jegou. Training data-efficient image transformers and distillation through attention. In International Conference on Machine Learning, volume 139, pages 10347-10357, July 2021.
[32] A. Van Den Oord, O. Vinyals, et al. Neural discrete representation learning, Advances in neural information processing systems, 30, 2017.
[33] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and
I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[34] Y. Wang, S. Joty, M. Lyu, King, C. Xiong, and S. Hoi. Vd-bert: A unified vision and dialog transformer with bert. arXiv preprint arXiv:2004. 20iv:2004.13278, 2020.
[35] T. Xiao, I. Radosavovic, T. Darrell, and J. Malik. Masked visual pre-training for motor control. arXiv preprint urXiv:2203.06122 2022.
[36] D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Reinforcement learning with prototypical representations. In International Conference on Machine Learning, pages 11920-11931. PMLR, 2021.
[37] D. Yarats, D. Brandfonbrener, H. Liu, M. Laskin, P. Abbeel, A. Lazaric, and L. Pinto. Don't change the algorithm, change the data: Exploratory data for offline reinforcement learning. arXiv preprini arXiv: Arxiv-2201. 13425, 2022.
[38] A. Zhan, P. Zhao, L. Pinto, P. Abbeel, and M. Laskin. A framework for efficient robotic manipulation. arXiv preprint arXiv:2012.074
2020.
[39] Q. Zheng, A. Zhang, and A. Grover. Online decision transformer. In International Conference on Machine Learning, 2022.
11


Reviews:
[{"review_id": "l0w7_fuOiNN", "reviewer": "Reviewer_PTCG", "summary": "This work proposes to use masked autoencoding to learn a self-supervised encoder-decoder transformer model from sampled trajectories. The tokens that are masked and then predicted are a sequence of (interleaved) states and actions. The trained model is then used for goal reaching (both in an open- and closed-loop setting), “skill prompting” and offline RL. Experiments on DM control suite tasks show promising results vs. imitation learning, and an autoregressive baseline. ", "questions": "- For the goal reaching task, the trajectory length must be fixed in advanced. How is this dealt with in practice? Is there a “skip this step” token?\n- How does trajectory length interact with training? Does the model generalize to trajectory lengths it has not seen during training? If not can that be fixed with an appropriate positional encoding, e.g. ALiBi (Press et al., “Train Short, Test Long …”)?\n- In the multi-goal setting, how does MaskDP perform if future goals are not available, e.g. if they are unmasked gradually? This baseline should help confirm that foresight is being leveraged to “produce a better plan”.\n- For the offline RL task, is it only the encoder that is being finetuned (as typical in vision tasks) or is it both encoder and decoder?", "limitations": "I'm concerned that the effect of expert (data-generating) policy on downstream tasks seems a bit too important to be left to future work. If it is left to future work, the authors should include at least some preliminary results in the next revision. ", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Positive\n\n- The method is straightforward, and the downstream use of the pre-trained model is crafty. Skill-prompting appears to solve an interesting inference problem. Insights on the appropriate masking strategy in this domain appears useful.\n- The experiments are well designed, and the performance is measured online (as opposed to just offline). The baselines are appropriate, and highlight the masking aspect while controlling for confounders.\n- The presentation is nice, and the paper is generally well written.\n\n### Not Positive\n\n- The application appears to be somewhat limited in scope (only robotics environments), and it is not clear if and how it might extend to pixels (though a VQGAN like approach can be promising).\n- The influence of the policy used to generate pre-training data is not studied but left to future work.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "f74lytXhqBC", "reviewer": "Reviewer_ufWH", "summary": "This paper presents a new way of pretraining RL models such that they can be directly applied to Goal Reaching and Skill Prompting tasks, or finetuned for Offline RL. In contrast to previous works that pre-train these models using an autoregressive objective, the authors propose using the masked autoencoder (MAE) objective.  The bidirectional transformer model is pre-trained through reconstructing masked state and action tokens.  The authors demonstrate how such a model can be deployed to the three types of downstream tasks using distinct masking strategies.  Although the method is a simple extension of MAE pre-training, which has recently shown state-of-the-art performance in certain CV and NLP tasks, this paper shows extensive empirical comparisons against baselines.  In particular, offline RL results on Walker are comparable to SoTA in ExoRL [31].  Ablation study is also provided to justify some of their design choices such as the mixed mask ratio strategy.\n\n[31] D. Yarats, D. Brandfonbrener, H. Liu, M. Laskin, P. Abbeel, A. Lazaric, and L. Pinto. Don’t\n380 change the algorithm, change the data: Exploratory data for offline reinforcement learning.", "questions": "I wonder if you could replace the word “zero-shot” performance with another word since goal reaching tasks are instances of the pre-training masked autoencoding task.\n", "limitations": "Yes, the authors have addressed the potential broader impact.\n", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:  The biggest strength of this paper is to investigate the use of MAE for unsupervised pretraining in RL.  Although I might have missed related works, this seems to be the first work that provides practical implementations of the method and comparisons against relevant GPT baselines that use the next state/action prediction loss.  Extensive experiments do show MaskDP outperform autoregressive baselines. \nOverall, the paper is clearly written.  It is interesting to see how MAE performs in unsupervised RL given its success in CV and NLP domains.  \n\nWeaknesses: As acknowledged by the authors themselves, MaskDP still uses some expert data, which are not “unlabeled” as in CV and NLP settings.   This is okay as the method can be compared against other methods using demonstrations.  Unfortunately, I’m not very aware of the state-of-the-art results on these tasks using demonstrations.  Although the walker result is compared against SoTA, I cannot comment on the significance of this offline RL result section.  Is there a reason why only the Walker domain is compared against SoTA results?\n\n-----\nAs the authors addressed my concerns, I've updated my score.", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "FwpngiXK8N", "reviewer": "Reviewer_cGtJ", "summary": "The paper proposes to adapt the Masked Autoencoder (MAE) from computer vision into decision transformer for RL demonstration modeling. The paper propose several downstream application after the unsupervised pre training: Goal reaching, prompting, and fine tuning. The experiments are done in the deep mind suit with Mojuco. The authors also provides analysis on masking ratio and scalability.", "questions": "Can you also add experiment on robotics domain like Metaworld?", "limitations": "Not that I can see", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\nThis paper is well written and easy to follow. While the idea of MAE is not novel, the application of MAE into RL is novel and the proposed downstream prompting and goal reaching are novel, which I think will have large impact. The experiment results are convincing and the choice of baselines are reasonable, showing the effectiveness of mixed masking ratio and versatilty of the large pretrained models.\n\nWeakness:\nI couldn’t find any major flaw, but one thing I would love to see is how is this work compared with PromptDT [1]. While I understand these can be concurrent work, but it could be better if there is discussion / comparison between these two very similar papers.\n\n[1] https://arxiv.org/abs/2206.13499", "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": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept