Uni[MASK]: Unified Inference in Sequential Decision Problems
Authors: Micah Carroll
Orr Paradise
Jessy Lin
Raluca Georgescu
Mingfei Sun
David Bignell
Stephanie Milani
Katja Hofmann
Matthew Hausknecht
Anca Dragan
Sam Devlin
Venue: NeurIPS 2022 Accept (2022)
Keywords: Multi-task Learning
Unsupervised Learning
Reinforcement Learning
Deep Learning

Abstract: Randomly masking and predicting word tokens has been a successful approach in pre-training language models for a variety of downstream tasks. In this work, we observe that the same idea also applies naturally to sequential decision making, where many well-studied tasks like behavior cloning, offline RL, inverse dynamics, and waypoint conditioning correspond to different sequence maskings over a sequence of states, actions, and returns. We introduce the UniMASK framework, which provides a unified way to specify models which can be trained on many different sequential decision making tasks. We show that a single UniMASK model is often capable of carrying out many tasks with performance similar to or better than single-task models. Additionally, after fine-tuning, our UniMASK models consistently outperform comparable single-task models.

Uni[MASK]
Unified Inference in Sequential Decision Problems
Micah Carroll', Orr Paradise', Jessy Lin', Raluca Georgescu', Mingfei Sun', David Bignell Stephanie Milani Katja Hofmann", Matthew Hausknecht', Anca Dragan', and Sam Devlin
'CC Berkeley
2 Microsoft Research
CMU
Abstract
Randomly masking and predicting word tokens has been a successful approach in pre-training language models for variety of downstream tasks. In this work, we observe that the same idea also applies naturally to sequential decision mak- ing, where many well-studied tasks like behavior cloning, offline reinforcement learning, inverse dynamics, and waypoint conditioning correspond to different sequence maskings over sequence of states, actions, and returns. We introduce the Uni[MASK] framework, which provides unified way to specify models which can be trained on many different sequential decision making tasks. We show that a single Uni[MASK] model is often capable of carrying out many tasks with perfor- mance similar to better than single-task models. Additionally, after fine-tuning, our Uni[MASK] models consistently outperform comparable single-task models. Our code is publicly available here.
1 Introduction
Masked language modeling [11] is a key technique in natural language processing (NLP). Under this paradigm, models are trained to predict randomly-masker subsets of tokens in sequence. For example, during training, a BERT model might asked to predict the missing words in the sentence "yesterday __ cooking to predict the next token conditioned only on the left context, bidirectional models trained on this objective learn to model both the left and right context to represent each word token. This leads to richer representations that can then be fine-tuned to excel on a variety of downstream tasks [11].
". Importantly, while unidirectional models like GPT [33] are trained
Our work investigates how masked modeling can be powerful idea in sequential decision problems. Consider a sequence of states and actions a collected across timesteps 81, 01,.... 8T, @r. If we consider each state and action as tokens of sequence (analogous to words in NLP) and mask the last action (81, a1, 82, 022 53, then predicting the missing token 03 amounts to a Behavior Cloning prediction with two timesteps of history [32], given that this masking corresponds to the inference IP( (aa) 8 81:3, 41:2). From this perspective, training model to predict missing tokens from all maskings of the form ($1, ($1,a1,.... .i, St. for all t [1...., T' corresponds training Behavior Cloning (BC) model.
In this work, we introduce the Uni[MASK] framework: Unified Inferences in Sequential Decision Problems via [MASK]ings, where inference tasks are expressed as masking schemes. In this framework, commonly -ttudiie tasks such as goal or waypoint conditioned BC [12, 36), offline reinforcement learning (RL) [25], forward or inverse dynamics prediction [18, 9, 6]. initial-state inference [38], and others are unified under simple sequence modeling paradigm. In contrast standard approaches that train a model for each inference task, we show how this framework naturally lends itself to multi-task
36th Conference on Neural Information Processing Systems (NeurIPS 2022).

---

training; single UnI[MASK] model can be trained to perform a variety of tasks out-of-the-bon by appropriately selecting sequence maskings training time.
We test this framework in a Gridworld navigation task and a continuous control environment. First, train Uni[MASK] model by sampling from the space of all possible maskings training time (random masking) and show how this scheme enables single Uni[MASK] model to perform BC, reward-conditioning, waypoint-conditoonnng and more by conditioning on the appropriate subsets of states, actions, and rewards. We then systematically analyze how the masking schemes seen at training time affect downstream task performance. Training on random masking generally does not compromise single-task performance. and in fact can outperform models that only train on the task of interest. In the continuous control environment, we confirm that a model trained with random masking and fine-tuned on BC or RL tends to outperform models specialized to those tasks.
Our results suggest that expressing tasks sequence maskings with the Uni[MASK] framework may be a promising unifying approach to building general-purpose models capable of performing many inference tasks in an environment [2], or simply offer an avenue for building better-performing single-task models via unified multi-task training.
In summary, contributions are:
1. We propose a new framework, Uni[MASK], that unifies inference tasks in sequential decision problems as different masking schemes in sequence modeling paradigm.
2. We demonstrate how randomly sampling masking schemes at training time produces single multi-inference-aass model that can do BC, reward-conditionnng dynamics modeling, and more out-of-the-Ae
3. We test how training on many tasks affects single-task performance and show how fine- tuning models trained with random masking consistently outperforms single-task models.
4. We show how the insights we have gained while developing our choice of Uni[MASK] architecture can be used to improve other state-of-the-arr methods.
2 Related Work
Transformer models. The great successes of transformer models [41] in other domains such as NLP [11, 33, 3] and computer vision [13, 19] motivates our work. Using transformers in RL and sequential decision problems has proven difficult due to the instability of training [30], but recent work has investigated using transformers in model-based RL [6], motion forecasting [29], learning from demonstrations [34], and teleoperation [10]. We focus on developing a unifying framework interpreting tasks in sequential decision problems maskings.
The utility of masked prediction. Work in both NLP [11] and vision [5, 19] have explored how masked prediction is useful as self-supervision task. In the context of language generation, [39] provides a framework for thinking about different masking schemes. Recent work has also explored how random masking can be used to do posterior inference in probabilistic program [43].
Sequential decision-making as sequence modeling. Previous and concurrent work [7, 21, 24] shows how to use GPT-style (causally-masked) masked) transformers to directly generate high-reward tra- jectories in an offline RL setting. We expand our focus many tasks that a sequence modeling perspective enables, including but not restricted to offline RL. Although previous work has cast doubt on the necessity of using transformers to achieve good results in offline RL [14), we note that offline RL [25] is just one of the various tasks we consider. Concurrent work generalizes the left-to-right masking in the transformer to condition on future trajectory information for tasks like state marginal matching [17] and multi-agent motion forecasting [29]. In contrast, we systematically investigate how a single bidirectional transformer can be trained to perform arbitrary downstream tasks in more complex settings than motion forecasting i..., we also consider agent actions and rewards in addition to states. The main thing that sets us apart from these works is a systematic view of all tasks that can be represented by this eqquence-modeling perspective, and detailed investigation of how different multi-task training regimes compare.
Prior work on tasks in sequential decision problems. While we use masked prediction as the self-supervision objective, previous work on self-supervised learning for RI. has investigated other auxiliary objectives, such as state dynamics prediction [37] or intrinsic motivation [31]. Typically,

---

Reward-conditoned (offline-RL)
Behavioral Cloning
Goal-conditoneed
Waypoint-conditonee
Past inference
Future inference
Forward Dynamics
Inverse Dynamics
t-0
L-T-1
I-T
[-]
1-2
[-]
[+]
Figure I: Uni[MASK] framework: Representing arbitrary tasks as masking schemes. For each task, we show the inputs the model colors) and the outputs the model must predict (translucent colors). For example, in future inference, the model must predict all future states and actions conditioned on the initial states and actions. Here we only display one input masking scheme for each task, but many tasks are fully represented by multiple masking schemes. For example, BC has up to T different masking schemes, one for each possible history length (although in practice one would generally use the model with sliding window).
to accomplish the tasks consider, prior work relies single-task models: for example, goal- conditioned imitation learning [12]. RL [22], maypoint-conditionn [36], property-conditionin [45, 17]. or dynamics model learning [18, 9]. Other work has focused on training models perform different "tasks" such as different games in Atari [24] or different environments and multi-modal prediction tasks [35]. In contrast, are interested in performing different inference tasks in single environment, such as RI. and forward dynamics modeling, using sequence modeling as a unifying framework.
3 The Uni[MASK] Framework
We introduce the Uni[MASK] framework. In Section 3.1 we propose a unifying interpretation of inference tasks in sequential decision problems masking schemes. In Section 3.2 we describe different ways of training Uni[MASK] models, and provide hypotheses about their efficacy.
We consider trajectories as sequences of states, actions, and optionally property tokens (e.g. reward): {(s0, Po.P ..... (ST,0r, pr))
Motivated by canonical problems in decision-making that involve reward, in most of our analysis return-to-go (RTG) as the property (the of rewards from timestep & to the end of the episode): that is, we set pt Ri where R, = Σr= Tv.. However, any property of the decision problem can be considered a "property token," " including specific environment conditions being satisfied, the style of the agent, the performance of the agent (e.g. the reward obtained in the timestep). In order to train on tasks requiring specific properties, one must have labels for them obtained either programmatically or through human annotators. We demonstrate how our model can be conditioned on non-reward property in Appendix F.
3.1 Tasks Masking Schemes
In the UnI[MASK] framework, we formulate tasks in sequential decision problems as input masking schemes. Formally, masking scheme specifies which input tokens are masked (determining what tokens are shown to the model for prediction) and which outputs of the model are masked before computing losses (determining which outputs the model should learn predict). For example, the masking scheme for BC unmasks (conditions on) Suz and aur-1, and the model must predict az
In Figure 1, illustrate how to unify commonly-studices tasks such as BC, goal and waypoint conditioned imitation, offline RL (reward-conditioned imitation). and dynamics modeling under our proposed representation of tasks masking schemes. We describe the masking scheme for each of these tasks in detail in Appendix B.
While reward-to- to-go (or other trajectory statistics) are not necessary, we formulate the most general form to showcase how one can easily condition on additional available properties of a trajectory. Using reward-to-ge also enables compare our method with previous offline- RL work [7].
3

---

Outpets
Eddectional Tranaformer
Inguts
Figure 2: The Uni[MASK] model takes in snippet of trajectory which is masked according to a masking scheme before inference time. For each input possible masking, there are (many) corresponding tasks of predicting the missing inputs, Above show input masking corresponding to conditioning on both reward-to-go and final (goal) state; we highlight the output corresponding to predicting the agent's next action, i.e. performing the inference P(a2 $02,T 40:1, Ro).
3.2 Model Architecture & Training Regimes
For our main experiments, we instantiate our Uni[MASK] framework using the BERT architecture [11] adapted to the sequential decision problem domain, consisting of positional encoding layer and stacked bidirectional transformer encoder (self-attention) layers (see Figure 2). One key difference with the original BERT architecture is that we stack the state, action, and property (e.g. RTG) tokens for each timestep into a single vector. While prior work in sequential maaisig had used timestep encoding [7] (which can be thought of as concatenating each observation with its environment timestep), we found traditional positional encoding [41, 11] to reduce overfitting. For cowaid-conditioned tasks, in each context window we only feed the first RTG token into the model along with the number of timesteps remaining in the horizon. This information is sufficient for reward-co weward-conditionigg at inference time, and we found that it outperformed the standard approach of feeding in the RTG token at every timestep [7, 21]. See Appendix D) for more model details, and Appendix for experiments with an alternative instantiation of Uni[MASK] with a feedforward neural network architecture.
3.2.1 Training regimes
We experiment with four ways to train a Uni[MASK] model on masked prediction, illustrated in Figure and described below.
single-task. Training on just one of the masking scheme described in Section 3.1.
multi-task. Training a single model on multiple masking schemes: each trajectory snippet is masked according to one of the schemes from Section 3.1 (chosen at random).
Intuition: Could allow single model perform well multiple tasks. Additionally, it might outperform s single- task on individual tasks, as the model could learn richer representa- tions of the environment from the additional masking schemes.
random-mask. Training a single model on a fully randomized masking scheme. For each trajectory snippet, first, a masking probability Pmask [0, 1] sampled uniformly random; then each state and action token is masked with probability Pmask; lastly, the first RTG token is masked with probability 1/2 and subsequent RTG tokens are always masked (see Appendix C for details).
Intuition:
Could allow single model to perform well on any sequence inference lask without the need specify the tasks of interest at training time. The model may learn richer repre- sentations than those of mu l. ti - task as it must reason about all aspects of the environment.
f finetun. Fine-tune a model pre- pre-trainn trained in random mask on a specific masking scheme.
Intuition: Performing fine-tuning could allow the model benefit from the improved representa- tions obtained from random- mask,, while specializing to the single task at hand.
3.2.2 Hypotheses
Based the intuition of the strengths of each training regime, we formulate the following hypotheses:
HI. First training on multiple inference tasks will lead to better performance on individual tasks than only training that inference task: (multi-tatt task,, random-mask, fik finetune) Lanle-task.

---

H2. Randomized mask training outperforms training on a specific set of tasks: random- mask multi-task.
H1 tests whether models learn richer representations by training multiple inference tasks. H2 tests a stronger claim: whether training on all possible tasks by randomly sampling masking at training time is better than selecting of specific maskings.
4 A Unified Model for Any Inference Task
We first demonstrate how random-mask enables a single Uni[MASK] model to perform arbitrary infer- ence tasks at test-time Gridworld environment, without the need for task-specific output heads or training schemes that are customized for the downstream task. We then show that random-man does not compromise performance most specific tasks of interest. Models trained with random-mas achieve comparable or better performance to single- maakee and wulli-task-modeen and fact consistently outperform after additional fine-tuning on the task of interest (finetune).
Environment Setup. We design fully observable to fixed goal location behind a locked door with the MiniGrid environment framework [8]. The agent and key positions are randomized in each episode. The agent receives reward of for each timestep it moves closer to the goal, -1 if it moves away from the goal, and 0 otherwise. We train Uni[MASK] models on training trajectories of sequence length 10 from noisy-rational agent [46]. More detailed information about the environment is in Appendix E.
Gridworld in which the agent should move
4.1
One Model to Rule Them All
As shown in Figure single Uni[MASK] model trained with random mask can used for arbitrary inference tasks by conditioning on specific sets of tokens. Unless otherwise indicated, we take the highest probability action from the model a, arg max,; P(a/ $o, ................................ .8), and then query the environment dynamics for the next state $t+1- The model can be used for imitation, reward- and condiconditionnng. or as a forward or inverse dynamics model (when querying for state predictions, as in the backwards inference task). If trajectories are labeled with properties at training time, the model can also be used for property-conditionin conditioned global properties of the trajectory, such as whether the trajectory passes through certain position at any timestep.
Figure 5, we show how the model can also be
Qualitatively, these results suggest that the model generalizes across masking schemes, since seeing the exact masking corresponding particular task training time is exceedingly rare (out of 2 x 2d 2 possible state, action, and RTG maskings for sequence of length T)).
Teekeve:
Tenteve:
fvaliwes
Finetuna
:
Finatune
Single-tast task
Mult ti task
Random-mon manddmm-ask
Figure 3: The four training regimes considered in this work: single-task, multi-tane random-man and f finetune.

---

Belavional Cloning
Gool-condilooee
Rewand-ooniifonee
Waypsiat-cuniiiumd
Backwards inferencc
10
1-1
5-
Conditoneed on (14))
Condriceudd goal
Cunditioned om R 2
Cordijone waypotnte
Conditioned Sa
Figure A Uni[MASK] model trained with random masking queried on various inference tasks. (1) Behavioral cloning: generating expert-like trajectory given an initial state. (2) Goal- conditioned: reaching an alternative goal. (3) Reward-conditioned: generating a trajectory that achieves particular reward. (4) Naypoint-condittonee reaching specified waypoints (or subgoals) at particular timesteps, e.g. going down on the first timestep instead immediately picking up the key. (5) Backwards inference: generating likely history conditioned in final state (by sampling actions and states backwards). Trajectories are shown with jitter for visual clarity.
4.2 Future State Predictions
Uses of the mandom-masse -tranned Uni[MASK] model are not limited to rolling out new trajectories (requesting inferences about the agent's 's next action). One can also request inferences for states and actions further into the future: e.g., "where will the agent be in 3 timesteps?". Given fixed initial set of observed states, we visualize the distribution of predicted states at each timestep in Figure particular timestep marginalizes over missing actions; for example, [P(s 80; 83, 86) models the possibility that the agent chooses either up left as the first action. Qualitatively, the state predictions suggest that the model accurately captures the environment dynamics and usual agent behavior; e.g. it correctly models that the agent has equal probability of going up and right
Since do not roll out actions, querying the model for the predicted state distribution
(leading the distribution over states at = 4), and that the agent must be at position (2, 1) at t = 5 to reach the door at t 6.
4.3 Measuring Single-Tass Task Performance
Next, investigate how random-mash model performs individual tasks, in comparison to single-task models trained exclusively on the evaluated task. we care about a single task (e.g. goal-conditioned imitation), should train model simply that task? Or there be advantages to training a general model first, and then fine-tuning it to the task of interest? For this set of experiments, we primarily consider validation loss as our measure of performance. Validation loss
Conditioned on property raaces
Property-coddliione
plsy Fyy
plsi 99.99))
Figure 5: Other types of property-conditioning. If the training dataset has additional property la- bels (e.g., whether the trajectory passes by the top left corner of the grid any timestep), the model can roll out trajectories conditioned on whether the property is exhibited.
pls, 99.
p/s (5,,15
Figure 6: Predicted state distributions. The
model conditioned on states at 0,3,6.

---

1.5
1
1.037
1.212
1.614
1.602
1.575
1000
5.644
Cleseg
1.04
1.022
1.252
1.674
1.589
1.584
1000
5.61
Grdttiendd
Gad
1.175
1.213
1.104
1.534
1.653
1.687
1000
5.968
1.4
Рerpart
1.259
1.301
1.187
1.249
1.682
1.75
1000
6.059
Training task
1.423
1.469
1.645
2.126
1.061
1.681
1000
3.633
1.3
1.025
1.056
1.23
1.652
1.81
1.045
1000
4.652
Dysariie
1.433
1.484
1.7
2.22
4.208
3.97
2.616
9.48
1.2
1.401
1.439
1.684
2.197
2.629
2.141
1000
1.05
Dysantio
Nata
1.017
1.034
1.124
1.434
1.102
1.072
1000
1.482
1.1
1.09
1.045
1.039
1.044
1.005
1.013
539
1.151
Random M.
1.03
1.0
Chosag
Cosliiined Colliion Gonlioeed
tafermncc
Islorne
Dysamso
Dynanio
Evaluation task
Figure 7: Task-specific validation losses (normalized column-wise). Each row corresponds to the performance of a single model evaluated in various ways, except for the last row-for which each cell is fine-tuned on the respective evaluation task. Loss values are averaged across six seeds and then divided by the smallest value in each column. Thus, for each evaluation task (i.e., column), the best method has value 1; value of 1.5 corresponds to loss that is 50% higher than the best model in the column. Note that the performance of multi-task model on the forward dynamics task is particularly poor since the environment is deterministic: we should expect overfitting (with single-task model) to perform the best. See Appendix for more details.
provides a general way to evaluate how well models fit the distribution of trajectories and transitions, which is what are concerned with for most inference tasks: i.c., "how well can the network predict the true state or action in the data?".
In Figure 7, report validation loss if the model is trained evaluated on another task. As expected, models trained on one masking (e.g. BC) perform well when queried on the task they were trained on (as seen the diagonal), but poorly when queried with another task (e.g. past inference).
task (or multiple tasks) and
First, we find that random mask training (but not multi-tan training) outperforms single-tast training on half of the tasks considered, showing that even if one interested in a single inference task, training many more tasks can sometimes improve performance. Specializing model trained with random mask via finetuning (finetune) leads to the best performance, outperforming single-task all tasks except behavior cloning. This means that even if one is interested in a single inference task, first training on multiple tasks generally improves performance. Overall, these results do not fully support H1, given multi-tasee poor performance and random-mast performance which is not consistently better than single-task.
We also find that random- mask training leads to lower loss values on almost all evaluation tasks relative to multi-tashe supporting H2: training on additional inference tasks beyond the specific ones of interest can augment performance.
Trajectory Generation in Complex Environment
In addition to Gridworld, we test our method in a partially observable, continuous-state and continuous-action environment, with larger trajectory horizon (200 timesteps).
5.1 Environment Setup
We adapt the Mujoco-physics Maze2D environment [16] (see Appendix for figures), in which point-mass object is placed at a random location in maze, and the agent is rewarded for moving

---

towards a randomly generated target location (making this task "goal- conditionnd by default"). We make this task harder by removing the agent's velocity information from each timestep's observation and increasing the amount of initial position randomization. These changes make the environment partially observable, forcing models trained on this data to implicitly infer the agent's veloeity from observed context.
Expert dataset. We want our expert data have some suboptimality so that reward-conditioning can be tested for setter-taanddemonstaator performance. We generate a dataset of expert trajectories by rolling out D4RL's PD controller (which is non-Markovian), and add noise the actions with zero-mean and 0.5 variance (which are then clipped to have each dimension between -1, 1). We generate 1000 trajectories of 200 timesteps, of which 900 are used for testing and 100 for validation. For more details on our adapted Maze2D environment and design decisions, see Appendix H.
5.2 Models Trained
For the Maze2D evaluations, we focus on test-time reward performance on behavior cloning and offline RL (reward-conditionnng) across various architectures and training regimes.
We consider Uni[MASK] models trained with the different training regimes: single- task, multi-task, random- and finetune. We additionally consider other architectures, such as ford-fowward// NN and Decision Transformer (DT) baselines [7]. We found that several of our design decisions for Uni[MASK] models using positional encoding instead of timestep encoding, inputting the return-to-go token at the first timestep with the number of timesteps in the horizon - also improved GPT-based models like DT. We call our improved baseline Decision-GPT (for implementation details, see Appendix G). We train our Decision-GPP model with the s single- Saask training regime. The only meaningful difference between Decision-Gi G and single taak Uni[MASK] model is whether the model is GPT- or BERT-based
For each architecture and applicable training regime, we train separate models to perform behavior cloning and offline RL rreward-oondiiionnng
The only exceptions are Uni[MASK] models trained with multi-task (trained to perform BC and RC) and random- aandommmask. We train two sets of such models, for context lengths of and 10 meaning that during both training and evaluation, the models will respectively only be able to see the last 5 or 10 timesteps of the agent's interaction with the environment.
5.3 Results
We report reward evaluation results for 1000 rollouts in the Maze environment with standard errors across 5 seeds in Table 1.
The value of pre-training and fine-tuning for UnI[MASK] models. We find that fine-tuning is critical for good performance in more complex environments. multi-task performs more-or-lewe comparably to single-tard taak in behavior cloning and reward conditioning; however, random -mak in this setting obtains significantly lower rewards (counter to H2). This suggests that multi- -task training can be effective in mostly maintaining reward performance while increasing the breadth of functionality, but training on too many tasks can hurt out-of-the-boo performance. However, f finetune recovers the performance loss, again out- nut-performing single-task (providing qualified support for H1). Surprisingly, for a context length of ten, fine-tuning multi-tare does not improve performance much fine-tuning the randomly masked model, suggesting that specifically training on random masking might provide benefits for adapting models to individual downstream tasks.
How do Uni[MASK] models compare to other architectures? For context length five, multi-task with finetuning and finetune Uni[MASK] models perform better than all baselines we consider. However, increasing the context length to ten, we see that Uni[MASK] models performs poorly across the board, with the finetuned conditions outperformed by our Decision-GrP baseline. We speculate that this might be related the documented difficulty of using BERT-like architectures (as that of Uni[MASK] models) for sequence generation [42,28].
that
Isolating the effect of GPT vs. BERT. In order to investigate the effect of using GPT-like architec- tures instead BERT-like ones, we can consider the comparison between single-task Uni[MASK] and our Decision-Gri GPT baseline: the main difference between these two models is only whether one
8

---

uses BERT or GPT as the backbone of the architecture." 2 We find that while using GPT seems to yield similar (or worse) performance BERT for context length five, using GPT seems to give an advantage for longer sequence lengths. In particular, note that larger context length enables GPT to increase performance, while performance worsens for single- task Uni[MASK]. This suggests that if one were able to use GPT architecture and train it with random masking and fine-tuning, it might be possible to get the best of both worlds.
Table 1: Maze2D Results. Comparing among Uni[MASK] models, we isolate the benefit of finetune: this training regime tends to perform best across tasks and sequence lengths. Com- paring single-task Uni[MASK] to our Decision-GPT model, we can isolate the effect of using a BERT-like architecture struggle to maintain the same generation quality. Every entry in the table corresponds to a separate model, except for the cells denoted with which use the same model across tasks (but not sequence lengths).
GPT-like architecture: for larger context lengths, BERT-like models
Context Length 5
Context Length 10
Model
BC
RC
BC
RC
Uni[MASK] Models
Uni[MASK] single-task
2.66
2.04
2.47 0.04
2.41 0.05
Uni[MASK] - multi-task (BC & RC)
2.01
2.01
2.39 0.03
2.39 0.03*
Uni[MASK] - multi-task + finetune
2.01
2.04
2.42
2.42 0.03
Uni[MASK] random-mask
2.19
0.09
0.07
Uni[MASK] finetune
0.03
2.01
2.55 0.03
2.61 # 0.03
Other architectures
Feedforward Neural Network
1.68 0.07
1.53
1.83 + 0.06
1.88 # 0.06
Decision Transformer [7] Our Decision-GPP model
0.07
1.49
1.58 0.06
1.70 0.07
2.66 + 0.01
2.32
2.74
2.73 # 0.02
6 Limitations and Future Work
Comparison to other specialized models. We show that Uni[MASK] outperforms feedforward networks, Decision Transformer models, and for short sequence lengths also our own improved GPT-based baseline. However, we do not compare our models directly to different models in prior work that are specialized for specific tasks (e.g. coal-conditioning models, etc.). While this is a limitation of our work, it is also not our main focus: propose unifying framework for variety of tasks in sequential decision problems, and extensively analyze how different training regimes affect performance.
Longer context lengths. One limitation in our experimentation is the relatively short context lengths used. We found that longer context lengths negatively affect the Uni[MASK] models' performance. In part, this could be addressed by designing masking schemes tailored to specific test-time tasks (see Appendix C), using principled masking schemes [26]. However, this degradation may be attributed to our use of a BERT-like (rather than GPT-like) architecture, which seems less compatible with longer sequence lengths. A clear avenue of future work would therefore be to get the "best of both worlds": long sequences and benefits of random- mando pre-training by using a GPT-like architectures, with our act like a bidirectional model. Recent methods in NLP might offer useful starting point [1, 15], as has been explored by concurrent work to ours [27].
and finetune training regimes. This requires finding ways to make GPT
Comparison to other applications of masked prediction and sequence models for sequential decision making, In concurrent work, MaskDP [27] has also applied masked prediction to sequential decision-making. Similarly to Uni[MASK], MaskDP pre-trains a bidirectional transformer to predict randomly-masked token sequences corresponding to states and actions in Markovian decision process. The main difference between our works is that we are more interested in comparing the performance between different training regimes, and testing the performance limits of having single
We additionally use input- saput-stacking for Uni[MASK] see Appendix D but in preliminary experiments we found this to not affect performance.

---

set of weights to perform a large variety of tasks out of the box. MaskDP instead focuses on getting the best performance possible on smaller subset of classic tasks (e.g. having separate architecture choices for RL). Future work could more systematically investigate the differences in our methods, e.g. how the MaskDP encoder-decoder architecture fares on multi-task performance as measured in our work. In addition, other architectural choices could be explored: in order to speed up training time efficiency, one could try to substitute BERT for XLNet or NADE-style approaches [44, 40]. Finally, another exciting direction for future work is determining whether the benefits obtained from random- mask (or even multi-task) apply to other types inferences more generally (e.g. Bayes Networks); alternatively, even trivially extending the approach to multi-agent settings (for which soken-stacking could prove more valuable), could enable interesting masking-enablees enabled queries [29].
Conclusion
Broader impacts. The prospect of very large "foundation models" [2] becoming the norm for sequential problems (in addition to language) raises concerns, in that it de-democratizes development and usage [23]. We use significantly smaller models and computational power than similar works, leaving open the option to have more modestly-sized environment-specifcc foundation models. However, acknowledge that this works still encourages this trend.
Summary. In this work we propose Uni[MASK], a framework for flexibly defining and training models which: 1) are naturally able to represent any inference task and support multi-task training in sequential decision problems, 2) match or surpass the performance the corresponding single-task models after multi-tase pre-training, and almost always surpasses them after fine-tuning.
Acknowledgments and Disclosure of Funding
We'd like to thank Miltos Allamanis, Panagiotis Tigas, Kevin Lu, Scott Emmons, Cassidy Laidlaw, and the members of the Deep Reinforcement Learning for Games team (MSR Cambridge), the Center for Human-Compatible AI, and the InterAct Lab for helpful discussions at various stages of the project. We also thank anonymous reviewers for their helpful comments. This work was partially supported by Open Philanthropy and NSF CAREER.
References
[1] Armen Aghajanyan, Bernie Huang, Candace Ross, Vladimir Karpukhin, Hu Xu, Naman Goyal, Dmytro Okhonko, Mandar Joshi, Gargi Ghosh, Mike Lewis, and Luke Zettlemoyer. CM3: A causal masked multimodal model of the internet. CoRR, abs/2201.07522 2022. URL
https://arxis: org/ /t///221.00000
[2] Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S. Bemsttein, Jeannette Bohg. Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson, Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri S. Chatterji, Annie S. Chen, Kathleen Creel, Jared Quincy Davis, Dorottya Demszky, Chris Donahue, Moussa Doumbouya, Esin Durmus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah D. Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitt, Daniel Ho, Jenny Hong, Kyle Hsu, Jing Huang, Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Geoff Keeling, Fereshte Khani, Omar Khattab, Pang Wei Koh, Mark S. Krass, Ranjay Krishna, Rohith Kuditipudi, and et al. On the opportunities and risks of foundation models. CoRR, fts/2108.00253. 2021. URL https: / /arxiv.vvvv..aa org/ h//g//bs/2// /bs/ 07258.
[3] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learn- ers. Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Annual
10

