Uni [MASK] : Unified Inference in Sequential Decision Problems
Anonymous Author(s)
AffiliationAddress email
Abstract
Randomly masking and predicting word tokens has been a successful approach inpre-training language models for a variety of downstream tasks.
In this work, weobserve 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 asequence of states, actions and returns.
We introduce the Uni [MASK] framework,which provides a unified way to specify models which can be trained on manydifferent sequential decision making tasks.
We show that a single Uni [MASK]model is often capable of carrying out many tasks with performance similar toor better than single-task models.
Additionally, after fine-tuning, our Uni [MASK]models consistently outperform comparable single-task models
Introduction
Masked language modeling [11] is a key technique in natural language processing (NLP).
Underthis paradigm, models are trained to predict randomly-masked subsets of tokens in a sequence.
Forexample, during training, a BERT model might be asked to predict the missing words in the sentence“yesterday I cooking a ”.
Importantly, while unidirectional models like GPT [32] are trainedto predict the next token conditioned only on the left context, bidirectional models trained on thisobjective learn to model both the left and right context to represent each word token.
This leads toricher representations that can then be fine-tuned to excel on a variety of downstream tasks [11].
Our work investigates how masked modeling can be a powerful idea in sequential decision problems.
Consider a sequence of states s and actions a collected across T timesteps s 1 ,a 1 ,...,s T ,a T .
Ifwe consider each state and action as tokens of a sequence (analogous to words in NLP) and maskthe last action ( s 1 ,a 1 ,s 2 ,a 2 ,s 3 , ) , then predicting the missing token a 3 amounts to a Behavior
Cloning prediction with two timesteps of history [31], given that this masking corresponds to theinference P ( a 3 ∣ s 1 ∶ 3 ,a 1 ∶ 2 ) .
From this perspective, training a model to predict missing tokens fromall maskings of the form ( s 1 ,a 1 ,...,s t , ,..., ) for all t ∈ [ 1 ,...,T ] corresponds to training
a
Behavior Cloning (BC) model.
In this work, we introduce the Uni [MASK] framework: Uni fied Inferences in Sequential Decision
Problems via [MASK] ings, where inference tasks are expressed as masking schemes .
Through thisframework, commonly-studied tasks such as goal or waypoint conditioned BC [13, 35], offlinereinforcement learning (RL) [25], forward or inverse dynamics prediction [20, 9, 5], initial-stateinference [37], and others can be unified under a simple sequence modeling paradigm.
In contrastto standard approaches that train a single-task model for each inference task, we show how thisframework naturally lends itself to multi-task training: a single Uni [MASK] model can be trained toperform a variety of tasks out-of-the-box by appropriately selecting sequence maskings at trainingtime.
We test this framework in a gridworld navigation task and a continuous control environment.
First, we train a Uni [MASK] model by sampling from the space of all possible maskings at training time
(random masking) and show how this scheme enables a single Uni [MASK] model to perform BC, reward-conditioning, waypoint-conditioning 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 outperforms models specialized to those tasks
Our results suggest that expressing tasks as 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 -performingsingle task models via unified multi-task training.
In summary, our contributions are:
We propose a new framework, Uni [MASK] , that unifies inference tasks in sequential decisionproblems as different masking schemes in a sequence modeling paradigm.2.
We demonstrate how randomly sampling masking schemes at training time produces a single multi-inference-task model that can do BC, reward-conditioning, dynamics modeling
andmore out-of-the-box.
We test how training on many tasks affects single-task performance, and show how finetuning models trained with random masking consistently outperforms single-task models.
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-art methods.
Related Work
Transformer models.
Our work is motivated by the huge success of transformer models [38] inother domains such as NLP [11, 32, 3] and computer vision [14, 21].
Using transformers in RL andsequential decision problems has proven difficult due to the instability of training [29], but recentwork has investigated how to use transformers in model-based RL [6], motion forecasting [28],learning from demonstrations [33], and tele-operation [10].
The utility of randomized masking.
In addition to being used as one of the main training objectivesfor BERT (the “cloze task”, [11]), the flexibility afforded by randomized masking in bidirectionalmodels has been utilized in other previous works applied to language [19, 27] and vision [4] – mostlyfor the purpose of speeding up auto-regressive decoding, which is not our focus.
Sequential decision-making as sequence modeling.
[7] and [22] focus on RL and show how one canuse GPT-style (causally-masked) Transformer models to directly generate high-reward trajectories inan offline RL setting.
Unlike this line of work, we focus on many tasks that a sequence modelingperspective enables one to do, rather than just offline-RL.
While some previous work has cast doubton the necessity of using transformers to achieve good results in offline RL [15], we note that offline RL [25] is just one of the various tasks we consider.
Concurrent work to ours generalizes the leftto-right masking in the Transformer to condition on future trajectory information for tasks such asstate marginal matching [18] and multi-agent motion forecasting [28].
In contrast, we systematicallyinvestigate how a single bidirectional Transformer can be trained to perform arbitrary downstreamtasks (in more complex settings than motion forecasting—i.e., we also consider agent actions andrewards in addition to states).
The main thing that sets us apart from these works is a systematic viewof all tasks that can be represented by this sequence-modeling perspective, and a detailed investigationof how different multi-task training regimes compare.
Self-supervised learning for sequential decision problems.
Training on random maskings of one’sdata can be considered as a form of self-supervised learning.
Previous work on self-supervisedlearning is mostly focused on improving RL by training models on auxiliary objectives such asstate dynamics prediction [36] or intrinsic motivation [30].
Typically, to accomplish the tasks weconsider, prior work relies on single-task models: for example, goal-conditioned imitation learning[12], RL [23], waypoint-conditioning [35], property-conditioning [40, 18], or dynamics model learning [20, 9].
Another recent work [34] focuses on scaling up transformer models to be able tohandle multi-modal input from a large variety of different environments (which they call “tasks” –which is different from our usage of the word! ).
We instead demonstrate how sequence modeling canbe a unifying framework for formulating and performing any inference task within an environmentwith a single model.
The Uni [MASK] Framework
We introduce the Uni [MASK] framework.
In Section 3.1 we propose a unifying interpretation ofinference tasks in sequential decision problems as masking schemes.
In Section 3.2 we describingdifferent ways of training Uni [MASK] models, an provide hypotheses about their efficacy.
We model trajectories as sequences of states, actions, and (optionally) property tokens: {+
 [Equation]+}
