Bootstrapped Model Predictive Control
Authors: Yuhang Wang
Hanwei Guo
Sizhe Wang
Long Qian
Xuguang Lan
Venue: ICLR 2025 Poster (2025)
Keywords: reinforcement learning
model-based reinforcement learning
model predictive control
expert iteration
continous control
Primary area: reinforcement learning

Abstract: Model Predictive Control (MPC) has been demonstrated to be effective in continuous control tasks. When a world model and a value function are available, planning a sequence of actions ahead of time leads to a better policy. Existing methods typically obtain the value function and the corresponding policy in a model-free manner. However, we find that such an approach struggles with complex tasks, resulting in poor policy learning and inaccurate value estimation. To address this problem, we leverage the strengths of MPC itself. In this work, we introduce Bootstrapped Model Predictive Control (BMPC), a novel algorithm that performs policy learning in a bootstrapped manner. BMPC learns a network policy by imitating an MPC expert, and in turn, uses this policy to guide the MPC process. Combined with model-based TD-learning, our policy learning yields better value estimation and further boosts the efficiency of MPC. We also introduce a lazy reanalyze mechanism, which enables computationally efficient imitation learning. Our method achieves superior performance over prior works on diverse continuous control tasks. In particular, on challenging high-dimensional locomotion tasks, BMPC significantly improves data efficiency while also enhancing asymptotic performance and training stability, with comparable training time and smaller network sizes. Code is available at https://github.com/wertyuilife2/bmpc.

Published as a conference paper at ICLR 2025
BOOTSTRAPPED MODEL PREDICTIVE CONTROL
Yuhang Wang, Hanwei Guo, Sizhe Wang, Long Qian, Xuguang Lan" National Key Laboratory of Muchin- Hybrid Augmented Intelligence National Engincering Research Center for Visual Information and Application Institute of Artificial Intelligence and Robotics, Xi'an an Jiaotong University, Xi'an, China (wertyuilife, clannnagisa, wangs zhe, qianlongym]@stu xjtu. edu, cn. xglan@mail xjtu. edu. cn
3XXiv:2503.18871V/ [cs.LG] 3 Apr 2025 2025
ABSTRACT
Model Predictive Control (MPC) has been demonstrated to be effective in con- tinuous control tasks. When a world model and a value function are available, planning a sequence of actions ahead of time leads to better policy. Existing methods typically obtain the value function and the corresponding policy in a model-free manner. However, plex tasks, resulting in poor policy learning and inaccurate value estimation. To address this problem, we leverage the strengths of MPC itself. In this work, we introduce Bootstrapped Model Predictive Control (BMPC), a novel algorithm that performs policy learning in a bootstrapped manner. BMPC learns network pol- icy by imitating an MPC expert, and in turn, uses this policy to guide the MPC process. Combined with model-based TD-learnnng our policy learning yields bet- ter value estimation and further boosts the efficiency of MPC. We also introduce a lazy reanalyze mechanism, which enables computationally efficient imitation learning. Our method achieves superior performance over prior works on diverse continuous control tasks. In particular, on challenging high-dimensional loco- motion tasks, BMPC significantly improves data efficiency while also enhancing asymptotic performance and training stability, with comparable training time and smaller network sizes. Code is available at https///gihubb.oo///evyllllliiie
find that such an approach struggles with com-
INTRODUCTION
Model-basee reinforcement learning (MBRL) algorithms that incorporate online planning often referred to as plan-based methods- demonstrated superior performance and data efficiency across range of domains, including chess (Silver al. 2016 2017, Schrittwieser et al., 2020). games (Ye et al. 202T), continuous control (Sikchi et al., 2022 Hansen et al. (023). and the rea- soning of large language models (Zhao et al. [2024] Putta et al. 2024). By leveraging future states and rewards predicted by a world model, planning algorithms can evaluate actions online and with greater accuracy, resulting in a better and more robust policy. This represents a key advantage of basel planning, in contrast to model-free algorithms that learn neural network directly through trial and error.
In the field of continuous control, model predictive control (MPC) has proven to be effective planning approach (Lowrey et al.) [2018] Hafner et al., 2019b Hansen al. 2022 Schubert et al.) 2023). A notable example is TD-MPC2 (Hansen et al. 2023), y an MPC-based MBRL algorithm with a robust world model, which demonstrates strong performance across a diverse range of con- tinuous control tasks. Similar to existing methods (Bhardwaj al. 2020 Sikchi et al.) (022), TD-MPC2 learns a network policy and a value function in model-free manner. During inference, it uses policy-guides MPC for online planning, integrating the world model and the value function. However, our experiments reveal that despite having high-quality interaction data from MPC, the model-free policy learning struggles with challenging control tasks. The struggle in policy learning further indicates poor value learning, which can lead to inaccurate value estimation during MPC and degrade the overall performance of the MPC policy.
"Corresponding author.

---

Published as a conference paper at ICLR 2025
Avg.
Avg. high-dim
1007
1007-
MPC Policy
Policy
MPPI
World Model
Environment
Epicadi roward
002
Export imitation
202-
202-
min
Buffer
p(+))
Frwironment steps (x10
Environment steps (X]07)
Model-based TD-Learning
Lazy Reanalyze X(6e)
|MPP (MPC policy)
TD MP-MPC2 IMPC policy)
BMPC (network policy!
TO-MPC2 fnetwork policy)
Overview. (left) BMPC learns a network policy through expert imitation and lazy rean- alyze mechanism, planning during inference using guided MPC, and performs model-based value learning in an on-policy manner. (right) Averaged evaluation performance of the network policy compared to the MPC policy in BMPC and TD-MPC2 on DMControl tasks. BMPC achieves better policy learning, which further boosts the performance of MPC. Mean and 95% Cls over 5 seeds.
Figure 1:
To address this problem, inspired by expert iteration algorithms (Anthony et al. 2017 Silver et al., 2017]. we propose Bootstrapped Model Predictive Control (BMPC), which performs policy learn- ing in bootstrapped manner. We first execute MPC guided by the action sequences generated by network policy, yielding a bootstrapped MPC expert. The network policy then updated by imitat- ing this expert, thus achieving policy improvement. By employing this iterative process, we leverage the capabilities of MPC planning to boost the efficiency of policy learning. For value learning, we compute TD-tare ID-ttrgets online using the world model to mitigate off-policy issues. For imitation target generation, we adopt a lazy reanalyze mechanism to maintain an expert action dataset, thereby sup- porting expert imitation from the MPC policy in computationally efficient manner. An overview of BMPC is presented in Figure
Through experiments, we show that learning network policy through expert imitation can better leverage the strengths of MPC than learning a policy in model-free manner, thus leading to better value estimation and MPC performance. Our method, BMPC, achieves superior sample efficiency over prior data-efficient RL methods across 42 continuous control tasks in DMControl (Tassa et al., 2018) and HumanoidBench (Sferrazza et al.) 2024). with comparable training time and smaller net- work sizes. In particular, in challenging high-dimensional locomotion tasks, BMPC significantly improves data efficiency while also enhancing asymptotic performance and training stability. Ad- ditionally, our lazy reanalyze approach reduces the proportion of reanalyzed samples required by expert iteration algorithms (Ye et al. 2021 Wang et al., 2024) from 99% to 0.8%, while maintaining comparable policy learning efficiency. This avoids the need for extensive re-planning, significantly reducing the computational cost of BMPCC
2 RELATED WORK
Model-based reinforcement learning. MBRL focuses on using model of the environment to help an agent make decisions, which typically involves learning a dynamic model and a reward model from data. To elaborate further, dyna-style approaches (Sutton 1991; Janner et al. 2019, Hafnrr et al. 2019a Okada & Taniguchi, 2022, Robine et al.) 2023 Hafner et al. 2023a) use the model simulate additional experiences based on real data, which improves the sample efficiency of the algorithm. In contrast, plan-based methods leverage the model for planning, resulting in better policies and further enhancing the sample efficiency of reinforcement learning. In the case of plan- based MBRL, if the model and value function are sufficiently accurate, planning alone can lead to highly effective policy (Hafner et al. 2019b; Hansen et al. 2022 Schubert et al., (203)).
Expert iteration. Typically, planning algorithms need to roll out large number of trajectories using the model to explore the solution space, which can be computationally intensive. Expert iteration methods (Anthony al. 2017| Silver et al. (017)) improve the efficiency of planning by employing a network policy to guide the search direction. The planning algorithm, guided by the network policy, can be considered an expert, allowing the network policy to learn from it and thus achieve
2

