{"id":"011a503f-2a6f-47ca-b0a9-632fa412646f","arxiv_id":"2412.05675","paper_version":2,"verdict":"CONDITIONAL","confidence":"MODERATE","novelty_score":7.0,"correctness_risk":"medium","formal_verification":"none","parameter_count":6,"one_line_summary":"M3PC runs model predictive control at test time on a pretrained masked trajectory Transformer, improving offline RL returns and enabling goal reaching without extra model training.","lead":"A pretrained trajectory model can be used as its own planner at test time: the agent proposes several action options, predicts what will happen for each using the model's masked reconstruction abilities, and picks the one with the best predicted outcome. This improves offline decision-making without retraining the model and adds a way to reach specified goals.","discovery_kind":"new_method","skeptic_critique":{"model":"deepseek-v4-flash","headline":"The key comparison (M3PC-M vs. BTM) uses a BTM pretrained with a modified NLL+entropy objective and a unified random-mask pattern; the claimed 6.0% gain may conflate planning with the changed pretraining objective, so the central 'same model, test-time only' claim is not cleanly isolated.","rationale":"I agree with the reader that the paper is largely sound, the offline result is clean in its own terms, and the code release is a real positive. The reader's weakest assumption concerns the reliability of the model's predicted states/rewards/returns for ranking candidate actions; that is a legitimate generalization concern about planning quality, but the paper provides substantial indirect evidence (ablation of specialized vs unified pretraining, Figure 6; the online planning ablation, Figure 7; consistent gains across six datasets) that the predictions are useful enough for ranking. My concern lands on a different, more fundamental point: the experimental setup does not isolate the claimed mechanism. The 'no additional parameter training' claim is conditional on a pretrained model that was trained with a modified objective (NLL action head plus entropy constraint), and the baseline BTM is that same modified model. A reader cannot tell from the paper whether the 6.0% offline gain comes from test-time MPC or from the change in the pretraining loss (and the uncertainty-aware action distribution it induces). This is not an accusation of dishonesty; it is a missing control. The paper's own ablation in Figure 7 includes a comparison of planning versus no planning for the online setting, but it does not hold the training objective fixed in the way needed for the offline headline claim. The reader's stated weakest assumption is about generalization of the learned world model; my concern is about identifiability of the treatment effect. Both support CONDITIONAL rather than ACCEPT, and neither is severe enough to push toward REJECT. A single retraining experiment with the original MSE objective would settle whether the headline mechanism is the actual cause of the improvement.","tokens_in":17773,"tokens_out":1945,"duration_ms":17073,"concrete_test":"Retrain the BTM with the original deterministic MSE reconstruction loss (no NLL action head, no entropy constraint) using the same unified random-mask pretraining, then run the full M3PC-M action-selection pipeline at test time with this model (using its MSE-based action proposals, rolled out with [FD] and [RP] masks) on the six D4RL datasets. If the average gain over the MSE-BTM baseline approximately reproduces the 6.0% improvement, the test-time MPC mechanism is the source of the gain. If the gain largely disappears, the reported improvement is attributable to the modified pretraining objective rather than to inference-phase MPC alone.","verdict_should_be":"CONDITIONAL","load_bearing_attack":"The central claim is that test-time MPC alone, with no additional parameter training, turns a pretrained masked trajectory model into a better decision-maker. The cleanest support is Table 1: M3PC-M shares weights with BTM and outperforms it by 6.0% on average. But the BTM used throughout is not trained with a plain masked autoencoding MSE objective; Section 4 replaces the action reconstruction loss with an NLL loss over an uncertainty-aware action distribution and imposes a minimum-entropy constraint (Eq. 2), solved via a Lagrangian (Eqs. 5-8). The 'same pretrained model' is therefore a model pretrained with this modified objective, and the baseline BTM is also this modified model used with the [RCBC] mask alone. That baseline is not the standard BTM/UniMask agent, and the paper never reports the gain of M3PC-M over a BTM trained with the original MAE-style objective but evaluated with uncertainty-aware sampling and the entropy constraint removed. The ablation in Figure 7 compares M3PC against BTM+MSE+noise and shows that planning helps, but that ablation is in the online finetuning setting and uses different exploration mechanisms; it does not isolate the offline contribution of the changed pretraining loss. Thus the 6.0% offline gain could partly reflect the modified action head rather than the MPC selection loop. The claim 'without any additional parameter training' is true in the narrow sense of the released checkpoints, but the paper's phrasing implies the same pretrained trajectory model is enhanced purely at inference, and the experimental design does not separate the test-time mechanism from the training-time objective change.","agreement_with_reader":"partial"},"referee_report":{"model":"deepseek-v4-flash","summary":"This paper proposes M3PC, a test-time Model Predictive Control framework for pretrained bidirectional masked trajectory models (BTMs). The core idea is to use the same pretrained model with different mask patterns to (i) propose action sequences via an uncertainty-aware action head, (ii) predict future states, rewards, and returns via forward-dynamics and reward/return masks, and (iii) select actions by a TD(λ)-style utility over the predicted outcomes. The authors present three sets of results: offline RL on D4RL (M3PC-M improves over the same-weight BTM by about 6% average, and M3PC-Q with an IQL critic achieves higher scores), offline-to-online finetuning with a 200K step budget (outperforming ODT), and goal-reaching experiments on D4RL and RoboMimic using a backward M3PC variant with path inference and inverse dynamics masks. The paper also includes ablations on planning, uncertainty-aware action reconstruction, entropy constraint, and the decay parameter λ, along with an inference-time comparison against TT and TD-MPC.","tokens_in":18118,"tokens_out":9344,"duration_ms":80593,"significance":"The strongest contribution is the clean within-model comparison in Table 1: M3PC-M and the BTM baseline share identical pretrained weights, so the reported 6% average improvement is directly attributable to the test-time MPC selection loop rather than to additional training. This supports the paper's central claim that a single masked trajectory Transformer can serve as both policy and world model for test-time planning. The O2O results show substantial gains over ODT, and the goal-reaching demonstrations suggest a genuinely new capability for backward planning with masked models. The paper ships code and reports ablations for key components, which aids reproducibility. The main weaknesses are that the uncertainty-aware pretraining objective is introduced as part of the method but not ablated in the offline setting, and the 'without any additional parameter training' statement does not apply to the M3PC-Q variant used in the O2O experiments.","major_comments":[{"comment":"The central comparison M3PC-M vs. BTM is internally clean because both share the same pretrained weights, but those weights come from a model trained with the modified uncertainty-aware action reconstruction objective (NLL loss plus entropy constraint, Eqs. (1)-(2)) rather than the standard masked-autoencoding MSE objective used by prior BTMs (e.g., Wu et al. 2023). The paper does not report an offline ablation that isolates the effect of the modified pretraining from the effect of test-time MPC; the only planning-vs-no-planning ablation (Figure 7) is in the online finetuning setting and uses different exploration mechanisms. As a result, the abstract's phrasing 'inference-phase MPC significantly improves the decision-making performance of a pretrained trajectory model without any additional parameter training' overstates the evidence: the 6.0% gain is demonstrated only for the authors' uncertainty-aware pretrained model, not for a standard masked trajectory model. Please either add an offline ablation (e.g., BTM trained with MSE + RCBC vs. the same model with MPC) or explicitly scope the claim to the uncertainty-aware BTM used throughout the paper.","section":"§4, Eq. (1)-(2); Table 1"},{"comment":"The offline-to-online results in Table 2 use the M3PC-Q variant, which requires training a separate IQL critic (Eq. (9)) and thus involves additional parameter training. The paper's general claim of 'without any additional parameter training' applies only to M3PC-M in the offline setting; the O2O experiments are a different regime where an extra value estimator is trained. This distinction should be stated clearly in the abstract and in the O2O section, and the table header should indicate that the reported M3PC numbers are M3PC-Q.","section":"§5, Table 2; §4, Eq. (9)"}],"minor_comments":[{"comment":"Standard deviations are reported only for M3PC-M and M3PC-Q; the baselines (BC, TD3+BC, IQL, DT, TT, BTM) are listed without variance, which makes it difficult to assess whether the gains over BTM on datasets such as halfcheetah-m-r (41.8 vs. 41.1) are statistically meaningful.","section":"Table 1"},{"comment":"The caption should explicitly state that the planning ablation is performed in the online finetuning phase, not in the offline setting, so that readers do not infer that the offline contribution of planning versus the modified pretraining objective has been isolated.","section":"Figure 7"},{"comment":"The statement that 'M3PC requires only two prediction steps for planning at each timestep' should mention that each step processes N candidate sequences in parallel; otherwise the reader may underestimate the computational cost when N=625 candidates are sampled.","section":"§4, Algorithm 1"},{"comment":"The table header uses 'M3PC (Ours)' without specifying whether the offline and online numbers come from M3PC-M or M3PC-Q; given that the offline values match M3PC-Q in Table 1, the variant should be identified explicitly.","section":"Table 2"},{"comment":"The goal-reaching experiments rely on hand-crafted subgoal trajectories that are manually designed for each task; this dependence should be acknowledged as a limitation, since the approach may require nontrivial engineering to apply to new tasks.","section":"§5, Goal Reaching; Appendix A"}],"recommendation":"major_revision","confidential_remarks":"The paper is already published at ICLR 2025, so the requested revisions are for a journal version or an extended rebuttal. The main concern is not the validity of the within-model comparison but the scope of the headline claim: the 'without any additional parameter training' statement is true for M3PC-M but not for M3PC-Q, and the effect of the modified pretraining objective is not separately ablated offline. These issues are addressable with a targeted ablation or careful rewording, and the strong empirical setup and code release make the paper worth publishing after revision."},"author_rebuttal":null,"desk_editor":{"model":"deepseek-v4-flash","letter":"The core idea here is genuinely useful: take a masked bidirectional trajectory Transformer and, instead of using one mask pattern for one task, combine several at test time to run a lightweight MPC loop. The same network proposes actions (RCBC), predicts forward states (FD), and predicts rewards/returns (RP); a softmax over sampled candidate sequences then picks the action. That is new relative to MTM/UniMask, which assign a single mask per task. The paper also adds an uncertainty-aware action reconstruction (NLL plus entropy constraint) and a backward variant (PI plus ID) for goal reaching. Code is out, and the evaluation covers D4RL, RoboMimic, and a real robot. Credit where due: the central offline comparison, M3PC-M vs. BTM, shares weights and improves consistently across six D4RL datasets. The O2O gains are substantial (26% over ODT total, though with variance). The goal-reaching demos, while qualitative, show the single-mask baseline fails and backward M3PC tracks the requested kinematics.\n\nThe soft spots are real but not fatal. The stress-test concern holds: the BTM used everywhere is not a standard MAE-trained BTM. Section 4 replaces the action reconstruction with an NLL loss over a stochastic action distribution and adds a minimum-entropy constraint solved by a Lagrangian. So the 'same pretrained model' is one with a modified action head. The 6.0% gain of M3PC-M over BTM isolates the effect of planning conditional on that modified pretraining, but the paper never reports the gain of M3PC over a BTM trained with the original MSE objective and evaluated deterministically. The only ablation touching this, Figure 7, is in the online finetuning setting and confounds exploration strategy with planning. Thus the abstract's 'without any additional parameter training' is true for the released checkpoints, but the implication that test-time MPC alone upgrades a vanilla masked trajectory model is not established. Second, the O2O headline uses M3PC-Q, which relies on a separately trained IQL critic; that is additional training, even if not of the Transformer. The paper is transparent about the variant in the method, but both main tables feature M3PC-Q prominently. Third, goal reaching depends on hand-crafted subgoal trajectories and is evaluated qualitatively; the out-of-distribution visualization is nice, but there is no quantitative success metric.\n\nOn balance, the core idea is sound and the experiments are honest enough. The strongest claim is overbroad, but the contribution stands: a test-time planning scheme that clearly helps a pretrained masked trajectory model, provided the model is trained the way they train it. I would send this to a serious referee; it deserves careful review, and the referee should push on the pretraining-objective confound and the M3PC-Q variant before accepting the headline.","headline":"Useful test-time MPC for masked trajectory models, but the 'no additional training' headline oversells the modified pretraining objective and the M3PC-Q variant.","tokens_in":18660,"tokens_out":3243,"would_cite":true,"duration_ms":30967,"reading_group":"yes","serious_thinker":"yes","would_accept_peer_review":true},"rs_alignment":null,"lean_confirmation":null,"pith_extraction":{"msc":[],"pacs":[],"model":"deepseek-v4-flash","headline":"A single masked trajectory Transformer can be turned into a test-time planner by having it score its own sampled action sequences with its own predicted futures, improving offline RL returns by 6.0% on average and enabling zero-shot goal…","keywords":["offline reinforcement learning","model predictive control","masked trajectory model","bidirectional transformer","test-time planning","offline-to-online RL","goal reaching","uncertainty-aware action reconstruction"],"falsifier":"Collect a suite of candidate action sequences from M3PC's uncertainty-aware policy head on a D4RL task, record the model's predicted TD(λ) utility for each, then execute each sequence in the environment and compare predicted utility to realized return; if the ranking correlation is near zero or negative on out-of-distribution candidates, the planning stage cannot be the source of the reported gains. A simpler check: run M3PC on a dataset deliberately corrupted by adding noise to rewards, and see whether the gains over the direct policy vanish.","tokens_in":17546,"feed_emoji":"🎯","tokens_out":4905,"duration_ms":41356,"temperature":0.7,"pith_summary":"This paper tries to establish that a single bidirectional Transformer trained once with masked autoencoding on offline trajectories can be turned into a better decision-maker at test time by using its own predictive heads inside a model predictive control loop, with no additional parameter training. The key move is to stop treating the pretrained model as a fixed policy and instead let it propose candidate action sequences, roll them forward into predicted states and rewards, score them, and pick the best. The authors show this improves the same model's offline RL returns by 6.0% on average across D4RL tasks, makes its offline-to-online finetuning more sample-efficient and 26% higher in total score than ODT, and unlocks zero-shot goal reaching including out-of-distribution states. A sympathetic reader should care because it suggests the versatility already latent in masked trajectory models can be converted into planning ability at inference time, without the cost of training a separate world model.","feed_headline":"One masked trajectory model out-plans its own policy at test time","feed_subtitle":"By scoring its own sampled action sequences with its own predicted futures, M3PC gains 6% offline and 26% vs ODT online.","key_machinery":"The machinery is the mask ensemble of a bidirectional trajectory model plus a TD(λ) utility scoring rule. A BTM is a Transformer whose tokens are states, actions, rewards, and return-to-go, pretrained to reconstruct randomly masked trajectories; at test time M3PC composes task-specific masks: the RCBC mask (actions given states and returns) generates uncertainty-aware candidate actions, the FD mask (future states given actions) rolls candidates forward, the RP mask (rewards and returns) scores them, and the PI/ID masks in backward M3PC infer waypoint states and the actions to reach them. The TD(λ) utility $U = (1-\\lambda)\\sum_{n=0}^{T-t-1}\\lambda^n G_{t:t+n} + \\lambda^{T-t}G_{t:T}$ blends short-term predicted rewards with long-term return estimates to rank candidates, and a softmax over utilities selects the action; parallel bidirectional prediction keeps the planning cost nearly constant in the horizon.","core_discovery":"The central claim is that a pretrained bidirectional trajectory model (BTM) contains, in one set of weights, the components of a model-based planner: a policy head that reconstructs actions with uncertainty, a forward dynamics head that predicts future states, and heads that predict rewards and returns. By composing different mask patterns at inference time, M3PC samples action sequences from the policy head, rolls them out with the dynamics head, assigns each a TD(λ)-style utility that blends short-term rewards with long-term return estimates, and executes the best action. The paper reports that this inference-phase MPC consistently beats the same model used as a direct policy, outperforming all offline RL baselines it compares against when guided by an IQL-trained Q-value estimator, and that the same procedure serves as an efficient online finetuning strategy and a backward goal-reaching mechanism via path inference and inverse dynamics masks.","pith_inferences":["If the model's predicted futures remain accurate for action sequences that stray from the offline data, M3PC's gains should grow with model scale and planning horizon; the paper leaves this open and notes its experiments used only a fixed small Transformer.","The same mask-composition recipe could be applied to other bidirectional sequence models—for example, extending M3PC to pixel observations by masking latent visual tokens rather than state vectors, a direction the paper flags as future work.","The TD(λ) decay parameter trades dynamics error against value-estimation error, so task-specific tuning of λ (which the paper finds mildly sensitive) could yield further gains on stochastic or long-horizon environments.","A direct test: compare M3PC's utility ranking against actual environment returns for sampled candidate actions; if the ranking correlation is low, the planning gain is coming from action resampling rather than from accurate prediction."],"forward_implications":["The same pretrained BTM, without finetuning, improves its offline RL normalized return by 6.0% on average across the six D4RL tasks, and M3PC-Q (guided by an IQL value estimator) beats all compared offline RL baselines.","With a 200K online interaction budget, M3PC's total score is 26% higher than ODT and 31% higher than IQL, and its improvement over finetuning is 123% larger than ODT's.","Backward M3PC reaches goal states unseen in the offline data—split, flip, and wiggle motions—where a single goal-reaching mask fails.","In manipulation, M3PC attains higher success rates than behavior cloning and IQL on RoboMimic tasks including a real-world can-picking task (0.70 vs 0.50 success rate for DT).","Because candidate actions are sampled from an uncertainty-aware distribution with an entropy lower bound, the planner also serves as an exploration strategy during online finetuning, collecting higher-return trajectories while maintaining diversity."],"supporting_citations":[{"why":"Establishes masked trajectory models as unified predictors and supplies the unified pretraining approach M3PC builds on.","marker":"(Wu et al., 2023)"},{"why":"Shows a unified bidirectional model can be applied zero-shot to downstream decision tasks with task-specific masks, the basis for the RCBC mask.","marker":"(Carroll et al., 2022)"},{"why":"Online Decision Transformer; provides the entropy-constrained action reconstruction idea and the O2O baseline M3PC compares against.","marker":"(Zheng et al., 2022)"},{"why":"Decision Transformer; defines the return-conditioned sequence modeling paradigm and is a baseline.","marker":"(Chen et al., 2021)"},{"why":"Trajectory Transformer; the prior sequence-modeling planner using beam search that M3PC contrasts with.","marker":"(Janner et al., 2021)"},{"why":"IQL; supplies the transition-wise value estimator used by M3PC-Q guidance.","marker":"(Kostrikov et al., 2021)"}],"fun_headline_variants":["Test-time MPC turns a masked trajectory model into its own planner","No retraining needed: M3PC plans better with the same model","Masked trajectory model plans at test time via MPC, no tuning","M3PC: sample, predict, score — the best action wins at test time","Offline RL boost: test-time MPC on a pretrained trajectory model"],"cache_read_input_tokens":3200,"weakest_assumption_plain":"The planner ranks sampled action sequences by the model's own predicted states, rewards, and returns, so the whole method stands or falls on whether those predictions stay accurate for action sequences that may differ from anything in the offline dataset.","fun_headline_variants_meta":{"raw":{"variants":["Test-time MPC turns a masked trajectory model into its own planner","No retraining needed: M3PC plans better with the same model","Masked trajectory model plans at test time via MPC, no tuning","M3PC: sample, predict, score — the best action wins at test time","Offline RL boost: test-time MPC on a pretrained trajectory model"]},"model":"deepseek-v4-flash","effort":"low","cost_usd":0.00022,"raw_usage":{"total_tokens":1445,"prompt_tokens":941,"completion_tokens":504,"prompt_tokens_details":{"cached_tokens":384},"prompt_cache_hit_tokens":384,"prompt_cache_miss_tokens":557,"completion_tokens_details":{"reasoning_tokens":408}},"tokens_in":557,"tokens_out":504,"duration_ms":4714,"temperature":1.0,"reasoning_tokens":408,"cache_read_input_tokens":384,"cache_creation_input_tokens":0},"cache_creation_input_tokens":0},"created_at":"2026-08-11T20:28:31.990357+00:00","model_set":{"reader":"deepseek-v4-flash"},"falsifier":"Collect a suite of candidate action sequences from M3PC's uncertainty-aware policy head on a D4RL task, record the model's predicted TD(λ) utility for each, then execute each sequence in the environment and compare predicted utility to realized return; if the ranking correlation is near zero or negative on out-of-distribution candidates, the planning stage cannot be the source of the reported gains. A simpler check: run M3PC on a dataset deliberately corrupted by adding noise to rewards, and see whether the gains over the direct policy vanish.","supporting_citations":[{"cited_title":"All inference times were benchmarked on a single NVIDIA RTX 3090 GPU","cited_arxiv_id":null,"evidence_quote":"Trajectory Transformer; the prior sequence-modeling planner using beam search that M3PC contrasts with."}],"review_version":1}