---

Conference on Neural Information Processing Systems 2020, NeurIPS 2020, December 6- 12. 2020, virtual, 2020. URL htwplaa--bbbttttt....aa...
[4] Micah Carroll, Orr Paradise, Jessy Lin, Raluca Georgescu, Mingfei Sun, David Bignell, Stephanie Milani, Katja Hofmann, Matthew Hausknecht, Anca Dragan, and Sam De- vlin. Codebase for "Uni[MASK]: Unified Inference in Sequential Decision Problems". https://gi https:://gthub..... com/ mi cawcarroll/ /uniMASK 2022. URL https: /ttt:////tthbbb com/ mi micahcarrolll/mmiM..
[5] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William Freeman. MaskGIT: Masked generative image transformer. In IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 11305-11315. IEFE, 2022. doi: 10. (099/CVPR22688202222) 103. URL https::/t///tt org/ 1109/ CVPR52688 2022. 01103.
[6] Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. TransDreamer: Reinforcement learning with transformer world models. CoRR, abs/2202.09481, 2022. URL hattps://arxtt org/ abs/2200 /bs/2220.09481..
[7] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Michael Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision Transformer: Rein- forcement learning via sequence modeling. In Marc' Aurelio Ranzato, Alina Beygelz- imer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan, editors, Ad- vances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 15084-15097, 2021. http: http::::://///ceeepeeeeeee
"E4886662a0dd1007726661676666668 Abtt://foe66AAbbatt.t....a Abstract........................ html.
[8] Maxime Chevalier-Boisvert Lucas Willems, and Suman Pal. Minimalistic gridworld environ- ment for openai gym. https: https://githhtt wtt.h//w.....mm..imm//////////// wwwpee////mmmmnniiinnn.... 2018.
[9] Paul F. Christiano, Zain Shah, Igor Mordatch, Jonas Schneider, Trevor Blackwell, Joshua Tobin, Pieter Abbeel, and Wojciech Zaremba. Transfer from simulation to real world through learning deep inverse dynamics model. CoRR, abs/1610.03518, 2016. URL http: //arxiv../////////a org/ /bs/ 1610. 03518.
[10] Henry M. Clever, Ankur Handa, Hammad Mazhar, Kevin Parker, Omer Shapira, Qian Wan, Yashraj S. Narang, Iretiayo Akinola, Maya Cakmak, and Dicter Fox. Assistive Tele-op: Lever- aging transformers collect robotic task demonstrations. CoRR, abs/2112.05122,, 2021. URL https: /ttpss:///xxvvvv/ org/ abs/2112. 05129.
[11] Jacob Devlin, Ming-Wes Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers language understanding. In Jill Burstein, Christy Doran, and Thamar Solorio, editors, Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume I (Long and Short Papers), pages 4171- 4186. Association for Computational Linguisties. 2019. doi: 10.18653/v1/n19. 1423. URL https:///oi... org/10. 10. 11865//1/n19-//111
[12] Yiming Ding, Carlos Florensa, Pieter Abbeel, and Mariano Phielipp. Goal-conditioned imitation learning. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Flo- rence d' d'Ach Buc, Emily B. Fox. and Roman Gamett, editors, Advances in Neural Information Processing Systems 32: Anmual Conference on Neural Information Pro- cessing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 15298-15309, 2019. URL https: httec:p://ppocceeiiii /ttpp//rdinn..nner.... nwwrins...... http://.....rrii...cpaa////////a pcc/pappr/2019//as) GGdaa60eea06115555800648833 www.:e/aaaaa-A-Aaaatt.eaaa. Abstract.htm
[13] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. image is worth 16x16 words: Transformers for image recognition at scale. In 9th International Conference on Learning Representations,

