The Surprising Difficulty of Search in Model-Based Reinforcement Learning

TL;DR

This work questions the assumption that better dynamics models automatically enable superior search in model-based RL, showing that search can harm performance due to distribution shift in value learning. The authors analyze why search fails even with accurate models, identifying overestimation caused by querying a value function with data collected under a search policy. They propose MRS.Q, which adds a robust, ensemble-based pessimism by using the minimum over 10 value-functions during both targets and MPC evaluation, along with short-horizon planning and SEM embeddings. Across 50+ tasks in Gym, DM Control, and HumanoidBench, MRS.Q achieves state-of-the-art results, demonstrating that addressing distribution shift and overestimation is key to unlocking the benefits of search. The work provides practical guidance for integrating search into RL, emphasizing how the value-learning pipeline should be coupled with planning to avoid degrade-and-hold scenarios.

Abstract

This paper investigates search in model-based reinforcement learning (RL). Conventional wisdom holds that long-term predictions and compounding errors are the primary obstacles for model-based RL. We challenge this view, showing that search is not a plug-and-play replacement for a learned policy. Surprisingly, we find that search can harm performance even when the model is highly accurate. Instead, we show that mitigating distribution shift matters more than improving model or value function accuracy. Building on this insight, we identify key techniques for enabling effective search, achieving state-of-the-art performance across multiple popular benchmark domains.

Figures (9)

• Figure 1: Absorbing N-chain. The absorbing N-chain is a simple environment with $N+1$ states and $A$ actions. The action $a_0$ moves the agent forward and any other action $\{a_i\}$ sends the agent to a self-looping absorbing state $s_\text{absorb}$. The reward is $0$ everywhere, other than the final state $s_{N-1}$ of the N-chain.
• Figure 2: Probability of finding a non-zero value trajectory. We analyze how the probability of discovering a non-zero value trajectory varies with the search horizon length (left) and the action space size (right) in the absorbing N-chain (\ref{['fig:nchain']}) according to $1 - (1 - \frac{1}{A^n})^m$. This demonstrates that the probability decays rapidly as either parameter increases.
• Figure 3: Value accuracy when using MPC. We report the percent error between value function estimates and true discounted returns at the end of training, where positive values indicate overestimation. Error bars represent standard deviation across 10 seeds. MR.Q+MPC replaces the behavior policy with MPC when acting in the environment. Reflecting \ref{['table:acc_perf']}, we mark environments where MPC improved MR.Q performance with + (Ant) and those where MPC harmed TD-MPC2 performance with - (dog-stand, Ant, Hopper, Humanoid, Walker2d). Adding MPC alone significantly increases overestimation compared to vanilla MR.Q, suggesting a correlation between overestimation and performance. In contrast, our approach (MRS.Q) uses MPC while significantly reducing overestimation bias.
• Figure 4: Change in selected action. We measure how the action selected by MPC or the policy network changes during learning. To do so, we measure the mean-squared difference of the selected action at intervals of 1000 time steps during learning. All actions are scaled to a range of $[-1,1]$. Across a wide range of environments, we see that actions selected by MPC change significantly more than actions selected by the policy network.
• Figure 5: Aggregate learning curves. Our approach MRS.Q offers significant performance benefits over state-of-the-art model-based RL methods. The shaded area captures a 95% stratified bootstrap confidence interval across 10 seeds. Gym scores are normalized using TD3 performance values fujimoto2018addressing (see \ref{['appendix:envs']} for details).