mlsanityData to 14 Oct 2026

UniMASK: Unified Inference in Sequential Decision Problems

arXiv 2211.10869 ↗First posted 20 Nov 2022cs.LGcited by 24

Not on the map. Why?

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.

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

36 of 50 references.

References · 36

  1. 2017Attention Is All You Needcited by 60,361
  2. 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
  3. 2020Language Models are Few-Shot Learnerscited by 32,419
  4. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 30,186
  5. 2021Masked Autoencoders Are Scalable Vision Learnerscited by 7,403
  6. 2021On the Opportunities and Risks of Foundation Modelscited by 4,033
  7. 2019XLNet: Generalized Autoregressive Pretraining for Language Understandingcited by 2,811
  8. 2020Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problemscited by 1,920
  9. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,546
  10. 2018World Modelscited by 1,513
  11. 2020D4RL: Datasets for Deep Data-Driven Reinforcement Learningcited by 1,427
  12. 2017Curiosity-driven Exploration by Self-supervised Predictioncited by 1,194
24 more
  1. 2022MaskGIT: Masked Generative Image Transformercited by 1,099
  2. 2022A Generalist Agentcited by 808
  3. 2021Offline Reinforcement Learning as One Big Sequence Modeling Problemcited by 644
  4. 2022InCoder: A Generative Model for Code Infilling and Synthesiscited by 498
  5. 2020Planning to Explore via Self-Supervised World Modelscited by 368
  6. 2019Stabilizing Transformers for Reinforcement Learningcited by 270
  7. 2022Multi-Game Decision Transformerscited by 214
  8. 2019BERT has a Mouth, and It Must Speak: BERT as a Markov Random Field Language Modelcited by 185
  9. 2021RvS: What is Essential for Offline RL via Supervised Learning?cited by 185
  10. 2019PRECOG: PREdiction Conditioned On Goals in Visual Multi-Agent Settingscited by 179
  11. 2019Goal-conditioned Imitation Learningcited by 162
  12. 2022TransDreamer: Reinforcement Learning with Transformer World Modelscited by 143
  13. 2022CM3: A Causal Masked Multimodal Model of the Internetcited by 142
  14. 2021Generalized Decision Transformer for Offline Hindsight Information Matchingcited by 106
  15. 2016Neural Autoregressive Distribution Estimationcited by 95
  16. 2016Transfer from Simulation to Real World through Learning Deep Inverse Dynamics Modelcited by 54
  17. 2020PMI-Masking: Principled masking of correlated spanscited by 54
  18. 2019Preferences Implicit in the State of the Worldcited by 34
  19. 2019A Generalized Framework of Sequence Generation with Application to Undirected Sequence Modelscited by 17
  20. 2021Hyperparameter Selection for Imitation Learningcited by 17
  21. 2021Teaching Autoregressive Language Models Complex Tasks By Demonstrationcited by 15
  22. 2019Learning Calibratable Policies using Programmatic Style-Consistencycited by 11
  23. 2021Assistive Tele-op: Leveraging Transformers to Collect Robotic Task Demonstrationscited by 11
  24. 2022Foundation Posteriors for Approximate Probabilistic Inferencecited by 6

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2022: 2; 2023: 9; 2024: 8; 2025: 3; 2026: 2.

24 papers.

Sort:

By area:

  1. 2023Foundation Models for Decision Making: Problems, Methods, and Opportunitiescited by 164
  2. 2024Real-World Robot Applications of Foundation Models: A Reviewcited by 61
  3. 2023A Survey on Transformers in Reinforcement Learningcited by 50
  4. 2022Masked Autoencoding for Scalable and Generalizable Decision Makingcited by 45
  5. 2023Masked Trajectory Models for Prediction, Representation, and Controlcited by 42
  6. 2022On Transforming Reinforcement Learning by Transformer: The Development Trajectorycited by 36
  7. 2023Emergent Agentic Transformer from Chain of Hindsight Experiencecited by 33
  8. 2023ACT: Empowering Decision Transformer with Dynamic Programming via Advantage Conditioningcited by 25
  9. 2023Large Sequence Models for Sequential Decision-Making: A Surveycited by 19
  10. 2024Inference via Interpolation: Contrastive Representations Provably Enable Planning and Inferencecited by 17
  11. 2023Towards Solving Fuzzy Tasks with Human Feedback: A Retrospective of the MineRL BASALT 2022 Competitioncited by 15
  12. 2023CEIL: Generalized Contextual Imitation Learningcited by 13
12 more
  1. 2024Position: Foundation Agents as the Paradigm Shift for Decision Makingcited by 9
  2. 2024Tackling Data Corruption in Offline Reinforcement Learning via Sequence Modelingcited by 7
  3. 2024Predictive Coding for Decision Transformercited by 6
  4. 2025Multimodal Diffusion Forcing for Forceful Manipulationcited by 4
  5. 2023MaskMA: Towards Zero-Shot Multi-Agent Decision Making with Mask-Based Collaborative Learningcited by 3
  6. 2024Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learningcited by 3
  7. 2024Efficient Offline Reinforcement Learning: First Imitate, then Improvecited by 0
  8. 2024M3PC: Test-time Model Predictive Control for Pretrained Masked Trajectory Modelcited by 0
  9. 2025FOSSIL: Harnessing Feedback on Suboptimal Samples for Data-Efficient Generalisation with Imitation Learning for Embodied Vision-and-Language Taskscited by 0
  10. 2025Planning as Descent: Goal-Conditioned Latent Trajectory Synthesis in Learned Energy Landscapescited by 0
  11. 2026Objective-Behavior Alignment: Diagnostics for MORL Policy Selectioncited by 0
  12. 2026PermVLA: Factorization Order as a Regularizer for VLA Learningcited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2022Towards Flexible Inference in Sequential Decision Problems via Bidirectional Transformers
  2. 2025Learning Unmasking Policies for Diffusion Language Models
  3. 2020Learning to Generalize for Sequential Decision Making
  4. 2024UNO Arena for Evaluating Sequential Decision-Making Capability of Large Language Models
  5. 2026UnMaskFork: Test-Time Scaling for Masked Diffusion via Deterministic Action Branching
  6. 2026Where-to-Unmask: Ground-Truth-Guided Unmasking Order Learning for Masked Diffusion Language Models
  7. 2025On the Limits of Test-Time Compute: Sequential Reward Filtering for Better Inference
  8. 2026CARM: Cancellation-Aware Response Masking for LLM Reinforcement Learning