# Masked Trajectory Models for Prediction, Representation, and Control

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

## Key facts

- Cited by 42 papers in the map (data through 14 Oct 2026).
- Placed in Decision transformers and sequence modeling, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 47 of its 66 references that are in the map.
- First posted on arXiv 4 May 2023.

## Placed in

- Offline RL and imitation learning from data / [Decision transformers and sequence modeling](https://mlsanity.com/offlinerl/decision-transformers-and-sequence-modeling/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Addressing Function Approximation Error in Actor-Critic Methods](https://mlsanity.com/paper/1802.09477/) (2018)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Curiosity-driven Exploration by Self-supervised Prediction](https://mlsanity.com/paper/1705.05363/) (2017)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Exploration by Random Network Distillation](https://mlsanity.com/paper/1810.12894/) (2018)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Learning Complex Dexterous Manipulation with Deep Reinforcement Learning and Demonstrations](https://mlsanity.com/paper/1709.10087/) (2017)
- [R3M: A Universal Visual Representation for Robot Manipulation](https://mlsanity.com/paper/2203.12601/) (2022)
- [Self-supervised Visual Feature Learning with Deep Neural Networks: A Survey](https://mlsanity.com/paper/1902.06162/) (2019)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [VIMA: General Robot Manipulation with Multimodal Prompts](https://mlsanity.com/paper/2210.03094/) (2022)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [COMBO: Conservative Offline Model-Based Policy Optimization](https://mlsanity.com/paper/2102.08363/) (2021)
- [Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videos](https://mlsanity.com/paper/2206.11795/) (2022)
- [Behavior Transformers: Cloning k modes with one stone](https://mlsanity.com/paper/2206.11251/) (2022)
- [Masked Visual Pre-training for Motor Control](https://mlsanity.com/paper/2203.06173/) (2022)

## Built on by (most cited first, 30 of 42)

- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [GAIA-1: A Generative World Model for Autonomous Driving](https://mlsanity.com/paper/2309.17080/) (2023)
- [Unified Video Action Model](https://mlsanity.com/paper/2503.00200/) (2025)
- [Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers](https://mlsanity.com/paper/2409.20537/) (2024)
- [Foundation Policies with Hilbert Representations](https://mlsanity.com/paper/2402.15567/) (2024)
- [Real-World Robot Applications of Foundation Models: A Review](https://mlsanity.com/paper/2402.05741/) (2024)
- [Zero-Shot Whole-Body Humanoid Control via Behavioral Foundation Models](https://mlsanity.com/paper/2504.11054/) (2025)
- [PlayFusion: Skill Acquisition via Diffusion from Language-Annotated Play](https://mlsanity.com/paper/2312.04549/) (2023)
- [RealGen: Retrieval Augmented Generation for Controllable Traffic Scenarios](https://mlsanity.com/paper/2312.13303/) (2023)
- [From LLMs to Actions: Latent Codes as Bridges in Hierarchical Robot Control](https://mlsanity.com/paper/2405.04798/) (2024)
- [Emergent Agentic Transformer from Chain of Hindsight Experience](https://mlsanity.com/paper/2305.16554/) (2023)
- [MUVO: A Multimodal Generative World Model for Autonomous Driving with Geometric Representations](https://mlsanity.com/paper/2311.11762/) (2023)
- [Inverse Dynamics Pretraining Learns Good Representations for Multitask Imitation](https://mlsanity.com/paper/2305.16985/) (2023)
- [Latent Diffusion Planning for Imitation Learning](https://mlsanity.com/paper/2504.16925/) (2025)
- [ACT: Empowering Decision Transformer with Dynamic Programming via Advantage Conditioning](https://mlsanity.com/paper/2309.05915/) (2023)
- [Learning from Imperfect Demonstrations with Self-Supervision for Robotic Manipulation](https://mlsanity.com/paper/2401.08957/) (2024)
- [Goal-Conditioned Predictive Coding for Offline Reinforcement Learning](https://mlsanity.com/paper/2307.03406/) (2023)
- [Emergence of Abstract State Representations in Embodied Sequence Modeling](https://mlsanity.com/paper/2311.02171/) (2023)
- [Position: Foundation Agents as the Paradigm Shift for Decision Making](https://mlsanity.com/paper/2405.17009/) (2024)
- [ACT-JEPA: Novel Joint-Embedding Predictive Architecture for Efficient Policy Representation Learning](https://mlsanity.com/paper/2501.14622/) (2025)
- [PASTA: Pretrained Action-State Transformer Agents](https://mlsanity.com/paper/2307.10936/) (2023)
- [Predictive Coding for Decision Transformer](https://mlsanity.com/paper/2410.03408/) (2024)
- [Masked Path Modeling for Vision-and-Language Navigation](https://mlsanity.com/paper/2305.14268/) (2023)
- [Investigating Pre-Training Objectives for Generalization in Vision-Based Reinforcement Learning](https://mlsanity.com/paper/2406.06037/) (2024)
- [Multimodal Diffusion Forcing for Forceful Manipulation](https://mlsanity.com/paper/2511.04812/) (2025)
- [Unveiling Markov Heads in Pretrained Language Models for Offline Reinforcement Learning](https://mlsanity.com/paper/2409.06985/) (2024)
- [Learning Versatile Skills with Curriculum Masking](https://mlsanity.com/paper/2410.17744/) (2024)
- [Task Tokens: A Flexible Approach to Adapting Behavior Foundation Models](https://mlsanity.com/paper/2503.22886/) (2025)
- [Self-supervised Pretraining for Decision Foundation Model: Formulation, Pipeline and Challenges](https://mlsanity.com/paper/2401.00031/) (2023)
- [Bridging MARL to SARL: An Order-Independent Multi-Agent Transformer via Latent Consensus](https://mlsanity.com/paper/2604.13472/) (2026)

## Abstract (arXiv)

We introduce Masked Trajectory Models (MTM) as a generic abstraction for sequential decision making. MTM takes a trajectory, such as a state-action sequence, and aims to reconstruct the trajectory conditioned on random subsets of the same trajectory. By training with a highly randomized masking pattern, MTM learns versatile networks that can take on different roles or capabilities, by simply choosing appropriate masks at inference time. For example, the same MTM network can be used as a forward dynamics model, inverse dynamics model, or even an offline RL agent. Through extensive experiments in several continuous control tasks, we show that the same MTM network -- i.e. same weights -- can match or outperform specialized networks trained for the aforementioned capabilities. Additionally, we find that state representations learned by MTM can significantly accelerate the learning speed of traditional RL algorithms. Finally, in offline RL benchmarks, we find that MTM is competitive with specialized offline RL algorithms, despite MTM being a generic self-supervised learning method without any explicit RL components. Code is available at https://github.com/facebookresearch/mtm
