0.0.1 Introduction↩︎

Model-based reinforcement learning (MBRL) has long been regarded as a promising paradigm for sample-efficient decision-making in complex environments. By learning an explicit dynamics model, agents can simulate future trajectories, plan ahead, and make informed decisions without requiring millions of real-world interactions [1].

Despite its conceptual appeal, MBRL has often struggled to meaningfully outperform its model-free counterparts [2]–[4]. A common explanation for this shortfall is model accuracy and compounding errors [5]–[8]. When a model is used to simulate trajectories over long horizons, small prediction errors accumulate, leading to increasingly unreliable state estimates. This phenomenon has motivated a substantial body of research focused on improving model accuracy [9]–[11], uncertainty-aware models [12]–[15], and long-horizon prediction [16]–[21]. The implicit assumption underlying these efforts is that better models will directly translate to better planning and, consequently, better performance.

In this work, we challenge this assumption by advancing a core hypothesis: model accuracy alone does not determine the effectiveness of search. We explore this hypothesis through the lens of MR.Q [4], a state-of-the-art algorithm that learns a dynamics model purely for representation learning, without incorporating search.

We begin in 0.0.4.1 by establishing a counterintuitive result: search can fail even with a perfect dynamics model and value function. Beyond a critical planning horizon, the search space expands dramatically, rendering exhaustive search computationally infeasible and high-value trajectories vanishingly rare under sampling-based approaches. Critically, this limitation is intrinsic to the search process; improving model quality cannot resolve it.

We next ask whether these theoretical limitations appear in practice. Our empirical analysis in 0.0.4.2 corroborates this theoretical limitation. Despite MR.Q’s architectural similarity to TD-MPC2 [22], a leading search-based MBRL method, we find that naively incorporating search harms MR.Q’s performance. This occurs even though MR.Q’s learned dynamics are comparably accurate to those of TD-MPC2, which benefits substantially from search. These results confirm that model accuracy alone does not determine whether search will be effective, but they also raise a natural question: if not model accuracy, what does determine the success of search?

We argue that the primary obstacle lies in the interaction between search and value learning. Because exhaustive search is infeasible, practitioners rely on sampling-based methods guided by a learned value function. Due to the computational cost of search, this value function is typically trained using rollouts from a learned, non-search policy network [22]. In 0.0.4.3, we show that this creates a mismatch between data collected with search and the distribution assumed during training, leading to overestimation bias that degrades both value function quality and overall performance [23].

Rather than trying to build better models, we directly address the overestimation bias introduced by search by using the minimum over an ensemble of value functions to produce pessimistic estimates for out-of-distribution actions. We instantiate this approach in a new algorithm: Model-based Representations for Search and Q-learning (MRS.Q). We evaluate MRS.Q on over 50 tasks spanning multiple benchmark domains using a single fixed set of hyperparameters. Our results demonstrate consistent improvements over both state-of-the-art model-free and model-based methods, validating that addressing overestimation, rather than solely improving model quality, is key to unlocking the benefits of search.

Our findings show the bottleneck is not just the model, but how we use it. Progress lies not only in building better dynamics models, but also in understanding and managing the fundamental tensions between value learning and search.

0.0.2 Related Work↩︎

MBRL. MBRL approaches can be broadly categorized by how they utilize their model. The simplest approach trains a model-free RL algorithm on synthetic samples generated by the model [14], [24]–[26]. A related strategy integrates the learned model into value learning by generating multi-step returns [11], [16], [27]–[30]. The PILCO family of methods [12], [31], [32] takes a different approach, using gradients from the model to directly improve the policy.

Search in MBRL. A more sophisticated approach uses the world model for search. One common strategy, rooted in optimal control, is model predictive control (MPC), which repeatedly searches through the model to identify the best action at each time step [10], [13], [15], [33]–[41]. For discrete action spaces, Monte Carlo tree search offers a related alternative [42]–[45]. The current state of the art for MPC-based methods is TD-MPC2 [22], [46], which uses short-horizon search with a value function.

Challenges in MBRL. MBRL faces several fundamental challenges. Compounding error in auto-regression predictions limits the ability to plan over long horizons [5]–[8]. The objective mismatch problem [47]: models are trained to optimize immediate predictions but are ultimately evaluated by policy performance. [48] showed that oracle models offer diminishing performance gains. MR.Q [4] demonstrates that models can improve representation learning without search, raising the question of when and how search adds value. Our empirical observations reinforce these concerns, suggesting that how the model is used matters more than its accuracy. Recent extensions of TD-MPC2 have also highlighted a performance gap between the learned policy network and MPC, applying regularization to align the policy with MPC-selected actions [49]–[51]. [49] demonstrate that MPC causes the value function to overestimate the performance of its policy network. Our experiments reveal that this failure mode extends beyond TD-MPC2, and that the value function can also overestimate the online, MPC-based policy.

