Masked Trajectory Models for Prediction, Representation, and Control

2023/05/04 by Philipp Wu, Wu, Philipp, Arjun Majumdar +11 · 7 citations
Computer Science · #Anomaly Detection Techniques and Applications #Artificial Intelligence (cs.AI) #Bayesian Modeling and Causal Inference #FOS: Computer and information sciences #Machine Learning (cs.LG) #Time Series Analysis and Forecasting

paper · pdf · doi:10.48550/arxiv.2305.02968

openalex publication_date 2023/05/04 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28

Abstract

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

Cited by

• Beyond Single-Episode Optimization: Sliding-Window Aware Generative Auto-Bidding for Long-Term Advertising Effectiveness
• Bridging MARL to SARL: An Order-Independent Multi-Agent Transformer via Latent Consensus
• Planning as Descent: Goal-Conditioned Latent Trajectory Synthesis in Learned Energy Landscapes
• MDG: Masked Denoising Generation for Multi-Agent Behavior Modeling in Traffic Environments
• Multimodal Diffusion Forcing for Forceful Manipulation
• GPS-MTM: Capturing Pattern of Normalcy in GPS-Trajectories with self-supervised learning
• Double Check My Desired Return: Transformer with Target Alignment for Offline Reinforcement Learning

Related

• M³PC: Test-time Model Predictive Control for Pretrained Masked Trajectory Model
• Traj-MAE: Masked Autoencoders for Trajectory Prediction
• RePreM: Representation Pre-training with Masked Model for Reinforcement Learning
• Masked Autoencoding for Scalable and Generalizable Decision Making
• Forecast-MAE: Self-supervised Pre-training for Motion Forecasting with Masked Autoencoders