Technical Summary: The Surprising Difficulty of Search in Model-Based Reinforcement Learning

Problem Statement

Model-based reinforcement learning (MBRL) is traditionally viewed as a promising paradigm for sample-efficient decision-making, relying on learned dynamics models to simulate future trajectories and plan ahead. The prevailing assumption in the field is that the primary obstacles to MBRL performance are model accuracy and compounding errors over long horizons. Consequently, much research has focused on improving dynamics models and uncertainty estimation.

However, this paper challenges the assumption that better models automatically lead to better planning. The authors investigate the effectiveness of search (specifically Model Predictive Control, or MPC) as a drop-in replacement for learned policies. They hypothesize that the difficulty of search in MBRL is not solely due to model imperfections but stems from fundamental issues in how search interacts with value learning, specifically regarding distribution shift and overestimation bias.

Methodology

Theoretical and Empirical Investigation

The authors first establish that search can fail even with perfect models. Using an "absorbing N-chain" MDP, they demonstrate that under uniform random search, the probability of finding a high-value trajectory decays exponentially as the planning horizon increases, rendering exhaustive search infeasible in large action spaces.

To isolate the effects of search from model quality, the authors compare two architecturally similar algorithms:

1. MR.Q: A state-of-the-art method that uses a model-based objective purely for representation learning, without search.
2. TD-MPC2: A leading search-based MBRL method that uses a latent world model for MPC.

They experimentally show that naively adding MPC to MR.Q degrades performance, even though MR.Q learns dynamics models with accuracy comparable to (or better than) TD-MPC2. Conversely, removing MPC from TD-MPC2 also alters its performance profile. This asymmetry confirms that model accuracy alone does not determine the success of search.

Root Cause Analysis: Overestimation Bias

The authors identify overestimation bias as the primary bottleneck. In standard MBRL setups:

• The value function is trained using rollouts from a learned policy network ($\pi$).
• During execution, search (MPC) selects actions based on the value function ($\pi_{MPC}$).
• This creates a distribution shift: the value function is queried on actions ($\pi_{MPC}$) that differ from the distribution it was trained on ($\pi$).
• This mismatch leads to the value function overestimating the returns of search-selected actions, particularly for out-of-distribution (OOD) actions, which degrades overall performance.

Proposed Solution: MRS.Q

Building on these insights, the authors propose MRS.Q (Model-based Representations for Search and Q-learning). MRS.Q modifies MR.Q to enable effective search by directly addressing the overestimation bias rather than solely improving model accuracy. Key modifications include:

1. Search-Based Action Selection: MRS.Q uses Model Predictive Path Integral (MPPI) control for action selection with a short horizon (3 steps), similar to TD-MPC2.
2. Aggressive Minimization over an Ensemble: To mitigate overestimation bias, MRS.Q employs an ensemble of 10 value functions. Crucially, it takes the minimum value across the entire ensemble during both:
   • Value target computation (Bellman updates).
   • Trajectory evaluation during MPC (selecting the best action sequence).
     This contrasts with TD-MPC2, which randomly samples a pair of value functions for updates and uses the mean for MPC evaluation.
3. Hyperparameter Adjustments:
   • Removal of Exploration Noise: Since MPC inherently induces action variation, additive Gaussian exploration noise is removed.
   • Simplicial Embeddings (SEM): Applied to the state encoder and dynamics model output to stabilize multi-step rollouts.
   • Increased Terminal Loss Weight: To improve the accuracy of episode termination predictions, which is critical for value estimation.

Key Contributions

1. Challenging the "Better Model" Dogma: The paper demonstrates that search can fail even with highly accurate or perfect dynamics models. The effectiveness of search is determined more by the interaction between search and value learning than by model fidelity alone.
2. Identifying Distribution Shift as the Bottleneck: The authors show that the distribution shift introduced when using search-gathered data to train a value function trained on non-search data leads to severe overestimation bias.
3. MRS.Q Algorithm: They introduce MRS.Q, which integrates search into a representation-learning framework by using an ensemble of value functions with aggressive minimization. This approach effectively mitigates overestimation bias.
4. State-of-the-Art Performance: MRS.Q achieves superior performance across multiple benchmarks compared to both model-free and model-based baselines, including TD-MPC2 and its recent extensions (TD-M(PC)2, BMPC, BOOM).

Results

The authors evaluated MRS.Q on over 50 tasks across three benchmark domains: Gym, DeepMind Control Suite (DMC), and HumanoidBench.

• Performance: MRS.Q outperforms all baselines on three of the four benchmark subsets (Gym, HumanoidBench-No Hand, HumanoidBench-Hand) and remains competitive on DMC.
• Comparison with TD-MPC2: Despite using the same search procedure (MPPI), MRS.Q significantly outperforms TD-MPC2, particularly in high-dimensional tasks (e.g., HumanoidBench with hand features) and Gym environments.
• Ablation Studies:
  • Ensemble Size: Performance improves as the ensemble size increases up to 10 value functions; using only 2 (as in standard Double Q-learning) is insufficient to mitigate the bias.
  • Minimization Strategy: Taking the minimum over the full ensemble during MPC trajectory evaluation is critical. Using the mean (as in TD-MPC2) results in lower performance.
  • Exploration Noise: Removing additive exploration noise improves performance, confirming that MPC provides sufficient exploration.
• Computational Cost: MRS.Q incurs a higher computational cost than MR.Q (approx. 53 hours vs. 10 hours for 1M steps on HalfCheetah), primarily due to the MPC search process rather than the ensemble of value functions.

Significance and Claims

The paper claims that the bottleneck in MBRL is not just the quality of the dynamics model, but how the model is used. The authors argue that progress in MBRL requires understanding and managing the fundamental tensions between value learning and search.

• Paradigm Shift: The work suggests that simply building better models is insufficient. Instead, addressing the overestimation bias introduced by the distribution shift between search and training is the key to unlocking the benefits of search.
• Generalizability: The technique of using an ensemble minimum to mitigate search-induced overestimation is shown to be effective not only for MRS.Q but also improves TD-MPC2 when applied to it.
• Practical Impact: By enabling more effective use of learned models, MRS.Q offers a path toward more sample-efficient RL, which is critical for applications in robotics and autonomous systems where data collection is costly or safety-critical.

The authors remain modest regarding limitations, noting that the computational cost of maintaining 10 value functions and running MPC is substantial, and the aggressive minimization strategy yields smaller gains on tasks with dense rewards and no early termination (like DMC), suggesting the method is most beneficial where overestimation is severe.