Overestimation bias. Overestimation bias is a well-known problem in RL, stemming from the maximization of an approximate value function [52]–[55]. In offline RL, overestimation bias is caused by errors induced by the distribution shift between the offline dataset and the learned policy [23]. A similar dynamic arises when search selects actions outside the distribution seen during value function training [49]–[51]. A common strategy for addressing overestimation bias is through the minimum of value functions [55]–[57]. In our work, we show that this technique is particularly effective for mitigating the overestimation induced by search.

0.0.3 Background↩︎

Reinforcement learning (RL). We follow the standard formulation [58], where an agent interacts with an environment modeled as a Markov Decision Process (MDP) defined by the tuple \((S, A, p, R, \gamma)\), with state space \(S\), action space \(A\), dynamics function \(p\), reward function \(R\), and discount factor \(\gamma\). The objective is to learn a policy \(\pi\) mapping state \(s \in S\) to action \(a \in A\), that maximizes the expected discounted return \(\mathbb{E}_\pi\left[\sum_{t=0}^{\infty} \gamma^t r_t \right]\) of rewards \(r_t \sim R(s_t, a_t)\). The value function \(Q^\pi(s,a) = \mathbb{E}_\pi\left[\sum_{t=0}^{\infty} \gamma^t r_t \right | s_0=s, a_0=a]\) measures the expected return when starting in state \(s\) and taking action \(a\).

Model-based reinforcement learning (MBRL) approaches learn an explicit dynamics model of the environment that predicts the next state \(s_{t+1}\) given the current state \(s_t\) and action \(a_t\). In this paper, we focus on methods that use these learned models to plan via search. Search refers to the process of evaluating candidate action sequences through model rollouts to select actions. Model predictive control (MPC) is an online search framework that optimizes actions over a finite horizon using model rollouts, executing only the first action before re-planning at each time step.

