UniMASK: Unified Inference in Sequential Decision Problems — part 1/4

Cleaned-text code points [0, 15540). Bibliography omitted here; source.md includes it.
# Uni [MASK] : Unified Inference in Sequential Decision Problems

###### 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 reinforcement learning, inverse dynamics, and waypoint conditioning correspond to different sequence maskings over a sequence of states, actions, and returns. We introduce the Uni [MASK] 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 Uni [MASK] model is often capable of carrying out many tasks with performance similar to or 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-masked subsets of tokens in a sequence. For example, 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 [ 33 ] are trained 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 ] .

Our work investigates how masked modeling can be a powerful idea in sequential decision problems. Consider a sequence of states s s and actions a a collected across T T timesteps s 1 , a 1 , … , s T , a T s_{1},a_{1},\dots,s_{T},a_{T} . If we consider each state and action as tokens of a sequence (analogous to words in NLP) and mask the last action ( s 1 , a 1 , s 2 , a 2 , s 3 , ) (s_{1},a_{1},s_{2},a_{2},s_{3},\rule{7.11317pt}{0.42677pt}) , then predicting the missing token a 3 a_{3} amounts to a Behavior Cloning prediction with two timesteps of history [ 32 ] , given that this masking corresponds to the inference ℙ ( a 3 | s 1 : 3 , a 1 : 2 ) \mathbb{P}(a_{3}|s_{1:3},a_{1:2}) . From this perspective, training a model to predict missing tokens from all maskings of the form ( s 1 , a 1 , … , s t , , … , ) (s_{1},a_{1},\dots,s_{t},\rule{7.11317pt}{0.42677pt},\dots,\rule{7.11317pt}{0.42677pt}) for all t ∈ [ 1 , … , T ] t\in[1,\dots,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 . In this framework, commonly-studied 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 a simple sequence modeling paradigm. In contrast to standard approaches that train a model for each inference task, we show how this framework naturally lends itself to multi-task training: a single Uni [MASK] model can be trained to perform a variety of tasks out-of-the-box by appropriately selecting sequence maskings at training time.

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 tends to outperform 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-performing single-task models via unified multi-task training.

In summary, our contributions are:

1. We propose a new framework, Uni [MASK] , that unifies inference tasks in sequential decision problems 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, and more out-of-the-box.

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-art 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 as maskings.

The utility of masked prediction. Work in both NLP [ 11 ] and vision [ 5 , 19 ] have explored how masked prediction is useful as a 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 a probabilistic program [ 43 ] .

Sequential decision-making as sequence modeling. Previous and concurrent work [ 7 , 21 , 24 ] shows how to use GPT-style (causally-masked) transformers to directly generate high-reward trajectories in an offline RL setting. We expand our focus to 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.e., 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 sequence-modeling perspective, and a 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 RL has investigated other auxiliary objectives, such as state dynamics prediction [ 37 ] or intrinsic motivation [ 31 ] . Typically, to accomplish the tasks we consider, prior work relies on single-task models: for example, goal-conditioned imitation learning [ 12 ] , RL [ 22 ] , waypoint-conditioning [ 36 ] , property-conditioning [ 45 , 17 ] , or dynamics model learning [ 18 , 9 ] . Other work has focused on training models to perform different “tasks” such as different games in Atari [ 24 ] or different environments and multi-modal prediction tasks [ 35 ] . In contrast, we are interested in performing different inference tasks in a single environment, such as RL 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 as 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 ): τ = { ( s 0 , a 0 , p 0 ) , … , ( s T , a T , p T ) } \tau=\{(s_{0},a_{0},p_{0}),\ldots,(s_{T},a_{T},p_{T})\} . 1 1 1 While reward-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-go also enables us to compare our method with previous offline-RL work [ 7 ] .

Motivated by canonical problems in decision-making that involve reward, in most of our analysis we use return-to-go (RTG) as the property (the sum of rewards from timestep t t to the end of the episode): that is, we set p t = R ^ t p_{t}=\hat{R}_{t} where R ^ t = ∑ t ′ = t T r t ′ \hat{R}_{t}=\sum_{t^{\prime}=t}^{T}r_{t^{\prime}} . However, any property of the decision problem can be considered a “property token,” including specific environment conditions being satisfied, the style of the agent, or 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 a non-reward property in Appendix F .

### 3.1 Tasks as Masking Schemes

In the Uni [MASK] framework, we formulate tasks in sequential decision problems as input masking schemes. Formally, a 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 to predict). For example, the masking scheme for BC unmasks (conditions on) s 0 : t s_{0:t} and a 0 : t − 1 a_{0:t-1} , and the model must predict a t a_{t} .

In Figure 1 , we illustrate how to unify commonly-studied tasks such as BC, goal and waypoint conditioned imitation, offline RL (reward-conditioned imitation), and dynamics modeling under our proposed representation of tasks as masking schemes. We describe the masking scheme for each of these tasks in detail in Appendix B .

### 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 a 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 decision-making 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 reward-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-conditioning 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 F 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 3 and described below.

single-task . Training on just one of the masking scheme described in Section 3.1 .

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 a single model to perform well on multiple tasks. Additionally, it might outperform single-task on individual tasks, as the model could learn richer representations of the environment from the additional masking schemes.

Training a single model on a fully randomized masking scheme. For each trajectory snippet, first, a masking probability p mask ∈ [ 0 , 1 ] p_{\textrm{mask}}\in[0,1] is sampled uniformly at random; then each state and action token is masked with probability p mask p_{\textrm{mask}} ; lastly, the first RTG token is masked with probability 1 / 2 1/2 and subsequent RTG tokens are always masked (see Appendix C for details). Intuition: Could allow a single model to perform well on any sequence inference task without the need to specify the tasks of interest at training time. The model may learn richer representations than those of multi-task as it must reason about all aspects of the environment.

Fine-tune a model pre-trained in random-mask on a specific masking scheme. Intuition: Performing fine-tuning could allow the model to benefit from the improved representations obtained from random-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 \left\{\text{\hyperlink{tr:multi}{{multi-task}}, \hyperlink{tr:random}{{random-mask}}, \hyperlink{tr:finetune}{{finetune}}}\right\}>\text{\hyperlink{tr:single}{{single-task}}} .

Randomized mask training outperforms training on a specific set of tasks: random-mask > multi-task \text{\hyperlink{tr:random}{{random-mask}}}>\text{\hyperlink{tr:multi}{{multi-task}}} .

H1 tests whether models learn richer representations by training on 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 a set 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 inference tasks at test-time on a 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-mask does not compromise performance on most specific tasks of interest. Models trained with random-mask achieve comparable or better performance to single-task and multi-task -models, and in fact consistently outperform after additional fine-tuning on the task of interest ( finetune ).