Any property of the decision problem can be considered a “property token”, including specificenvironment conditions being satisfied, the style of the agent, or the performance of the agent (e.g. thereward obtained in the timestep).
Motivated by canonical problems in decision making that involvereward, in almost all of our analysis and experiments we consider the return-to-go property (the sumof rewards from timestep t to the end of the episode): that is, we set p t
= ˆ R t where ˆ R t =
∑ Tt ′ =
t r t ′ .
3.1 Tasks as Masking Schemes
In the Uni [MASK] framework, we formulate tasks in sequential decision problems as input maskingschemes.
Formally, a masking scheme specifies which input tokens are masked (determining whattokens are shown to the model for prediction) and which outputs of the model are masked beforecomputing losses (determining which outputs the model should learn to predict).
For example, themasking scheme for BC unmasks (conditions on) s
∶ t and a 0 ∶ t − 1 , and the model must predict a t .
In Figure 1, we illustrate how commonly-studied tasks such as BC, goal and waypoint conditionedimitation, offline RL (reward-conditioned imitation), and dynamics modeling can be unified underthis representation of tasks as masking schemes.
We describe the masking scheme for each of thesetasks in detail in Appendix A.1.
3.2 Model Architecture & Training Regimes
Throughout the body of this, we instantiate our Uni [MASK] framework with the BERT architecture[11] consisting stacked bidirectional transformer encoder (self-attention) layers and positional encod-
3.2.1 Training regimes
We experiment with four training regimes , illustrated in Figure 3 and described below.single-task .
Training on just one of the masking scheme described in Section 3.1.
Intuition: A single, unified interface for all tasks (data pipeline, neural architecture, etc.),where only the masking schemes needs to be varied for different tasks.
This is a feature of Uni [MASK] as a whole, so we think of single-task as a baseline training regime.multi-task .
Training a single model on multiple masking schemes.
Intuition: A single model should perform well on any of the multiple tasks on which it wastrained.
Additionally, it might outperform single-task on individual tasks, as the modelmay learn richer representations of the environment from the additional masking schemes.random- mask .
Training a single model on a randomized masking scheme.
Intuition: A single model should
perform well on any sequence inference task, without requiring to specify the tasks of interest at training time.
The model may learn representations even richer than that of multi-task as the model must reason about all components of the environment. finetune .
Fine-tune a model pre-trained in random-mask on a specific masking scheme.
Intuition: A fine-tuned model benefits from the improved representations obtained fromrandom-mask , while specializing to the single task at hand.
3.2.2 Hypotheses
Based on the intuition of the strengths of each training regime, we formulate the following hypotheses:
H1.
First training on multiple inference tasks will lead to better performance on individual tasks than
only training on that inference task: { multi-task , random-mask , finetune } > single-task .
H2.
Random masks training outperforms training on a specific set of tasks: random-mask > multi- task . So, H1 tests whether models indeed learn richer representations by training on multiple inferencetasks.
H2 tests a stronger claim: whether training on all possible tasks by randomly samplingmaskings at training time is better than selecting a set of specific maskings 2 Sample a random trajectory snippet and apply a masking (chosen ).Sample
a random trajectory snippet and mask each token at random.
A Unified Model for Any Inference Task
We first demonstrate how random-mask enables a single Uni [MASK] model to perform arbitrary inference tasks at test-time on a gridworld environment, without the need for task-specific output heads ortraining schemes that are customized for the downstream task.
We then show that random-mask doesnot compromise performance on most specific tasks of interest.
Models trained with random-maskachieve comparable or better performance to single-task and multi-task -models, and in factconsistently outperform after additional fine-tuning on the task of interest ( finetune ).
Environment Setup.
We design a fully observable 4 × 4 gridworld in which the agent should moveto a fixed goal location behind a locked door with the MiniGrid environment framework [8].
Theagent and key positions are randomized in each episode.
The agent receives a reward of 1 for eachtimestep 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 T = 10 from a noisy-rational agent[41].
More detailed information about the environment is in Appendix A.4.
4.1 One Model to Rule Them All
As shown in Figure 4, a single Uni [MASK] model trained with random-mask can be used for arbitraryinference tasks by conditioning on specific sets of tokens.
Unless otherwise indicated, we take the Conditioned on property = {+reaches top left corner s 
 t= t= t= t= ... 1 p+}
