# M3PC: Test-time Model Predictive Control for Pretrained Masked Trajectory Model

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

## Key facts

- No citing papers in the map yet (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 43 of its 48 references that are in the map.
- First posted on arXiv 7 Dec 2024.

## 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 43)

- [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)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [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)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [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)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reduction](https://mlsanity.com/paper/1906.00949/) (2019)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [Offline Reinforcement Learning as One Big Sequence Modeling Problem](https://mlsanity.com/paper/2106.02039/) (2021)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [Temporal Difference Learning for Model Predictive Control](https://mlsanity.com/paper/2203.04955/) (2022)
- [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)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [Cal-QL: Calibrated Offline RL Pre-Training for Efficient Online Fine-Tuning](https://mlsanity.com/paper/2303.05479/) (2023)
- [Offline-to-Online Reinforcement Learning via Balanced Replay and Pessimistic Q-Ensemble](https://mlsanity.com/paper/2107.00591/) (2021)
- [Online Decision Transformer](https://mlsanity.com/paper/2202.05607/) (2022)
- [Model-Based Offline Planning](https://mlsanity.com/paper/2008.05556/) (2020)
- [Plan Online, Learn Offline: Efficient Learning and Exploration via Model-Based Control](https://mlsanity.com/paper/1811.01848/) (2018)
- [Policy Expansion for Bridging Offline-to-Online Reinforcement Learning](https://mlsanity.com/paper/2302.00935/) (2023)

## Abstract (arXiv)

Recent work in Offline Reinforcement Learning (RL) has shown that a unified Transformer trained under a masked auto-encoding objective can effectively capture the relationships between different modalities (e.g., states, actions, rewards) within given trajectory datasets. However, this information has not been fully exploited during the inference phase, where the agent needs to generate an optimal policy instead of just reconstructing masked components from unmasked ones. Given that a pretrained trajectory model can act as both a Policy Model and a World Model with appropriate mask patterns, we propose using Model Predictive Control (MPC) at test time to leverage the model's own predictive capability to guide its action selection. Empirical results on D4RL and RoboMimic show that our inference-phase MPC significantly improves the decision-making performance of a pretrained trajectory model without any additional parameter training. Furthermore, our framework can be adapted to Offline to Online (O2O) RL and Goal Reaching RL, resulting in more substantial performance gains when an additional online interaction budget is provided, and better generalization capabilities when different task targets are specified. Code is available: https://github.com/wkh923/m3pc.