Research map / Reinforcement learning / Offline RL

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

ICLR 2025 · Reinforcement learning › Offline RL · approach: Decision Transformer sequence models

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes M^3PC, a test-time model predictive control method that uses a single pretrained bidirectional masked trajectory model as both policy and world model via mask ensembles, improving offline RL, offline-to-online finetuning, and zero-shot goal reaching on D4RL and RoboMimic.

Problem

Masked trajectory models pretrained with random masking are typically used at inference only as policies (reconstructing actions), leaving their predictive capacity for decision-making unexploited; the paper targets better test-time decision-making, online finetuning efficiency, and goal reaching wi

Approach

Use different mask patterns of one pretrained bidirectional trajectory model (RCBC for action proposals, forward dynamics and reward/return prediction masks for rollout and evaluation, path-inference and inverse-dynamics masks for backward goal-reaching) to perform sample-based MPC at test time, with uncertainty-aware action distributions trained via an entropy-constrained NLL objective.

Key results

• Offline D4RL: M^3PC-Q total 429.8 vs BTM 372.8 and IQL 401.0 across six datasets (M^3PC-M 395.9).
• Online finetuning with 200K samples: M^3PC total 530.8 vs ODT 422.7 and IQL 405.5 (reported as 26% higher than ODT).
• O2O comparison at 100K budget: M^3PC average 71.7 to 76.8, comparable to OFF2ON (73.5) and PEX (80.2).

Measured effects reported in the paper

Surprising findings

• Single-mask goal-reaching (RCBC-style goal mask) fails to reach out-of-distribution goals, while backward M^3PC succeeds.
• Two separately pretrained specialized policy and world models do not improve over the unified pretrained BTM under MPC.

Limitations

• Limited to state-based observations; does not handle pixel observations.
• Unclear whether increasing Transformer capacity improves test-time planning.
• Backward goal-reaching depends on a limited planning horizon and requires hand-crafted subgoal trajectories.

Experimental setup

Datasets: D4RL hopper-medium D4RL hopper-medium-replay D4RL walker2d-medium D4RL walker2d-medium-replay D4RL halfcheetah-medium D4RL halfcheetah-medium-replay RoboMimic Can-Pair RoboMimic Square-MH RoboMimic Lift-MG RoboMimic Can-Lim (customized) RoboMimic Can-Real (real-world)

Benchmarks: D4RL locomotion RoboMimic

Metrics: normalized return total normalized score goal-reaching success

Baselines: BC TD3+BC IQL DT TT BTM ODT AWAC OFF2ON PEX TD-MPC

Builds on: Decision Transformer Trajectory Transformer ODT IQL TD-MPC

Closest papers by meaning

• Overcoming Slow Decision Frequencies in Continuous Control: Model-Based Sequence Reinforcement Learning for Model-Free Control (ICLR 2025 · similarity 0.86)
• Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models (NeurIPS 2025 · similarity 0.85)
• Bootstrapped Model Predictive Control (ICLR 2025 · similarity 0.85)
• Are Expressive Models Truly Necessary for Offline RL? (AAAI 2025 · similarity 0.85)
• TD-MPC2: Scalable, Robust World Models for Continuous Control (ICLR 2024 · similarity 0.85)
• Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning (ICML 2025 · similarity 0.85)