Property-conditioned
If the training dataset is additionally labeled with property indicators (e.g. whether the trajectory passes by the top left corner of the grid at any timestep ), the model can roll out trajectories conditioned on whether the property is exhibited.
t= p(s1 | s0, s3, s6) p(s2 | s0, s3, s6) t= p(s4 | s0, s3, s6) p(s5 | s0, s3, s6) t= highest probability action from the model a t = arg max a ′ t P ( a ′ t
∣ s 0 ,a 0 ,...,s t ) , and then query theenvironment dynamics for the next state s t
+ 1 .
The model can be used for imitation, reward- andgoal-conditioning, or as a forward or inverse dynamics model (when querying for state predictions, asin the backwards inference task).
If trajectories are labeled with properties at training time, the modelcan also be used for property-conditioning.
In Figure 5, we show how the model can be conditionedon global properties of the trajectory (as opposed to conditioning on states or actions at specifictimesteps), such as whether the trajectory passes through a certain position at any timestep.
Qualitatively, these results suggest that the model is capable of generalizing across masking schemes,since seeing the exact masking corresponding to a particular task at training time is exceedingly rare(out of 2 T × 2 T × 2 possible state, action, and reward maskings for a sequence of length T ).
4.2 Future State Predictions
Uses of the random-mask -trained Uni [MASK] model are not limited to rolling out new trajectories(requesting inferences about the agent’s next action).
One can also request inferences for states andactions further into the future: e.g., “where will the agent be in 3 timesteps ?”.
Given a fixed initialset of observed states, we visualize the distribution of predicted states at each timestep in Figure 6.
Since we do not roll out actions, querying the model for the predicted state distribution at a particulartimestep marginalizes over missing actions; for example, P ( s 1 ∣ s 0 ,s 3 ,s 6 ) models the possibility thatthe agent chooses either up or left as the first action.
Qualitatively, the state predictions suggest thatthe model accurately captures the environment dynamics and usual agent behavior; e.g. it correctlymodels that the agent has equal probability of going up and right at t = 3 (leading it to the distributionover states at t = 4 ), and that the agent must be at position (2, 1) at t = 5 in order to reach the door att = 6 .
4.3 Measuring Single-Task Performance
Next, we investigate how random-mask performs on individual tasks, in comparison to asingle-task trained exclusively on the evaluated task. .
For this set of experiments, we primarily consider validation loss as our measure of performance.
Validation loss provides a generalway to evaluate how well models fit the distribution of trajectories and transitions, which is primarilywhat we are concerned with for most inference tasks: i.e., how well can the network predict the truestate or action in the data (e.g. next-action prediction for behavior cloning)?
In Figure 7, we report validation loss if the model is trained on one task (or multiple tasks) andevaluated on another task.
As expected, models that have been trained on one masking (e.g. BC;predicting the next action conditioned on the past) perform poorly when queried with another task(e.g. predicting the previous states conditioned on the final state; past inference).
We find that many-task models ( multi-task and random-mask ) perform comparably (or better) tosingle-task .
Specializing a model trained on multiple tasks via finetuning ( finetune ) leads to thebest performance, outperforming single-task on all tasks except BC and forward dynamics.
These Additionally, we find that random-mask leads to lower loss values on almost all evaluation tasksrelative to multi-task , supporting H2 : training on additional inference tasks beyond the specificones of interest can further augment performance.
Trajectory Generation in a Complex Environment
In addition to gridworld, we also test our method in a partially observable, continuous-state andcontinuous-action environment, with a larger trajectory horizon (200 timesteps).
5.1 Environment Setup
We adapt the Mujoco-physics Maze2D environment [17] (see Figure 14), in which a point-massobject is placed at a random location in a maze, and is rewarded for moving towards a randomlygenerated target location (making this task “goal-conditioned by default”).
Notably, we make this task harder than the original by: 1) removing the velocity information of theagent from each timestep’s observation, and 2) increasing the amount of initial position randomization.
These changes make the environment partially observable, forcing models trained on this data toimplicitly infer the agent’s velocity from observed context.
Expert dataset.
We want our expert data to have some suboptimality so that reward-conditioningcan be tested for better-than-demonstrator performance.
We generate a dataset of expert trajectoriesby rolling out D4RL’s PD controller (which is non-Markovian), and additionally add noise to theactions 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 forvalidation.
More details on our adapted Maze2D environment and design decision can be found in Appendix A.7.
5.2 Models Trained
For the Maze2D evaluations, we focus on test-time reward performance on behavior cloning andoffline RL (reward-conditioning) across various architectures and training regimes.
We consider Uni [MASK] models trained with the different training regimes: single-task ,multi-task , random-mask , and finetune .
We additionally consider other architectures, such as afeed-forward NN and Decision Transformer (DT) baselines [7].
We found that several of our designdecisions for Uni [MASK] models – using positional encoding instead of timestep encoding, inputtingthe 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 A.7).
We train our Decision-GPT model with the single-task training regime.
Theonly meaningful difference between Decision-GPT and a single-task Uni [MASK] model is whetherthe model is GPT- or BERT-based.
For each architecture and applicable training regime, we train separate models to perform behaviorcloning and offline RL (reward-conditioning).
The only exceptions are Uni [MASK] models trainedwith multi-task (trained to perform BC and RC) and random-mask .
We train two sets of suchmodels, for context lengths of 5 and 10 – meaning that during both training and evaluation, themodels will respectively only be able to see the last 5 or 10 timsteps of the agent’s interaction withthe environment.
5.3 Results
We report reward evaluation results for 1000 rollouts in the Maze environment with standard errorsacross 5 seeds in Table 1.
The value of pre-training and fine-tuning for Uni [MASK] models.
We find that fine-tuning iscritical for good performance in more complex environments.
We see that multi-task on performscomparably to single-task in behavior cloning and reward conditioning; however, random-maskin this setting suffers in terms of reward-performance.
This suggests that multi-task trainingcan be effective in increasing reward performance, but training on too many tasks can hurt out-ofthe-box performance.
However, finetune recovers the performance loss, again out-performingsingle-task (providing qualified support for H1 ).
Surprisingly, fine-tuning multi-task does notimprove performance as much as fine-tuning the randomly masked model, suggesting that specificallytraining on random masking provides benefits for adapting models to individual downstream tasks(providing qualified support for H2 ).
How do Uni [MASK] models compare to other architectures?
For context length five, we see thatfinetune Uni [MASK] models perform better than all baselines we consider.
However, increasingthe context length to ten, we see that Uni [MASK] models performs poorly across the board, withfinetune outperformed even by our Decision-GPT baseline.
We speculate that this might be relatedto the documented difficulty of using BERT-like architectures (as that of Uni [MASK] models) forsequence generation [39].
Isolating the effect of GPT vs. BERT.
In order to investigate the effect of using GPT-like architectures instead of BERT-like ones, we can consider the comparison between single-task Uni [MASK]and our Decision-GPT baseline: the main difference between these two models is only whether oneuses BERT or GPT as the backbone of the architecture.
4 We find that while using GPT seems toyield similar performance to BERT for context length five, using GPT seems to give an advantagefor longer sequence lengths.
In particular, note that a larger context length enables GPT to increaseperformance, while this is not the case for single-task Uni [MASK] .
This suggests that if one wereable to use a GPT architecture and train it with random masking and fine-tuning, it might be possibleto get the best of both worlds.
Limitations and Future Work
Comparison to other specialized models.
We show that Uni [MASK] outperforms feedforwardnetworks, 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 priorwork that are specialized for specific tasks (e.g. goal-conditioning models, etc.).
While this is alimitation of our work, it is also not our main focus: we propose a unifying framework for a variety of [Table] decision problems [Table]
analyze how different training regimes affectperformance.
Longer context lengths.
One limitation in our experimentation are the relatively short context lengthsused.
As we saw in our experiments, using longer context lengths starts affecting Uni [MASK] models’performance negatively.
In part, this could be addressed by designing masking schemes tailed tospecific test-time tasks (see Appendix A.2), or using principled masking schemes [26].
However, thisdegradation may be attributed to our use of a BERT-like (rather than GPT-like) architecture, whichseems less compatible with longer sequence lengths.
A clear avenue of future work would thereforebe to get the “best of both worlds”: long sequences and benefits of random-mask pre-training byusing a GPT-like architectures, with our random-mask and finetune training regimes.
This whichwould require finding ways to make GPT act like a bidirectional model – for which recent methods in NLP might offer a useful starting point [1, 16].
Other future work.
Another exciting direction for future work is trying to see whether the benefitsobtained from random-mask (or even multi-task ) would also apply to other types of inferences(e.g. Bayes Networks more generally); alternatively, even trivially extending the approach to multiagent settings (for which token-stacking could prove more valuable), could enable many interestingmasking-enabled queries [28].
Conclusion
Broader impacts.
The prospect of very large “foundation models” [2] becoming the norm forsequential problems (in addition to language) raises concerns, in that it de-democratizes developmentand usage [24].
We do use significantly smaller models and computational power than similarwork, leaving open the option to have more modestly-sized environment-specific foundation models.
However, we acknowledge that this works still encourages this problematic trend.
Summary.
In this work we propose Uni [MASK] , a framework for flexibly defining and trainingmodels which: 1) are naturally able to represent any inference task and support multi-task trainingin sequential decision problems, 2) match or surpass the performance of single-task models aftermulti-task pre-training, and almost always surpasses them after fine-tuning.
We believe our approachprovides an elegant conceptual formulation and exciting experimental results which warrant furtherinvestigation by the research community.
Checklist
For all authors
(a) Do the main claims made
in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See Section 6.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] See
Section 7.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[N/A](b)
Did you include complete proofs of all theoretical results?
[N/A]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main exper-
imental results (either in the supplemental material or as a URL)?
[Yes] The code islinked at the top of the Appendix.(b) Did you specify all the training details (e g., data splits, hyperparameters,
how they
were chosen)?
[Yes] See the relevant Appendix sections, for each environment and
model.
(c) Did you report error bars (e.g., with respect to the random seed after running exper-
iments multiple times)?
[Yes] Yes, with the exception of Figure 7.
More details in Appendix A.
(d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)?
[Yes] See Appendix.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators?
[Yes]
(b) Did you mention the license of the assets?
[Yes] See the relevant Appendix sections.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]Yes, our codebase is linked in the appendix as a URL.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount
spent on participant compensation?
[N/A]