Masked Trajectory Models for Prediction, Representation, and Control

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

• The paper introduces a masked trajectory model that leverages self-supervised learning to reconstruct state-action sequences.
• It demonstrates versatility by adapting a single model for behavior cloning, forward and inverse dynamics, and return-conditioned control tasks.
• Experimental results show that MTM is competitive with leading offline RL algorithms on diverse continuous control benchmarks.

Overview of Masked Trajectory Models for Prediction, Representation, and Control

The paper "Masked Trajectory Models for Prediction, Representation, and Control" presents a framework called Masked Trajectory Models (MTM) aimed at advancing sequential decision making tasks. This work explores the intersection of self-supervised learning and reinforcement learning, focusing on trajectory modeling. MTM is designed to handle state-action sequences and is capable of reconstructing trajectories from masked versions of the same sequences. By employing random masking patterns during training, MTM learns multifaceted representations which can be used for a variety of inference tasks.

Key Components and Capabilities

The primary contribution of the paper is the introduction of a self-supervised learning paradigm using MTM, which leverages transformer architectures akin to those used in vision and NLP for sequence modeling tasks. The model exploits masked prediction as a training mechanism, forcing it to develop robust representations. This allows MTM to serve as multiple types of models such as forward dynamics, inverse dynamics, or even as an agent in offline reinforcement learning (RL) environments by modifying the masking patterns at inference time.

A salient aspect of MTM lies in its versatility. With the same learned weights, the model can perform several tasks, including:

• Behavior Cloning (BC): Learning to mimic expert behavior using state-action demonstrations.
• Return Conditioned Behavior Cloning (RCBC): Inferring actions that achieve specified returns, pertinent in offline RL.
• Inverse Dynamics (ID): Inferring actions necessary to transition between states, valuable for state-based imitation.
• Forward Dynamics (FD): Predicting future states given current states and actions, useful in model-based RL.

Experiments and Results

The authors evaluate MTM across a range of continuous control benchmarks, notably from the D4RL and Adroit suites, as well as DM-Control datasets. The experiment results underline the efficacy of MTM in offline RL where it shows competitiveness with or even outperforms specialized algorithms like CQL and IQL in certain environments, without integrating explicit RL components.

Furthermore, MTM's distinctive ability to operate on heteromodal datasets—datasets with incomplete or varying modalities—highlight its robustness and broader applicability. This is demonstrated by training MTM on datasets with mixed-modal data, such as state-only or state-action sequences, enhancing its performance on tasks using incomplete data.

Practical and Theoretical Implications

Practically, MTM poses significant implications for designing generalist models that simplify the learning pipelines traditionally needing separate components. Its demonstrated versatility across multiple tasks with a single network reduces the model complexity and training times when handling large-scale decision-making problems.

Theoretically, the work posits new directions in learning paradigms for RL and control tasks, where self-supervised learning objectives can sufficiently lead to high-quality representations and tasks' performance without explicit rewards-based optimization.

Future Directions

The versatility and data efficiency exhibited by MTM highlight its potential for further exploration. Future research could address scaling the model to tasks involving longer trajectory sequences, enhancing real-time inference capabilities, and exploring more complex data modalities, including those found in video streams. Moreover, integrating MTM with online learning frameworks could further refine its performance by enabling faster adaptation during active interactions with environments.

In conclusion, the paper provides significant insights into the deployment of masked prediction objectives in RL contexts and paves the way for creating robust, general-purpose frameworks adaptable to a wide array of decision-making scenarios.