M^3PC: Test-time Model Predictive Control using Pretrained Masked Trajectory Model
Authors: Kehan Wen
Yutong Hu
Yao Mu
Lei Ke
Venue: ICLR 2025 Poster (2025)
Keywords: Offline-to-Online Reinforcement Learning
Model-based Reinforcement Learning
Masked Autoencoding
Robot Learning
Primary area: reinforcement learning
TL;DR: Enhance Transformer for RL by employing the Model itself for test-time MPC, achieving better performance in offline RL and offline-to-online RL for both simulated and real-world robotic tasks, with additional goal-reaching capabilities.

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. 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 capacity 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 given, and better generalization capabilities when different task targets are specified. Code is available: \href{https://github.com/wkh923/m3pc}{\texttt{https://github.com/wkh923/m3pc}}.

Published as a conference paper at ICLR 2025
M TEST-TIME MODEL PREDICTIVE CONTROL FOR PRETRAINED MASKED TRAJECTORY MODEL
Kehan Wen'1, Yutong Hu Yao Mu Lei Ke
'TTH Zurich, "KU Leuven, "Hong Kong University, "Carnegie Mellon University
ABSTRACT
Recent work in Offline Reinforcement Learning (RL) has shown that a unified Transformer trained under masked auto-encoding objective effectively cap- ture 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 op- timal 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 Pre- dictive Control (MPC) at test time to leverage the model's s own predictive capa- bility to guide its action selection. Empirical results on D4RL, and RoboMimic show that our inference-phase MPC significantly improves the Docision-making performance of a pretrained trajectory model without any additional parame- ter training. Furthermore, (020) RL and Goal Reaching RL, resulting in more substantial performance gains when additional online interaction budget is provided, and better gen- eralization capabilities when different task targets are specified. Code is available: https://gii weiggiithuuee thub. com 9w3 923333p22
framework can be adapted to Offline Online
INTRODUCTION
The Masked Modeling paradigm has simple, self-supervised training objective: predicting ran- domly masked subset of the original sequence. It has become a powerful technique for generation representation learning for sequential data, e.g., language tokens (Devlin et al. 2018) image patches (He et al., 2022). Unlike autoregressive models like GPT (Brown et al. 2020). which con- dition only on the past context in the "left", bidirectional models trained with this objective learn to model the context from both sides, leading to richer representations and deeper understandings of the data's 'ss underlying dependencies.
Given that a sequential Decision-making trajectory inherently involves sequence of states s and actions a, and other optional augmented properties like return-to-ge (RTG) (Chen al. 2021 or approximate state-action value v (Yamagata et al. 2023) across T timesteps, the mask modeling paradigm can be adapted easily for sequential decision-make makiin tasks. For example, in the case of Reinforcement Learning, the policy output P(als) at each time step can be regarded as predicting a masked action a conditioned on given states s. Moreover, recent works (Carroll et al.] 20022] Liu et al. [2022] Wu et al. (202)) have demonstrated that a unified bidirectional trajectory model (BTM) pretrained with a highly random masking pattern can be applied zero-shot in various downstream tasks. Different reconstruction tasks can be deliberately created by applying appropriate masks to different modalities- whether states, actions, or rewards- during inference time.
However, the inherent flexibility and versatility of models trained with random masking techniques have not been fully exploited in deployment settings. Previous research has highlighted the mul- titasking capabilities of Bidirectional Trajectory Models (BTMs) by assigning one single specific mask pattern to individual tasks, such as the RCBC mask commonly used in offline RI. after the pretraining phase. Our findings, in contrast, suggest that integrating multiple capabilities such as short-term reward and long-term return prediction, along with forward dynamics, could significantly
Equal contribution. Corresponding authors. This work was done ETH Zurich.

---

Published as a conference paper at ICLR 2025
past
Teture.
+23%
Behavior Cloning
offline dataset:
zero-shof:
+12%
walk
split
Offine RL
Online Finetuning
improvements
Existing DT 81M
MPPC (Ours)
DTTIOONTI BTM (SPCCOOUS)
(a)
(b)
(c)
Figure 1: Benefits of equipping pretrained bidirectional trajectory model with our test-time PCP. (a) Instead of generating actions solely based on history context, we leverage the full ca- pacity of the masked pretrained model to predict future outcomes (e.g. states, rewards, returns) as test-time self-enhanced decision making approach. Such a MPC framework can be used to achieve higher return inference time to reach given goal state (in dashed square block) even unseen during offline training.
(b) Forward PC achieves better offline learning performances, using the same model without any finetuning, and gains better 020 improvement when online finetuning is allowed after offline pretrain. (c) Backward PC unlocks zero shot goal reaching capability. Given desired state, the walker agent split its legs to large degree without any prior experience.
enhance cetiiiion-mainng These capabilities allow the agent to explicitly evaluate action candidates and determine the optimal one, rather than merely relying on implicit mappings from expected re- turns to policies.
Building on these insights, we introduce the MC framework: Enhancing Decision-Making by using the Masked Model itself for test-time Model Predictive Control. Our framework decomposes decision-making tasks into series of simpler steps in a typical sample-based MPC style: sam- pling potential actions, inferring possible future states, evaluating these actions based on predicted outcomes, and selecting the final optimal action. We then demonstrate how pretrained model, equipped with our adaptation and ensemble of masks. can efficiently and effectively handle these subtasks. Our empirical results demonstrate that, by using PC for final decision decision-making, making, the same pretrained model can get substantial decision quality improvement in offline RL and goal-reaching RI., outperforming traditional single-mask models. Furthermore, MCPC supports sample-efficient online finetuning - a capability rarely seen in previous sequential modeling agents. By fully lever- aging the potential of pretrained BTM, MC evolves the model from multitasking framework into an inference phase self-enhancing, and a finetuning phase self-improving generalist agent. We summarize results in Figure and highlight contributions as:
We present PC, a novel framework that utilizes mask ensembles address complex decision- making tasks, effectively leveraging the multitasking abilities of a pretrained bidirectional tra- jectory model (BTM).
We demonstrate that MC not only improves the test-time performance of the same pretrained BTM in offline RL by 6.0%, but also enables efficient finetuning through online interactions with environments, outperforming specialized anfline-to-online (020) RL algorithms, such as ODT, by 26.0%.
We show that MC can be adapted for goal-reaching tasks, effectively guiding agents spec- ified goal states-even when these states are out-of-distribution relative to the datasets used for pretraining.
RELATED WORK
Transformers for Sequential Decision Making, The Transformer (Vaswani et 2017) architec- ture has been extensively applied in sequential decision decision-making tasks such as reinforcement learning (RL) CChnn et al. 2021 Janner et al. 2021 Wang et al. (022) and imitation learning (IL.) (Reed) et al. 2022] Shahullah et al. 2022: Brohan et al. 2022 Baker et al. 2022, Jia et al.) (023). Repre- sentative work such as Decision Transformer (DT) (Chen et al. 2021 and its variants (Zheng et al.)
2022 Yamagata et al. 2023) learn return-con return-conditioned policy using causal-masked Transformer.
Recent studies (Carroll et al. 2022 Liu et al. 2022 Wu et al. 2023) utilize a bidirectional Trans- former to model trajectories, highlighting the model's 's versatility enhanced by the mask prediction training objective. These studies focus on the potential of trajectory Transformers to unify various decision-making tasks, typically employing a unique mask pattern tailored to each specifie down-
2

---