---

Published as a conference paper at ICLR 2025
policy improvement. By combining these two aspects, expert iteration can bootstrap the efficiency of planning and the capability of the network policy. Although tree-search-busee expert iteration has achieved strong performance and data efficiency across range of domains (Wang et al. 2024),
in the domain of continuous control, pure MPC-based MBRL. without expert iteration remains the superior approach (Hansen et al. 2023).
Imitating and enhancing MPC. Methods that aim to mimic or enhance MPC have gained signif-
icant attention in RL and robotics. Guided policy search methods (Levine & Koltun 2013 Levine et al. 2016 Zhang et al., 2016 Sun et al. 2018) help RL policies explore effectively in challenging tasks by using an external plan-based policy, such as MPC, to guide the learning of neural net- & Boots
work policy. Pan et al. (2017) (010)
et al. (2023) propose methods
where a neural network learns by mimicking an MPC controller, either to achieve faster policy or to develop more efficient neural planner. Alternatively, there are approaches that attempt to learn residual policy on top of MPC (Silver et 2208820188 [Sacks et al.) 2024). which serves as another way of bootstrapping MPC. Additionally, Power& & Berenson (2022); Sacks & Boots (0233) aim to improve MPC via learned sampling distributions. Unlike MBRL, the MPC policy in the aforementioned methods is derived from a real, carefully designed model, rather than learned model. Moreover, these approaches do not adopt the iterative policy optimization scheme like expert iteration.
BACKGROUND
Problem formulation. We address reinforcement learning problems in continuous action spaces, which are modeled as an infinite-horizon Markov Decision Process (MDP). This MDP is defined by the tuple A,,A, P. R, 7) where S R and A Rm state and action spaces, s E S are states, a E A are actions, P S x A + S is the state transition function, R S x A + IR the reward function, and 7 is the discount factor. The objective in reinforcement learning is to derive a policy π S + A that maximizes the expected discounted cumulative reward, expressed as E, EE= r're), where r;- = R(s,, π (s:)).
TD-MPC2. TD-MPC2 (Hansen et al. 2023) is a plan-based MBRL algorithm that learns a world model, a Q-function, and a corresponding policy, which are then used for MPC to derive an plan- based policy π. The model components of TD-MPC2 can be described by a tuple (h, d. R, p), where z = h((s) is the encoder that maps the observation s into a latent space vector z, z' dd(z. a) is the latent-space dynamics model, i = R(z, a) is the reward prediction function, q = Q(z, a) is the Q-value prediction function, and à = p(z) is the prior neural network policy. In this paper, we omit the representation of task embedding inputs of TD-MPC2 as we do not focus on its multi-task capabilities. Similar to model-free approaches like SAC (Haamoja et al. 2018). TD-MPC2 learns the Q-function through iterations of the Bellman equation, and the network policy p is optimized by maximizing the Q-value with entropy regularization, which can be formalized as the following update rules:
Φ + arg min (s..)))) [CE(Qo a), + (z'.po((((
(1)
0 < arg max E (o + 3H (-(pp (-12))] z = h(s), z' = h(ss)
(2)
where H is the entropy p. p. 33 is hyperparameter for loss balancing, 0, Φb, 6- denote the parameters of the neural networks for p. Q, and the target Q-network, respectively. B represents the replay buffer. CE is the cross-entropy, used because TD-MPC2 formulates value prediction as a discrete regression problem. For simplicity, we omit the temporal expansion of the latent vector in update rules for TD-MPC2; for the full temporally weighted objectives, see Hansen et al. 2023).
MPC with a policy prior. During inference, TD-MPC2 performs MPC planning guided by the prior policy pe. Specifically. TD-MPC2 leverages Model Predictive Path Integral (MPPI) (Williams et al.) 2015) as its underlying MPC algorithm. MPPI models the action sequence (a,, as+1, ..., a++ of length H as being drawn from a time-dependent multivariate Gaussian with diagonal covariance, parameterized as (n, a). where μ4, E fm During the planning process, MPPI iteratively sam- ples sequences from (H, o²), estimates their values by rolling out trajectories with the model, and updates (μ,, o)) based on a weighted average of the top-k sequences, thus maximizing the expected estimated value of action sequence, which is expressed as:
E 1.0 aut (1, [Q(ze. -t)
11 0" = arg max
(3)
3

---

Published as a conference paper at ICLR 2025
Dog
Run
1000
500
500
800
400
400
(p 1WA0
600
200
Actinn voose offernteeeuur qurc pwwcyeeppeootoeieee neteurk pulcy)
0
0.5
2.0
1.5
2.0
05
1.0
15
2.0
0.5
1.0
1.5
(x1.6))
Enwiromment I=30%)
Envionment (K10")
policy)
Inebwork policy!
Walker Run
Humanold Run
Dog Run
Figure 2: Performance gap between TD-MPC2 policies. (left) Evaluation performance of the net- work policy compared to the MPC policy in TD-MPC2. The network policy struggles with complex tasks like Dog Run and Humanoid Run. Mean and 95% CIs over 5 seeds. (right) Distributions of action value differences during MPPI over environment steps.
H Σ
Q ==)+))==**** 1(( ( a.+
7 R(Zzth ax+h)
(4)
To integrate prior policy guidance, a portion of the sampled sequences drawn from the distribution pg. Ultimately, the first action the optimized distribution a, VVu is selected for execution
during inference. For more details on the planning procedure, see Hansen et al. (2022).
In this paper, adopt the same world model architecture and MPPI algorithm as TD-MPC2, and we use similar mathematical notations to describe our algorithm components.
METHOD
4.1 INSUFFICIENT POLICY LEARNING IN MODEL-PREES APPROACH
model-free approach is adopted TD-MPC2 learn both the Q-function and max-Q network policy. During inference, it leverages MPC to plan actions, guided by the network policy. Since the MPC policy-based based on the learned model and online planning-typically outperforms the network policy, it provides higher-quality samples for model-free learning, thereby improving the efficiency of both policy and value learning. However, we find that in challenging environments, even with high-quality samples from the MPC policy, the model-free approach still struggles. leading performance gap between the MPC policy and the network policy. This gap indicates inaccurate value estimation, further degrading the planning performance of MPC.
Figure shows the evaluation performance of both the network policy and MPC policy during training. three locomotion tasks in DMControl Tassa al.) 2018]: Walker Run, Humanoid Run, and Dog Run. In all three tasks, a policy performance gap evident, though the extent of the gap varies. In simpler task Walker Run, the network policy performs comparably the MPC policy. However, in more complex tasks like Humanoid Run and Dog Run, the network policy struggles to improve, while the MPC policy maintains high performance.
To better understand the root cause of this performance gap, we analyze the planning process of t. MPPI. We represent the prior action sequences generated by the network policy as { H Tim1> where n is the number of prior action sequences, and the final action sequence selected by MPPI appe
as a To quantify the optimization achieved by MPPI, we compute the difference between the value the final selected sequence and the average value of the prior sequences:
1 72
AQ(Z:) Q(zi, a
Σ Q(z,, a n
(5)
where Q denotes the value estimated by the model in Equation 4 Overall, AQ(Z,) reflects the improvements that MPC makes to action value time-step t. As illustrated in Figure 2 plot the distributions of AQ over training steps within an episode. The figure shows that improvements
4

---

Published as a conference paper at ICLR 2025
brought by MPC continuously increasing in Humanoid Run and Dog Run, whereas in Walker Run, the network policy already performs well, leaving little room for MPC to further optimize.
These experiments demonstrate that, although TD-MPC2 exhibits strong performance in high- dimensional locomotion tasks, much of its performance is due to the use of online planning. In contrast, network policy learning struggles considerably with these tasks. Consequently, that this inefficient policy learning indicates poor value learning, reducing the data efficiency of the algorithm. Specifically, because the value network is learned represent the value of the corre- sponding network policy, this gap implies that the terminal value used by MPC is derived from a weaker policy, which will result in inaceurate value estimation during planning. natural solution, which we propose, is to leverage the imitation of the MPC expert to achieve policy learning.
argue
4.2 BOOTSTRAPPED MODEL PREDICTIVE CONTROL
We propose BMPC, a plan-based MBRL algorithm based on TD-MPC2's world model. BMPC learns neural network policy by imitating an MPC expert, and in turn, uses this policy guide the MPC process. The world model is leveraged in BMPC to perform on-policy TD-leaming of a value network, which is used for terminal value calculation during MPC. Additionally, we introduce lazy reanalyze mechanism to maintain expert dataset for more computationally efficient imitation learning. The algorithm for BMPC training is presented in Algorithm
Policy learning through expert imitation. BMPC uses prior policy po to guide the planning process of MPC. other words, we can describe this as MPC bootstrapping Po into an expert policy, denoted as T( 2, po). Thus, we learn the policy po by imitating which can be formalized as the following objective:
H1 x
(0)
E (8.a)o.
(3 |z:)]]
(6)
0=
Zg = h((so:: Zet1 = d(zz, a.), S EMA(Per(KL(T,,p) 95) - (T) 5), 0.99)
where H is the entropy, KT. is the Kullback-Leibler divergence, 20: are latent vectors rolled out through model h and d. 33 and \ are hyperparameters for loss balancing and temporal weighting, respectively. Empirically, when action space is large, imitating the action distribution æ is more ef- ficient than imitating the exact actions x, especially when both the student and expert policies' distributions belong to the same parametric family. Since the MPC policy is parameterized as a mul- tivariate Gaussian, we choose to parameterize the neural network policy as well, allowing us to compute the KL divergence in closed form. As the KL divergence between multivariate Gaussian distributions can vary significantly across tasks and action spaces, affecting training stability, we normalize the KL loss using moving percentiles S to keep the loss value within an acceptable range. This method also commonly used balance the policy loss and entropy loss (Hanner et al.) 2023a Hansen et al. 2023).
multivariate Gaussian
Model-based TD-learning Since we do not employ SAC-style max-Q approach for policy im- provement, we opt to learn state value function Vg instead of a state-action value function Qo- We construct n-step TDDttarett V using the latest model, policy, and target value network. The value network learns to minimize the cross-entropy loss with respect to the discretized TD-target:
H
(v(0)= =
E
Σ
x (GB(V)(2).VV
20 -h((s): 21-1 1 = d(zr, ax)
(5.0)a.H B
v(2)) = N (z(4N)
* *Riz'a Po
1441 d(zi,,P((22)
where N is the TD horizon, Zo.H are latent vectors rolled out through model h and d. V is the TD-target computed using the model d, R and the policy pe in an on-policy manner. In practice, we found that N is more suitable choice, likely because the world model of TD-MPC2 is trained with a short horizon (H = 3, A = 0.5), limiting its ability to predict rewards over long sequences. As result, setting N too large would lead to excessive compounding errors.
5

---

Published as a conference paper at ICLR 2025
Indeed, using the original Q-iteration method for value learning is also feasible choice. However, due to the changes in the policy learning approach, this option introduces certain off-po issues, which can lead to lower data efficiency and unstable training in tasks where the policy varies sig- nificantly during training. We compare the results of the two value learning approaches Section 5.1
Lazy reanalyze. In practice, it is costly to compute the policy objective direetly as it requires re-planning for all samples during every update, which is infeasible for MPC algorithms. Instead, we choose to maintain imitation targets in the replay buffer through lazy reanalyze, thus resulting in a surrogate policy objective 8
During every k-th network update, we draw b samples from the batch, re-plan them, and obtain fresh imitation targets, i.e., the mean and standard deviation of the action distribu- tion ne = T(- ((:)h(s)), pa). These targets are then placed back into the replay buffer. This re- analyzing process is performed independently of the training process. Thus, we can approx-
Algorithm BMPC training
Require:
Initialize pe, Vg. h, d, R randomly.
B, k: replay buffer, lazy reanalyze interval. while not converged do
1:
2:
// CCCllctt experience
3:
for step O...T do
4:
Tt x x-I(h(s:,, at Tτ.
imately regard the replay buffer as expert
5:
dataset, and directly sample state-action pairs from it for supervised learning. To increase ex- ploration in MPC planning, we additionally add noise the prior policy during re-planning.
6:
(s.+1. + env env.step(a.)
7:
B BU (st, a re, St41, Te)
8:
// Update networks
9:
for num updates per episode do
This approach is inspired by the reanalyze pro- posed in [Schrittwieser et al. (2020). common
10:
{st, ar, Ft. S1+1, We ha++ ~ B
11:
Update h, d, as in TD-MPC2.
method in sample-efficien expert iteration al-
12:
Update Vo via Equation 7
gorithms (Ye et al., 2021 Wang al.) 2024).
13:
Update pe via Equation 8
The key difference is that reanalyze performs
14:
Update Vo- via EMA.
re-planning during every network update, with
15:
Lazy reanalyze if update step % k == then
99%-100% of samples re-planned (i.e., reana- 16:
lyze ratio). These reanalyzed samples are im- mediately used for imitation learning and then 18:
17:
TEt+H <- T-110(5a4+111)))))) update
B
Te:++
discarded. In contrast, lazy reanalyze performs far fewer re-plannings and places the reanalyzed samples back into the buffer for reuse.
Specifically, the reanalyze interval l and the reanalyze batch size b are hyperparameters, where choose k 10 and b 20, with batch size of 256 for network updates. Under this setup, lazy reanalyze achieves a computational cost equivalent to reanalyze ratio of 0.8%, which is over 100 times lower than the typical reanalyze ratio of 99% (Ye et al. 2021 Wang al. 2024). When combined with batched MPPI planning on GPU, lazy reanalyze introduces only a 10% -00 increase in training wall-time.
The surrogate policy objective with lazy reanalyze can be formalized as:
C Clazy
E
(8)
(s,
where me is the expert action distribution maintain in the replay buffer.
EXPERIMENTS
In this paper, propose BMPC to more effectively leverage the strengths of MPC in continuous control tasks. Our approach integrates expert imitation for policy learning, performs model-based TD-learning for value learning, and introduces lazy reanalyze to better utilize re-planning results. Through our experiments, we aim to answer the following key questions:
How does BMPC perform as data-efficient continuous control algorithm compared to the current ofate-of-the-ar methods?
Does BMPC lead better policy learning, and how this be further leveraged?

---

Published as a conference paper at ICLR 2025
How does lazy reanalyze affect the performance and training time of BMPC?
To ensure direct comparison, we evaluate BMPC on 28 DMControl TTassa et al. 2018) tasks used in the TD-MPC2 (Hansen al., 2023 | and 14 tasks from HumanoidBench (Sterrazza et al. 2024) locomotion suite. The tasks covers a diverse range of continuous control challenges, including sparse reward, locomotion with high-dimensional state and action space (up to A E I. Visualizations of the tasks can be found in Appendix D To avoid ambiguity, we use the term "environment step" in our experiments, where an environment step refers to the number of inference steps multiplied by the action repeat (2 in DMControl, 1 in HumanoidBench).
Baselines. We select three state-of-ther trttthe-arttt data-efficient RL methods as baselines: (1) SAC (Haarnoja al. 2018), model-free actor-critic algorithm rooted in maximum entropy reinforce- ment learning; (2) DreamerV3 (Hafner et al., 2023a), a Dyna-style MBRI. algorithm that trains a model through reconstruction loss and learns a model-free policy from trajectories imagined by the model; (3) TD-MPC2 (Hansen et al., 2023). an MPC-based MBRL algorithm that learns a network policy and Q-function in model-free manner, and performs MPC planning at inference based on them and the model. Both MBRL methods learn an implicit model to roll out sequences in latent space and predict rewards.
For SAC, we use the results from TD-MPC2 and HumanoidBench; For DreamerV3, we use its default settings on DMControl tasks, corresponding to a network size of 12M, and use the results from Humanoidßench:; For TD-MPC2, we use its default configuration, corresponding to network with 5M parameters. For BMPC, we adopt a network almost identical to TD-MPC2, except that we replace the Q-network with V-network. Notably, as we do not use Q-based method, we find that BMPC's performance is less dependent on ensemble networks, allowing us to reduce the default 5 ensemble value networks to 2. This results in smaller network for BMPC, with only 3M parameters. We use the same hyperparameters for BMPC across all tasks, see Table 2] and detailed baseline configurations are provided in Appendix B
5.1 RESULTS
Benchmark performance. We first compare BMPC against the baselines across all 28 DMControl tasks. Due to limited space, we only present results for 10 selected tasks, as shown in Figure 3 The training curves for all tasks are provided in Appendix C Our results show that BMPC consistently achieves either superior or comparable performance relative to the baselines on most tasks. Notably, on -ichdimennsioall locomotion tasks, such as Dog and Humanoid, BMPC shows significant im- provements in data efficiency, despite having fewer learnable parameters. The results on DMControl indicate that BMPC maintains the performance of TD-MPC2 across wide range of control tasks, while significantly enhancing performance in tasks where model- approach struggles.
We further compare BMPC with the baselines on the high-dimensional tasks, as shown in Figure The environment steps are extended from 1M to 4M for a comprehensive comparison. In addition to improved data efficiency at short training lengths, BMPC also outperforms the baselines in terms of asymptotic performance and training stability, as indicated by the confidence intervals of the curves. Finally, the average dieps-to-solve (the number of steps required to achieve 795, which is 90% of the asymptotic reward) across the 7 tasks for BMPC is 90k steps, while for TD-MPC2 it is 360k steps, indicating a 300% increase data efficiency.
On the HumanoidBench locomotion suite, which requires the agent to control more complex embodiment-a a Unitree robot with a large action space (AA E 1 BMPC maintains supe- rior performance compared baselines, further demonstrating its advantage on challenging high- dimensional tasks. We present results for all tasks in Figure
Leveraging better policy learning. In BMPC, we obtain network policy through expert imitation. To verify whether this network policy is better and how we can leverage it, we design ablation experiments as shown in Figure 6 and Figure 7a
Figure 6 illustrates the performance gap between the network policy and MPC policy for BMPC and TD-MPC2 DMControl, similar to Figure Training curves for all tasks are provided in Appendix Surprisingly, through expert imitation, BMPC enables its network policy to perform nearly on
Excluding custom tasks created for multitask training.

---

Published as a conference paper at ICLR 2025
Avg. 28 tasks
high-dim tasks
Cheetah Run
Cartpole Swingup Sparse
600-
400
200
Deg Stand
Dog Trol
Finger Turn Hard
Hopper Sland
400-
200
Humanoid Stand
Humanoid Walk
Quadruped Walk
Walker Stand
1000
800
Episode
200
0
0/6
0.6
0/6
02
1.3
Environment stops (x10*) Environment (x10") Environment stops (x110))
Environment stops (tops (x100")
BMPC (ours)
TD-MPC2
SAC
DreamerV3
Figure 3: DMControl tasks. Comparing BMPC to baselines on DMControl tasks. In the top left, present the average performance of 7 high-dimensional locomotion tasks and all 28 tasks. Mean and 95% Cls over 5 seeds? Training curves for all tasks are provided in Appendix
Average
Dog Stand
Dog Trat
Dog Walk
1000
Episode neward
800
800
M
200-
Dog Run
Humanold Stand
Humanoid Walk
Humanold Run
8d0
Episode
400-
200
1.0
2.0
3.0
1.0
2.0
3.0
4.0
1.0
2.0
3.0
4.0
1.0
2.0
3.0
4.0
Environment steps (X109)
Environment steps (x10%)
Environment steps (x109) Environment steps Coppsxx10 1x10
BMPC (ours)
TD-MPC2
DreamerV3
Figure 4: High-dimensional locomotion tasks. Comparison of BMPC with baselines on the most challenging high-dimensional locomotion tasks; The environment steps are extended 4M for a comprehensive comparison. In the top left, we present the results averaged over all 7 tasks. Mean and 95% Cls over seeds.
par with its MPC policy. In contrast, TD-MPC2 exhibits a substantial performance gap between its network policy and MPC policy, particularly in challenging tasks such Dog and Humanoid. This indicates that BMPC's network policy can serve as a viable final inference strategy without the need for online planning, which suitable for real-time control tasks that demand low latency.
We further conduct ablation studies to show how this improved network policy contributes to BMPC's performance. We introduce three BMPC variants for comparison: (1) Variant /: based on TD-MPC2, we use expert imitation to additionally learn a network policy, which used to guide MPPI planning, while value learning still relies the original policy; (2) Variant based Vari- ant 1, we use both two network policies to guide MPPI planning simultaneously; (3) Variant 3: we replace TD-MPC2's policy learning approach with expert imitation. thus changing the policy used in value learning, while still learning value based on Q-iteration; (4) BMPC: based on Variant 3, we adopt model-based on-policy TD-learning. For further details of these variants, see Appendix
Except SAC, which seeds.
8

---

Published as a conference paper at ICLR 2025
Avg. 13 tasks
Balance Hard
Balance Simple
Crawl
Hurdle
1000
1000
1003
1000
1000
800
800
800
600
Episode
800
400-
402-
400-
400-
200
200-
200
200-
Pole
Reach
Hard
800
800
11200
800
800
600
8400
600
600-
200
200
2000
200
200
Sit Simple
Slide
Stair
Stand
Walk
1000
1000
1000
1000
1.000
return
800-
600
600
600
600
600
1.0
400
400
400
0.5
1.0
15
2.0
0.5
16
1/5
0.5
1.0
1.5
2.0
0.5
3.0
1.5
2
0.5
1/5
20
Environment, steps (x110))
Environment steps (x10%)
Environment steps (x10f)
Environment steps (x110))
Environment steps (x110%
BMPC (ours)
TD-MPC2
DreamerV3
SAC
Figure 5: HumanoidBench locomotion suite. Comparing BMPC to baselines on HumanoidBench locomotion suite. In the top left, we present the average performance of all 13 tasks except for Reach due to the different reward scales. Mean and 95% Cls over seeds.
Waller Run
Deg Run
Hurraneid Run
1000
800
700
800
560
Ray
Actoe deturecceseeie padcy nehause Doblyy)
70g
styal
05
15
Envtronment: steos (X101)
Enwroomenn stepe 1-30%)
Envfronmmnn steas 1x10%
policy)
policy)
BMPC Instwork policy)
TD MPC2 (network policy)
Walker
Hum Humanold Run
Figure 6: Performance gap between BMPC policies. (left) Evaluation performance of the network policy compared the MPC policy in BMPC and TD-MPC2, The performance gap of BMPC is barely noticeable. Mean and 95% CIs over 5 seeds. Curves for all tasks are provided in Appendix (right) Distributions of action value differences during MPPI over environment steps for BMPC.
Figure Seguree7|showwe shlss the results of BMPC variants on the Dog Run and Humanoid Run. We find that even with better policy, using it guide MPC does result in improved performance, shown by the results of Variant / and Variant 2. The key to performance improvement lies in using this policy learn better value function, as indicated by the results of Variant 3 and BMPC. Additionally, based value learning avoids off- off-policy issues, leading to improved data efficiency.
Lazy reanalyze ablation. BMPC maintains an expert dataset through lazy reanalyze achieve computationally efficient expert imitation. We explore the relationship between the frequency of re- planning and performance by evaluating lazy reanalyze interval i E (10, 40, 80, oo), corresponding to reanalyze ratios of 0.8%, 0.2%, 0.1%, and 0%, respectively. The results, shown in Figure 7bby indicate that the reanalyze interval considerably affects BMPC's performance. However, as decreases, the impact becomes less pronounced. Beyond = 10, further increasing the reanalyze frequency does not significantly improve performance but results in substantial computational over- head. Thus, set k 10 the default value. It is worth noting that in all our experiments, the replay buffer size 1M, which is quite large for an expert dataset, but even with low- frequency reanalyzing, the freshness of the data is sufficient to support expert imitation.

---

Published as a conference paper at ICLR 2025
Dog Run
Humanoid Run
Dog Run
Humanoid Run
630
303
530
Epooode towarrs
403
430
300
320
200
230
100
100
022
04
0.5
0.5
20
0.2
08
08
1.0
022
044
0.6
09
10
0.2
0.4
0.6
10
Environment steps (X10f)
Environment staps (x10f)
Environment steps (x105)
Environment steps (X10f)
Veriant
Variant
km10
Km40
km80
km
Performance of BMPC variants
(b) Lazy reanalyze ablation
Figure 7: Ablations. (a) Performance of BMPC variants with different ways to leverage the network policy, demonstrating the importance value learning, (b) Performance of BMPC with different lazy reanalyze interval k; Mean and 95% CIs over seeds.
Training wall-time. We compare the training wall-time and -ime-to-oolve of BMPC and TD-MPC2 in Walker Walk and Dog Walk, where time-to-solve is defined as the time required to achieve rewards of 899 and 872, respectively (90% of the asymptotic reward), as shown in Table 1 The experiments are conducted using single RTX 3090 GPU. BMPC and TD-MPC2 have similar training times per 500k steps, although lazy reanalyze increases training time by approximately 20%, this increase is offset by the reduced size of the network. Due to its high data efficiency, the time-to-solve BMPC is 2 times shorter than TD-MPC2 on Dog Walk.
Notably, our BMPC implementation uses only single thread for training and does not fully parallelize the lazy reanalyze process. Since lazy reanalyze operates independently of net- work training, it could be parallelized using separate thread further reduce training time. Additionally, BMPC's network policy performs nearly on par with MPC planning on DMCon- trol tasks, making it feasible to use the network policy for inference without planning, which would significantly reduce inference time,
Table 1: Wall-time. Time-to-solee and time per 500k environment steps for the Walker Walk and Dog Walk. Mean of runs.
Walker Walk
Dog Walk
Woll-imee())
TD-MPC2
BMPC
TD-MPC2
BMPC
time-to-solve
0.40
0.43
2.03
0.87
Iv500k steps
7.32
8.47
8.71
CONCLUSIONS AND FUTURE DIRECTIONS
In this paper, we introduce BMPC, plan-based MBRL algorithm that leverages the strong perfor- mance of MPC achieve better policy and value learning in continuous control tasks. Our approach demonstrates superior performance, particularly in complex locomotion tasks, while maintaining comparable training time and a smaller network size compared state-of-the-arr methods.
Due to the current world model's lack of long-horizon prediction capability, BMPC is limited to a TD-horizon 1. As the model capability improves, there is potential to extend the TD-horizon for enhanced value learning. Moreover, BMPCC can further benefit from advanced expert iteration techniques, such as the plan-based value estimation proposed in Wang et al. (2024).
It is also worth exploring the combination of expert imitation with the max- Q gradient for joint pol- icy improvement, integrating the strengths of both approaches, such as by combining loss functions or integrating independent value functions for planning.
Finally, BMPC can also be applied in multi-task and offline settings. Although experiments in these areas have yet to be conducted, we believe BMPC can further harness the capabilities of MPC.
REPRODUCIBILITY STATEMENT
We have open sourced our work, available at httpps//github..om. www.b//oo/wwtttllee///mmmm By running the code with the default configuration, the results presented in this paper can be reproduced.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
This work was supported in part by NSFC under grant No.62125305, No.U23A20339, No.62088106 I02, No.62203348.
REFERENCES
Thomas Anthony, Zheng Tian, and David Barber. Thinking fast and slow with deep learning and tree search. Advances in neural information processing systems, 30, 2017.
Mohak Bhardwaj, Sanjiban Choudhury, and Byron Boots. Blending mpe & value function approxi- mation for efficient reinforcement learning. arXiv preprint arXiv:2012.05909 2020.
Adam Fishman, Adithyavairavan Murali, Clemens Eppner, Bryan Peele, Byron Boots, and Dieter Fox. Motion policy networks. In Conference on Robot Learning, pp. 967-977. PMLR, 2023.
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 arXiv: 191.0199,,.01633 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pp. 2555-2565 PMLR, 2019b.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104:: 2023a.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.041044 2023b.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203..m 049555 2022.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpe2: Scalable, robust world models for contin- uous control. arXiv preprint arXiv:2310.1680 16828, 2023.
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.
Sergey Levine and Vladlen Koltun. Guided policy search. International conference on machine learning, pp. 1-9. PMLR, 2013.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuo- motor policies. Journal of Machine Learning Research, 17(39):1-40, 2016.
Kendall Lowrey, Aravind Rajeswaran, Sham Kakade, Emanuel Todorov, and Igor Mordatch. Plan online, learn offine: Efficient learning and exploration via model-based control. arXiv preprint arXiv: 18/ 1.01848, 2018.
Masashi Okada and Tadahiro Taniguchi. Dreamingv2: Reinforcement learning with discrete world models without reconstruction. In 2022 IEEF/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 985-9911 IEEE, 2022.
Yunpeng Pan, Ching-An Cheng, Kamil Saigol, Keuntaek Lee, Xinyan Yan, Evangelos Theodorou, and Byron Boots. Agile autonomous driving using end-to-end deep imitation learning. arXiv preprint arXiv:1709.07174 2017.
Yunpeng Pan, Ching-An Cheng, Kamil Saigol, Keuntack Lee, Xinyan Yan, Evangelos A Theodorou, and Byron Boots. Imitation learning for agile autonomous driving. The International Journal of Robotics Research, 39(2-3):286-303 2020.
11

---

Published as a conference paper at ICLR 2025
Thomas Power and Dmitry Berenson. Variational inference mpe using normalizing flows and out- of-distribution projection. arXiv preprint arXiv:220S. 023.00..466..22 2022.
Pranav Putta, Edmund Mills, Naman Garg, Sumeet Motwani, Chelsea Finn, Divyansh Garg, and Rafael Rafailov. Agent q: Advanced reasoning and learning for autonomous ai agents. arXiv preprint arXiv:2408.077 199, 2024.
Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world mod- els are happy with 100k interactions. arXiv preprint arXiv:2303.07109 2023.
Jacob Sacks and Byron Boots. Learning to optimize in model predictive control. In 2022 Interna- tional Conference on Robotics and Automation (ICRA), pp. 10549-10556. IEEE, 2022.
Jacob Sacks and Byron Boots. Learning sampling distributions for model predictive control. In Conference on Robot Learning, pp. 1733-1742. PMLR, 2023.
Jacob Sacks, Rwik Rana, Kevin Huang. Alex Spitzer, Guanya Shi, and Byron Boots. Deep model predictive optimization. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 16953. IEEE, 2024.
Julian Schritwieser,, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-009 2020.
Ingmar Schubert, Jingwei Zhang, Jake Bruce, Sarah Bechtle, Emilio Parisotto, Martin Riedmiller, Jost Tobias Springenberg, Arunkumar Byravan, Leonard Hasenclever, and Nicolas Heess. A generalist dynamics model for control. arXiv preprint ar arvivv2300. 10912005.0010,, 2023.
Carmelo Sferrazza, Dun-Ming Huang, Xingyu Lin, Youngwoon Lee, and Pieter Abbeel. Humanoid- bench: Simulated humanoid benchmark for whole-body locomotion and manipulation. arXiv preprint arXiv:2403. 2024.
Harshit Sikchi, Wenxuan Zhou, and David Held. Learning off-policy with online planning. In Conference on Robot Learning, pp. 1622-1633, PMLR, 2022.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489, 2016.
David Silver, Julian Schritwieser,, Karen Simonyan, loannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al. Mastering the game of go without human knowledge. nature, 550(7676):354-35) 2017.
Tom Silver, Kelsey Allen, Josh Tenenbaum, and Leslie Kaelbling, Residual policy learning. arXiv preprint /8/2:v88//200633 2018
Liting Sun, Cheng Peng, Wei Zhan, and Masayoshi Tomizuka. A fast integrated planning and control framework for autonomous driving via imitation learning. In Dynamic Systems and Con- trol Conference, volume 51913, pp. V003T37A012. American Society of Mechanical Engineers, 2018.
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160-163, 1991.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.00640 2018
Shengjie Wang, Shaohuai Liu, Weirui Ye, Jiacheng You, and Yang Gao. Efficientzero v2: Mastering discrete and continuous control with limited data. arXiv preprint arXiv:2403.00564 2024.
Grady Williams, Andrew Aldrich, and Evangelos Theodorou. Model predictive path integral control using covariance variable importance sampling. arXiv preprint urXiv:/509.000 149, 2015.
12

---

Published as a conference paper at ICLR 2025
Weirui Ye, Shaohuai Liu, Thanard Kurutach, Picter Abbeel, and Yang Gao. Mastering atari games with limited data. Advances in neural information processing systems, 34:25476-25444 2021.
Tianhao Zhang, Gregory Kahn, Sergey Levine, and Pieter Abbeel. Learning deep control policies for autonomous aerial vehicles with mpc-guided policy search. In 2016 IEEE international con- ference on robotics and aulomation (ICRA), pp. 528-535. IEEE, 2016.
Zirui Zhao, Wee Sun Lee, and David Hsu. Large language models commonsense knowledge for large-scale task planning. Advances in Neural Information Processing Systems, 36, 2024.
A IMPLEMENTATION DETAILS
Exploration in lazy reanalyze. The network policy of BMPC uses neural network to obtain the mean and log standard deviation Gaussian distribution. The mean is derived by applying the tanh function to squash the output. The log standard deviation is first squashed the range [-1,1] using the tanh function, and then mapped linearly to the range [log std min, log std max).
During lazy reanalyze, Specifically, the default values are log_std iog_st__min min_mmin---- and log dog_std_mm 1. When reanalyzing, log_std_min and log_std ma Mathematically, this adjustment is equivalent to log_stdreanalyse log_std x 0.75 + 0.25.
increase the value of log std min to enhance exploration in MPC.
This mechanism does increase overall performance but helps prevent BMPC from prematurely converging to local optima due to iterative policy learning. larger policy std during lazy reanalyze may further improve the exploration of MPC, but this would likely require increasing the number of MPPI iterations to ensure convergence, which we opt not to do.
Policy loss. In our experiments, we try both log probability loss in Wang et al. (2024) and KL loss for expert imitation, observing substantial differences in the policy std. We find that network policy accurately mimics the std of the MPC policy through KL loss: while using log probability loss tends to result in large policy std, affecting the optimality of the algorithm. We speculate this may be due to the use of lazy reanalyze, leading the expert action dataset to encompass behaviors of different MPC policies over longer period, necessitating larger std for the network policy. Alternatively, it may result from the MPC policy's s tendency to produce "shaky actions") leading to larger policy std when using log probability loss.
Hyperparameters. We use the same hyperparameters for all tasks. BMPC is based on the world model and MPPI of TD-MPC2. For the hyperparameters in these components, default values as those in TD-MPC2 (Hansen et al., (023). The hyperparameters of BMPC are detailed in Table 2
the same
"see hww.fh/w/ggtticcooooocccccooeene
13

---

Published as a conference paper at ICLR 2025
Table 2: BMPC Hyperparameters. We use the same hyperparameters for all tasks.
Hyperparameter
Value
Entropy loss coef.
x 10 4
Batch size
256
TD horizon (N)
Number ensemble value networks
Lazy reanalyze interval (k)
Lazy reanalyze batch size (b)
20
Planning horizon
Re-planning horizon
Policy log std. min.
-3
Policy log std. max.
Policy log std. min. (re-planning)
-2
Policy log std. max. (re-planning)
14

---

Published as a conference paper at ICLR 2025
BMPC variants. We conduct ablation studies by introducing three BMPC variants in the Experi- ments section. Below, we provide a detailed explanation these variants along with their respective design choices.
Key Facts: The design of BMPC and its variants guided by the following components:
A network policy can be learned using:
(1a) Max-Q gradient-based learning;
(1b) Expert imitation.
A network policy is used for:
(2a) Computing the TD-target during value learning;
(2b) Guiding the planning process of MPPI.
The value function can be learned using:
(3a) Off Offl-o TD-learning (Q QQiitrationn;
(3b) On-policy model-based TD-learning.
BMPC Variants: We define the following three BMPC variants based on different combinations of the above components:
Variant 1:
Learn network policy A using (1a) and network policy B using (1b).
Use policy A for (2a) and policy B for (2b).
Learn the value function using (3a).
Variant 2:
Leern network policy A using (1a) and network policy B using (1b).
-se policy A for (2a) and both policies A and B for (2b).
Learn the value function using (3a).
Variant 3:
Learn network policy A using (1b).
Use policy A for both (2a) and (2b). Learn the value function using (3a).
BMPC: For comparison, our proposed BMPC approach is defined follows:
Learn network policy A using (lb).
Use policy A for both (2a) and (2b).
Learn the value function using (3b).
Additional Remarks: For all variants, we use 5 ensemble Q-networks instead of 2 as in BMPC. For Variant 2, we generate 24 guiding trajectories from each policy (A and B), and increase the number of clite trajectories in the MPPI process from 64 88.
15

---

Published as a conference paper at ICLR 2025
B BASELINES DETAILS
We report our implementation details of baselines this section.
SAC. We use the results from TD-MPC2 and HumanoidBench for DMControl and Humanoid- Bench, respectively. For the hyperparameters of SAC, see Hansen et al. (2023) and Sferrazza et al. 2024).
DreamerV3. For DMControl, we use the latest implementation" 0 referencing Hafner et al. (0023a) which differs from the older version Hafner proprio tasks, corresponding to a network size of 12M and a UTD ratio of 512. We discover that the performance of DreamerV3 diverges from what is reported in Hansen et al. (202)) For example, the performance on the Fish Swim has improved, while the performance on the Walker Run has decreased. This is likely because the newer version of DreamerV3 changes the policy learning approach in continuous action space, from stochastic backpropagation to the Reinforce. However, since both results underperform relative to TD-MPC2, this does not affect the overall comparison. For comprehensive list of hyperparameters, please refer the original paper (Hafner et al. 2023a) For HumanoidBench, we use the results from HumanoidBench repository
(2023b). We use the default settings on DMC
TD-MPC2. For both DMControl and HumanoidBench, use the latest code with its default hyperparameters" For a comprehensive list of hyperparameters, please refer to their original pa- per(Hansen et al. 2023
C ALL TRAINING CURVES
Average
Dog Run
Dog Stand
Dog Trot
1000
Episode reward
800-
400
200
Dog Walk
Humanoid Run
Humanoid Stand
Humanoid Walk
1000
Episode reward
600
400
200
0
1.0
2.0
3.0
4.0
1.0
2.0
3.0
4.0
1.0
2.0
3.0
4.0
1.0
2.0
3.0
4.0
Environment steps (x10*)
Environment steps (x10*)
Environment steps (x10%)
Environment steps (x 10
BMPC (MPC policy)
TD-MPC2 (MPC policy)
BMPC (network policy)
TD-MPC2 (network policy)
Figure 8: Performance of different policies on hign-dimensional locomotion tasks. Evaluation performance of the network policy compared to the MPC policy in BMPC and TD-MPC2 the high-dimensional locomotion tasks. The environment steps are extended to 4M for comprehensive comparison. In the top left, present the average performance. Mean and 95% CIs over 5 seeds.
"We results and code in htttts///attaa.commiiclllsseeeea
We use results in http:///iillcccomcccaaaaeazzaaan hench We use the code in http:///hhtub.oommmmmi///m/emmei
16

---

Published as a conference paper at ICLR 2025
Avg. 28 taskss
Avg. high-dImm tasks
Acrobot Swingup
Cartpole Balance
Carpooe Balance Sparse
1600
Episode reward
600
200
Cartpole Swingup
Carppoe Swingup
Cheetah
Dog
1.000
Episade reward
800
400
Dog Stand
Dog Trot
Dog Walk
Finger Spin
Finger Turn Easy
Eptcode reward
800-
400
Finger Turn Hard
Fish Swim
Hopper Hop
Hopper Stand
Humanoid Run
1.000
Eplsode reward
600
200
Humanoid Stand
Humanoid Walk
Pendulum Swingup
Quadruped Run
Quadruped Walk
1.000
Egisode reward
600
200
Reacher Easy
Walker
Walker Stand
1000
Episede reward
800
600
400
200
044
Environment steps (x105) Enwironment (K10f) Environment (x109)
Environment steps (x105) Environment (x10%)
BMPC (ours)
TD-MPC2
SAC
DreamerV3
Figure 9: All DMControl tasks. Comparing BMPC to baselines on DMControl tasks. In the top left, we present average performance of high-dimensional locomotion tasks and tasks. Mean and 95% Cls over 5 seeds" ?7
FExcept SAC, which seeds.
17

---

Published as a conference paper at ICLR 2025
Avg. 28 tasks
Avg. high-dIm tasks
Acrobot Swingup
Cartpole Balance
Cartpole Balance Sparse
1000
Episode reward
600
200-1
Cartpole
Caripole Swingup Sparse
Cheetah Run
Catch
Dog
Episode reward
800
400
Dog Stand
Dog Trot
Dog Walk
Finger Spin
Finger Turn Easy
1000
Episade reward
800
400-
Finger Turn Hard
Fish Swimn
Happer Hop
Happer Stand
Humanaid Run
1000
Egisade reward
600
200
Swingup
Quadruped Run
Quadruped
1000
Episode reward
600
400
200
Reacher Easy
Reacher
Walker
Walker Stand
Walker Walk
Episode rewerd
800
400
1000
0.2
0/4
Environment stapps (x105)
Environment stops (x10f)
Environment steps (x109) Environment stops (X109)
Environment stops (x10%)
BMPC (MPC
TD-MPC2 (MPC policy)
BMPC (network policy)
TD-MPC2 (network policy)
Figure 10: Performance of different policies on DMControl. Evaluation performance of the network policy compared to the MPC policy in BMPC and TD-MPC2. In the top left, present the average performance of 7 high-dimensional locomotion tasks and all 28 tasks. Mean and 95% Cls over 5 seeds.
18

---

Published as a conference paper at ICLR 2025
D TASK VISUALIZATION
Figure 11: DMControl tasks visualization. Images of all the embodiments we control in the DM- Control tasks. The tasks include controlling them to run, walk, jump, balance, reach, and perform actions like swing-up and spin, covering diverse range of continuous control scenarios.
Figure 12: HumanoidBench locomotion suite visualization. Images of the Unitree robot we con- trol in the HumanoidBench locomotion suite. The tasks include running, walking, crawling, bal- ancing, sitting, reaching, and performing actions like walking on stairs or walking while avoiding collisions with poles, which cover diverse range of robotic locomotion scenarios.
19


Reviews:
[{"review_id": "sGIGzVTOQM", "reviewer": "Reviewer_4Ngs", "summary": "The paper proposes Bootstrapped Model Predictive Control (BMPC), a model-based RL method that combines MPC with a policy, where the policy gives MPC better prior distribution of action to sample over, while the MPC uses a refined plan to train this policy. The authors propose lazy reanalyze method to make imitation of MPC policy more efficient, by eliminating the frequency of replanning & immediate update. The authors study the performance of proposed BMPC on a variety of continuous control tasks, compared against SOTA MBRL methods as well as model-free methods.", "questions": "In terms of the trade-off between exploration & exploitation, do you have to carefully tune/have a curriculum on the sigma std for the MPC?\n\nWhat's the effect of EMA in Eqn(7)? Is it mainly for stabilization? How grounded is this design choice?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is clearly written and the method is well-motivated: the policy can make MPC search more efficient while MPC's high-quality results can help with policy learning. \n\nThe result and analysis are solid and convincing. The authors benchmarked on a good number of environments & standard benchmarks.\n\nThe good ablations answer a lot of my questions, and I expect all readers to find them very informative.", "weaknesses": "The outstanding concern is the legitimacy of letting MPC & policy bootstrapping each other. For example, it's unclear what influence this has on exploration. Will the added policy hurt the exploration of MPC? Could you highlight how exploration is achieved in your model or is it just like in TD-MPC2?\n\nThe authors chose Dreamer v3 as the MBRL baseline in the main evaluations. Since Dreamer v3 is specifically designed for visual observation & multi-tasking, I wonder whether the authors tried tuned it for state-based evaluations.", "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": "xiJgU7buxP", "reviewer": "Reviewer_tbYo", "summary": "This paper proposes Boostrapped Model Predictive Control (BMPC), a model-based RL method that improves upon TD-MPC. TD-MPC consists of two policies: a network policy trained with RL, and an MPC policy guided by the network policy and its value function. The authors found the network policy in TD-MPC to perform much worse than the MPC policy. This is attributed to the network policy being trained in a model-free manner, which is less sample-efficient than MPC. The poor network policy then seeps into the value function, which in turn impacts the performance of the MPC policy. To mitigate this issue, the authors propose to train the network policy by directly imitating the MPC policy, while still using the network policy and its value function to guide the MPC policy. This effectively \"distills\" the MPC policy into a neural network, avoiding the policy mismatch in TD-MPC. However, this kind of imitation introduces significant overhead into policy learning, as each training step requires replanning on all samples. To address this, the authors propose Lazy Reanalyze, which replans on a small set of transitions every few training steps. With Expert Imitation and Lazy Reanaylze, BMPC achieves better asymptotic performance in fewer number of environment steps compared to TD-MPC on high-dimension DM Control tasks, all without introducing significant overhead.", "questions": "1. Are there other reasons behind TD-MPC's bad value function? If it's just the sample inefficiency model-free policy optimization, then the network policy should eventually catch up to the MPC policy. But this is not the case as shown in Figure 2 -- the asymptotic performance of the network policy is worse.\n2. What is the tradeoff of pure MPC policy distillation compared to model-free policy optimization? Are there domains or settings where BMPC performs worse than TD-MPC? If so, what are the reasons?\n3. Can you describe the three variants in Figure 6 (a) in more detail? I would also suggest adding the details of the variants in the appendix.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. The paper is motivated by a thorough analysis of the TD-MPC algorithm, which crystalizes into a key insight: the performance of TD-MPC is impacted by a mismatch between the network policy and the MPC policy. The proposed algorithm then follows as an intuitive solution. The overall flow of the paper is excellent, making it a pleasant read.\n2. The empirical results of BMPC demonstrate a significant performance gain over TD-MPC, especially in complex high-dimensional locomotion tasks. In most tasks, BMPC is more sample-efficient than TD-MPC. And in a few tasks, it also achieves higher asymptotic performance.\n3. The ablation experiments supplement the main experiment well, providing insights into the effect of hyperparameters and key design decisions. This adds to the credibility of the method\n4. Overall, I believe this paper proposes a simple method that brings a new perspective to model-based RL. Therefore, I recommend acceptance.", "weaknesses": "1. The expert imitation procedure introduces overhead into the training pipeline, as each training step requires replanning. Although this is sidestepped by Lazy Reanalyze, it remains a fundamental limitation of the method.\n2. The experiments are run with a small number of seeds (3 seeds). \n3. The experiments succinctly prove the point that the authors try to make. That said, it would strengthen the paper to include experiments across more diverse domains (those in TD-MPC 2).", "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": "XE1XWTSEyl", "reviewer": "Reviewer_qEhR", "summary": "Current state-of-the-art in MBRL learn a policy and Q-function in a model-free manner, then couple it with online planning via MPC to improve the policy. However, the authors of this paper find that on complex tasks, there is a substantial gap between the policy and policy+MPC performance. They propose a number of modifications, including an expert iteration-like scheme, for improving the final policy. This involves alternating between improving the policy with MPC online and imitating the updated policy as an expert offline. They find that not only does this improve performance, but it also closes the gap between the network policy and combining it with MPC. With their algorithm, BMPC, the policy learned via expert iteration actually nearly matches the performance of coupling it with online re-planning. Therefore, it is a viable final policy which amortizes the online planning with MPC.", "questions": "- How is the MPC policy computed for imitation learning without rerunning the planning algorithm every batch? Is the mean and covariance of the MPC policy stored in the replay buffer as well?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- Improving the performance of MBRL algorithms by combining model-based and model-free components is a timely and important problem.\n- The paper is well organized and overall clearly written. It does a good job explaining the novelty and results and provides enough information to support its claims.\n- The analysis of the performance gap between the model-free policy and bootstraped MPC policy in TD-MPC2 is particularly thorough and very informative, doing a good job to motivate the paper.\n- The evaluations are very thorough, considering all 28 environments in the DMControl benchmark tasks and an adequate selection of baselines, such as SAC, DreamerV3, and TD-MPC2. There is also a nice outline of questions to be considered in the evaluation section. And the results do indicate that BMPC can outperform these baselines, suggesting it is a promising method for combining the advantages of model-free and model-based RL.\n- The lazy reanalyze method proposed in the paper is a nice solution to recomputing expert actions while still making use of a replay buffer. It's also nice that this reanalyze interval is a free hyperparameter rather than coupled directly to the number of updates to the policy.", "weaknesses": "- Some details are missing from the exposition, such as how the MPC policy is computed each batch without having to re-run the planning algorithm. One possibility is that the mean and covariance of the Gaussian is stored in the replay buffer. However, the paper would benefit from making this more concrete.\n- It would be worth mentioning other work which attempt to imitate an expert MPC controller [1-4]. Alternatively, there are also approaches which attempt to learn residual on MPC [5-6], which is another way of bootstrapping MPC. And there are approaches which attempt to improve MPC via learned sampling distributions [7, 8]. Although not expert iteration, they are definitely related approaches which are important to mention.\n\nReferences:\n[1] Pan et al., \"Agile Autonomous Driving using End-to-End Deep Imitation Learning,\" RSS 2018.\n[2] Pan et al., \"Imitation learning for agile autonomous driving,\" IJRR 2020.\n[3] Sacks et al., \"Learning to optimize in model predictive control,\" ICRA 2022.\n[4] Fishman et al., \"Motion policy networks,\" CoRL 2023.\n[5] Sacks et al., \"Deep Model Predictive Optimization,\" ICRA 2024.\n[6] Silver et al., \"Residual Policy Learning,\" arXiv 2018.\n[7] Power et al., \"Variational Inference MPC using Normalizing Flows and Out-of-Distribution Projection,\" RSS 2022.\n[8] Sacks et al., \"Learning Sampling Distributions for Model Predictive Control,\"  CoRL 2022.", "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": "rj3RuSDEpl", "reviewer": "Reviewer_Ayiq", "summary": "This paper presents BMPC, a plan-based MBRL algorithm that integrates expert imitation for policy learning, performs model-based TD-learning for value learning, and introduces lazy reanalyze to utilize re-planning results better.", "questions": "1. The paper replaces the Q-network with a V-network. What are the specific implications of using either Q or V for the training of the algorithm?\n\n2. In line 260, the paper states that \"setting N too large would lead to excessive compounding errors,\" yet it seems that increasing N would be less biased for the V-value estimation. Could you clarify this contradiction?\n\n3. In line 325, could you provide a detailed explanation of \"environment step\"?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well-written and easy to follow.\n\nFor high-dimensional locomotion tasks, the paper presents ablation experiments comparing MPC policy with network policy. The results indicate that learning a network policy through expert imitation significantly enhances the performance of MPC.\n\nThe approach is evaluated against several baseline agents on DeepMind Control Suite environments.", "weaknesses": "1. The paper relies on TD-MPC2 and introduces modifications such as expert imitation and lazy reanalysis. However, these improvements seem overly engineering-focused and do not provide significant insights into the research direction.\n\n2. It is unclear how the expert policy π is initially established as the prior policy. Could you describe how expert policy π evolves during training?\n\n3. In the experimental section, the authors use only 3 seeds as following TD-MPC2. However, I believe that this is not enough. Could you provide additional results with more seeds as DreamerV3 if possible?\n\n4. The related work section provides a rather simplistic overview of MPC research. I recommend enhancing this section for greater comprehensiveness.", "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)