M$^3$PC: Test-time Model Predictive Control for Pretrained Masked Trajectory Model

TL;DR

This work tackles leveraging a pretrained masked Bidirectional Trajectory Model for decision making in offline and online RL. It introduces M$^3$PC, a test-time Model Predictive Control framework that uses mask ensembles to perform forward and backward planning (via RCBC/FD/RP and PI/ID masks) and selects actions through a parallel, model-driven evaluation of future outcomes, all without further training. Key contributions include uncertainty-aware action reconstruction, an entropy-constrained objective, and TD($\lambda$)-style utilities for reward-maximizing planning; the approach yields strong improvements on D4RL and RoboMimic, enables efficient online finetuning, and demonstrates zero-shot goal-reaching capabilities. The results suggest practical impact for deploying pretrained multitask trajectory models as generalist agents with robust generalization to unseen tasks and goals.

Abstract

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.

Figures (16)

• Figure 1: Benefits of equipping pretrained bidirectional trajectory model with our test-time M$^3$PC. (a) Instead of generating actions solely based on history context, we leverage the full capacity of the masked pretrained model to predict future outcomes (e.g. states, rewards, returns) as a test-time self-enhanced decision making approach. Such a MPC framework can be used to achieve higher return at inference time or to reach a given goal state (in dashed square block) even unseen during offline training. (b) Forward M$^3$PC achieves better offline learning performances, using the same model without any finetuning, and gains better O2O improvement when online finetuning is allowed after offline pretrain. (c) Backward M$^3$PC unlocks zero shot goal reaching capability. Given a desired state, the walker agent can split its legs to a large degree without any prior experience.
• Figure 2: Model overview. The bidirectional trajectory model is pre-trained using MAE loss that aims to reconstruct the whole MDP trajectory taken a [Random] masked trajectory. After pretraining, the model show multiple capabilities by applying different test-time masks. E.g., Return-Conditioned Behaviour Clone [RCBC] Mask: Predict actions given states, expected return and context trajectory. Reward and Return Prediction [RP] Mask: Predict intermediate rewards and future return given states and actions. Forward Dynamics [FD] Mask: Predict future states given current state and future actions. Inverse Dynamics [ID] Mask: Infer actions needed taken to perform a given state path. As a pretrained masked transformer can always reconstruct the full trajectory, for those MDP-elements that are not related to the given task, e.g., the rewards during [RCBC], we omit and mark them as gray.
• Figure 3: Leverage the Masked Model itself for test-time Model Predictive Control. Our pipeline utilizes BTM's versatile inference capabilities to enhance decision making. (a) Forward M$^3$PC. We employ [RCBC], [FD] and [RP] masks to build an MPC pipeline for planning, prediction, and action resample. (b) Backward M$^3$PC. Given a goal state that we finally want to reach, we first use Path Inference [PI] mask to infer the waypoint-states, followed by a Inverse Dynamic [ID] mask to get the action sequence conditioned on those waypoints, and finally execute the first one.
• Figure 3: Offline Results on RoboMimic. Success rate of various offline pretrained agents in manipulation tasks. We report the mean of 5 seeds (50 trials for simulator and 20 trials for real world). We exclude the BC and IQL from real-world implementation due to their poor performance in the corresponding simulated tasks.
• Figure 4: Exploration Rollout Statistics. Results from two example runs of the Hopper task on the medium dataset using the same offline pretrained BTM agent. One run employs Gaussian noise for exploration, while the other utilizes M$^3$PC. The red line represents the offline result. Compared to naive Gaussian noise exploration, M$^3$PC significantly improves the agent's exploration quality by generating more high-return trajectories while maintaining stochasticity, including some mid-level or low-return trajectories.