TD-MPC2. TD-MPC2 [22] is a MBRL algorithm that learns a latent world model that is used for MPC. The state embedding \(\mathbf{z}_s\) is learned together with the dynamics function, reward model, and value function by co-optimizing loss functions for each: \[\label{eqn:TDMPC295update} \begin{align} \Loss ( \mathbf{z}_{s}, &F, R, Q) = \Loss_\text{Dynamics}( F(\mathbf{z}_s, a) - \mathbf{z}_{s'} ) \\ +~&\Loss_\text{Reward}( R(\mathbf{z}_s, a) - r ) \\ +~&\Loss_\text{Value} ( Q(\mathbf{z}_s,a) - ( r + \y Q(\mathbf{z}_{s'}, \pi(\mathbf{z}_{s'})) ) ). \end{align}\tag{1}\] During MPC, trajectories \(\tau = (\tilde{\mathbf{z}}_0, a_0, \tilde{\mathbf{z}}_1, a_1, ..., \tilde{\mathbf{z}}_N)\) are generated by recursively rolling out the model from an embedding \(\mathbf{z}_{s_0}\) of an initial observed state \(s_0\): \[\tilde{\mathbf{z}}_t = F(\tilde{\mathbf{z}}_{t-1}, a_{t-1}) \quad\text{ where }\quad \tilde{\mathbf{z}}_0 = \mathbf{z}_{s_0}.\] The value \(V\) of a trajectory \(\tau\) is computed by summing a final value \(Q\) with predicted rewards along the model-predicted embeddings over a finite horizon \(N\): \[V(\tau) = \sum_{t=0}^{N-1} \y^t R(\tilde{\mathbf{z}}_t, a_t) + \y^N Q(\tilde{\mathbf{z}}_{N}, \pi(\tilde{\mathbf{z}}_{N})). \label{eqn:tdmpc295eval95traj}\tag{2}\]

TD-MPC2 uses a learned policy network \(\pi\), trained to maximize value, in its value updates, but selects actions during environment interaction using Model Predictive Path Integral (MPPI) control [59]. At each time step, candidate action trajectories are initially sampled from a mixture of random and policy-derived actions, then evaluated following 2 . The highest-valued trajectories inform the next round of candidate sampling, and this process repeats. The final trajectory is selected through value-weighted sampling, and its first action is executed.

MR.Q. MR.Q [4] is a model-free RL method that uses a model-based objective to learn a state-action embedding \(\mathbf{z}_{sa}\) which has an approximately linear relationship to the value function \(Q^\pi(s,a)\). Unlike most model-based methods, which use search or model rollouts, MR.Q’s objective is solely used to learn useful representations. More formally, MR.Q’s learning objective co-optimizes the embedding \(\mathbf{z}_{sa}\) with weights capturing the reward and dynamics of the environment: \[\begin{align} \Loss ( \mathbf{z}_{s}, W_p, W_r) =~&\Loss_\text{Dynamics}( \mathbf{z}_{sa}^\top W_p - \mathbf{z}_{s'}) \\ +~&\Loss_\text{Reward}( \mathbf{z}_{sa}^\top W_r - r ). \end{align}\] The value function is trained independently with a standard Bellman update [60], [61], taking the minimum between two value functions in the target [55]: \[\label{eqn:MRQ95value} \Loss_\text{Value} \lp r + \y \min (Q_1, Q_2)(\mathbf{z}_{s'a'}) - Q(\mathbf{z}_{sa}) \rp,\tag{3}\] where the action \(a'\) is sampled from the policy  \(\pi(\mathbf{z}_{s'})\). The state embedding \(\mathbf{z}_s\) is learned upstream, end-to-end with \(\mathbf{z}_{sa}\), i.e., \(g(\mathbf{z}_s, a) = \mathbf{z}_{sa}\). Since \(\mathbf{z}_{sa}\) is a non-linear function of \(\mathbf{z}_s\) and \(a\), the resulting dynamics and reward function are functionally similar to TD-MPC2’s. Crucially, although MR.Q leverages a model-based objective, it is strictly used to learn the embeddings rather than forward dynamics predictions or search.

Adding and removing MPC. The architectural similarity between TD-MPC2 and MR.Q allows us to isolate the effect of search by adding MPC to MR.Q and removing it from TD-MPC2. For clarity, let \(\pi_\text{MPC}\) denote the implicit policy induced by MPC, which selects actions via MPPI using the learned model and value function (2 ). Methods with MPC select actions according to \(\pi_\text{MPC}\), while methods without MPC select actions directly from the learned policy \(\pi\).

Distribution shift. Distribution shift arises in value estimation when the behavior policy used for data collection differs from the target policy \(\pi\) used in the value update ([eqn:TDMPC2_update,eqn:MRQ_value]). For example, if MPC selects actions during data collection but the value function is trained with actions from the learned policy \(\pi\), this mismatch may cause the value function to be queried on actions it was not trained on, ultimately leading to unreliable value estimates [23].

0.0.4 The Challenge of Search↩︎

In this section, we investigate the fundamental challenges limiting the effectiveness of search in MBRL. Conventional knowledge suggests that model accuracy is the key predictor of MBRL performance. We challenge this assumption, demonstrating that it does not strictly hold. While model accuracy remains important, our results suggest that there are other underappreciated factors at play.

We present three key findings:

1. Naive search can fail with a perfect model. We show theoretically that even with perfect model and value functions, uniform random search can fail with high probability over longer horizons due to the exponential growth of the search space (0.0.4.1).

2. Model accuracy alone does not determine whether search improves performance. Contrary to expectations, adding MPC to MR.Q [4] degrades its performance, despite MR.Q having a more accurate model than TD-MPC2 [22], a method where MPC proves beneficial (0.0.4.2).

3. Search induces overestimation bias. We show that simply replacing the learned policy with MPC when collecting data introduces overestimation bias that degrades value function quality (0.0.4.3).

Together, these findings motivate the need for methods that explicitly address these challenges rather than focusing solely on improving model accuracy.

0.0.4.1 Search with a perfect model

We begin by examining the theoretical limits of search under ideal conditions. Perhaps surprisingly, search can fail with high probability even with perfect dynamics and a perfect value function. To illustrate this, consider the N-chain MDP defined in 1. In this environment, a single action \(a_0\) advances the agent along the chain, while all other actions lead to an absorbing state. The reward is zero everywhere except at the final state. We define search failure as the inability to discover any trajectory with non-zero value, and analyze a uniform random search that samples action sequences, simulates trajectories via the model, and selects the trajectory with the highest value.

In this N-chain environment, regular states have zero reward but non-zero value, since the final state remains accessible. In contrast, the absorbing state has both zero reward and zero value. Consequently, if the agent selects any action other than \(a_0\), the trajectory will end in the absorbing state and will necessarily have zero value. Therefore, even with access to the ground-truth value function, an agent can only select action \(a_0\) at each time step to discover a trajectory with non-zero value.

Under random search, the probability of finding such a trajectory is \(1 - (1 - \frac{1}{A^n})^m\), where \(A\) is the number of actions, \(n\) is the search horizon, and \(m\) is the number of sampled trajectories. Intuitively, only 1 of \(A^n\) possible action sequences consists entirely of \(a_0\), so the probability of randomly sampling a successful trajectory decreases exponentially with horizon length.

In 2 we show that this probability degrades rapidly as the search horizon increases. For example, with \(A = 10\) actions and \(m = 1000\) trajectories, a search horizon of \(n=3\) gives a probability of approximately \(0.63\) of discovering a non-zero valued trajectory. However, extending the search horizon to \(n=10\) causes the probability to fall to \(10^{-7}\).

The failure in this example could be avoided by acting according to the value function rather than relying on search. More broadly, this illustrates that depending solely on imagined trajectories may be insufficient when the search space is large, particularly in domains with continuous states and actions. While modern search-based methods such as TD-MPC2 [22] use policy-guided sampling to narrow the search space, they still face combinatorial challenges as the horizon increases, particularly when the policy is imperfect or the reward landscape requires precise action sequences. This helps explain why MBRL methods often employ short search horizons (three time steps), both to constrain the search space and to reduce compounding model error.

0.0.4.2 Does model accuracy determine performance?

Having established that search can fail even with perfect models in theoretical settings, we now examine the role of model accuracy in practical settings with learned components. We compare two recent state-of-the-art methods, MR.Q [4] and TD-MPC2 [22], both with and without MPC. We focus on these methods because the similarities in architecture and learning objectives allow us to isolate the effect of search from confounding factors.

In 1 we evaluate each method’s learned dynamics model accuracy and task performance using the following metrics:

• Dynamics error. The mean-squared error between predicted and actual next-state embeddings, weighted by \(\y^t\) over the three-step search horizon. To ensure comparable error magnitudes, we modify MR.Q’s state encoder by replacing its final activation function with Simplicial Embeddings (SEM) [62], matching TD-MPC2’s architecture and ensuring that differences in embedding scale do not inflate one method’s error.

• Unroll error. The absolute error between the value prediction used in MPC and the ground-truth trajectory value. We compute the value prediction by unrolling the dynamics over the search horizon, and summing the discounted predicted rewards with the value estimate of the final predicted state (2 ). The ground-truth value is computed similarly, unrolling trajectories with the simulator and summing the true discounted reward over the full episode length.

• Performance \(\Delta\). The performance difference with and without MPC (i.e., MR.Q with MPC vs. MR.Q, and TD-MPC2 vs. TD-MPC2 without MPC). Positive values indicate that MPC improved performance.

Dynamics and unroll errors directly measure model quality. Unroll error, in particular, measures the quality of the signal that MPC uses to select actions (2 ), as it reflects both dynamics and value prediction accuracy. Performance \(\Delta\) shows the change in performance using the identical MPC procedure with each algorithm.

Despite comparable model accuracy between methods, using MPC gives divergent outcomes, improving TD-MPC2 while degrading MR.Q. 1 quantifies this surprising asymmetry. Although both methods exhibit similar dynamics and value prediction errors, TD-MPC2’s performance improves substantially with MPC across nearly all DM Control Suite (DMC) tasks [63], while MR.Q’s consistently decreases. Given their architectural and loss function similarities, we would expect comparable error rates to translate into comparable search benefits.

The substantial gap in MPC performance confirms that model accuracy alone does not determine whether search will help, suggesting other factors are at play.

0.0.4.3 Why does search harm MR.Q?

Given that model accuracy does not fully explain these results, we now investigate what other factors are responsible. Following [49]–[51], we hypothesize that distribution shift and its downstream effects on value estimation are the key bottleneck to search performance.

Consider the standard value function update used by both TD-MPC2 and MR.Q, based on a value target that evaluates a learned policy network \(\pi\): \[\label{eqn:td95update} Q(s,a) \approx r + \y Q(s', \pi(s')).\tag{4}\] This value update evaluates the learned policy \(\pi\) rather than the action originally taken in the environment. When MPC is used to collect data, there is a mismatch in which the value function is trained on actions from \(\pi_\text{MPC}\) but queried on actions sampled from \(\pi\). Because \(\pi \neq \pi_\text{MPC}\), these queries are out-of-distribution, which prior work has shown leads to approximation errors and overestimated values [23].

One might resolve this by using search to compute value targets directly, but doing so is prohibitively expensive. Every sample in a mini-batch requires a target, multiplying the number of costly search calls by the batch size and update-to-data ratio (e.g., \(256\times\) with default hyperparameters).

To test this hypothesis, we measure the accuracy of value functions learned by MR.Q, MR.Q with MPC, TD-MPC2, and our proposed approach, MRS.Q (introduced in the following section) in 3. Value error is computed by taking the difference between the learned value estimate and the true discounted return of the behavior policy: the learned policy network for MR.Q, and the MPC policy \(\pi_\text{MPC}\) for the others. Unlike prior work [49] which examined overestimation bias of the learned policy \(\pi\), our experiments focus on the behavior policy (largely the MPC policy \(\pi_\text{MPC}\)).

3 shows that simply changing the behavior policy from the learned policy (MR.Q) to MPC (MR.Q+MPC) introduces significant overestimation bias. Moreover, 3 reveals a clear relationship between value error and the performance change \(\Delta\) from using MPC (1):

• MR.Q+MPC: Switching from the learned policy to MPC introduces overestimation in nearly every environment. MPC improves performance only in Ant, where overestimation is the lowest.

• TD-MPC2: Compared to MRQ+MPC, overestimation is considerably lower overall. However, it remains high in environments where TD-MPC2’s learned policy outperforms vanilla TD-MPC2 (dog-stand, Ant, Hopper, Humanoid, Walker2d).

These findings suggest that (1) MPC introduces significant overestimation bias, and (2) although the design choices in TD-MPC2 help mitigate this bias, they are insufficient as a standalone solution.

0.0.5 Effective Search in Model-based RL↩︎

Having identified overestimation bias as a key challenge limiting search effectiveness, we now present our approach. Our algorithm, Model-based Representations for Search and Q-learning (MRS.Q), builds on MR.Q with targeted modifications that enable effective use of search.

MRS.Q introduces three key modifications to MR.Q: (1) search-based action selection with a short horizon, (2) aggressive minimization over an ensemble of value functions, and (3) minor hyperparameter changes. Action selection. Following TD-MPC2, MRS.Q selects actions using Model Predictive Path Integral (MPPI) control [59] when acting in the environment. As such, we also use the same MPPI hyperparameters as TD-MPC2. Critically, this means that we adopt the same short MPC horizon of \(3\) steps, using the value function to estimate returns beyond this window. As discussed in 0.0.4.1, naive search can fail with high probability when planning over long horizons due to the expanded search space. A short horizon keeps the search space tractable and limits compounding error in the model rollout.

Minimum over value functions. Recent work addressing distribution shift in MBRL [49]–[51] has focused on constraining the policy to match actions selected by search. However, this approach is problematic: search-selected actions shift rapidly during training as the model evolves.

In 4, we compare the rate of change of actions selected by search versus the learned policy network over the course of learning. Throughout training, search-selected actions exhibit substantially more variation than policy network actions, suggesting that constraining the policy to match such an unstable target may destabilize learning or introduce noise into training.

Rather than constraining or modifying the learned policy, we address distribution shift by directly reducing overestimation in the value function. We achieve this through aggressive minimization across an ensemble of 10 value functions. While taking the minimum over a pair of value functions has become a standard technique [55], this approach naturally extends to larger ensembles: \[\label{eqn:min} Q(s,a) \approx r + \y \min_{i \in \{1,2,..,10\}} Q_i(s', \pi(s')).\tag{5}\] Crucially, we apply this minimization wherever the value function is evaluated with the policy: not only when computing value targets (as above), but also when computing final values during MPC trajectory evaluation (2 ). Applying the minimum during trajectory evaluation prevents the search procedure from favoring potentially inflated value estimates over model-predicted rewards, while ensuring consistency in how the value function is used throughout the algorithm.

TD-MPC2 also employs multiple value functions (\(N=5\)), however, rather than using the full ensemble, it randomly samples a pair, taking the minimum for value updates [64] and the mean during MPC. As shown in 3, using the minimum across the full ensemble significantly reduces overestimation compared to MR.Q (\(N=2\)) naively combined with MPC. We further defend this choice in the ablation studies in 0.0.6.2.

0pt

Hyperparameter changes. In addition to the ensemble, we make only three hyperparameter changes to MR.Q:

• No exploration noise. As shown in 4, search produces greater variation in action selection than the policy network. This variation in action selection implicitly induces exploration, making additive Gaussian noise unnecessary. We therefore remove it entirely (see 0.0.6.2).

• Simplicial embedding (SEM). Following TD-MPC2, we apply SEM [62], [65] to both the state encoder and the output of the dynamics model. SEM normalizes latent representations onto the probability simplex, which stabilizes multi-step dynamics rollouts by reducing drift in predicted states. Because SEM produces sparse representations with smaller loss magnitudes, we increase the dynamics loss weight from \(1 \rightarrow 20\) to match the value used by TD-MPC2 and balance the loss terms.

• Increased terminal loss weight. Since predicting when episodes terminate is essential for accurate value estimation, we increase the weight of the terminal loss function from MR.Q’s default of \(0.1 \rightarrow 1\).

All other hyperparameters and design choices follow the original MR.Q defaults. We highlight these minor adjustments to show that enabling effective search in MR.Q does not require major modifications to hyperparameter settings.

0.0.6 Results↩︎

We evaluate MRS.Q on a diverse set of continuous control tasks to demonstrate the performance benefits of MRS.Q.

Benchmarks. We evaluate on over 50 tasks spanning three benchmark domains: MuJoCo [66], [67], DeepMind Control Suite (DMC) [63], and HumanoidBench [68]. For HumanoidBench, we split the results into two subsets, one which uses hand features, matching the experimental setting used by BMPC and BOOM [50], [51], and one which does not, matching the experimental setting used by SimbaV2 [69].

Baselines. We evaluate MRS.Q against several methods, including MR.Q [4] (both with and without MPC), TD-MPC2 [22], three state-of-the-art extensions of TD-MPC2 that regularize the policy to mitigate distribution shift: TD-M(PC)\(^2\) [49], BMPC [50], and BOOM [51], and the state-of-the-art model-free method SimbaV2 [69]. All methods are trained for 1M environment steps (full details in 1.0.1.5).

Evaluation protocol. For each task, we train agents for 1M environment steps and 10 seeds. We evaluate agents every 5k steps for 10 episodes without exploration noise. We report mean performance with 95% confidence intervals.

0.0.6.1 Main Results

We evaluate aggregate performance across all benchmarks, as shown in 5 and 2. MRS.Q outperforms all baselines on three out of four benchmarks and remains competitive on the fourth (DMC). This demonstrates the effectiveness of our proposed modifications. While naively combining MPC with MR.Q degrades performance, our approach enables MRS.Q to fully leverage MPC.

Although TD-MPC2 was specifically designed for search and both methods use the same search procedure, MRS.Q achieves substantially stronger results, particularly on Gym and the higher-dimensional HumanoidBench tasks with hand features. This demonstrates that MRS.Q successfully unlocks MR.Q’s potential for search.

Furthermore, MRS.Q surpasses the TD-MPC2 extensions (TD-M(PC)\(^2\), BMPC, BOOM) that regularize the policy towards search-selected actions. This suggests that simply imitating search behavior is insufficient to address distribution shift, likely due to the rapid evolution of selected actions during search (see 4).

0pt

0.0.6.2 Ablation Studies

We ablate key components of MRS.Q to understand their individual contributions. 3 summarizes the results.

Number of value functions. We vary the number of value functions used in the ensemble. We find that performance improves as the ensemble size increases up to approximately 10 value functions, after which gains plateau. As shown in 3, using the default 2 value functions is insufficient to address the overestimation induced by MPC, while an ensemble of 10 provides a good trade-off between bias reduction and computational cost.

Exploration noise. Re-adding exploration noise with MR.Q’s default of \(\N(0,0.2^2)\) induces a small performance drop across the four benchmarks. As illustrated in 4, using MPC already induces substantial change in selected actions, making additional exploration noise on top of MPC unnecessary and often harmful.

SEM. Replacing SEM in the state encoder with MR.Q’s default ELU activation function [70] reduces performance on Gym and DMC, confirming that stabilizing latent dynamics is important for effective multi-step planning. However, this replacement slightly improves performance on HumanoidBench, suggesting that SEM is not the primary driver of performance gains in MRS.Q.

Minimum over ensemble. Instead of taking the minimum over the ensemble of value functions, we match TD-MPC2 by taking the minimum over two randomly sampled ensemble members [64]. We find there is a significant performance loss in Gym and HumanoidBench - Hand. This suggests that overestimation bias mitigation is particularly important in high-dimensional tasks with early termination. Conversely, we remark this technique is not necessary in the DMC benchmark, possibly due to smoother dense rewards, and the lack of early termination.

Minimum in MPC. Our minimization strategy replaces all evaluations of the value function with a minimum over the entire ensemble. To measure the impact of minimization within MPC, we experiment with using the mean of the ensemble when evaluating trajectories during MPC, i.e., in 2 , as done in TD-MPC2. We find that this gives a moderate performance drop across the four benchmarks.

Minimum TD-MPC2. To evaluate the generalizability of our approach, we apply it to TD-MPC2 by taking the minimum over an ensemble of 10 value functions. We find that this improves performance across all four benchmarks.

0.0.7 Conclusion↩︎

In this work, we highlight a critical bottleneck in MBRL methods: the search process itself and the distribution shift it introduces. When value functions are trained using a non-search policy, but using search-gathered data, the resulting distribution shift leads to overestimation bias that undermines the benefits search is meant to provide. Our method, MRS.Q, adds search to MR.Q [4] and addresses this overestimation bias by taking the minimum over an ensemble of value functions. This targeted intervention gives consistent improvements over MR.Q and state-of-the-art model-based and model-free RL baselines on more than 50 tasks across popular benchmark domains.

More broadly, our results challenge a foundational assumption in MBRL: better models lead to better performance. Our analysis demonstrates that search can fail despite highly accurate, or even perfect, dynamics models, and that naively incorporating search into state-of-the-art algorithms can actively degrade performance. These findings show that model accuracy alone is insufficient and equal attention should be given to how models are used. By directly addressing the challenges of search, we can unlock performance gains that better models alone cannot provide.

Limitations. Despite MRS.Q’s strong performance, several limitations remain. First, the computational cost of maintaining 10 value functions and running MPC at each step is substantial. Although this cost is comparable to other MBRL methods such as TD-MPC2 [22], it represents a significant increase over the search-free MR.Q (see 1.0.2). Second, the aggressive minimization strategy provides smaller gains on tasks with dense rewards and no early termination (e.g., DMC), suggesting that the approach is most beneficial in settings where overestimation is severe. Lastly, because we adopt TD-MPC2’s search method, our approach inherits its short planning horizon, and our evaluation is limited to continuous control benchmarks. Nevertheless, our work offers novel insights into MBRL and highlights potential pitfalls of incorporating search into existing algorithms.

Impact Statement↩︎

Our contributions to sample-efficient MBRL have potential applications in robotics and autonomous systems, where data collection is costly, time-consuming, or safety-critical. By enabling more effective use of learned models, MRS.Q could reduce the real-world samples needed to train physical systems where extensive trial-and-error is impractical. While we evaluate only in simulation, the sample efficiency and performance gains demonstrated here may help accelerate the deployment of RL in real-world settings.

1 Appendix↩︎

1.0.1 Experimental Details↩︎

1.0.1.1 Hyperparameters

1.0.1.2 MPC Procedure

For our search algorithm, we follow the TD-MPC2 [22] version of Model Predictive Path Integral (MPPI) control [59]. Let \(f(s) = \mathbf{z}_s\) be the state encoder, let \(g(\mathbf{z}_s, a) = \mathbf{z}_{sa}\) be the state-action encoder, and let \(h(\mathbf{z}_{sa})\) be the dynamics prediction model. Let \(t_\text{episode}\) denote the current time step of the episode.

1.0.1.3 Network Architecture

This section describes MRS.Q’s network architecture using PyTorch code blocks [73]. The network definitions follow MR.Q’s architecture, with minor modifications to incorporate SEM [62].

Preamble

import torch
import torch.nn as nn
import torch.nn.functional as F

zs_dim = 512
za_dim = 256
zsa_dim = 512

def ln(x):
    return F.layer_norm(x, (x.shape[-1],))

def SEM(self, x):
    shape = x.shape
    x = x.reshape(*shape[:-1], -1, 8)
    x = F.softmax(x, dim=-1)
    return x.reshape(*shape)

State Encoder \(f\) Network

The state encoder is a three-layer MLP with a hidden dimension of 512. Each hidden layer is followed by LayerNorm and an ELU activation. The final layer applies LayerNorm with a learnable affine transformation, followed by SEM.

The resulting state embedding \(\mathbf{z}_s\) is trained end-to-end with the state-action encoder.

self.zs1 = nn.Linear(state_dim, 512)
self.zs2 = nn.Linear(512, 512)
self.zs3 = nn.Linear(512, zs_dim)
self.ln = nn.LayerNorm(zs_dim)

self.activ = F.elu

def forward(self, state):
    zs = self.activ(ln(self.zs1(state)))
    zs = self.activ(ln(self.zs2(zs)))
    return SEM(self.ln(self.zs3(zs)))

State-Action Encoder \(g\) Network

The action input passes through a linear layer followed by an ELU activation. The processed action is then concatenated with the state embedding and fed through a three-layer MLP with a hidden dimension of 512. LayerNorm and ELU activations follow the first two layers.

The resulting state-action embedding \(\mathbf{z}_{sa}\) is passed to a linear layer that predicts the next state embedding, reward, and terminal signal. The final layer of the dynamics model applies LayerNorm with a learnable affine transformation, followed by SEM.

self.za = nn.Linear(action_dim, za_dim)
self.zsa1 = nn.Linear(zs_dim + za_dim, 512)
self.zsa2 = nn.Linear(512, 512)
self.zsa3 = nn.Linear(512, zsa_dim)

self.ln = nn.LayerNorm(zs_dim)
self.dynamics_model = nn.Linear(zsa_dim, zs_dim)
self.reward_model = nn.Linear(zsa_dim, 65)
self.terminal_model = nn.Linear(zsa_dim, 1)

self.activ = F.elu

def forward(self, zs, action):
    za = self.activ(self.za(action))
    zsa = torch.cat([zs, za], 1)
    zsa = self.activ(ln(self.zsa1(zsa)))
    zsa = self.activ(ln(self.zsa2(zsa)))
    zsa = self.zsa3(zsa)
    return (
        SEM(self.ln(self.dynamics_model(zsa))), 
        self.reward_model(zsa), 
        self.terminal_model(zsa), 
        zsa
    )

Value \(Q\) Networks

The value network is a four-layer MLP with a hidden dimension of 512. LayerNorm and ELU activations follow the first three layers. An ensemble of 10 value networks is used, each sharing the same architecture and forward pass.

self.l1 = nn.Linear(zsa_dim, 512)
self.l2 = nn.Linear(512, 512)
self.l3 = nn.Linear(512, 512)
self.l4 = nn.Linear(512, 1)

self.activ = F.elu

def forward(self, zsa):
    q = self.activ(ln(self.l1(zsa)))
    q = self.activ(ln(self.l2(q)))
    q = self.activ(ln(self.l3(q)))
    return self.l4(q)

Policy \(\pi\) Network

The policy network is a three-layer MLP with a hidden dimension of 512. LayerNorm and ReLU activations follow the first two layers, with a tanh function applied to the output.

self.l1 = nn.Linear(zs_dim, 512)
self.l2 = nn.Linear(hdim, 512)
self.l3 = nn.Linear(512, action_dim)

self.activ = F.relu

def forward(self, zs):
    a = self.activ(ln(self.l1(zs)))
    a = self.activ(ln(self.l2(a)))
    return torch.tanh(self.l3(a))

1.0.1.4 Environments

All main experiments for MRS.Q and baseline methods used 10 random seeds, while ablation studies used 5 seeds. Evaluations were performed every 5k time steps, taking the average performance across 10 episodes.

Gym [66], [67], [74]. We evaluate on the same five environments used by MR.Q [4], all based on the -v4 version with no preprocessing applied. Following MR.Q, we report aggregate scores using TD3-normalized performance, with TD3 reference scores obtained from TD7 [75]: \[\text{TD3-Normalized}(x) := \frac{x - \text{random score}}{\text{TD3 score} - \text{random score}}.\]

DM Control Suite (DMC) [63]. We evaluate on the same 28 environments used by MR.Q, with an action repeat of 2 following prior work [4], [22].

HumanoidBench [68]. We evaluate on the same set of 14 -v0 environments for both HumanoidBench experiment configurations. Following prior work [50], we omit the reach task from aggregate scores due to its different reward scale. Experiments without hands use the h1 environments (e.g., h1-walk-v0), while experiments with hands use the h1-hand environments (e.g., h1hand-walk-v0). For action repeat settings, MRS.Q and the model-based baselines (TD-MPC2, BMPC, BOOM) use an action repeat of 1, following author-provided hyperparameters. The model-free baselines (MR.Q, SimbaV2) use an action repeat of 2; MR.Q+MPC also uses an action repeat of 2 to ensure a fair comparison with MR.Q.

1.0.1.5 Baselines

All experiments were run for 10 seeds over 1M environment steps and use default author-provided hyperparameters for all tasks.

MR.Q [4]. Results were obtained from the authors’ GitHub repository (https://github.com/facebookresearch/MRQ), commit 280d9c0263964463522d16ec2daee57b1b7bf087, except for HumanoidBench, where we re-ran the authors’ code. For HumanoidBench, we use an action repeat of 2, as this achieved a slightly higher performance than an action repeat of 1.

TD-MPC2 [22]. Results were obtained by re-running the authors’ code (https://github.com/nicklashansen/tdmpc2, commit 8bbc14ebabdb32ea7ada5c801dc525d0dc73bafe). Unlike the MR.Q paper, which used separate TD-MPC2 codebases for terminal and non-terminal environments, we use this single codebase for all experiments, as it is designed to handle both settings.

TD-M(PC)\(^2\) [49]. Results were obtained by re-running the authors’ code (https://github.com/DarthUtopian/tdmpc_square_public), commit d1c2632c36effd2f7b661bfe5f822a3db8054d40. To enable termination prediction for Gym tasks, we extend the dynamics model to include termination signals, following the TD-MPC2 implementation. This modification applies only to Gym tasks with termination conditions.

BMPC [50]. Results were obtained by re-running the authors’ code (https://github.com/wertyuilife2/bmpc), commit 6746ea7d265a8b82c8347fd7e894373ce00333fa. To enable termination prediction for Gym tasks, we extend the dynamics model to include termination signals, following the TD-MPC2 implementation. This modification applies only to Gym tasks with termination conditions.

BOOM [51]. Results were obtained by re-running the authors’ code (https://github.com/molumitu/BOOM_MBRL), commit 9b3156d1fadda5ca5318274c404b694ad7b0786f. To enable termination prediction for Gym tasks, we extend the dynamics model to include termination signals, following the TD-MPC2 implementation. This modification applies only to Gym tasks with termination conditions.

SimbaV2 [69]. Results were obtained by re-running the authors’ code (https://github.com/DAVIAN-Robotics/SimbaV2), commit 86899c277cdc697b2b02d827243de1ea93f20a1d.

1.0.1.6 Model and Value Accuracy Experiments

To compute the error terms reported in 1 and 3, we evaluate error terms on state-action pairs sampled according to a fully trained agent. Specifically, for each algorithm, we use the final policy to collect 50 state-action pairs along a trajectory at regular intervals of 20 time steps, resetting the environment as needed.

We obtain ground-truth value estimates by resetting the internal simulator to each recorded state, executing the initial action, and then following the agent’s policy until episode termination. We repeat this process 100 times per state-action pair to ensure reliable estimates.

1.0.2 Computational Cost↩︎

We evaluate the computational cost of MRS.Q against our baselines using consistent hardware, with results reported in 5. For a fair comparison, torch.compile is disabled for all approaches. All experiments are conducted on HalfCheetah-v4 for 1M time steps.

The results show that although MRS.Q incurs significantly higher computational cost than MR.Q, this overhead is primarily due to MPC rather than the ensemble of 10 value functions.

1.0.3 Full Results↩︎

1.0.3.1 Gym

0pt

1.0.3.2 DMC

0pt

1.0.3.3 HumanoidBench (No Hand)

0pt

1.0.3.4 HumanoidBench (Hand)

0pt

1.0.4 Ablation Results↩︎

0pt

0pt

0pt

0pt

0pt

0pt

0pt

0pt

References↩︎