M3PC: Test-time Model Predictive Control for Pretrained Masked Trajectory Model
On the map: Offline RL and imitation learning from data → Decision transformers and sequence modeling
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.
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.
Builds on
43 of 48 references.
Closest ancestors · 7
Same direction, newest first
- 2023Masked Trajectory Models for Prediction, Representation, and Control4 May
- 2022Masked Autoencoding for Scalable and Generalizable Decision Making23 Nov
- 2022Q-learning Decision Transformer: Leveraging Dynamic Programming for Conditional Sequence Modelling in Offline RL8 Sep
- 2022Bootstrapped Transformer for Offline Reinforcement Learning17 Jun
- 2022Online Decision Transformer11 Feb
- 2021Offline Reinforcement Learning as One Big Sequence Modeling Problem3 Jun
- 2021Decision Transformer: Reinforcement Learning via Sequence Modeling2 Jun
Other references · 36
- 2017Attention Is All You Needcited by 60,361
- 2018BERT: Pre-training of Deep Bidirectional Transformers for Language Understandingcited by 40,399
- 2020Language Models are Few-Shot Learnerscited by 32,419
- 2021Masked Autoencoders Are Scalable Vision Learnerscited by 7,403
- 2018Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actorcited by 4,885
- 2022RT-1: Robotics Transformer for Real-World Control at Scalecited by 2,061
- 2020Conservative Q-Learning for Offline Reinforcement Learningcited by 1,809
- 2019Dream to Control: Learning Behaviors by Latent Imaginationcited by 1,653
- 2021Offline Reinforcement Learning with Implicit Q-Learningcited by 1,293
- 2018Off-Policy Deep Reinforcement Learning without Explorationcited by 1,241
- 2021A Minimalist Approach to Offline Reinforcement Learningcited by 899
- 2019Stabilizing Off-Policy Q-Learning via Bootstrapping Error Reductioncited by 840
24 more
- 2022A Generalist Agentcited by 808
- 2018Recurrent World Models Facilitate Policy Evolutioncited by 785
- 2018Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Modelscited by 774
- 2019When to Trust Your Model: Model-Based Policy Optimizationcited by 773
- 2020AWAC: Accelerating Online Reinforcement Learning with Offline Datasetscited by 671
- 2020MOPO: Model-based Offline Policy Optimizationcited by 636
- 2022Temporal Difference Learning for Model Predictive Controlcited by 397
- 2022Video PreTraining (VPT): Learning to Act by Watching Unlabeled Online Videoscited by 345
- 2022Behavior Transformers: Cloning k modes with one stonecited by 304
- 2021Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemblecited by 302
- 2023Cal-QL: Calibrated Offline RL Pre-Training for Efficient Online Fine-Tuningcited by 237
- 2021Offline-to-Online Reinforcement Learning via Balanced Replay and Pessimistic Q-Ensemblecited by 214
- 2020Model-Based Offline Planningcited by 129
- 2018Plan Online, Learn Offline: Efficient Learning and Exploration via Model-Based Controlcited by 113
- 2023Policy Expansion for Bridging Offline-to-Online Reinforcement Learningcited by 104
- 2018Temporal Difference Models: Model-Free Deep RL for Model-Based Controlcited by 89
- 2022When does return-conditioned supervised learning work for offline reinforcement learning?cited by 76
- 2021Model-Based Offline Planning with Trajectory Pruningcited by 30
- 2022UniMASK: Unified Inference in Sequential Decision Problemscited by 24
- 2023GOPlan: Goal-conditioned Offline Reinforcement Learning by Planning with Learned Modelscited by 20
- 2023Chain-of-Thought Predictive Controlcited by 18
- 2023Plan To Predict: Learning an Uncertainty-Foreseeing Model for Model-Based Reinforcement Learningcited by 17
- 2023Goal-Conditioned Predictive Coding for Offline Reinforcement Learningcited by 10
- 2021The Value of Planning for Infinite-Horizon Model Predictive Controlcited by 8
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026Inference Time Policy Optimization for Offline RL with Differentiable World Models
- 2023Multi-Objective Decision Transformers for Offline Reinforcement Learning
- 2025Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models
- 2025A New Perspective on Transformers in Online Reinforcement Learning for Continuous Control
- 2025Offline RLAIF: Piloting VLM Feedback for RL via SFO
- 2024Trajectory-Oriented Policy Optimization with Sparse Rewards
- 2026Shortcut Trajectory Planning for Efficient Offline Reinforcement Learning
- 2023Model-Based Reinforcement Learning with Multi-Task Offline Pretraining