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

Cleaned-text code points [15540, 31494). Bibliography omitted here; source.md includes it.

Environment Setup. We design a fully observable 4 × 4 4\times 4 Gridworld in which the agent should move to a 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 a reward of 1 1 for each timestep it moves closer to the goal, − 1 -1 if it moves away from the goal, and 0 0 otherwise. We train Uni​ [MASK] models on training trajectories of sequence length T = 10 T=10 from a noisy-rational agent [ 46 ] . More detailed information about the environment is in Appendix E .

### 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 arbitrary inference tasks by conditioning on specific sets of tokens. Unless otherwise indicated, we take the highest probability action from the model a t = arg ⁡ max a t ′ ⁡ ℙ ⁡ ( a t ′ ∣ s 0 , a 0 , … , s t ) a_{t}=\arg\max_{a_{t}^{\prime}}\mathbb{P}(a_{t}^{\prime}\mid s_{0},a_{0},\ldots,s_{t}) , and then query the environment dynamics for the next state s t + 1 s_{t+1} . The model can be used for imitation, reward- and goal-conditioning, 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-conditioning. In Figure 6 , we show how the model can also be conditioned on global properties of the trajectory, such as whether the trajectory passes through a certain position at any timestep.

Qualitatively, these results suggest that the model generalizes 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 2^{T}\times 2^{T}\times 2 possible state, action, and RTG maskings for a sequence of length T 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 and actions further into the future: e.g., “where will the agent be in 3 timesteps ?”. Given a fixed initial set 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 particular timestep marginalizes over missing actions; for example, ℙ ⁡ ( s 1 ∣ s 0 , s 3 , s 6 ) \mathbb{P}(s_{1}\mid s_{0},s_{3},s_{6}) models the possibility that the agent chooses either up or 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 at t = 3 t=3 (leading it to the distribution over states at t = 4 t=4 ), and that the agent must be at position (2, 1) at t = 5 t=5 to reach the door at t = 6 t=6 .

### 4.3 Measuring Single-Task Performance

Next, we investigate how a random-mask model performs on individual tasks, in comparison to single-task models trained exclusively on the evaluated task. If we care about a single task (e.g. goal-conditioned imitation), should we train a model simply on that task? Or can 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 provides a general way to evaluate how well models fit the distribution of trajectories and transitions, which is what we are concerned with for most inference tasks: i.e., “how well can the network predict the true state or action in the data?”.

In Figure 7 , we report validation loss if the model is trained on one task (or multiple tasks) and 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 on the diagonal), but poorly when queried with another task (e.g. past inference).

First, we find that random-mask training (but not multi-task training) outperforms single-task training on half of the tasks considered, showing that even if one is interested in a single inference task, training on many more tasks can sometimes improve performance. Specializing a model trained with random-mask via finetuning ( finetune ) leads to the best performance, outperforming single-task on 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-task ’s poor performance and random-mask ’s 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-task , supporting H2 : training on additional inference tasks beyond the specific ones of interest can augment performance.

## 5 Trajectory Generation in a Complex Environment

In addition to Gridworld, we test our method in a partially observable, continuous-state and continuous-action environment, with a larger trajectory horizon (200 timesteps).

### 5.1 Environment Setup

We adapt the Mujoco-physics Maze2D environment [ 16 ] (see Appendix H for figures), in which a point-mass object is placed at a random location in a maze, and the agent is rewarded for moving towards a randomly generated target location (making this task “goal-conditioned 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 velocity from observed context.

Expert dataset. We want our expert data to have some suboptimality so that reward-conditioning can be tested for better-than-demonstrator performance. We generate a dataset of expert trajectories by rolling out D4RL’s PD controller (which is non-Markovian), and add noise to the actions with zero-mean and 0.5 0.5 variance (which are then clipped to have each dimension between − 1 , 1 -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-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 a feed-forward 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-GPT model with the single-task training regime. The only meaningful difference between Decision-GPT and a single-task 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 (reward-conditioning). The only exceptions are Uni​ [MASK] models trained with multi-task (trained to perform BC and RC) and random-mask . We train two sets of such models, for context lengths of 5 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-less comparably to single-task in behavior cloning and reward conditioning; however, random-mask 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-box performance. However, finetune recovers the performance loss, again out-performing single-task (providing qualified support for H1 ). Surprisingly, for a context length of ten, fine-tuning multi-task does not improve performance as much as 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, we see that 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-GPT baseline. We speculate that this might be related to the documented difficulty of using BERT-like architectures (as that of Uni​ [MASK] models) for sequence generation [ 42 , 28 ] .

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 one uses BERT or GPT as the backbone of the architecture. 2 2 2 We additionally use input-stacking for Uni​ [MASK] – see Appendix D – but in preliminary experiments we found this to not affect performance. We find that while using GPT seems to yield similar (or worse) performance to BERT for context length five, using GPT seems to give an advantage for longer sequence lengths. In particular, note that a 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 a GPT architecture and train it with random masking and fine-tuning, it might be possible to get the best of both worlds.

## 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. goal-conditioning models, etc.). While this is a limitation of our work, it is also not our main focus: we propose a unifying framework for a 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 ), or 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-mask pre-training by using a GPT-like architectures, with our random-mask and finetune training regimes. This requires finding ways to make GPT act like a bidirectional model. Recent methods in NLP might offer a useful starting point [ 1 , 15 ] , as has been explored by concurrent work to ours [ 27 ] .

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 a 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 a single set of weights to perform a large variety of tasks out of the box. MaskDP instead focuses on getting the best performance possible on a 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 of inferences more generally (e.g. Bayes Networks); alternatively, even trivially extending the approach to multi-agent settings (for which token-stacking could prove more valuable), could enable interesting masking-enabled queries [ 29 ] .

## 7 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-specific foundation models. However, we 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 of the corresponding single-task models after multi-task 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.

## Appendix A Code

Our codebase can be found in [ 4 ] . Our code uses assets from gym-minigrid [ 8 , Apache License 2.0] , Decision Transformer [ 7 , MIT License] , and D4RL [ 16 , Apache License 2.0] .

## Appendix B Training regimes

Each batch is made of many randomly sampled trajectory snippets. Across the tasks depicted in Figure 1 , for each snippet τ t : t + k \tau_{t:t+k} we describe the input masking and predicted outputs in detail:

• Behavioral Cloning. Select i ∈ [ 0 , k ] i\in[0,k] uniformly. Feed s t : t + i , a t : t + i − 1 s_{t:t+i},a_{t:t+i-1} to the network (include no actions if i = 0 i=0 ), with all other tokens masked out. Have the network only predict the next missing action a i a_{i} .