---

ICLR 2021. Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL https: //openreview, wtt..//oooor//aaeiicooaa....
[14] Scott Emmons, Benjamin Eysenbach, Ilya Kostrikov, and Sergey Levine. Rvs: What is essential for offline RL via supervised learning? In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://openrevier net/forumidd-aaeeeoo S87XXAXXI
[15] Daniel Fried, Armen Aghajanyan, Jessy Lin, Sida Wang, Eric Wallace, Freda Shi, Ruiqi Zhong, Wen-tau Yih, Luke Zettlemoye, and Mike Lewis. Incoder: A generative model for code infilling and synthesis. CoRR, abs/2204.05999 2022. doi: 10448550/2X1..2244.59999 URL https://dde
org/10. 48550/ arxiv 2204. 05999.
[16] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL; datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020. URL https:
//arxiv. org/ abs///0//./0//////////////////// 07219.
[17] Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu. Generalized decision transformer for offline hindsight information matching. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL
[18] David Ha and Jürgen Schmidhuber. World models. CoRR, abs/ htbs/1803.10122, 2018. URL httt://arxivy wwwpe/iw..rr//b///00003 10122.
[19] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li. Piotr Dollár, and Ross B. Girshick. Masked autoencoders are scalable vision learners. In JEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2022, New Orleans, LA, USA, June 18-24, 2022, pages 15979- URL https://doi. 10g/10.
15988. IEEE, 2022. doi: 10. 1110//CV/5528882222.001111 1109/ C009//VP852688 2022 01553.
[20] L' L'conard Hussenot, Marcin Andrychowicz. Damien Vincent, Robert Dadashi, Anton Raichuk, Lukasz Stafiniak, Sertan Girgin, Raphaël Marinier, Nikola Momchev, Sabela Ramos, Manu Orsini, Olivier Bachem, Matthieu Geist, and Olivier Pietquin. Hyperparameter selection for imitation learning. In JCML, 2021.
[21] Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as big sequence modeling problem. In Marc'Aurelia Aurelio Ranzato, Alina Beygelz- imer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan. editors, Ad- vances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pages 1273-1286, 2021. URL https: httpt:p::/ppoocediigg weteeeedinns.nnurippp.. wwwrips.ccpaapeeee 220por/222///// //hahh/ 09fe htw.:/077444000---bbaarra Abstract. html.
[22] Leslie Pack Kaelbling. Learning to achieve goals. In Ruzena Bajesy, editor, Proceedings of the 13th International Joint Conference Artificial Intelligence. Chambéry, France, August 28 September 3, 1993, pages 1094-1099. Morgan Kaufmann, 1993.
[23] Pratyusha Kalluri. Don't ask artificial intelligence good or fair, ask how it shifts power. Nature, 583(7815): July 2020. doi: 10. 1011//431116-02000 000-5880000000-----2 URL https: 10w010033111000 020 -20-02003-
[24] Kuang-Huei Lee, Ofir Nachum, Mengjiao Yang, Lisa Lee, Daniel Freeman, Winnie Xu, Sergio Guadarrama, Ian Fischer, Eric Jang, Henryk Michalewski, and Igor Mordatch. Multi-game decision transformers, CoRR, abs/2205. 15241, 2022. doi: 10.485550/rxxv.22255. 12..2225...2..1... URL https://doi.
org/10. 48550/ arxiv. 2205 15241.
[25] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. CoRR, abs/2005.016 .146,, 2020. URL. https:/ /aaxxiv.v... org/ abs///bs/.0../////////// 01643.
12

