# UniMASK: Unified Inference in Sequential Decision Problems

arXiv: https://arxiv.org/abs/2211.10869 · HTML: https://mlsanity.com/paper/2211.10869/

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 50 references that are in the map.
- First posted on arXiv 20 Nov 2022.

## Builds on (most cited first, 30 of 36)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [On the Opportunities and Risks of Foundation Models](https://mlsanity.com/paper/2108.07258/) (2021)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [MaskGIT: Masked Generative Image Transformer](https://mlsanity.com/paper/2202.04200/) (2022)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [InCoder: A Generative Model for Code Infilling and Synthesis](https://mlsanity.com/paper/2204.05999/) (2022)
- [Planning to Explore via Self-Supervised World Models](https://mlsanity.com/paper/2005.05960/) (2020)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Multi-Game Decision Transformers](https://mlsanity.com/paper/2205.15241/) (2022)
- [BERT has a Mouth, and It Must Speak: BERT as a Markov Random Field Language Model](https://mlsanity.com/paper/1902.04094/) (2019)
- [RvS: What is Essential for Offline RL via Supervised Learning?](https://mlsanity.com/paper/2112.10751/) (2021)
- [PRECOG: PREdiction Conditioned On Goals in Visual Multi-Agent Settings](https://mlsanity.com/paper/1905.01296/) (2019)
- [Goal-conditioned Imitation Learning](https://mlsanity.com/paper/1906.05838/) (2019)
- [TransDreamer: Reinforcement Learning with Transformer World Models](https://mlsanity.com/paper/2202.09481/) (2022)
- [CM3: A Causal Masked Multimodal Model of the Internet](https://mlsanity.com/paper/2201.07520/) (2022)
- [Generalized Decision Transformer for Offline Hindsight Information Matching](https://mlsanity.com/paper/2111.10364/) (2021)
- [Neural Autoregressive Distribution Estimation](https://mlsanity.com/paper/1605.02226/) (2016)
- [Transfer from Simulation to Real World through Learning Deep Inverse Dynamics Model](https://mlsanity.com/paper/1610.03518/) (2016)
- [PMI-Masking: Principled masking of correlated spans](https://mlsanity.com/paper/2010.01825/) (2020)
- [Preferences Implicit in the State of the World](https://mlsanity.com/paper/1902.04198/) (2019)

## Built on by (most cited first, 24 of 24)

- [Foundation Models for Decision Making: Problems, Methods, and Opportunities](https://mlsanity.com/paper/2303.04129/) (2023)
- [Real-World Robot Applications of Foundation Models: A Review](https://mlsanity.com/paper/2402.05741/) (2024)
- [A Survey on Transformers in Reinforcement Learning](https://mlsanity.com/paper/2301.03044/) (2023)
- [Masked Autoencoding for Scalable and Generalizable Decision Making](https://mlsanity.com/paper/2211.12740/) (2022)
- [Masked Trajectory Models for Prediction, Representation, and Control](https://mlsanity.com/paper/2305.02968/) (2023)
- [On Transforming Reinforcement Learning by Transformer: The Development Trajectory](https://mlsanity.com/paper/2212.14164/) (2022)
- [Emergent Agentic Transformer from Chain of Hindsight Experience](https://mlsanity.com/paper/2305.16554/) (2023)
- [ACT: Empowering Decision Transformer with Dynamic Programming via Advantage Conditioning](https://mlsanity.com/paper/2309.05915/) (2023)
- [Large Sequence Models for Sequential Decision-Making: A Survey](https://mlsanity.com/paper/2306.13945/) (2023)
- [Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inference](https://mlsanity.com/paper/2403.04082/) (2024)
- [Towards Solving Fuzzy Tasks with Human Feedback: A Retrospective of the MineRL BASALT 2022 Competition](https://mlsanity.com/paper/2303.13512/) (2023)
- [CEIL: Generalized Contextual Imitation Learning](https://mlsanity.com/paper/2306.14534/) (2023)
- [Position: Foundation Agents as the Paradigm Shift for Decision Making](https://mlsanity.com/paper/2405.17009/) (2024)
- [Tackling Data Corruption in Offline Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2407.04285/) (2024)
- [Predictive Coding for Decision Transformer](https://mlsanity.com/paper/2410.03408/) (2024)
- [Multimodal Diffusion Forcing for Forceful Manipulation](https://mlsanity.com/paper/2511.04812/) (2025)
- [MaskMA: Towards Zero-Shot Multi-Agent Decision Making with Mask-Based Collaborative Learning](https://mlsanity.com/paper/2310.11846/) (2023)
- [Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning](https://mlsanity.com/paper/2409.06985/) (2024)
- [Efficient Offline Reinforcement Learning: First Imitate, then Improve](https://mlsanity.com/paper/2406.13376/) (2024)
- [M3PC: Test-time Model Predictive Control for Pretrained Masked Trajectory Model](https://mlsanity.com/paper/2412.05675/) (2024)
- [FOSSIL: Harnessing Feedback on Suboptimal Samples for Data-Efficient Generalisation with Imitation Learning for Embodied Vision-and-Language Tasks](https://mlsanity.com/paper/2510.11307/) (2025)
- [Planning as Descent: Goal-Conditioned Latent Trajectory Synthesis in Learned Energy Landscapes](https://mlsanity.com/paper/2512.17846/) (2025)
- [Objective-Behavior Alignment: Diagnostics for MORL Policy Selection](https://mlsanity.com/paper/2606.21321/) (2026)
- [PermVLA: Factorization Order as a Regularizer for VLA Learning](https://mlsanity.com/paper/2610.04659/) (2026)

## Abstract (arXiv)

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 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. Our code is publicly available at https://github.com/micahcarroll/uniMASK.