Published as a conference paper at ICLR 2025
stream task. Building upon these insights, our work proposes hamnessing the functional versatility of pretrained models to enhance pocision-making More specifically, we investigate whether utilizing two more mask patterns can lead to improved decision decision-making making within single downstream task.
Offline RL with Online Finetuning. Traditional off-poliiy RL algorithms often suffer from boot- strapping error accumulation (Fujimoto al.) 2019] Nair al., 2020). To mitigate these issues, most offline RL algorithms employ regularization techniques to mitigate errors caused by out-of- distribution actions (Fujimoto al. 2019 Nair et al. 2020 Kumar et al., 2020;; Kostrikov et al.) 2021] An et al.) 2021] Kumar et al. 2019). However, finetuning an offline RL algorithm can be chal- lenging due to its inherent conservatism and the offline-to-online data distribution shift (Nair al.) 2220 Yu & Zhang 2023). Many techniques such as value calibration (Nakamoto et al.) 2024). bal- anced replay (Lee et al. 2022) and policy expansion (Zhang al.] (023) have been investigated to improve the online sample efficiency. In parallel, some work (Chen et al. 2021 Zheng et al. 2022) following supervised learning (SL.) paradigm can naturally ensure -idisttibbtionn learning but also suffer from poor online sample efficiency (Brandfonbrener et al. 2022). Our approach adheres to the SL paradigm while incorporating DP-based module to improve online sample efficiency.
Model-based RL. Learning dynamics model of the environment can be used for policy learn- ing (Pong et al., 2018 Ha & Schmidhuber 2018; Hafner et al., 2019) or planning (Silver et al.)
2008 Walsh et al. 2010: Zhang et al., [2019] Yu et a 20200. Recent work has explored the feasi-
bility of MPC in online RL (Chua et al. [0088 Janner et al. 2019 Wu et al.) [022] Lowrey et al.
2018 Hatch & Boots [021 Hansen al. (2022) Similar planning methods have also been tailored
for offline RL using techniques such as behavior cloning regularization (Argenson & Dulac-Amold
2023). Instead of maintaining sep- arate world and policy models, Trajectory Transformer (TT) (Janner et al., 202I) frames RL as a sequential modeling problem and performs beam search planning based on return heuristics. Our work follows a similar paradigm but leverages a bidirectional Transformer and mask autoencoding to enable a more flexible and computationally efficient planning process.
2020 and trajectory pruning (Zhan al., 2021 Wang et al.
PRELIMINARY
We consider the environment as a Markov Decision Process (MDP), formally defined by the tuple M = (S, A. P, R. Po). In this notation, S represents the state space, and A represents the action space. The transition probability distribution, P(s++1| 8t, at), defines the likelihood of transitioning from state s; to state 8:41 given action a, The reward function, R(st, at), assigns a reward for each action taken a particular state. The discount factor, denoted by 7, quantifies the preference for immediate rewards over future rewards. The maximum episode length, also referred to as the horizon of the MDP, is denoted as H.
episoee
Additional notations are introduced to adapt RL to sequential modeling. We denote the train- ing data distribution as T, which may be dynamic when the agent interacts with the environ- ment. A trajectory т, consisting of T states, actions, RTGs and rewards is represented as ($1,9) 91, T1, ... ST, gr, aT, rT). Note that some other properties can also be directly or indirectly accessed from the training data such next-states ( s/y), estimated values (UI, uT) for state-action pairs, but we do not model these modalities on the Transformer.
4 METHOD
This section details how we leverage a bidirectional trajectory model's versatile prediction capabili- ties within the PC framework enhance an agent's decision-making. In a typical MPC process, the system repeatedly solves an optimization problem to identify the best sequence of actions over a finite horizon by evaluating the outcomes of these actions and then executing the action at the current timestep. The following subsections describe our approach to adapting a BTM to perform these MPC steps: First, we enable the BTM to reconstruct actions with uncertainty, allowing us to sample from distribution of action proposals. Next, we demonstrate how to use different mask- ing patterns for forward or backward prediction for MDP sequence elements. These predictions serve references for evaluating the expected outcomes of action proposals, which we use to de- termine the optimal action to execute. As illustrated in Figure by breaking down pocision-making into these structured steps and using the BTM for versatile predictions, our MCPC framework en-
3

---

Published as a conference paper at ICLR 2025
trajectory output
JRCBCJ
[RP]
untokenize
Bidirectional Transformer
trujectery imput
modality embeding:
Emneedin
[FD]
[ID]
imestep embeding
Miskel Token
tokenizer Natten
training-time mask pattern
model IO structure
test-time capabilities
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 pre- training. the model show multiple capabilitiess 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 per- form a given state path. As pretrained masked transformer can always reconstruct the full trajec- tory, for those MDP-elements that are related to the given task, e.g., the rewards during [RCBC], we omit and mark them as gray.
hances the agent's s ability beyond simply imitating behaviors observed in offline data, e.g.. achieving higher reward incomes or diverse goals which typically fall in offline RL and goal reaching domains, respectively.
Bidirectional Trajectory Model. We illustrate the model architecture and how it process a masked MDP trajectory as Figure 2 To perform masked trajectory modeling, we first flatten and tokenize the different elements of the raw trajectory sequence. This tokenization involves three components: modality encoder that lifts elements from the raw modality space to a common represen- tation space, along with the addition of timestep embeddings and modality-type embeddings. These components collectively enable the Transformer to distinguish between different sequence elements.
We employ sicoder-decoder architecture with both the encoder and decoder being bidirectional Transformers. The tokenized and flattened trajectory fed into the Transformer encoder, where only unmasked tokens are processed. The decoder then processes the full trajectory sequence, leveraging values from the encoder when available or substituting a mask token when not. The decoder is trained to reconstruct the original sequence, including the unmasked tokens.
Training-phar Mask Pattern. Inspired by previous work (Wu et al. 2023 Zeng et al. 2024 we employ two-step masking pattern for training. Firstly, randomly mask proportion of elements in the trajectory т. Secondly, we mask all elements to the right randomly chosen position. By learning predict the mask elements, the model is trained to handle temporal dependencies and infer outcomes based solely on past events.
Uncertainty-Awar Action Reconstruction. To equip the agent with robust decissonnmakinng capa- bilities beyond mere imitation, our method employs uncertainty-awar action reconstruction rather than predicting the masked action deterministically. The primary focus of MAE lies perfectly reconstructing each token of the sequence, typically by optimizing a Mean Squared Error (MSE) loss. This inherently leads deterministic action reconstruction, which limits the agent's ability to account for uncertainties associated with the actions.
To address this limitation, we propose reconstructing an uncertainty-awant action distribution A by minimizing a Negative Log Likelihood (NLL) loss J(B) denoted by
((0) = - Σ - log Pa (a(oa\Masked(())
BraT
(I)
Inspired by ODT (Zheng et al. 2022), additionally impose lower bound on trajectory-level
action entropy H. to encourage the agent's online exploratory behavior. The overall constraint problem is formally defined
4

---

Published as a conference paper at ICLR 2025
future a with uncertuinty
combine short long-term reward
(a)
predict 8.r4 given
(b)
ID(() utility
sample's 4
Co A cooo
predict
future s,7, given
execusie
Figure 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 PºP. We employ [RCBC]. [FD] and [RP] masks to build an MPC pipeline for planning, predic- tion, and action resample. (b) Backward MCC. Given goal state that we finally want to reach, we first use Path Inference [PI] mask to infer the waypoint- followed by a Inverse Dynamic [ID) mask to get the action sequence conditioned those waypoints, and finally execute the first one.
1 m(0) subject to HJ 2 3, H. 1 BrTT
Σ H [Pa((a)aasked((]]]
where H [-] denotes the Shannon entropy of the distribution, 8 represents the predefined target en- tropy. avoid explicitly handling the inequality constraint, we solve the Lagrangian dual problem of Equation 2 Implementation details are provided in Appendix
Forward MCPC for Reward Maximization. A bidirectional trajectory model agent has demon- strated zero-shot ability in offline RL tasks when equipped with [RCBC] mask shown in previous work (Carroll et al., 2022, Wu et al. (023). By predicting actions conditioned on states and RTGs, the agent generates actions by imitating trajectories with similar RTGs in offline data. The performance of this imitative behavior is inherently upper- goueded by the best trajectory in the offline data.
To address this limitation, we propose refining the decision decision-making process by implementing an explicit wavarddmaximization procedure using the forward dynamics function and the return and ward prediction functions provided by the unified trajectory model. Typically, this process divides decision-making into three substeps: generating action proposals, rolling the future, and select- ing action proposals based on their potential utilities. Suppose we have access to both intermediate and long-term reward estimation for candidate action sequence arTr represented by Tt.T and ge.,, respectively. We define the TD(A)-style utility U for this candidate action as follows
T-1-1
n-
U U=
Σ
X"Grato + XT XT-GGTT where Gettn = Σ Fretk + " gttns
k=0
where the decay parameter À determines the weights of longer horizon estimates that contribute to the final result which can help trade off the errors from dynamies predictions and value estimates. We construct categorical distribution P using softmax for proposal selection:
P cxp(EU") E, exp(EU0)) Vi E [1, [I,...,N],
(4)
where & denote the softmax temperature. Notably, MC requires only two prediction steps for planning at each timestep. Leveraging the bidirectional nature of Transformers and the masked au- toencoding paradigm, MC can predict all future actions given current states and all future states given future actions parallel. This parallel prediction capability mitigates the computational cost's linear growth with respect to the planning horizon which is commonly observed in planning algo- rithms such as beam search in TT (Janner et al. 2021) or CEM in TD-MPC (Hansen et al.) 2022). We detail the decision-making process for reward-man wwwedddmaxiiiziiiio in Algorithm Since RTG value a trajectory-wise Monte Carlo estimation, becomes uninformative when datasets' behavior poli- cies are diverse. We optionally extend PC by replacing RTG guidance with transition -isse value for a better heuristic. this case, we calculate this value with a standalone value estimator updated in dynamic programming way proposed in IQI. (Kostrikov et al. 2021).
5

---

Published as a conference paper at ICLR 2025
Using the "Utility' metric to estimate future actions before they taken, forward MC can also adapt to an exploration strategy in the subsequent online finetuning phase, where equation 3 are used again. During the offline-to-online process, instead of executing the expectation in categorical distribution equation 4 the MCPC agent samples actions from the candidate set according to the possibilities proportional to their utility. This introduces stochasticity, maintaining overall superior actions while ensuring diversity in the experience collected during exploration, thereby balancing the exploration and exploitation.
Algorithm Forward M PC for Reward Maximization
1: Input: Current state St, past trajectory Tet, discount factor y, decay parameter X, number of candidates N, softmax temperature &
2: Initialize: Proposal action set A, Utilily set
3: Output: Selected action a
4: A < Initialize an empty list for candidate actions
5: ou. t Predict uncertainty-awaat
action distribution sequence using [RCBC] mask as Fig.
6: for 1 to N do
7: a.r f Sample candidate action sequence from distribution ou:
8: 1:T <- Roll out the candidate sequence with [FD) mask as Fig.
9: Ter PeT < Simulate intermediate rewards and long-term rewards using [RP] mask as Fig.
10: U +- Calculate expected utility
> using Equation
11: Append a;, Us to A,,, respectively.
12: end for
13: P Construct candidate selection distribution
b using Equation
14:
15: return <- [ ~ P] if online, else a +- Eip [A]]
Backward PC for Goal Reaching. The ability of a BTM to infer past tokens conditioned on future events sets apart from GPT-based models. This feature is particularly advantageous for implementing MPC from reverse or "backward" perspective when the objective is to achieve a specified goal state. Unlike the goal-reaching mask proposed in previous works (Liu et al. 2022] Carroll et al. (022). which masks all elements along the trajectory except the current and final states to reconstruct the action at the current timestep, we leverage the BTM's bidirectional conditioning capability to inpaint a transition path that guides action selection. We refer to this method as back- ward M MCC
Specifically, the backward MCPC approach uses Path Inference (PI) mask (illustrated in Figure ((b) to guide the model predicting sequence intermediate states leading to the goal. Once path is established, the model employs an Inverse Dynamics (ID) mask deduce the necessary actions to transition between consecutive states along the predicted path. This approach eliminates the need generate a large number of candidates and roll out each one, which inherently demands significant computational resources. Instead, it implicitly performs the same function as traditional MPC by selecting the first action in sequence that most satisfies the given goal.
EXPERIMENTS
Our experiments aim to answer the following questions:
Q1: Can forward MCPC enable the (same) agent to achieve higher accumulated rewards in offline RL and subsequent online finetuning?
Q2: Can backward M enable the agent to perform diverse tasks given target states?
Q3: How does each algorithmic component contribute to M
Q4: Is the pretrained model capable enough perform MC in more complex environments that demand the knowledge of interaction with external objects, e.g. manipulation?
Tasks and Datasets. To answer these questions, we utilize D4RL and RoboMimic dataset suites. We apply three D4RL. locomotion domains (Hopper, Wa Walker2d, Hal FCheetah) with two

---

more competitive results, Published as a conference paper at ICLR 2025
dataset types for each task: medium (m) and medium-reples (m-x), used benchmark our proposed forward PC in offline RL and O20 settings. The RoboMimic encompasses three manip- ulation tasks (Can, Lift. Square). We utilize three official datasets (can-pair, ecuare-mllh 1 lif and two customized datasets (can-1 1 i,, can-real) )) to evaluate M*PC's potential real-world application, particularly in robotic manipulation tasks. Detailed descriptions of the tasks and datasets are provided in Appendix D
Table I: Offline Results on D4RL. Comparison of the average normalized return against several baseline methods without online finetuming. PC and M MCPC- shortened for method PC with (M)onte-carlo return estimation and -Q)ue estimation guidance heuristics, respec- tively. We report the mean and standard deviation of 5 seeds. The best result for each dataset is highlighted bold. Note that M shares the same weights as a pretrained BTM, but con- stantly outperforms BTM in all tasks due the test-time enhancement brought by MC
Dataset
BC
TD3+BC
IQL
DT
TT
BTM
M*PC-M
M¹PC-Q
hopper-m
53.5
60.4
63.8
65.1
61.1
64.3
70.7 26.2
73.645.0
walker2d-m
63.2
82.7
79.9
67.6
79.0
72.5
86.4
halfcheetah-m
42.4
48.1
47.4
42.2
46.9
43.0
43.9
51.240.7
hopper-m-r
29.8
64.4
92.1
81.8
91.5
75.3
80.4
78.3
walker2d-m-
21.8
85.6
73.7
82.1
82.6
76.6
78.2+10.2
92.2+2.1
halfcheetah-m-rr
35.7
44.8
44.1
48.3
41.9
41.1
41.8 10.5
48.2
Total
246.4
386.0
401.0
387.1
403.0
372.8
395.9
429.8
Q1: Offline RL. We present the offline results of PC with Monte Carlo return estimation guid-
ance (M*PC-M) and Q-value estimation guidance (M²PC-Q) in Table [ To evaluate the offline RL performance of our proposed method, we compare it against the following baselines: (1) BC: behavior cloning, which directly mimics the behaviors in the offline dataset; (2) TD3+BC (Fujimoto
& Gu b): an off-policy RL method incorporating a behavior cloning regularization term; (3)
IQL Kostrikov et al. 1)2 y model-free algorithm designed to avoid bootstrapping errors by learn- ing implicit f-functions; (4) DT (Chen et al., 2021): : sequence sroguence-modeling model free approach that predicts actions conditioned on expected returns; (5) TT (Janner et al., (202]): a sequence-more modeling
model based approach that utilizes beam search planning and (6) BTM: which shares the same pre- trained model as our method but applies only the [RCBC] mask for policy inference. The results demonstrate that MCPC significantly improves reward accumulation compared to BTM, consistently outperforming it across all datasets and domains, irrespective of the guidance heuristic used. This indicates that PC'' planning phase effectively refines the action proposals generated by BTM. Furthermore, as a generalist agent, PCPC-M performs competitively with specialized offline RL algorithms such TD3+BC and IQL. Notably, PCPC- achieves outperforming all baselines by a considerable margin.
Online Finetuning. Under the 020 setup, we compare method against IQL. and ODT (Zheng et al. 2022), a specially designed 020 method for DT. The full online training curves of each al- gorithm can be found in Appendix C In Table 2 we report the performance of each algorithm with a 200K online sample budget. To ensure a fair comparison, we use the best performance between ODT's original paper (Chen al. 202T) and our result running its open-sourced im- plementation. Our method outperforms the other two methods in all the environments except the hopper-medi um dataset. After fine-tuning, our total performance score is 31% higher than IQL and 26% higher than ODT, with improvements over finetuning 123% more substantial than those of ODT. We plot the normalized exploration rollout statistics of the M agent and the BTM agent in Figure 4 The results show that MCPC is more likely to collect trajectories of high quality while maintaining some randomness to cover diverse states. Additional results are provided Appendix
Q2: Goal Reaching. To assess whether our proposed method can effectively guide an agent to specified goal states, evaluate it on the following three tasks: (a) Halfcheetah flipping, (b) Walker performing splits, and (c) Hopper wiggling. Due to the limited planning horizon of our model, we provide a sequence of consecutive subgoals ensure that each goal-reaching task remains within

---

Published as a conference paper at ICLR 2025
Noise Explore Return Mean
100
MPPC Explore Return Mcan
100
60
60
40
40
0.5
1.5
Step
10
Step
-10
Figure 4: Exploration Rollout Statistics. Results from two example runs of the Hopper task on the medium um dataset using the same offline pretrained BTM agent. One run employs Gaussian noise for exploration, while the other utilizes MºPC. The red line represents the offline result. Compared to naive Gaussian noise exploration, 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.
Table 2: Online Finetuning Results on D4RL. Comparison of normalized returns before and after online finetuning, as well as the improvement achieved using a 200K online sample budget. We report the mean from five seeds. The best final result for each dataset is highlighted in bold and the greatest improvement is highlighted in green.
Dataset
IQL online
ODT online
PC (Ours) online
offline
6
offline
offline
6
hopper-m
63.8
66.8
+3.0
67.0
97.5
+30.6
73.6
93.9
+20.3
walker2d-m
79.9
80.3
+0.4
72.2
76.8
+4.6
86.442.6
91.947.8
+5.5
halfcheetah-m
47.4
47.4
+0.0
42.7
42.2
-0.6
51.220.7
69.3
+18.1
hopper-m-r
92.1
96.2
+4.1
86.6
88.9
+2.3
78.3
103.5
+25.2
walker2d-m-r
73.7
70.6
-3.1
68.9
76.9
+7.9
92.2+:244
105.2
+13.0
halfcheetah-m-4
44.1
44.1
+0.0
40.0
40.4
+0.4
48.2
67.0
+18.8
Total
401.0
405.5
+4.5
377.4
422.7
+45.3
429.8
530.8
+101.0
the model's planning horizon capacity, rather than directly providing the final desired goal state. Details regarding subgoal selection for each task provided in Appendix A
These tasks deviate from the reward mechanisms typically seen offline data but can be extrap- olated stitched together from offline trajectories. Our results, showcased in Figure 5 illustrate that backward PC enables the agent to generalize to diverse tasks rather than merely imitating offline experiences. This demonstrates the model's ability to adapt to new challenges by leveraging its knowledge of complex dynamics to achieve specific goals.
Figure 5: Demonstration for D4RL Goal Reaching. One evaluation visu- alization for (a) Halfcheetah flipping, (b) Walker doing splits, and (c) Hop- per wiggling predefined frequency. These behavior are all unseen in the of- fline dataset during pretraining, see Ap- pendix for more details.
(
(b)
(
Additionally, we evaluated the BTM's goal- ability using a single goal-reaching mask, sim- ilar to previous studies (Carroll et al. 2022 Liu al. 2022). This method involves keeping the
8

---

Published as a conference paper at ICLR 2025
current state and goal state unmasked while directly executing the inpainted action. However, as detailed in Appendix C this approach failed to consistently enable the agent to reach the goal state. This discrepancy underscores the effectiveness of our model-based approach.
100
Normalized D4RL Tes-fime Return
80
MPPC-S
BTM
MPCC
60
40
20
hopper-m
hopper-m-r
walker2d-m
walker2d-m-rr
haffeheetah-m
halfcheetthh-m-r
Figure 6: Offline RL comparison between unified and specialized model. We report the nor- malized average returns of BTM and PC unified pretrained agent compared specified pretrained agents, denoted by BTM-S and M respectively. The results represent the mean five seeds. The comparison suggests that the unified pretrained model leads to efficient representations and better performance with BTM and M.PC
Q3: Ablation Studies. We conduct ablation studies to evaluate the contribution of individual com- ponents to the success of our method. Specifically, we examine whether unifying the pretraining process using random masking technique enhances M performance. To this end, we pre- trained two separate models a policy model and world model-using the same training objec- tive and model structure as the BTM. However, these models were only applied with the [RCBC] mask and [FD] mask, respectively, during the training phase. These specialized models were then integrated to implement MPC. Our findings, illustrated in Figure 6 indicate that two specialized pretrained models do not improve decision quality compared to the unified pretrained BTM. More- over, implementing MPC with separate policy and world models does not significantly enhance decision-making compared to using only the specialized policy model. This suggests that the uni- fied pretraining approach benefits performance, as the bidirectional Transformer cohesively captures both policy behaviors and environmental dynamics, leading to more effective planning during MPC.
(0) A w/ planing (ours) (b) A w/o planing (c)
Figure 7: Ablation Study on Planning and Uncertainty-aware Action Recon- struction. We ablate sample-based planning, uncertainty-aware action reconstruction, and both components investigate their contri- butions to the algorithmic performance in the online finetuning phase. We report average results over six datasets. Mean of five seeds. The shaded area represents the averaged per- task standard deviation across random seeds.
A w/o planing (d) A w/ planing
Norallzzdd Score (%)
100
80
60
20
0
0.4 0.6 Environment Steps (x106)
0.8
We further justify some design choices in MOPC's online finetuning phase by comparing: (a) Our MCPC as in Algorithm [] which combines uncertainty-aware action distribution A reconstruction and planning-based action resampling (b) randomly sampling from A for exploration (c) the orig- inal BTM's method of action a reconstruction, trained with MSE loss, adding fixed action noise (O aT) for exploration, maintaining the same entropy level as PC (d) performing planning- based action resample using (c)'s decisions. We present the averaged finetuning process across D4RL datasets Figure 7 The results highlight the effectiveness of our key contributions. Specifi- cally, uncertainty-awar policy for exploration is crucial for maintaining online training stability,

---

Published as a conference paper at ICLR 2025
and forward planning significantly improves sample efficiency. Figure 7 also shows that the per- formance drops drastically when naively using the uncertainty- "mawwrre original BTM for explo- ration. Per-task training curves and additional ablation studies are provided in Appendix C
Table 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 imple- mentation due to their poor performance in the corresponding simulated tasks.
RCCB
Goal (seen)
Goal (umseen)
Dataset
BC
IQL
DT
PC
Figure 8: Skill Generalization in Can-Pick task. Simulated environments on the top and real-world environments on the bottom. The columns show the original behavior (left), behavior conditioned 0.77
Can-Pair
0.64 0.34 0.94 0.98
Square-MH
0.53 0.13 0.21 0.28
Lift-MG
0.65
0.29 0.93
Can-Lim
0.25 0.27 0.46
L..16
on the seen goal state (mid), and behavior condi- tioned on the unseen goal state (right).
Can-Real
0.50 0.70
Q4: Manipulation. In addition to the self-body motion control tasks explored in earlier experi- ments, we shift our focus to manipulation tasks to assess whether the proposed MCPC method can be effectively applied to robotics tasks requiring interaction with objects in the environment. We con- ducted experiments across three simulated tasks in RoboMimic-Can Square, and Lift-each with varying levels of complexity. Additionally, we utilized datasets of varying quality levels, including machine- -ennrreee (MG), mid-level buman-demonstrated (MH), and paired positive-and -ndieiaaive-eeeiiie (Pair) demonstrations. We also created customized simulated task named Can-Lim, variant of the Can-Pick task, in which the dataset was adapted to a scenario where the relative pose between the gripper and the can is unavailable. Finally, tested our method real-world Can-Pick task, referred to as Can-Real. The results are compared against several offline RL baselines, as shown in Table 3
To evaluate generalization capabilities, we conducted a goal-conditioned RL experiment on the Can- Picking task with the Paired dataset. This dataset contains 50% perfect demonstrations, where the agent successfully picks up the can and places it into the box in the right corner, and 50% negative demonstrations, where the can is thrown off the table, resulting in no reward. As illustrated in Figure 8 by specifying the final goal states, we can control the agent's 's behavior to either complete the original task or reproduce the throwing-away behavior. Furthermore, by specifying a final state numerically between two observed states in the dataset, the model generates actions that enable the agent to reach a previously unseen statinlacing the can into the box adjacent to the right one.
DISCUSSIONS AND LIMITATIONS
We propose MOPC, test-time MPC framework designed enhance the inference performance of masked Transformers pretrained under offline RL settings. MC offers the following benefits: (1) Improved Decision-Making without Further Training: During inference, 'C improves mecisinn-making with high computational efficiency. (2) Enhanced Finetuning Efficiency: With an additional online interaction budget, PC achieves better final performance and outperforms the previous sequential modeling 020 approach, ODT. This enhances the agent's continuous learning ability. (3) Generalization Ability: The framework demonstrates notable generalization capabili- ties, generating actions that effectively drive the agent toward unseen goal states in both simulated and woaldworld tasks.
While MCPC demonstrates promising results, there areas for further investigation: (1) Han- dling Pixel Observations: Currently, our framework is limited to environments with state-based observations. Future work will explore methods for handling pixel observations.
(2) Transformer Scalability: Our experiments employed a fixed-structure masked trajectory Transformer. remains unclear whether increasing the Transformer's capacity would lead to better test-time planning.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty- based offline re- inforcement learning with diversified g-ensemble. Advances in neural information processing systems, 744777677477 2021.
Arthur Argenson and Gabriel Dulac-Amold. Model-based offline planning. arXiv preprint MxXi::0008.05566 2020.
Bowen Baker, 1lge Akkaya, Peter Zhokov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune. Video pretraining (vpt): Learning to act by watching unlabeled online videos. Advances in Neural Information Processing Systems, 35:24639-24654, 2022.
David Brandfonbrener, Alberto Bietti, Jacob Buckman, Romain Laroche, and Joan Bruna. When does return- -ttun--oondiioone supervised learning work for offline reinforcement learning? Advances in Neural Information Processing Systems, 35:1542-1553, 2022.
Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint arXiv:2212.06817 2022.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Micah Carroll, Orr Paradise, Jessy Lin, Raluca Georgescu, Mingfei Sun, David Bignell, Stephanie Milani, Katja Hofmann, Matthew Hausknecht, Anca Dragan, et al. Uni [mask]: Unified inference in sequential decision problems. Advances in neural information processing systems, 35:35365- 35378, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 14411-10055977 2021.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in a handful of trials using probabilistic dynamics models. Advances neural information processing systems, 31, 2018.
Jacob Devlin, Ming- Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805, 2018.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132-20145, 2021.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pp. 2052-2062 PMLR. 2019.
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in neural information processing systems, 31, 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. PMLR, 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint 1rXiv:1912.01660 2019.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint MXXv:22033.04555 2022.
Nathan Hatch and Byron Boots. The value of planning for anfinite-horizon model predictive control. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pp. 7372-7378. IEEE, 2021.
11

---

Published as a conference paper at ICLR 2025
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked au- toencoders are scalable vision learners. In Proceedings of the IEEE/CVE conference on computer vision and pattern recognition, pp. 16000-16009, 2022.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. Advances in neural information processing systems, 32, 2019.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1273-122 1286, 2021.
Zhiwei Jia, Vineet Thumuluri, Fangchen Liu, Linghao Chen, Zhiao Huang, and Hao Su. Chain-of- thought predictive control. arXiv preprint arXiv:2304.00776 2023.
Iyya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. arXiv preprint 1.X1::2110006699 2021.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. Advances in neural information processing systems, 32, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 179-1199, 2020.
Seunghyun Lee, Younggyo Seo, Kimin Lee, Pieter Abbeel, and Jinwoo Shin. Offline-to-online reinforcement learning via balanced replay and pessimistic q-ensemble. In Conference Robot Learning, pp. 1702-1712. PMLR, 2022.
Fangchen Liu, Hao Liu, Aditya Grover, and Pieter Abbeel. Masked autoencoding for scalable and generalizable decision making. Advances in Neural Information Processing Systems, 12618, 2022.
Kendall Lowrey, Aravind Rajeswaran, Sham Kakade, Emanuel Todorov, and Igor Mordatch. Plan online, learn offline: Efficient learning and exploration via model-basee control. arXiv preprint arXiv:1811.01888, 2018
Ashvin Nair, Abhishek Gupta, Murtaza Dalal, and Sergey Levine. Awac: Accelerating online rein- forcement learning with offline datasets. arXiv preprint arXiv:2006.0933 2020.
Mitsuhiko Nakamoto, Simon Zhai, Anikait Singh, Max Sobol Mark, Yi Ma, Chelsea Finn, Aviral Kumar, and Sergey Levine. Cal-ql: Calibrated offline rl pre-training for efficient online fine- tuning, Advances in Neural Information Processing Systems, 36, 2024.
Vitchyr Pong, Shixiang Gu, Murtaza Dalal, and Sergey Levine Temporal difference models: Model- free deep for model-based control. arXiv preprint arXiv: 0908/, 2018
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Martn, Mai Gimenez, Yury Sulsky, Jackie Kay. Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint wxXiv2205..06175 2022.
Nur Muhammad Shafiullah, Zichen Cui, Ariuntuya Arty Altanzaya, and Lerrel Pinto. Behavior transformers: Cloning k modes with one stone. Advances in neural information processing sys- 2022.
tems, 35:22955-22545
David Silver, Richard S Sutton, and Martin Müller. Sample-based learning and search with perma- nent and transient memories. In Proceedings of the 25th international conference on Machine learning, pp. 968-975, 2008
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, łukasz. Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- fion processing systems, 30, 2017.
Thomas Walsh, Sergiu Goschin, and Michael Littman. Integrating basele-based planning and basel-based reinforcement learning. In Proceedings of the aaai conference on artificial intel- ligence, volume 24, pp. 612-617, 2010.
12

---

Published as a conference paper at ICLR 2025
Kerong Wang, Hanye Zhao, Xufang Luo, Kan Ren, Weinan Zhang, and Dongsheng Li. Bootstrapped transformer for offline reinforcement learning. Advances in Neural Information Processing Sys- tems, 35:34748-344766 2022.
Mianchu Wang, Rui Yang, Xi Chen, Hao Sun, Meng Fang, and Giovanni Montana. Goplan: Goal- conditioned offline reinforcement learning by planning with learned models. arXiv preprint arXiv:2310.20025. 2023.
Philipp Wu, Arjun Majumdar, Kevin Stone, Yixin Lin, Igor Mordatch, Pieter Abbeel, and Aravind Rajeswaran. Masked trajectory models for prediction, representation, and control. In Interna- tional Conference on Machine Learning, pp. 37607-37623. PMLR, 2023.
Zifan Wu, Chao Yu, Chen Chen, Jianye Hao, and Hankz Hankui Zhuo Plan to predict: Learning umcertainty-foreseeimn model for model-based reinforcement learning. Advances in Neural Information Processing Systems, 15:45569-1 2022.
Taku Yamagata, Ahmed Khalil, and Raul Gantos-Rodrigued Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline rl. Inter- national Conference on Machine Learning, pp. 119989-3900) PMLR, 2023.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33:14129- 2020.
Zishun Yu and Xinhua Zhang. Actor-critic alignment for offline-to-online reinforcement learning. In International Conference on Machine Learning, pp. 40452-44474, PMLR, 2023.
Zilai Zeng, Ce Zhang, Shijie Wang, and Chen Sun. Gondicondi predictive coding for offline reinforcement learning. Advances Neural Information Processing Systems, 36, 2024.
Xianyuan Zhan, Xiangyu Zhu, and Haoran Xu. Model-based offline planning with trajectory prun- ing. arXiv preprint arXiv:2105.07351, 2021.
Haichao Zhang, We Xu, and Haonan Yu. Policy expansion for bridging offline-to-online reinforce- ment learning. arXiv preprint arXi::2302.00933 2023.
Marvin Zhang, Sharad Vikram, Laura Smith, Pieter Abbeel, Matthew Johnson, and Sergey Levine. Solar: Deep structured representations for model-based reinforcement learning. In International conference on machine learning. pp. 7444-7453. PMLR, 2019.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In international conference on machine learning, pp. 17742-27059) PMLR, 2022.
13

---

Published as a conference paper at ICLR 2025
A
IMPLEMENTATION DETAILS
Loss Function Construction. We consider the Lagrangian of Equation 2 given by:
L(8, J(f) a(3 -H))
(5)
where a is non-negative Lagrange multiplier. The training objective then become
max L(0, o).
(6)
e20 8
We alternately optimize 0 and a as follows:
Optimizing 0 with fixed a, which involves:
min (J(0) - o7())
(7)
0
Optimizing a with fixed 0, formulated as:
min (Hg - 9). e20
(8)
This iterative training of 9 and a ensures compliance with the entropy constraint while optimizing the objective function J()).
Transition-wise Value Estimator. We choose IQL (Kostrikov al., 2021) algorithm train the value estimator because its Bellman updates do not require an explicit policy function. Typpcally, IQL simultaneously learns a critic network Qy and value network Vo with the losses defined by:
Jo(()) B(saan;s) [rr VV(ss') 2x(s,a))²)
(Ø) ((..a)~T [|e 1/ft--1eecee (e..
(s, (o()))
where is constant hyperparameter named expectile used to control the conservatism of the value estimation. The critic network Q will be applied estimate the long-term reward for given state-action pair in our approach. t is set to 0.7 for D4RL locomotion tasks and 0.9 for RoboMimic manipulation tasks.
Goal State Definitions in Goal Reaching Tasks, In the goal-reaching setup for Hopper, Walker, and HalfCheetah, we craft rough trajectories based on the specific anticipated dynamics of each agent. For the Hopper, a sinusoidal trajectory is designed for the foot joint to induce a wiggling motion, while the other two joints' initial positions are maintained. In the case of the Walker, a lin- early increasing trajectory for the thigh joint facilitates the splits, with the dynamics of other joints extracted from the offline trajectory which correspond to stepping behavior, providing rough guid- ance. For the HalfCheetah, the primary flipping motion is guided by linear trajectories that set the body height decrease and a full 180-degree rotation to simulate a flip. Complementing this, the dy- namies of the other joints and body movements are derived from offline datasets that capture detailed flipping steps, providing a coherent and realistic motion base. Subgoals are extracted from these tra- jectories specific intervals intervals-every every fifth, thirtieth, and every timestep, respectively- guide each model towards achieving the intended maneuvers, ensuring that while the main actions are precisely targeted, the full spectrum of body dynamics remains realistically integrated and synchronized with the models' overall movements.
In the Can Pick task, we deploy specific guidance trajectories for each distinct thhovig away and moving to nearer box. For the "throwing-away behavior, we directly select suitable trajectory from an offline dataset without any modifications, ensuring that the agent replicates proven effective throwing motion. For the "moving to nearer box" behavior, the process begins by selecting a "moving to correct box" trajectory from the offline dataset. To tailor this trajec- tory to the specific task, apply affine transformation to adjust the horizontal positions of the Franka robot's end effector and the object along the trajectory. This transformation proportionally reduces the distance the object needs to be moved, customizing the trajectory the current scenario. Subgoals are then extracted from these guidance trajectories at every state, providing detailed, step- by-step targets that guide the agent's 's actions towards successful task completion.
Hardware. The entire training process, including both pretraining and finetuning, is performed on NVIDIA 3090 GPUs. During the offline pretraining phase, we train the BTM model for 140K
14

---

Published as a conference paper at ICLR 2025
gradient steps, which takes approximately 4 hours per dataset single GPU. For the finetuning phase, we allow million online exploration steps for figure plot and report the performance with 0.2 million exploration steps. The finetuning phase including exploration and evaluation in simulator takes between 7 and 9 hours per dataset on a single GPU, while finetuning the pretrained trajectory model itself takes half of the total time.
15

---

Published as a conference paper at ICLR 2025
B HYPERPARAMETERS
Table 4: Hyperparameters.
Hyperparameter
Offline
Online
Training
Nonlinearity
GELU
GELU
Batch size
2048
512
Trajectory-segment length
8
Dropout
0.10
0.10
Learning rate
0.0001
0.0001
Weight decay
0.005
0.005
Target entropy 8
-3
-3
Scheduler
cosine decay
Warmup steps
40000
Training steps
140000
Evaluation
Context length
4
Bidirectional Transformer
of Encoder Layers
2
Decoder Layers
Heads
4
Embedding Dim
512
512
Mode Decoding Head
Number of Layers
2
2
Embedding Dim
512
512
Reward Maximization
Decay Parameter
0.6
0.6
Candidate Number N
625
625
Softmax temperature 5
1.0
1.0
16

---

Published as a conference paper at ICLR 2025
C
ADDITIONAL RESULTS
Online Finetuning Results. We report the per-task online training curves over 1 million online samples for our method and our reproductions of baseline methods in Figure 9 ablations in Figure
10
MPPC
IQL
GCT
8TM
hopper-medium
hosper-medlum-reply
140
140
Normalized Soore (%6)
120
Normalized (%)
Vhm s mm
n 11/w
0.2
0.4
0.6
0.8
0.2
0.4
0.8
0.8
Environment Steps (x100)
Environment Steps (x100)
hallcheetal-mmedlum
hellcheetal-mmddlun--eelllyy
100
100
Normalized Score (%)
80
80
Normalized
20
0.2
0.6
0.8
0.4
0.8
Environment Steps (x100)
Environment Steps (x10")
walker2d-medium
wallker2d-medium-rrplayyy
Normalized Score (%6
Nomalized
80
80
40
40
20
20
Environment Staps (x109)
Environment Steps (x10²)
Figure 9: D4RL Benchmark Comparison. Per-task Online Training Curves for PC and baseline methods. Mean of seeds. The shaded area represents the standard deviation across seeds.
We furthermore compete MC with some stronger. specialized 020 bascline methods with the 100k online sample budget practice: (1) AWAC (Nair et al., 2020 a representative 020 approach utilizing advantage-weighee cling framework for offline RL and online finetuning; (3) OFF2ON Lee et al. (2022), a CQL-based pessimistic Q-ensemble method that incorporates balanced replay to encourage near on-policy samples from the offline dataset; and (4) PEX (Kostrikov et al. 2021). an IQL-based algorithm cused policy expansion. We evaluate the baselines the D4RL locomotion datasets, with the results summarized 5 The results demonstrate that M SC achieves performance compara- ble SOTA specialized 020 methods such OFF2ON and PEX.
actor-critic; (2) ODT (Zhngg et al., 2022). unified sequential mod-
17

---

Published as a conference paper at ICLR 2025
Dataset
AWAC
ODT
OFF2ON
PEX
PC
hopper-m
57.8 > 55.1
73.4
97.5 80.2
56.5 87.5
73.6 81.3
walker2d-m
35.9 + 72.1
72.0-72..2
66.2 72.4
80.1 00.1-92..
86.4 + 74.9
halfCheetah-m
43.0 42,4
42.7 42.1
-9.3-559.6
50.8 118-600.9
51.2 + 64.0
hopper-mr
37.7 - 60.1
-0.4-78.5
28.2
31.5
78.3 78.6
walker2d-mr
24.5 79.8
79.8
44.2 71.8
17.7
89.2
80.1
92.2 98.8
halfCheetah-mr
40.5 - 41.2
32.4 39.7
42.1 60.0
45.5 51.3
62.7
Average
39.9 58.5
54.2
+ 61.9
48.5 + 73.5
57.4 80.2
76.8
Table 5: 020 Baseline Comparison Results. Comparison of normalized returns before and after online finetuning with 100K online sample budget. We report the mean of four seeds.
planing (ours)
M planing
w/o planing
/ planing
hopper-medium
hogper-mediun-rrollyy
Normalize Soore (66)
120
Nommaizee Soore (56)
120
80
80
60
40
to
0.8
Environment Steps (x10
Environment (x10")
hellchhettah-medlun
Maltchhettah-mmadllnn
100
100
Nommaized Soore (16
80
Nomalized Stoore (%)
80
40
0.2
0.4
0.8
0.8
0.2
0.4
0.8
0.8
Environment Steps Oc10*)
Enwironment Stepe (x10%)
welker2d-mediuu
welkeraad-medum-neplay
140
Nomralized Score (8)
120
Normalzed Score (96)
120
100
100
80
60
40
8
0.8
Environment Steps 0c10
Environment Steps (x109)
Ablation Studies for Algorithmic Components Contribution. Mean of 5 seeds.The shaded area represents the standard deviation across seeds.
Figure 10:
Inference Time.
We have introduced MOPC's computational efficiency due to the parallel predic- tion nature of the mask autoencoding paradigm in the methodology section. For completeness, we report the inference time of M*PC's planning overhead with respect to range of planning horizons (1 to 8) in Fig. II We additionally include two methods for references:
(1) TT (Janner et al. 2021). sequential modeling approach that employs beam search for test-time planning: (2) TD- MPC (Hansen et al. 2022), a representative model-based RL method combining MPC and temporal difference learning. All inference times were benchmarked single NVIDIA RTX 3090 GPU.
18

---

Published as a conference paper at ICLR 2025
Note that used the original implementations of the baseline methods, the number of parame- ters is not aligned across approaches. Results demonstrate that PC much more computational efficient compared to the sequential modeling approach TT. Furthermore, as the planning horizon increases, MCC even outperforms TD-MPC, despite the latter being a more lightweight model.
Compurational Time (mno, log scak)
MPC
Figure 11: Inference Time Comparison. MC much more computational efficient compared to sequential modeling approach TT and even out- perform lightweight model TD-MPC as planning horizon increases.
107
Planning Horizon (steps)
Ablation Study on Decay Parameter. Decay parameter play significant role in balanc- ing the weight of instant rewards and long-term value. We provide the training curves for À E {0.0, 0.1, 0.3, 0.5, 0.7. Figure 12 indicates our approach is not sensitive to the choice for since each choice outperforms the baseline (randomly sampling action from A for exploration) by large margin, and has minor difference learning speed (fine-tuning improvements happen slower when À 0.1 and long step stability (performance drops after 800k online steps when À 0.9. We choose À = 0.6 in all the experiments as an intermediate choice for balancing converge speed and online training stability.
Performance 9.1M
Performance
Performance
300
Pird Nomalbord Store
Find Nomeloord Score
200
15
100
20
044
0.4
0.
U.S
Baseline
A= 0.6
-BI 0.1M
-4- 0.5M Steps
-6- 1M Steps
Figure 12: Ablation Study for X Choices. Normalized score function of choice with 0. 0.1m, 0.5m, 1.0m online steps. The red star represents our default choice (0.6) while the grey line denotes baseline results (explore w/o planning). Mean of seeds.
Ablation Study Entropy Constraint. We also report the effects of entropy constraint we im- posed in Equation [ The results of offline results M³PC-M, M$PC-Q and online finetuning results M'PC-online summarized in Table 6 Empirical results show that entropy constraint does not have substantial influences on offline results but significantly boost the online sample efficiency.
Datasets
MCCC-M
PCCC
MOPC-online w/o
w/o
W
w/o
w
W
hopper-m
84.347.3
70.7+6.
81.643.5
73.6+5.6
94.9
93.9+15.8
halfchectah-ne m
43.8+0.6
43.9+3.9
50.0+0.3
51.2+0.7
71.5
69.3
walker2d-m
79.9+1.
80.9+2.5
80.7+7.2
86.4+2.6
68.3-25.0
91.9
hopper-mr
75. 75.1
80.415.2
76.8
78.3
88.7
103.5
walker2d-mm
78.5+16.0
78.2+10.2
94.0.a.s
92.211.4
108.1
105.241.0
halfcheetah-mr
40.0+1.0
41.840.5
48.0+0.8
48.2+0.4
70.2
67.0+7.1
Average
66.9
66.0
71.8
71.6
83.6
88.5
Table 6: Ablation Study on Entropy Constraint. Comparison of PCCM-M MºPC-Q, and online results w/o entropy constraint across D4RL datasets.
19

---

Published as a conference paper at ICLR 2025
Goal Reaching Results. We show more results in goal reaching tasks here. To demonstrate the extent to which our unseen goal is out of the distribution, we show together the PCA dimension- reduced results for all states in the offline dataset which the model was pretrained, and the states in the trajectory of reaching the given goal. As in Fig[T] The different tasks have different out- of-distribution cases: For the walker-split task, the agent starts with to a state never seen before (the angle of the hip-joint). For the cheetah-flip task, the initial state and goal state are both seen in the offline dataset, the normal state usually corresponds better rewards, while the flip-over state hardly leads to any reward, as the original task in the dataset is run fast. However, conditioned on the state given, the agent finds many unseen intermediate states to finally transit to flip-over state. For the Hopper-Wiggh task, the agent strings together a series of near-in-distributtin states to form a loop wiggling action, which is not seen the dataset.
state and finally reach
Hopper-Wizale
CheelatiFFip
Walker-Spii
States in offine catasets (normalized reward)
States nn online poal-reaching trajestory (normalized timesteps)
Figure 13: Visualization of states in different tasks after 2-dim PCA mapping.
Additionally, we show the goal states we take as input in order to reach the final behavior, and how well BTM with single Goal Reaching mask and backward PC can follow those states. We only plot the most representative dimension in the state vector for each task, respectively. E.g., Angle of the front tip (dim(1]) of cheetah, and angle of the thigh joint (dim[2]) of walker and angle of the top (dim[1]) of hopper. As in Fig 14 with only single mask, the agent can hardly achieve the goal, and the overall behavior resembles the behavior cloning result from the pretrain dataset. However, with backward MC the agent can successfully follow the kinematics guidance, although some do exactly satisfy the dynamics. Moreover, we show that the same pretrained model with backward MCPC PC can reach wiggling behavior of different frequencies in hopper environment, with proper goal states.
20

---

Published as a conference paper at ICLR 2025
Goal state
Naive Goal-reaching (single mask)
Goal-reaching with MPPC
Cheetah-Fiip
Walker-Spiit
0.00
-0.25
-1
0.50
-0.75
-2
1.00
-3
150
200
100
150
200
Hopper-wiggle frequency=2
0.3
0.1
0.0
0.1
-0.
Hopper-wiggle froquency=6
0.3
0.2
0.0
-0.1
-0.2
Hopper-wiggle frequency=0
0.3
0.2
0.1
-0.1
-0.2
-0.
200
400
600
800
1000
Figure 14: Comparison between Backward PC and a single Mask in Goal-Reaching Tasks. We present the goal states and resulting states after policy execution across three goal-reaching tasks, focusing on a single key dimension. The single Mask fails to guide the agent toward the goal states when the given current-goal state pairs are out of distribution.
21

---

Published as a conference paper at ICLR 2025
D TASKS AND DATASETS
The dataset utilization checklist is shown in Table
(b)
Tasks Setup. (a) Locomotion tasks in D4RL: halfcheetah, hopper, walker2d (from left right); (b) Manipulation tasks RoboMimic: lift, can, square (from left right), (c) Left view and front view of real-world manipulation task setup.
Figure 15:
D4RL. We consider three representative D4RL locomotion domains (Hopper, Walker, and Hal SalfCheetah), Each domain contains two datasets (medium, medium-reet replay) which have different data compositions. The medi um datasets contain 1M samples collected by a partially- trained SAC (Haarnoja et al. 2018) agent. The medi wawiiumm-lpppp samples in the replay buffer observed during training until the agent reaches the "medium" level. We both these two types of datasets offline RL and O20 RL.
dataset consists of recording all
RoboMimic. RoboMimic includes suite of manipulation task datasets designed for the Franka Panda robot, focusing on three specific tasks: Can, Square, and Lift. The dataset for pretrain- ing encompasses four distinct categories: (1) Multi-Human (MH), consisting of six sets with each containing 50 demonstrations by different pairs of demonstrators; (2) Machine Generated (MG), generated by Soft Actor-Critic (SAC) agent at various stages of its training, providing a spectrum of behaviors from early exploratory to more refined tactics; and (3) Paired, where single expe- rienced operator recorded two demonstrations for each of 100 initializations of the Can task-one demonstrating correct placement and the other tossing the object outside. detailed the state space and action space definition for each environment in Robomimic below, including customized en- vironments can-limit and can-real.
The Action Space and State Space for Manipulation. The action space for each timestep is 7-dimensional vector per arm, where the first six coordinates represent control signals the op- erational space control (OSC) space, and the last coordinate controls the opening and closing of the gripper fingers. The observation space includes 7-dimensional vector for the absolute end ef- fector position quaternion and 2-dimensional vector for the left and right finger relative poses of the gripper in addition to task-specified object observations. In the "Lift" task, object observations include 10-dimensional vector consisting of the absolute cube position and quaternion (7-dim), and the cube position relative to the robot end effector (3-dim). In the "Can" task, the object ob- servations are 14-dimensional vector, including the absolute can position and quaternion (7-dim), and the can's position and quaternion relative to the robot end effector (7-dim). For the "Square" task, object observations also form a 14-dimensional vector with the absolute square nut position and quaternion (7-dim) and their relative positions and quaternions (7-dim) to the robot end effec- tor. In the "Can-Limit" task, the object observations include only the absolute can position (3-dim), excluding relative position knowledge to align with goal-reaching tasks where precise relative poses are unnecessary. In the "Can-real" task, which is a real-world environment similar to Can Can-Limin Limit, ob- ject position is detected using vertically placed depth cameras, with actions output 20 Hz, and
22

---

Published as a conference paper at ICLR 2025
robot joint torques adjusted at 500 Hz to achieve the desired Cartesian poses based on the operational space controller.
Table 7: Dataset Utilization. We outline the dataset utilization for each experiment part here, a checkmark means the corresponding dataset is use for pretraining.
Dataset
Offline RL
Goal Reaching RL Online Finetuning
hopper-medi um-v2
hooper-mmedium-repplaayyy
walker22d-nedium-12
walkeer20-mmddiun--eeeyyy>>>
halffchettahhmmediin-2222
Can-Pair
Square-MH
Lift-MG
Can-Lim
Can-Real
23


Reviews:
[{"review_id": "qIn5ncy7ms", "reviewer": "Reviewer_cxyR", "summary": "This paper introduces M^3PC, an approach that leverages properly designed masking schemes to perform test-time MPC with masked trajectory models for decision making tasks. The proposed method enables action reconstruction with uncertainties for better robustness, as well as forward and backward prediction through different masking patterns for solving various downstream tasks. Evaluations are performed on both simulated motion control and real-robot manipulation tasks, across offline and offline-to-online setups. The efficacy and generalization capabilities of the method are supported by the extensive experimental results.", "questions": "1. Given that one benefit of utilizing sequence models for decision-making is their strong capability to model long-term dependencies, when performing backward M^3PC for goal-reaching tasks, I wonder how the method should be properly used in long-horizon cases where the number of state transitions needed to reach the goal state might be larger than the planning horizon/time budget of the sequence model. I also wonder if it’s possible to have some quantitative results on how the proposed method performs on one or two long-horizon goal-reaching tasks (e.g. AntMaze-Large).\n2. The overall framework is novel, but some design choices made in the framework are reminiscent of prior works, such as Dreamer [1], MTM [2], TT [3], UniMASK [4], and MaskDP [5]. I'm aware the authors have properly cited and briefly discussed these works, but it would be great if the authors could discuss the discrepancy between M^3PC and these works in detail in the related work or elsewhere, which might help highlight the unique contribution of M^3PC. For example, how is the sampling process of M^3PC different from that of TT, how is [PI] mask different from the goal-reaching mask used by MaskDP, and how are the applied return heuristics related to the ones used by prior works?\n3. Line 294: it seems [PI] mask is illustrated in Figure 3 (b) instead of Figure 2.\n\n[1] Hafner et al. Dream to Control: Learning Behaviors by Latent Imagination. ICLR 2020.\n\n[2] Wu et al. Masked Trajectory Models for Prediction, Representation, and Control. ICML 2023.\n\n[3] Janner et al. Offline Reinforcement Learning as One Big Sequence Modeling Problem. NeurIPS 2021.\n\n[4] Carroll et al. Uni[mask]: Unified inference in sequential decision problems.NeurIPS 2022.\n\n[5] Liu et al. Masked autoencoding for scalable and generalizable decision making. NeurIPS 2022.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. This work proposes a novel framework that combines the pre-trained Masked Trajectory Model with MPC. Allowing action reconstruction with uncertainty improves the robustness under stochastic settings.\n2. The authors provide extensive evaluations in both simulated and real-robot setups, demonstrating the effectiveness of the proposed method.\n3. The authors thoroughly discussed the limitations of the proposed method.\n4. This paper is well-written and easy to follow.", "weaknesses": "1. Perform sampling-based planning with large sequence models can usually be computationally intensive and time-consuming, given that during test time M^3PC only executes the first action of the plan in every interaction step.\n2. Trajectory Transformer (TT) [1] also performs planning using sequence models, even though it does not leverage any masking schemes or bidirectional transformer, including TT as one of the baselines would be a great bonus, especially when TT also proposes to leverage the Q function learned by IQL as return heuristics.\n3. Please see the first question below.\n\n[1] Janner et al. Offline Reinforcement Learning as One Big Sequence Modeling Problem. NeurIPS 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jlaUKN3GH0", "reviewer": "Reviewer_HJV8", "summary": "This paper introduces a unified trajectory model (Bidirectional Trajectory Model, BTM) trained through multiple auxiliary tasks. The pretrained BTM can serve both as a policy to output actions and as an environment model to predict future states. Building on this, the authors enhance policy capabilities using an MPC-based approach. The experiments demonstrate that the proposed algorithm significantly improves performance in offline RL and effectively handles offline-to-online RL and goal-conditioned RL tasks.", "questions": "See weaknesses.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This paper presents a versatile pretrained BTM that achieves four functionalities by designing different masking methods. The authors cleverly combine these functionalities to implement MPC and goal-reaching based on BTM.\n2. M$^3$PC shows substantial improvement over previous methods (DT, BTM, and ODT) in experiments. The authors also validate the method on a real robotic arm manipulation task, greatly enhancing the paper's significance.", "weaknesses": "I don't have major concerns. However, I would be grateful if the authors could clarify the following questions:\n\n1. How much additional time overhead does the MPC process introduce? Including a description of time overhead in the experiments would improve the comprehensiveness of the method evaluation.\n2. It would be beneficial to test the goal-reaching capability in the antmaze task, as it may better visualize the trajectories inferred by the model in these tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "3tXv44W7Sh", "reviewer": "Reviewer_CMiL", "summary": "The paper introduces M3PC, an approach that builds on bidirectional trajectory models trained with masked prediction. During inference, M3PC refines behavior using Model Predictive Control (MPC) applied to a pre-trained model. To support this, the model includes probabilistic action \"heads,\" allowing it to sample actions with uncertainty awareness. This setup enables forward MPC for reward maximization and reverse MPC for goal-reaching. In forward MPC, M3PC uses a utility function (a TD($\\lambda$)-like combination of local rewards and either Return-To-Go or Q-values) alongside the model's forward predictions to estimate the performance of an action sequence and propose refinements. In reverse MPC, the model first plans a trajectory toward the goal, then selects actions that follow this trajectory using inverse dynamics.\n\nA robust set of experiments demonstrates M3PC's effectiveness across standard offline reinforcement learning, offline to online settings, goal-reaching tasks, and a real robot experiment.", "questions": "- While reading the main part of the paper I assumed the \"goal reaching\" task only needs the final position but from the appendix it seems there are entity goal trajectories? I guess it makes sense given the limited horizon of the model, but I would appreciate a bit more elaboration clarity here. \n\n- How are the actions selected from predicting the inverse dynamics during backward planning (sampling or mean of the gaussian head)? \n\n- What is indicated by the shaded areas in the reward curves?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper stands out due to its simplicity and clever conceptual foundation, building effectively on a straightforward masked trajectory approach. The evaluations are thorough and demonstrate strong performance across various benchmarks, showcasing the approach's applicability to a wide set of problems. Additionally, the paper includes helpful ablation studies that validate the importance of key design choices.", "weaknesses": "The only weaknesses are some missing clarity, detail, and context in the experiment section, see the list below\n\n- Which task is used in Figure 7? It seems a bit like there is an \"introduction\" paragraph missing before the one starting in l451\n- The acronyms MPC-M and MPC-Q) appear a bit out of nowhere, I think I could piece it together, but introducing the two variants more explicitly would be helpful \n- Missing standard deviations in table 3 \n- Include (or replace IQL with) a \"conservative Q-Learning\" approach tailored to O2O for better context, (e.g. Cal-QL: Calibrated Offline RL Pre-Training for Efficient Online Fine-Tuning, Nakamoto et al 2023). Such approaches could also be included in the related work and given the standardized benchmark, results can be simply taken from the respective works. \n- ( I realize that this is a bit of a personal preference / due to my background so this is just a suggestion, not affecting my assessment: I would appreciate highlighting the robo-mimic results, in particular those including a real robot, and more discussion with regards to the effect of the different datasets used)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "afaAIprUfB", "reviewer": "Reviewer_6kkf", "summary": "This paper proposes a model-based planning method M$^3$PC for offline RL. M$^3$PC has two parts — a bi-directional trajectory model and model-predictive control (MPC). The trajectory model acts as a prior policy and an environment model, and test-time MPC incorporates the trajectory model to decide the most promising actions. The novelty of this paper mainly lies in the trajectory model, which predicts actions, states, and rewards using different masks. In these experiments, it shows good results in D4RL benchmarks, offline-to-online settings, and goal-reaching tasks.", "questions": "1. There is a typo in the caption of Figure 2: \"the model can show have multiple capabilities...\"\n\n2. The legend in Figure 7 is unclear. Could the authors link the curves with methods (a), (b), (c), and (d) in the text?\n\n3. The motivation for the constraint in Eq. (2) is not adequately explained. The trajectory-level entropy constraint could lead the policy to sample out-of-distribution actions, which is problematic in offline RL and may harm performance. It would be better to compare and explain the results for a version without the trajectory-level entropy.\n\n4. In the related work, some model-based offline planning methods have not been discussed, such as MBOP (Argenson and Dulac-Arnold, 2021), MOPP (Zhan et al., 2022), and GOPlan (Wang et al., 2024). Notably, MOPP and GOPlan consider uncertainty during planning and tend to remove/prune trajectories with high uncertainty. \n\n5. It would be useful to report the computation time for M$^3$PC. Typically, test-time planning methods require more computation than actor-critic methods, such as TD3-BC and IQL. Could the authors also add test-time planning baselines to Table 1?\n\n6. In the ablation study, the training details of the policy model and the world model should be provided. Q3 shows that a unified model performs better than the separated models, but it would be better to further demonstrate that the proposed unified model outperforms those in other papers. For example, the well-known world model RSSM (Hafner et al., 2019) could be compared. \n\n**References**\n\nArgenson, A., \\& Dulac-Arnold, G. (2021). Model-Based Offline Planning. ICLR.\n\nHafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., \\& Davidson, J. (2019). Learning Latent Dynamics for Planning from Pixels. ICML.\n\nWang, M., Yang, R., Chen, X., Sun, H., Fang, M., \\& Montana, G. (2024). GOPlan: Goal-conditioned Offline Reinforcement Learning by Planning with Learned Models. TMLR.\n\nZhan, X., Zhu, X., \\& Xu, H. (2022). Model-Based Offline Planning with Trajectory Pruning. IJCAI.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. This paper proposes four interesting masks for a bi-directional transformer, contributing to formulating offline RL as a sequence modeling problem. \n\n2. The presentation, with nice figures, is clear and informative.", "weaknesses": "This paper appears to have fewer novel contributions, and the experiments should be expanded. More details are provided in the Questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)