---

[26] Yoav Levine, Barak Lenz, Opher Lieber, Omri Abend, Kevin Leyton- Broown Moshe Tennen- holtz, and Yoav Shoham. PMI-Masking: Principled masking of correlated spans. 9th Inter- national Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. J
[27] Fangchen Liu, Hao Liu, Aditya Grover, and Pieter Abbeel. Masked autoencoding for scalable and generalizable decision making. In Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2022, NewrIPS 2022, November, 2022, New Orleans, LA, USA, 2022.
[28] Elman Mansimov, Alex Wang, and Kyunghyun Cho. generalized framework of sequence generation with application to undirected sequence models. CoRR, abs/1905. 2019. URL http://ariv..
org/
12790.
[29] Jiquan Ngiam, Benjamin Caine, Vijay Vasudevan, Zhengdong Zhang, Hao- Lewis Chiang, Jeffrey Ling, Rebecca Roelofs, Alex Bewley, Chenxi Liu, Ashish Venugopal, David Weiss, Benjamin Sapp. Zhifeng Chen, and Jonathon Shlens. Scene transformer: A unified multi-task model for behavior prediction and planning. CoRR, abs/2106.08417, 2021. URL https: //arxiv. org/ harg//bb//// abs/2106.... 08417.
[30] Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Çaglar Gülçehree Siddhant M. Jayakumar, Max Jaderberg, Raphaäl Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learning- In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 7487- 7498. PMLR, 2020. URL http: /ttp:://prree http://ppoeeeen.../////// mlr, prwss/esssseee. w11// http:ppatiiptarrr /ttess//111/parrsttt00..... htmml...ooo html.
[31] Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell. Curiosity-driven explo- ration by supe-sude prediction. In Doina Precup and Yee Whye Tch, editors, Proceedings of the 34th International Conference Machine Learning, ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 2778-2787. html.
PMLR, 2017. URL http: http://procceddnng/ mlr wwt.:ps/eo/ppptaaa..ai.e
[32] Dean Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural Comput., 3(8):88 97, 1991. doi: 10.1162/neco. 19teeecc..991.... 1991.3.1.88. URL https://doi. org/
10. 10..1122/ncco. 1991. 3. 1.8.88
[33] Alec Radford, Karthik Narasimhan, Tim Salimans, and Ilya Sutskever. Improv- ing language understanding by generative pre-training https://opent: wwtpppnai...///llpeeeeeeeeeeeeee com/blog/ language- unt.aaggse.suuprriiieeee...// 2018.
[34] Gabriel Recchia. Teaching autoregressive language models complex tasks by demonstration. CoRR, abs/2109.02 2021. URL https::////// /ttp://///xvvv/// http//aaiiiv...raaaaaaaa/ org/abs/2100 02102.
[35] Scott E. Reed, Konrad ZoIna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, Tom Eccles, Jake Bruce, Ali Razavi, Ashley Edwards, Nicolas Heess, Yutian Chen, Raia Hadsell, Oriol Vinyals, Mahyar Bordbar, and Nando de Freitas. A generalist agent. CoRR, abs/2205.06175. 2022. doi: 200550/rrxi..22050617777 URL https: //doi. org/ arg/10..0 488550 arXiv. 2205. 06175.
[36] Nicholas Rhinehart, Rowan McAllister, Kris Kitani, and Sergey Levine. PRECOG: PREdiction conditioned goals in visual multi- agent settings. In 2019 IEEE/CVF International Conference on Computer Vision, ICCV 2019, Seoul, Korea (South). October 27 November 2, 2019, pages 2821-2830. IEEE, 2019. doi: 10.1 (0.1009/CCC22219.000999 URL https:///oi. org/ 10//110.1000 1109/ ICCV 2019. 00291.
[37] Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 8583-8592, PMLR, 2020. URL http:/ btttt://proceeeee pttp//eeeccee..ee//oo// mlr, press/ /aaessvviil w111/veea12a/e sekar html.
13

---

[38] Rohin Shah, Dmitrii Krasheninnikov. Jordan Alexander, Pieter Abbeel, and Anca D. Dragan. Preferences implicit in the state of the world. In 7th International Conference on Learning Reprexentations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.
URL, http///pwreeeeeeeeeee/rrrrrae/ea
[39] Yi Tay. Mostafa Dehghani, Vinh Q. Tran, Xavier Garcia, Dara Bahri, Tal Schuster, Huaixiu Steven Zheng, Neil Houlsby, and Donald Metzler. Unifying language learning paradigms. CoRR, abs/2205.05131, 2022. doi: 10448550/a2xi.2200.55533 URL https:
//doi.org/10 4851000r.55/////x 2205.05131.
[40] Benigno Uria, Mare-Alexandre Côté, Karol Gregor, lain Murray, and Hugo Larochelle. Neural autoregressive distribution estimation. J. Mach. Learn. Res., 17:205:1-205:37, 2016. URL
http://jmlr... wwwooogg///aeeeeeeee papprs/vili 1t//v/111111222// html.
[41] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008 2017. URL https: /ttp:://///ooceeinggg// neur pourips.. issiipss. cc/ htt/ppppprr///// 2017 //ashh/ aateee4aaaaabbtraaccaaaaa html.
[42] Alex Wang and Kyunghyun Cho. BERT has a mouth, and it must speak: BERT as a markov random field language model. CoRR, abs/1902.04094, 2019. URL http://arxiv. org/ abbs/ 1902. 1902.04094.
[43] Mike Wu and Noah D. Goodman. Foundation posteriors for approximate probabilistic inference. CoRR, htss/2255.09735. 2022. doi: 10.48550//rxxv.2205.0077 URL https: org/10. 48550/ar h0t0:/rXiv..200 2205. 09735.
[44] Zhilin Yang. Zihang Dai, Yiming Yang, Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. XLNet: Generalized autoregressive pretraining for language understanding. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d' Alché-Buc, Emily B. Fox, and Roman Garnett, editors, Advances Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14 2019, Vancouver, BC, Canada, pages 5754-5764, 2019. URL https: httppppp: http///roceeee...e//e/ neun cc/paper/ 2019/hash/dece /ttpa////cerr555eeeee6666633eeee wwstrcccccc....tt.t.......c..... html..
[45]
Erie Zhan, Albert Tseng, Yisong Yue, Adith Swaminathan, and Matthew J. Hausknecht. Learn ing calibratable policies using programmatic style-consistency. In Proceedings of the 37th International Conference on Machine Learning, ICMI. 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 11001-11011. PMLR, 2020. URL http: / /ttp///reeeeen..ee//e/ mlr, wwwpppeessseeee //t//vv1lla///// 2tt99//2220aa... htmll.
[46] Brian D. Ziebart, Andrew L. Maas, Andrew Bagnell, and Anind K. Dey. Maximum entropy inverse reinforcement learning. In Dieter Fox and Carla P. Gomes, editors, Proceedings of the Twenty-Third AAAI Conference on Artificial Intelligence, AAAI 2008, Chicago, Illinois, USA, July 13-17, 2008, pages 1433-1438. AAAI Press, 2008. URL http: http:p://// //www, aaai. org/ Library/AAne Baaary/AAAA//A /200////00////// aaai - 227 .hhp.
14


Reviews:
[{"review_id": "nrpZFTMxu_F", "reviewer": "Reviewer_4afK", "summary": "This paper presents the Uni[MASK] framework, which generalizes masked language modeling (in natural language processing) to more general sequence decision problems. Given actions, states and optional property tokens (such as the return-to-go), multiple tasks such as behavioral cloning, reward-conditional offline RL, forward dynamics, etc., may be represented by specific masking schemes. The authors propose to train a model by using a randomized masking scheme (instead of single-task only or with a fixed set of tasks). Such a Uni[Mask] model may also be fine-tuned on a given task afterwards. A pre-trained Uni[MASK] model often performs comparably to single-task training, and fine-tuning further improves results.\n\n[Post author response] I updated the rating I assigned to this paper, in particular because of the correction to the originally misleading discussion.\n", "questions": "**Questions**\n\nMost important question: Could you please clarify exactly what you mean by \"BERT\" architecture vs \"GPT\" architecture, and how they interact with the masking schemes you propose? In NLP, BERT uses a transformer encoder (self-attention) with some masked inputs whereas GPT uses a transformer decoder (causal self-attention, i.e. not looking at the future).\n\nL160+: Should there be a reward for shorter paths? Otherwise, there could be useless \n\nL190: Why no test data?\n\nL212: Do you mean figure 15 in the appendix? In general, figures referred in the main text should also be within the paper itself.\n\n**Suggestion**\n\nPlease move the description of random masking (end of appendix 1) and appendix 2 to the main text. This is too important to leave out.\n\n**Other**\n\nL235: This points to the wrong appendix.", "limitations": "Some limitations are discussed.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\nThe proposed approach neatly formalizes multiple tasks under a general masking framework.\n\nTraining a Uni[MASK] model, then finetuning it on a given task, generally leads to better results than single-task training only.\n\nThe approach is tested in both discrete (MiniGrid) and continuous (Maze2D) environments.\n\n**Weaknesses**\n\nThe discussion around line 200 is misleading. In figure 7, multi-task performs worse than single-task (although sometimes negligibly), which doesn't support H1.\n\nThe distinction between BERT-like and GPT-like models should be described more clearly (see question below).\n\nThe MiniGrid environment appears to be very simple. I would be interested in seeing how well the approach scales with larger grids.", "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": "8w2CSxqqCFM", "reviewer": "Reviewer_YeV6", "summary": "Proposes an approach to using a single bidirectional transformer to accomplish many different reinforcement learning tasks, including behavior cloning, (sub)goal conditioning, and conditioning on properties more generally. It accomplishes this by creating a sequence of out of the (state, action, properties) for all timesteps, then either (1) training on a single type of mask, (2) training across many masking styles drawn from fixed set, and (3) randomly masking positions in the sequence. (2) and (3) can be finetuned in single task fashion. They evaluate on a toy gridworld and the mujoco 2d maze tasks. They find that Uni[MASK] performs better across both given a short truncated sequence, but it underperforms a GPT model for longer sequences. They attribute this to difficulties in training BERT models.", "questions": "\n\t1. Clarify which task is shown in table 1. The grid world or maze?\n\t2. Consider expanding on what position vs. timestep encoding is in the body of the paper (line 231)?\n\t3. The \"bert has no mouth and must speak\" paper was withdrawn\n\t4. Take a look at XLNet and NADE style masking, which would allow combining benefits of GPT and BERT\n\t5. One possibile reason the longer sequence is failing is because it is never exposed to many masked positions in a row at the end of the sequence. Randomized decoding as in xlnet might help here, or applying more structured masking (such as blocking out spans instead of ). See the uni\n\t\t1. Perhaps training on both multitask and \n\t6. What happens if you train on both multitask and random masking? \n\t7. Potential improvements\n\t\t1. Vary the sequence lengths instead of holding fixed\n\t\t\t1. This is more natural for GPT style models, but can be done with XLNet/NADE approaches as well\n\t\t\t2. How might Uni[MASK] handle the full length 200 timestep sequences instead of truncating to last 5/10?\n\t\t2. Ideally would have more complicated tasks than the two presented\n\t\t3. Consider studying how \n\t8. Relevant concurrent work: foundation posterior, which applies a similar approach to inference in probabilistic programs (i.e. generalized bayesian networks) in Stan.\n\t9. nits\n\t\t1. works -> work in line 300", "limitations": "There are no major concerns with this paper", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "\n\t- originality\n\t\t- Concurrent work suggests that the core ideas of this paper are in the air (see: foundation posterior, multigame decision transformer, The UL2 paper on unifying masking schemes).\n\t- quality\n\t\t- The paper is well executed. It has simple experiments which verify the core claims.\n\t\t- My main concern is that the length 10 context experiments underperforming. Would like to see more study of this instead of alluding to claims that BERT has this known failure mode.\n\t- clarity\n\t\t- The paper is well written and a pleasant read. It adequately cites prior work, and it does a great job of explaining and illustrating the core ideas.\n\t- significance\n\t\t- The paper is a solid contribution to the literature. Overall, its core findings and technique align with concurrent and prior 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": "-8Z5PQ1f9s7", "reviewer": "Reviewer_SqCy", "summary": "This submission proposed a new framework, Uni[MASK], which performs unified inference in sequential decision problems, using different masking schemes. \n\nThe authors demonstrated that randomly sampling masking schemes help the bidirectional transformer model being able to do behavior cloning, rewarding conditioning, dynamics modeling and more.\n\nThe experiments are first on grid-world and then extended to Mujoco-physics Maze2D environment. The comparison between single-task, multi-task, random-task, with finetune are through and solid, integrating the theory with practice, and quite convincing. The experiments in Mujoco-physics Maze2D environment shows that Uni[MASK] framework outperforms baselines, such as Feedforward Neural Network and Decision Transformer.\n\nI fully endorse that this submission is a wonderful work: well designed experiments, diplomatic and sound illustration and the novelty of unifying decision making through Uni[MASK] and bidirectional models.", "questions": "From line 285 to line 2899, \"A clear avenue of future work would therefore be to get the “best of both worlds”: long sequences and benefits of random-mask pre-training by using a GPT-like architectures, with our random-mask and finetune training regimes. This which would require finding ways to make GPT act like a bidirectional model – for which recent methods in NLP might offer a useful starting point\". In the rebuttal period, may I know the opinion from the authors, why making GPT act like a bidirectional model is important in the decision making scenario? The answers to this question would make me (maybe other readers) understand this paper better.", "limitations": "* Relative short context lengths. Experiments with longer context lengths would make this submission more convincing.\n* The experiments are mainly on the grid world (grid size is 4* 4). If the grid sizes increase more, this work would become more persuasive.", "rating": 9, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n* Uni[MASK] unifies inference task in sequential decision problems.\n* The illustration of motivation, theory and experiments design is convincing and sound, written and presented very well.\n* The authors demonstrate how randomly sampling masking schemes at training time produces a single multi-inference-task model that can do behavior cloning, reward-conditioning, dynamics modeling and more.\n* The authors test how training on many tasks affects single-task performance.\n* The authors show how fine-tuning models trained with random masking consistently outperforms single-task models.\n* The proposed Hypothesis H1 ( {multi-task, random-mask, finetune} > single-task) and H2 ( {random-mask > multi-task} ) make sense intuitively, as \"H1 test whether models indeed learn richer representations by training on multiple inference tasks\" and H2 test \"whether training on all possible tasks by randomly sampling maskings at training time is better than selecting a set of specific maskings\"\n* The results from well designed experiments strongly support hypothesis H1 and H2.\n* The authors extensively analyze how different training regimes (combination of single-task, multi-task, random mask, fintune) affect performance. The analysis is very convincing to me.\n* The authors also make experiments comparison with GPT-like architectures, and the analysis is also very convincing: \"We find that while using GPT seems to yield similar performance to BERT for context length five, using GPT seems to give an advantage for longer sequence lengths,\" and \"This suggests that if one were able to use a GPT architecture and train it with random masking and fine-tuning, it might be possible to get the best of both worlds.\"\n\nWeakness:\n* Relative short context lengths, and the experiments is mainly on the grid world (4 * 4), if the grid size increases, it might further improve this work.\n", "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": {}}]

Author rebuttal:

Decision: Accept